[EC Connector] Added Build Connector Worker Meta for EC Connector

.buildkite/test_areas/misc.yaml

@@ -127,6 +127,7 @@ steps:
     - pytest -v -s v1/test_kv_cache_spec_registry.py
     - pytest -v -s v1/cudagraph/test_cudagraph_manager.py
     - pytest -v -s -m 'cpu_test' v1/kv_connector/unit
+    - pytest -v -s -m 'cpu_test' v1/ec_connector/unit
     - pytest -v -s -m 'cpu_test' v1/metrics
 
 - label: Extract Hidden States Integration

tests/v1/ec_connector/unit/test_ec_output_aggregator.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for ECOutputAggregator."""
+
+import pytest
+
+from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata
+from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator
+from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
+from vllm.v1.outputs import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    ECConnectorOutput,
+    KVConnectorOutput,
+    ModelRunnerOutput,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+
+class FakeWorkerMeta(ECConnectorWorkerMetadata):
+    """Records merge order. `aggregate` returns a new object, as the base class
+    declares: an aggregator discarding the return value would lose the merge.
+    """
+
+    def __init__(self, saves: list[str]):
+        self.saves = saves
+
+    def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta":
+        return FakeWorkerMeta(self.saves + other.saves)
+
+
+def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput:
+    return ModelRunnerOutput(
+        req_ids=[], req_id_to_index={}, ec_connector_output=ec_output
+    )
+
+
+def test_aggregate_folds_every_rank_onto_output_rank():
+    """EC work done on any rank reaches the scheduler via output_rank's output.
+
+    The middle rank reports no worker metadata: it must neither seed nor clobber
+    the accumulator.
+    """
+    outputs = [
+        _worker_output(
+            ECConnectorOutput(
+                finished_sending={"mm0"},
+                ec_connector_worker_meta=FakeWorkerMeta(["mm0"]),
+            )
+        ),
+        _worker_output(ECConnectorOutput(finished_recving={"mm1"})),
+        _worker_output(
+            ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm2"]))
+        ),
+    ]
+
+    result = ECOutputAggregator().aggregate(outputs, output_rank=2)
+
+    assert result is outputs[2]
+    assert result.ec_connector_output.finished_sending == {"mm0"}
+    assert result.ec_connector_output.finished_recving == {"mm1"}
+    assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0", "mm2"]
+
+
+def test_aggregate_leaves_no_ec_output_when_no_worker_reported():
+    """Empty per-worker reports must not reach the scheduler as an empty object."""
+    outputs = [_worker_output(ECConnectorOutput()), _worker_output(ECConnectorOutput())]
+
+    result = ECOutputAggregator().aggregate(outputs, output_rank=0)
+
+    assert result is outputs[0]
+    assert result.ec_connector_output is None
+    assert ECOutputAggregator().aggregate([None], output_rank=0) is None
+
+
+def test_aggregate_does_not_write_through_the_shared_empty_output():
+    """A rank with nothing to report yields the shared empty output singleton.
+
+    Folding another rank's metadata onto it must not write through to the
+    module-level object, which every later step would then carry.
+    """
+    outputs = [
+        _worker_output(
+            ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm0"]))
+        ),
+        EMPTY_MODEL_RUNNER_OUTPUT,
+    ]
+
+    result = ECOutputAggregator().aggregate(outputs, output_rank=1)
+
+    assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None
+    assert result is not EMPTY_MODEL_RUNNER_OUTPUT
+    assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0"]
+
+
+def test_chaining_with_kv_aggregator_preserves_both_outputs():
+    """MultiprocExecutor chains both aggregators and keeps only the last result,
+    so each must merge onto the same output_rank output rather than replace it.
+    """
+    outputs = [
+        _worker_output(ECConnectorOutput(finished_sending={"mm0"})),
+        _worker_output(None),
+    ]
+    outputs[1].kv_connector_output = KVConnectorOutput(invalid_block_ids={7})
+
+    result = None
+    for aggregator in (
+        KVOutputAggregator(expected_finished_count=1),
+        ECOutputAggregator(),
+    ):
+        result = aggregator.aggregate(outputs, output_rank=1)
+
+    assert result is outputs[1]
+    assert result.kv_connector_output.invalid_block_ids == {7}
+    assert result.ec_connector_output.finished_sending == {"mm0"}

tests/v1/ec_connector/unit/test_worker_ec_connector.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the V2 GPU model runner's EC connector wrapper."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm.distributed.ec_transfer.ec_connector.base import (
+    ECConnectorBase,
+    ECConnectorMetadata,
+)
+from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
+from vllm.v1.worker.gpu.ec_connector import NO_OP_EC_CONNECTOR, ActiveECConnector
+
+pytestmark = pytest.mark.cpu_test
+
+WORKER_META = object()
+
+
+def _scheduler_output() -> SimpleNamespace:
+    return SimpleNamespace(
+        ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset()
+    )
+
+
+def _connector(
+    encoder_cache: dict | None = None,
+    is_producer: bool = True,
+    is_consumer: bool = False,
+) -> tuple[ActiveECConnector, MagicMock]:
+    fake = MagicMock(spec=ECConnectorBase)
+    fake.is_producer = is_producer
+    fake.is_consumer = is_consumer
+    fake.get_finished.return_value = (None, None)
+    fake.build_connector_worker_meta.return_value = WORKER_META
+    with patch("vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=fake):
+        return ActiveECConnector(SimpleNamespace(), encoder_cache or {}), fake
+
+
+@pytest.mark.parametrize(
+    ("is_producer", "is_consumer"), [(True, False), (True, True), (False, True)]
+)
+def test_saves_newly_added_caches_for_every_producer(is_producer, is_consumer):
+    """An ec_both node is also a producer: it must offload what it just computed."""
+    encoder_cache = {"mm_old": None}
+    connector, fake = _connector(encoder_cache, is_producer, is_consumer)
+
+    with connector.maybe_get_output(_scheduler_output()):
+        encoder_cache["mm_new"] = None
+
+    saved = [call.kwargs["mm_hash"] for call in fake.save_caches.call_args_list]
+    assert saved == (["mm_new"] if is_producer else [])
+    assert fake.start_load_caches.called == is_consumer
+
+
+def test_worker_meta_is_reported_on_context_exit():
+    """Reported in the finally block, so is_empty() sees it only after the exit."""
+    connector, fake = _connector()
+
+    with connector.maybe_get_output(_scheduler_output()) as output:
+        assert output.ec_connector_worker_meta is None
+
+    assert output.ec_connector_worker_meta is WORKER_META
+    assert fake.clear_connector_metadata.called
+
+
+def test_no_forward_reports_without_running_the_model():
+    connector, _ = _connector()
+
+    output = connector.no_forward(_scheduler_output())
+
+    assert output.ec_connector_output.ec_connector_worker_meta is WORKER_META
+
+    empty = NO_OP_EC_CONNECTOR.no_forward(_scheduler_output())
+    assert empty is EMPTY_MODEL_RUNNER_OUTPUT

tests/v1/executor/test_executor.py

@@ -9,6 +9,7 @@
 
 import pytest
 
+from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator
 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.sampling_params import SamplingParams
@@ -98,6 +99,7 @@ def collective_rpc(
         non_block: bool = False,
         unique_reply_rank: int | None = None,
         kv_output_aggregator: KVOutputAggregator = None,
+        ec_output_aggregator: ECOutputAggregator | None = None,
     ) -> Any | list[Any] | Future[Any | list[Any]]:
         # Drop marker to show that this was run
         with open(".marker", "w"):
@@ -110,6 +112,7 @@ def collective_rpc(
             non_block,
             unique_reply_rank,
             kv_output_aggregator,
+            ec_output_aggregator,
         )
 
 

tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py

@@ -55,6 +55,7 @@ def _run_engine_core_handshake(
     class _FakeScheduler:
         def __init__(self, **kwargs: Any) -> None:
             self.connector = connector
+            self.ec_connector = None
 
         def get_kv_connector(self) -> KVConnectorBase_V1:
             return connector

tests/v1/test_outputs.py

@@ -6,7 +6,13 @@
 import torch
 
 from vllm.platforms import current_platform
-from vllm.v1.outputs import LogprobsLists, LogprobsTensors
+from vllm.v1.outputs import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    ECConnectorOutput,
+    LogprobsLists,
+    LogprobsTensors,
+    ModelRunnerOutput,
+)
 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
 from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors
 
@@ -200,3 +206,14 @@ def test_slice_all_requests(self):
         assert len(sliced.logprob_token_ids) == 9  # All tokens
         assert sliced.logprob_token_ids == self.logprobsLists.logprob_token_ids
         assert sliced.cu_num_generated_tokens is None
+
+
+def test_with_ec_conn_output_copies_shared_empty_output():
+    """The shared empty output is copied, never written to."""
+    ec_output = ECConnectorOutput(finished_sending={"mm_hash"})
+
+    result = ModelRunnerOutput.with_ec_conn_output(EMPTY_MODEL_RUNNER_OUTPUT, ec_output)
+
+    assert result is not EMPTY_MODEL_RUNNER_OUTPUT
+    assert result.ec_connector_output is ec_output
+    assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -180,6 +180,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
         hidden_states=None,
         aux_hidden_states=None,
         finished_req_ids=set(),
+        ec_connector_output=None,
         routed_experts=None,
         num_tokens_across_dp=None,
     )

vllm/distributed/ec_transfer/ec_connector/base.py

@@ -20,6 +20,8 @@
 
         get_finished() - called with ids of finished requests, returns
             ids of requests that have completed async sending/recving.
+        build_connector_worker_meta() - builds metadata to be sent
+            back to the scheduler-side connector
 """
 
 import enum
@@ -56,6 +58,27 @@ class ECConnectorMetadata(ABC):  # noqa: B024
     pass
 
 
+class ECConnectorWorkerMetadata(ABC):
+    """
+    Abstract Metadata used to communicate back
+    Worker ECConnector -> Scheduler ECConnector.
+
+    Each worker can output its own metadata.
+    For a single engine step, all metadata objects returned by workers
+    will be aggregated using the `aggregate` method below, before
+    being passed to the Scheduler ECConnector.
+    """
+
+    @abstractmethod
+    def aggregate(
+        self, other: "ECConnectorWorkerMetadata"
+    ) -> "ECConnectorWorkerMetadata":
+        """
+        Aggregate metadata with another `ECConnectorWorkerMetadata` object.
+        """
+        pass
+
+
 class ECConnectorBase(ABC):
     def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole):
         self._connector_metadata: ECConnectorMetadata | None = None
@@ -190,6 +213,16 @@ def get_finished(
         """
         return None, None
 
+    def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None:
+        """
+        Build the ECConnector worker metadata for this engine step.
+
+        Returns:
+            ECConnectorWorkerMetadata: the worker metadata.
+            None if no worker metadata is available.
+        """
+        return None
+
     # ==============================
     # Scheduler-side methods
     # ==============================

vllm/distributed/ec_transfer/ec_connector/utils.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""EC connector helper utilities."""
+
+from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput
+
+
+class ECOutputAggregator:
+    """Merge every worker's EC connector output onto the single
+    ModelRunnerOutput that reaches the scheduler.
+
+    Mirrors KVOutputAggregator: only `output_rank`'s output is returned to the
+    scheduler, but the EC connector may have run on any rank.
+    """
+
+    def aggregate(
+        self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0
+    ) -> ModelRunnerOutput | None:
+        output = outputs[output_rank]
+        if not output:
+            return None
+
+        finished_sending = set[str]()
+        finished_recving = set[str]()
+        worker_meta = None
+        for model_runner_output in outputs:
+            assert model_runner_output is not None
+            ec_output = model_runner_output.ec_connector_output
+            if not ec_output:
+                continue
+
+            finished_sending |= ec_output.finished_sending or set()
+            finished_recving |= ec_output.finished_recving or set()
+
+            if meta := ec_output.ec_connector_worker_meta:
+                worker_meta = (
+                    meta if worker_meta is None else worker_meta.aggregate(meta)
+                )
+
+        aggregated = ECConnectorOutput(
+            finished_sending=finished_sending or None,
+            finished_recving=finished_recving or None,
+            ec_connector_worker_meta=worker_meta,
+        )
+        if aggregated.is_empty():
+            output.ec_connector_output = None
+            return output
+
+        # `output` is the shared empty output whenever `output_rank` had no work,
+        # so attach through the copy-on-write helper.
+        return ModelRunnerOutput.with_ec_conn_output(output, aggregated)

vllm/v1/engine/core.py

@@ -173,6 +173,8 @@ def __init__(
         )
         if self.scheduler.connector is not None:  # type: ignore
             self.model_executor.init_kv_output_aggregator(self.scheduler.connector)  # type: ignore
+        if self.scheduler.ec_connector is not None:  # type: ignore
+            self.model_executor.init_ec_output_aggregator()
 
         mm_registry = MULTIMODAL_REGISTRY
         self.mm_receiver_cache = mm_registry.engine_receiver_cache_from_config(