[EC Connector] Call to EC Connector update_connector_output from scheduler (#49579)
Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Signed-off-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
omerpaz95 committed
Aug 10, 2026 at 09:25 UTC
61c1dd0966fd8d91954003cf6b7ab746448a46f8
3 files changed
+43
tests/v1/core/test_async_scheduler.py
+1
@@ -272,6 +272,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
272
273
scheduler.perf_metrics = None
274
scheduler.connector = None
275
+ scheduler.ec_connector = None
276
scheduler.structured_output_manager = Mock()
277
scheduler.structured_output_manager.should_advance.return_value = True
278
scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = (
tests/v1/core/test_scheduler.py
+37
@@ -40,6 +40,7 @@ from vllm.v1.kv_cache_interface import (
40
)
41
from vllm.v1.outputs import (
42
DraftTokenIds,
43
+ ECConnectorOutput,
44
KVConnectorOutput,
45
ModelRunnerOutput,
46
make_empty_encoder_model_runner_output,
@@ -3237,6 +3238,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
3238
3239
scheduler.perf_metrics = None
3240
scheduler.connector = None
3241
+ scheduler.ec_connector = None
3242
scheduler.structured_output_manager = Mock()
3243
scheduler.structured_output_manager.should_advance.return_value = True
3244
scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = (
@@ -4771,6 +4773,41 @@ def test_scheduler_kv_connector_stats():
4773
assert final_stats == expected_data
4774
4775
4776
+def test_ec_connector_update_connector_output_called():
4777
+ """Test that worker-side EC connector output is forwarded to the
4778
+ EC connector's update_connector_output hook."""
4779
+ scheduler = create_scheduler(
4780
+ model="llava-hf/llava-1.5-7b-hf",
4781
+ use_ec_connector=True,
4782
+ ec_role="ec_consumer",
4783
+ )
4784
+ scheduler.ec_connector.update_connector_output = Mock()
4785
+
4786
+ scheduler_output = SchedulerOutput(
4787
+ scheduled_new_reqs=[],
4788
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
4789
+ num_scheduled_tokens={},
4790
+ total_num_scheduled_tokens=0,
4791
+ scheduled_encoder_inputs={},
4792
+ scheduled_spec_decode_tokens={},
4793
+ num_common_prefix_blocks=[],
4794
+ finished_req_ids=set(),
4795
+ free_encoder_mm_hashes=[],
4796
+ )
4797
+ ec_connector_output = ECConnectorOutput(finished_sending={"hash_test1"})
4798
+ model_runner_output = ModelRunnerOutput(
4799
+ req_ids=[],
4800
+ req_id_to_index={},
4801
+ ec_connector_output=ec_connector_output,
4802
+ )
4803
+
4804
+ scheduler.update_from_output(scheduler_output, model_runner_output)
4805
+
4806
+ scheduler.ec_connector.update_connector_output.assert_called_once_with(
4807
+ ec_connector_output
4808
+ )
4809
+
4810
+
4811
# ==============================================================================
4812
# Variable-length encoder cross-attention block allocation tests
4813
# ==============================================================================
vllm/v1/core/sched/scheduler.py
+5
@@ -1672,6 +1672,7 @@ class Scheduler(SchedulerInterface):
1672
pooler_outputs = model_runner_output.pooler_output
1673
num_nans_in_logits = model_runner_output.num_nans_in_logits
1674
kv_connector_output = model_runner_output.kv_connector_output
1675
+ ec_connector_output = model_runner_output.ec_connector_output
1676
cudagraph_stats = model_runner_output.cudagraph_stats
1677
1678
# Every GPU write enqueued by this and earlier steps has completed, so it is
@@ -1979,6 +1980,10 @@ class Scheduler(SchedulerInterface):
1980
if kv_connector_output:
1981
self._update_from_kv_xfer_finished(kv_connector_output)
1982
1983
+ # EC Connector: update state from worker-side EC connector output.
1984
+ if self.ec_connector is not None and ec_connector_output:
1985
+ self.ec_connector.update_connector_output(ec_connector_output)
1986
+
1987
# Worker-side KV connector stats from the model runner output.
1988
kv_connector_stats: KVConnectorStats | None = (
1989
kv_connector_output.kv_connector_stats if kv_connector_output else None