@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f1aeff6007

[Bugfix][EC Connector] Fail the request, not the engine, when a remote encoding cannot arrive (#55290)

Signed-off-by: Tianyu Guo <guoty@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>

Tianyu Guo committed Sep 9, 2026 at 11:13 UTC f1aeff60073dda82a423d7d081f132c8a458e16b
8 files changed +536 -104
tests/v1/ec_connector/integration/README.md
+20
@@ -58,6 +58,26 @@ EC_SHARED_STORAGE_PATH="/tmp/my_ec_cache" bash ./tests/v1/ec_connector/integrati
58
59 ## How It Works
60
61 +### NIXL EC failure isolation (1E + 1PD)
62 +
63 +```bash
64 +PATH="$PWD/.venv/bin:$PATH" CUDA_VISIBLE_DEVICES=0,1 \
65 + .venv/bin/python -m pytest \
66 + tests/v1/ec_connector/integration/test_nixl_failure.py -v -s
67 +```
68 +
69 +Requires two CUDA GPUs and NIXL. `MODEL` overrides the default
70 +`Qwen/Qwen2.5-VL-3B-Instruct`. The test starts real E/PD servers with CUDA
71 +graphs enabled and uses the proxy's rewrite helper. Only the failing request's
72 +control endpoint is replaced: a ZMQ peer returns `NACK_MISSING` and verifies
73 +that the consumer actually requested the encoding.
74 +
75 +Metadata-only requests must return a request-level error; requests retaining
76 +the image must succeed through local fallback. After each case, a fresh
77 +metadata-only image request must succeed through the real NIXL transfer path.
78 +No server restart, proxy retry, cache hit, or timeout race can mask the result.
79 +This is a correctness test, not a throughput benchmark.
80 +
81 ### Mooncake EC (1E + 1PD)
82
83 ```bash
tests/v1/ec_connector/integration/test_nixl_failure.py new
+202
@@ -0,0 +1,202 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""NIXL request failure isolation with real E/PD servers and an injected NACK.
4 +
5 +Requires two CUDA GPUs and NIXL. Set MODEL to override the small default model.
6 +"""
7 +
8 +import asyncio
9 +import copy
10 +import json
11 +import os
12 +import uuid
13 +from pathlib import Path
14 +
15 +import aiohttp
16 +import msgspec
17 +import pytest
18 +import zmq
19 +import zmq.asyncio
20 +from PIL import Image
21 +
22 +from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import (
23 + rewrite_for_decode,
24 +)
25 +from tests.utils import RemoteOpenAIServer
26 +from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import (
27 + XferAck,
28 + XferReq,
29 + XferStatus,
30 +)
31 +from vllm.multimodal.utils import encode_image_url
32 +from vllm.platforms import current_platform
33 +from vllm.utils.network_utils import get_open_port
34 +
35 +MODEL = os.environ.get("MODEL", "Qwen/Qwen2.5-VL-3B-Instruct")
36 +
37 +
38 +@pytest.fixture(scope="module")
39 +def servers():
40 + pytest.importorskip("nixl")
41 + if not current_platform.is_cuda() or current_platform.device_count() < 2:
42 + pytest.skip("Requires two CUDA GPUs")
43 + devices = os.environ.get("CUDA_VISIBLE_DEVICES", "0,1").split(",")
44 + common = [
45 + "--host",
46 + "127.0.0.1",
47 + "--dtype",
48 + "bfloat16",
49 + "--max-model-len",
50 + "4096",
51 + "--max-num-seqs",
52 + "4",
53 + "--gpu-memory-utilization",
54 + "0.7",
55 + "--no-enable-prefix-caching",
56 + ]
57 +
58 + def config(role):
59 + return json.dumps(
60 + {
61 + "ec_connector": "ECCPUConnector",
62 + "ec_role": role,
63 + "ec_connector_extra_config": {
64 + "ec_enable_nixl": True,
65 + "ec_cpu_bytes": 64 << 20,
66 + },
67 + }
68 + )
69 +
70 + with (
71 + RemoteOpenAIServer(
72 + MODEL,
73 + common
74 + + [
75 + "--mm-encoder-only",
76 + "--disable-hybrid-kv-cache-manager",
77 + "--ec-transfer-config",
78 + config("ec_producer"),
79 + ],
80 + env_dict={
81 + "CUDA_VISIBLE_DEVICES": devices[0],
82 + "VLLM_EC_SIDE_CHANNEL_HOST": "127.0.0.1",
83 + "VLLM_EC_SIDE_CHANNEL_PORT": str(get_open_port()),
84 + },
85 + ) as encoder,
86 + RemoteOpenAIServer(
87 + MODEL,
88 + common
89 + + ["--enable-mm-embeds", "--ec-transfer-config", config("ec_consumer")],
90 + env_dict={"CUDA_VISIBLE_DEVICES": devices[1]},
91 + ) as consumer,
92 + ):
93 + yield encoder, consumer
94 +
95 +
96 +async def _post(session, server, body):
97 + async with session.post(server.url_for("v1/chat/completions"), json=body) as r:
98 + return r.status, await r.json()
99 +
100 +
101 +async def _encode(session, encoder):
102 + image = Image.open(Path(__file__).with_name("hato.jpg")).convert("RGB")
103 + image.thumbnail((224, 224))
104 + item_uuid = uuid.uuid4().hex # Avoid both EC and processor cache hits.
105 + body = {
106 + "model": MODEL,
107 + "messages": [
108 + {
109 + "role": "user",
110 + "content": [
111 + {
112 + "type": "image_url",
113 + "image_url": {"url": encode_image_url(image)},
114 + "uuid": item_uuid,
115 + },
116 + {"type": "text", "text": "Describe the image briefly."},
117 + ],
118 + }
119 + ],
120 + "max_tokens": 16,
121 + "temperature": 0,
122 + "chat_template_kwargs": {"enable_thinking": False},
123 + }
124 + status, result = await _post(session, encoder, body)
125 + assert status == 200, result
126 + params = result["ec_transfer_params"]
127 + assert len(params) == 1, params
128 + ((mm_hash, handle),) = params.items()
129 + assert handle["metadata"] and handle["peer_host"] and handle["size_bytes"] > 0
130 + rewritten = rewrite_for_decode(
131 + body,
132 + {
133 + 0: {
134 + **handle["metadata"],
135 + "mm_hash": item_uuid,
136 + "ec_mm_hash": mm_hash,
137 + }
138 + },
139 + )
140 + assert rewritten["messages"][0]["content"][0]["type"] == "image_embeds"
141 + for request in (body, rewritten):
142 + request["ec_transfer_params"] = {item_uuid: copy.deepcopy(handle)}
143 + return body, rewritten, item_uuid
144 +
145 +
146 +@pytest.mark.asyncio
147 +@pytest.mark.parametrize("local_image", [False, True], ids=["metadata-only", "pixels"])
148 +async def test_missing_remote_encoding_is_request_scoped(servers, local_image):
149 + """A terminal NACK must preserve local fallback and subsequent serving."""
150 + encoder, consumer = servers
151 + timeout = aiohttp.ClientTimeout(total=120)
152 + async with aiohttp.ClientSession(timeout=timeout) as session:
153 + original, rewritten, mm_hash = await _encode(session, encoder)
154 + request = original if local_image else rewritten
155 +
156 + # Keep the real encoder's shape/metadata; replace only the control peer.
157 + # The socket stays bound until completion, so no port or timing race is
158 + # needed to make the consumer receive NACK_MISSING.
159 + with zmq.asyncio.Context() as ctx, ctx.socket(zmq.ROUTER) as router:
160 + router.setsockopt(zmq.LINGER, 0)
161 + port = router.bind_to_random_port("tcp://127.0.0.1")
162 + request["ec_transfer_params"][mm_hash].update(
163 + peer_host="127.0.0.1", peer_port=port
164 + )
165 +
166 + async def reject_read():
167 + identity, delimiter, payload = await router.recv_multipart()
168 + req = msgspec.msgpack.decode(payload, type=XferReq)
169 + assert req.mm_hash == mm_hash
170 + ack = XferAck(
171 + mm_hash=req.mm_hash,
172 + status=XferStatus.NACK_MISSING,
173 + session_id=req.session_id,
174 + )
175 + await router.send_multipart(
176 + [identity, delimiter, msgspec.msgpack.encode(ack)]
177 + )
178 +
179 + rejection = asyncio.create_task(asyncio.wait_for(reject_read(), 30))
180 + try:
181 + status, result = await _post(session, consumer, request)
182 + # A frontend rejection must not count as successful injection.
183 + await rejection
184 + finally:
185 + rejection.cancel()
186 + await asyncio.gather(rejection, return_exceptions=True)
187 +
188 + if local_image:
189 + assert status == 200, result
190 + assert result["choices"][0]["message"]["content"], result
191 + else:
192 + assert status == 500, result
193 + assert result["error"]["type"] == "InternalServerError", result
194 +
195 + async with session.get(consumer.url_for("health")) as health:
196 + assert health.status == 200, "The failed read killed EngineCore"
197 + # A fresh metadata-only image must cross the real NIXL data path after
198 + # the fault, not just pass a frontend health check or hit a local cache.
199 + _, healthy, _ = await _encode(session, encoder)
200 + status, result = await _post(session, consumer, healthy)
201 + assert status == 200, result
202 + assert result["choices"][0]["message"]["content"], result
tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py
+117 -35
@@ -1,13 +1,22 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import time
4 import uuid
5
6 +import pytest
7 +import torch
8 +
9 import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
10 from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
11 from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
12 ECSharedRegion,
13 )
14 from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler
15 +from vllm.multimodal.inputs import (
16 + MultiModalFieldElem,
17 + MultiModalKwargsItem,
18 + MultiModalSharedField,
19 +)
20
21 _N, _BS, _HID, _ES = 16, 64, 32, 2
22
@@ -18,10 +27,12 @@ class _Pos:
27
28
29 class _Feature:
21 - def __init__(self, mm_hash, length=1):
30 + def __init__(self, mm_hash, length=1, data=None):
31 self.mm_hash = mm_hash
32 self.identifier = mm_hash
33 self.mm_position = _Pos(0, length)
34 + self.data = data
35 + self.modality = "image"
36
37
38 class _Request:
@@ -94,6 +105,7 @@ def _consumer_sched(monkeypatch):
105 # directly since this helper builds gate-off then flips fields on.
106 s._hidden_dim = _HID
107 s._element_size = _ES
108 + s._metadata_resolver._cache["image"] = {"image_grid_thw"}
109 return s
110
111
@@ -150,13 +162,17 @@ def test_new_remote_read_defers_then_completes(monkeypatch):
162 s.shutdown()
163
164
153 -def test_no_params_is_ready(monkeypatch):
165 +@pytest.mark.parametrize("params", [None, {"x": 123}])
166 +def test_no_remote_announcement_is_ready(monkeypatch, params):
167 s = _consumer_sched(monkeypatch)
155 - assert s.ensure_cache_available(_Request([_Feature("x")], params=None), 0) is True
168 + assert s.ensure_cache_available(_Request([_Feature("x")], params=params), 0)
169 + assert s.take_unavailable_requests() == set()
170 s.shutdown()
171
172
159 -def test_tombstoned_read_discards_and_blocks_retry(monkeypatch):
173 +@pytest.mark.parametrize("payload", [None, "metadata", "pixels", "embeds", "shm"])
174 +def test_failed_read_falls_back_only_with_local_input(monkeypatch, payload):
175 + """A remote miss is fatal only when local model input is unavailable."""
176 s = _consumer_sched(monkeypatch)
177 fake = _FakeSession()
178
@@ -167,26 +183,71 @@ def test_tombstoned_read_discards_and_blocks_retry(monkeypatch):
183 return True
184
185 monkeypatch.setattr(s, "_start_xfer", _fake_start)
170 - req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
186 + fields = {
187 + "metadata": {"image_grid_thw": torch.tensor([1, 2, 2])},
188 + "pixels": {"pixel_values": torch.ones(1, 3)},
189 + "embeds": {"image_embeds": torch.ones(1, _HID)},
190 + "shm": {"address": 123, "monotonic_id": 1},
191 + }
192 + data = (
193 + MultiModalKwargsItem(
194 + {
195 + key: MultiModalFieldElem(value, MultiModalSharedField(batch_size=1))
196 + for key, value in fields[payload].items()
197 + }
198 + )
199 + if payload is not None
200 + else None
201 + )
202 + req = _Request([_Feature("h1", 1, data)], params=_params("h1", 1))
203
204 # Step 1: start the read.
205 assert s.ensure_cache_available(req, 0) is False
206 s.build_connector_meta(scheduler_output=None)
207
176 - # Producer rejected the read -> tombstoned.
208 fake._results.tombstoned.add("h1")
209 s._sessions[("h", 1)] = fake
210
180 - # Step 2: poll discards the entry and records a tombstone; the tombstone
181 - # is consumed by admit so the request proceeds to local compute.
182 - assert s.ensure_cache_available(req, 0) is True
211 + can_fallback = payload in ("pixels", "embeds", "shm")
212 + assert s.ensure_cache_available(req, 0) is can_fallback
213 assert s._cache.get("h1") is None
214 assert "h1" not in s._in_flight
185 - s.build_connector_meta(scheduler_output=None)
215 + assert s.take_unavailable_requests() == (set() if can_fallback else {"r1"})
216 + # Drained by the read, so the scheduler cannot abort it twice.
217 + assert s.take_unavailable_requests() == set()
218 + if can_fallback:
219 + # A later scheduling step must not retry the failed remote source.
220 + s.build_connector_meta(scheduler_output=None)
221 + assert s.ensure_cache_available(req, 0)
222 + assert fake.started == ["h1"]
223 + s.shutdown()
224
187 - # Step 3: tombstone was consumed, so a fresh transfer starts again.
188 - assert s.ensure_cache_available(req, 0) is False
189 - assert "h1" in s._in_flight
225 +
226 +@pytest.mark.parametrize("retry", [False, True], ids=["in-flight", "retryable"])
227 +def test_remote_wait_budget_survives_retries_and_long_steps(monkeypatch, retry):
228 + """Neither retries nor delayed scheduling may restart the request budget."""
229 + s = _consumer_sched(monkeypatch)
230 + fake = _FakeSession()
231 + s._sessions[("h", 1)] = fake
232 +
233 + def start(mm_hash, info, size):
234 + assert s._cache.alloc(mm_hash, 1) is not None
235 + fake.started.append(mm_hash)
236 + return True
237 +
238 + monkeypatch.setattr(s, "_start_xfer", start)
239 + now = [1000.0]
240 + monkeypatch.setattr(time, "monotonic", lambda: now[0])
241 + req = _Request([_Feature("h1")], params=_params("h1", 1))
242 + assert not s.ensure_cache_available(req, 0)
243 + for elapsed in (20, 40, 121):
244 + now[0] = 1000.0 + elapsed
245 + if retry:
246 + fake._results.retryable.add("h1")
247 + s.build_connector_meta(scheduler_output=None)
248 + assert not s.ensure_cache_available(req, 0)
249 + assert s.take_unavailable_requests() == ({"r1"} if elapsed > 60 else set())
250 + assert len(fake.started) == (3 if retry else 1)
251 s.shutdown()
252
253
@@ -225,17 +286,50 @@ def test_retryable_read_re_requests_without_admitting(monkeypatch):
286 s.shutdown()
287
288
228 -def test_size_mismatch_skips_transfer(monkeypatch):
289 +@pytest.mark.parametrize(
290 + "size_fields",
291 + [{"size_bytes": 999}, {"size_bytes": None}, {"size_bytes": "big"}, {}],
292 + ids=["mismatch", "null-size", "text-size", "no-size"],
293 +)
294 +def test_invalid_remote_size_fails_the_request(monkeypatch, size_fields):
295 + """Invalid remote sizes fail the request without raising in the scheduler."""
296 s = _consumer_sched(monkeypatch)
230 - # Advertised size disagrees with pos.length * hidden_dim * element_size.
231 - bad = {"h1": {"peer_host": "h", "peer_port": 1, "size_bytes": 999}}
232 - req = _Request([_Feature("h1", 1)], params=bad)
233 - assert s.ensure_cache_available(req, 0) is True
297 + announced = {"peer_host": "h", "peer_port": 1, **size_fields}
298 + req = _Request([_Feature("h1", 1)], params={"h1": announced})
299 + assert s.ensure_cache_available(req, 0) is False
300 + assert s.take_unavailable_requests() == {"r1"}
301 assert "h1" not in s._in_flight
302 assert s._cache.get("h1") is None
303 s.shutdown()
304
305
306 +def test_deferral_budget_is_per_request(monkeypatch):
307 + """Shared hashes have independent wait budgets and cancellation cleanup."""
308 + s = _consumer_sched(monkeypatch)
309 + monkeypatch.setattr(s._cache, "alloc", lambda key, n: None)
310 + params = _params("h1", 1)
311 + old = _Request([_Feature("h1", 1)], params=params, req_id="old")
312 + new = _Request([_Feature("h1", 1)], params=params, req_id="new")
313 +
314 + now = [1000.0]
315 + monkeypatch.setattr(time, "monotonic", lambda: now[0])
316 + assert s.ensure_cache_available(old, 0) is False
317 + now[0] += sched_mod._ADMIT_DEFER_TIMEOUT_S - 1
318 + assert s.ensure_cache_available(new, 0) is False
319 + assert s.take_unavailable_requests() == set()
320 +
321 + now[0] += 2
322 + # `old` is now past its budget; `new` is not.
323 + assert s.ensure_cache_available(new, 0) is False
324 + assert s.take_unavailable_requests() == set()
325 + assert s.request_finished(new) == (False, None)
326 + assert set(s._deferred_since) == {("old", "h1")}
327 + assert s.ensure_cache_available(old, 0) is False
328 + assert s.take_unavailable_requests() == {"old"}
329 + assert "h1" not in s._in_flight
330 + s.shutdown()
331 +
332 +
333 def test_metadata_only_entry_admits_without_transfer(monkeypatch):
334 # A producer that reported placeholder metadata but had no cache entry
335 # to transfer omits peer_host/peer_port/size_bytes entirely; the consumer
@@ -250,7 +344,7 @@ def test_metadata_only_entry_admits_without_transfer(monkeypatch):
344 s.shutdown()
345
346
253 -def test_orphan_not_ready_entry_falls_back_no_realloc(monkeypatch):
347 +def test_orphan_not_ready_entry_defers_no_realloc(monkeypatch):
348 # Post-quarantine, post-tombstone-consumed orphan: a not-ready cache
349 # entry exists for the mm_hash, but the hash is tracked in none of
350 # _in_flight / _step_completed / _tombstones. Its blocks are held by a
@@ -274,25 +368,13 @@ def test_orphan_not_ready_entry_falls_back_no_realloc(monkeypatch):
368 monkeypatch.setattr(s._cache, "alloc", _spy_alloc)
369 req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
370
277 - # Must not raise; request falls back to local compute (admitted True),
278 - # the orphan entry is untouched, and alloc is never called again.
279 - assert s.ensure_cache_available(req, 0) is True
371 + # Must not raise; the request is deferred while the DMA settles, the
372 + # orphan entry is untouched, and alloc is never called again.
373 + assert s.ensure_cache_available(req, 0) is False
374 assert calls == []
375 assert s._cache.get("h1") is entry
376 assert not entry.ready
283 - s.shutdown()
284 -
285 -
286 -def test_alloc_failure_falls_back_to_local(monkeypatch):
287 - s = _consumer_sched(monkeypatch)
288 -
289 - # Force the cache to reject the allocation.
290 - monkeypatch.setattr(s._cache, "alloc", lambda key, n: None)
291 - req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
292 -
293 - # _start_xfer returns False -> request admitted for local recompute.
294 - assert s.ensure_cache_available(req, 0) is True
295 - assert "h1" not in s._in_flight
377 + assert s.take_unavailable_requests() == set()
378 s.shutdown()
379
380
tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py
+1
@@ -63,6 +63,7 @@ def test_request_finished_producer_emits_params(monkeypatch):
63
64 delay, params = s.request_finished(_Request([_Feature("h1", length=2)]))
65 assert delay is False
66 + assert entry.evictable
67 assert params == {
68 "h1": {
69 "metadata": {},
tests/v1/ec_connector/unit/test_session.py
+5 -3
@@ -451,14 +451,16 @@ def test_consumer_session_poll_not_ready_nack_goes_to_retryable():
451 assert not results.completed
452
453
454 -def test_consumer_session_take_results_clears_state():
454 +def test_consumer_session_ack_timeout_reports_retry_once():
455 + """A delayed producer ACK is retryable, and its result is drained once."""
456 s = _make_consumer_session()
457 s.start_xfer("h1", [0, 1], deadline=time.monotonic() - 1)
458 s.poll([], time.monotonic())
459 r1 = s.take_results()
459 - assert "h1" in r1.tombstoned
460 + assert "h1" in r1.retryable
461 + assert not r1.tombstoned
462 r2 = s.take_results()
461 - assert not r2.tombstoned # cleared after first take
463 + assert not r2.retryable # cleared after first take
464
465
466 def test_consumer_session_on_peer_down_cancels_waiting_ack():
vllm/distributed/ec_transfer/ec_connector/cpu/connector.py
+4
@@ -104,6 +104,10 @@ class ECCPUConnector(ECConnectorBase):
104 request, num_computed_tokens
105 )
106
107 + def take_unavailable_requests(self) -> set[str]:
108 + assert self.connector_scheduler is not None
109 + return self.connector_scheduler.take_unavailable_requests()
110 +
111 def update_state_after_alloc(self, request: "Request", index: int) -> None:
112 assert self.connector_scheduler is not None
113 self.connector_scheduler.update_state_after_alloc(request, index)
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py
+174 -64
@@ -7,6 +7,7 @@ Owns the mmap region and the embedding cache, and handles the producer
7 for the ECCPUConnector.
8 """
9
10 +import time
11 from collections.abc import Collection
12 from typing import TYPE_CHECKING, Any
13
@@ -29,12 +30,20 @@ from vllm.logger import init_logger
30
31 if TYPE_CHECKING:
32 from vllm.config import VllmConfig
33 + from vllm.multimodal.inputs import MultiModalFeatureSpec
34 from vllm.v1.core.sched.output import SchedulerOutput
35 from vllm.v1.outputs import ECConnectorOutput
36 from vllm.v1.request import Request
37
38 logger = init_logger(__name__)
39
40 +# Total remote wait budget per request item, including all retries.
41 +_ADMIT_DEFER_TIMEOUT_S = 60.0
42 +
43 +
44 +class _RemoteUnavailable(Exception):
45 + """An item cannot be fetched within its request's remote wait budget."""
46 +
47
48 class ECCPUScheduler:
49 """Scheduler delegate for the ECCPUConnector."""
@@ -92,6 +101,11 @@ class ECCPUScheduler:
101 self._in_flight: set[str] = set()
102 self._tombstones: set[str] = set()
103 self._step_completed: set[str] = set()
104 + # Per-request clocks survive retries and are cleared on completion.
105 + self._deferred_since: dict[tuple[str, str], float] = {}
106 + # Requests needing a remote encoding that will not arrive. Drained by
107 + # take_unavailable_requests(); the scheduler aborts them.
108 + self._unrecoverable: set[str] = set()
109 self._peer_host: str | None = None
110 self._peer_port: int | None = None
111 # Model shape for size checks + compat hash; only set by
@@ -117,6 +131,7 @@ class ECCPUScheduler:
131 compute_ec_compatibility_hash,
132 )
133 from vllm.distributed.ec_transfer.ec_connector.cpu.session import (
134 + PRODUCER_PIN_LEASE_S,
135 ProducerSession,
136 )
137 from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
@@ -136,7 +151,9 @@ class ECCPUScheduler:
151 # its scheduler step, so its reply latency scales with the encoder's
152 # batch size: a deployment whose steps run longer must raise this.
153 self._ack_timeout_s = float(
139 - self._ec_config.get_from_extra_config("consumer_ack_timeout_s", 2.0)
154 + self._ec_config.get_from_extra_config(
155 + "consumer_ack_timeout_s", PRODUCER_PIN_LEASE_S
156 + )
157 )
158 self._compat_hash = compute_ec_compatibility_hash(
159 vllm_version=VLLM_VERSION,
@@ -203,79 +220,172 @@ class ECCPUScheduler:
220 def _nixl_consumer_admit(
221 self, request: "Request", num_computed_tokens: int
222 ) -> bool:
206 - """Admit a request only once all its remote encodings are cached.
223 + """Admit a request once its required remote encodings are ready.
224 +
225 + Each item has a per-request wait budget that survives transfer retries.
226 + Unavailable items fall back to local input; without it, the request is
227 + reported by take_unavailable_requests() for the scheduler to abort.
228
208 - Returns True when every encoding the request needs is either already
209 - ready in the local cache or has no remote source. Any encoding whose
210 - NIXL READ is still in flight (or was just started this step) defers
211 - the request, which the scheduler re-presents on a later step.
229 + Returns:
230 + True if no item needs to wait for a remote encoding. False if any
231 + item is pending or the request must be aborted.
232 """
213 - params: dict[str, dict[str, Any]] = (
214 - getattr(request, "ec_transfer_params", None) or {}
215 - )
216 - if not params:
233 + if not request.ec_transfer_params:
234 return True
235 + now = time.monotonic()
236 pending = False
237 for feature in request.mm_features:
220 - pos = feature.mm_position
221 - if pos.offset + pos.length <= num_computed_tokens:
222 - continue
223 - mm_hash = feature.identifier
224 - entry = self._cache.get(mm_hash)
225 - if entry is not None and entry.ready:
226 - # Local hit: upstream's update_state_after_alloc pins and
227 - # loads it through the same path as a natively cached entry.
228 - continue
229 - if mm_hash in self._in_flight:
230 - pending = True
231 - continue
232 - if mm_hash in self._step_completed:
233 - pending = True
234 - continue
235 - if mm_hash in self._tombstones:
236 - self._tombstones.discard(mm_hash)
237 - continue
238 - if entry is not None:
239 - # Present but not ready and not being fetched: its blocks are
240 - # held by a quarantined/settling DMA and cannot be reused.
241 - # Fall back to local compute this step; never re-alloc a
242 - # mm_hash already in the cache.
243 - continue
244 - info = params.get(mm_hash)
245 - # Absent entirely, or present with only placeholder metadata (no
246 - # cache entry existed on the producer side at request_finished
247 - # time): nothing to fetch, fall back to local compute.
248 - if not info or "peer_host" not in info:
249 - continue
250 - expected = pos.length * self._hidden_dim * self._element_size
251 - if int(info.get("size_bytes", -1)) != expected:
252 - logger.warning(
253 - "EC consumer: size mismatch mm_hash=%s; local encode", mm_hash
254 - )
255 - continue
238 try:
257 - started = self._start_xfer(mm_hash, info, expected)
258 - except Exception:
259 - logger.exception(
260 - "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash
261 - )
262 - continue
263 - if not started:
264 - continue
265 - self._in_flight.add(mm_hash)
266 - pending = True
239 + if not self._admit_item(request, feature, num_computed_tokens, now):
240 + pending = True
241 + except _RemoteUnavailable as error:
242 + if not self._fail_or_fallback(request, feature, str(error)):
243 + return False
244 return not pending
245
246 + def _admit_item(
247 + self,
248 + request: "Request",
249 + feature: "MultiModalFeatureSpec",
250 + num_computed_tokens: int,
251 + now: float,
252 + ) -> bool:
253 + """Check one encoding, starting a remote read if needed.
254 +
255 + Returns:
256 + True if already computed, cached, or not remotely sourced. False
257 + while waiting for a read or cache space, including across retries.
258 +
259 + Raises:
260 + _RemoteUnavailable: The remote source is invalid, the read failed,
261 + or the item's total wait budget expired.
262 + """
263 + pos = feature.mm_position
264 + mm_hash = feature.identifier
265 + if pos.offset + pos.length <= num_computed_tokens:
266 + self._deferred_since.pop((request.request_id, mm_hash), None)
267 + return True
268 + announced = (request.ec_transfer_params or {}).get(mm_hash)
269 + # Without a producer address there is nothing to fetch: the request
270 + # carries what the model needs and the encoder runs locally.
271 + # `ec_transfer_params` reaches us from the request, so its shape is
272 + # checked rather than assumed.
273 + remote: dict[str, Any] | None = (
274 + announced
275 + if isinstance(announced, dict) and "peer_host" in announced
276 + else None
277 + )
278 +
279 + entry = self._cache.get(mm_hash)
280 + if entry is not None and entry.ready:
281 + # Local hit: upstream's update_state_after_alloc pins and
282 + # loads it through the same path as a natively cached entry.
283 + self._deferred_since.pop((request.request_id, mm_hash), None)
284 + return True
285 + if remote is None:
286 + return True
287 + since = self._deferred_since.setdefault((request.request_id, mm_hash), now)
288 + if now - since > _ADMIT_DEFER_TIMEOUT_S:
289 + raise _RemoteUnavailable(f"the remote wait exceeded {now - since:.0f}s")
290 + if mm_hash in self._in_flight or mm_hash in self._step_completed:
291 + return False
292 +
293 + if mm_hash in self._tombstones:
294 + self._tombstones.discard(mm_hash)
295 + raise _RemoteUnavailable("the remote read failed")
296 +
297 + if entry is not None:
298 + # Present but not ready and not being fetched: its blocks are
299 + # held by a quarantined/settling DMA and cannot be reused.
300 + return False
301 +
302 + expected = pos.length * self._hidden_dim * self._element_size
303 + try:
304 + size = int(remote["size_bytes"])
305 + except (KeyError, TypeError, ValueError) as error:
306 + raise _RemoteUnavailable("the announced size was unusable") from error
307 + if size != expected:
308 + logger.warning(
309 + "EC consumer: size mismatch mm_hash=%s announced=%d expected=%d",
310 + mm_hash,
311 + size,
312 + expected,
313 + )
314 + raise _RemoteUnavailable("the announced size was wrong")
315 +
316 + try:
317 + started = self._start_xfer(mm_hash, remote, expected)
318 + except Exception as error:
319 + logger.exception(
320 + "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash
321 + )
322 + raise _RemoteUnavailable("the read could not be started") from error
323 + if started:
324 + self._in_flight.add(mm_hash)
325 + return False
326 +
327 + def _fail_or_fallback(
328 + self, request: "Request", feature: "MultiModalFeatureSpec", why: str
329 + ) -> bool:
330 + """Resolve an unavailable remote encoding and clear its wait budget.
331 +
332 + Returns:
333 + True if local media or embeddings are available, removing this
334 + item's remote announcement so subsequent steps do not retry it.
335 + False if only placeholder metadata (or no input) is available,
336 + recording the request for take_unavailable_requests() to drain.
337 + """
338 + mm_hash = feature.identifier
339 + self._deferred_since.pop((request.request_id, mm_hash), None)
340 + data = feature.data
341 + metadata = (
342 + self._metadata_resolver.fields_for(feature.modality) if data else set()
343 + )
344 + if data and any(
345 + value.data is not None for key, value in data.items() if key not in metadata
346 + ):
347 + # Passthrough embeddings bypass the processor cache. SHM address
348 + # items therefore refer to media that the worker can encode.
349 + request.ec_transfer_params = dict(request.ec_transfer_params or {})
350 + request.ec_transfer_params.pop(mm_hash, None)
351 + logger.warning(
352 + "EC consumer: request %s mm_hash=%s: %s; using local input",
353 + request.request_id,
354 + mm_hash,
355 + why,
356 + )
357 + return True
358 + self._unrecoverable.add(request.request_id)
359 + logger.error(
360 + "EC consumer: request %s needs remote encoding mm_hash=%s but %s. "
361 + "No local model input is available; failing the request.",
362 + request.request_id,
363 + mm_hash,
364 + why,
365 + )
366 + return False
367 +
368 + def take_unavailable_requests(self) -> set[str]:
369 + """Return and clear IDs of requests that cannot obtain their encodings.
370 +
371 + The scheduler must abort these requests rather than leave them waiting.
372 + Requests that can fall back to local input are not included, and each
373 + recorded failure is returned only once.
374 + """
375 + if not self._unrecoverable:
376 + return set()
377 + failed = self._unrecoverable
378 + self._unrecoverable = set()
379 + return failed
380 +
381 def _start_xfer(
382 self, mm_hash: str, info: "dict[str, Any]", size_bytes: int
383 ) -> bool:
384 """Allocate a not-ready cache entry and start a NIXL READ into it.
385
386 Returns True if the transfer was started. Returns False when the
275 - cache cannot accommodate the encoding, in which case the request
276 - falls back to local recomputation.
387 + cache cannot accommodate the encoding; admission will retry it.
388 """
278 - import time
389 from math import ceil
390
391 from vllm.distributed.ec_transfer.ec_connector.cpu.session import (
@@ -286,7 +396,7 @@ class ECCPUScheduler:
396 entry = self._cache.alloc(mm_hash, n_blocks)
397 if entry is None:
398 logger.debug(
289 - "EC consumer: cache full for mm_hash=%s (%d blocks); local encode",
399 + "EC consumer: cache full for mm_hash=%s (%d blocks); deferring",
400 mm_hash,
401 n_blocks,
402 )
@@ -319,8 +429,6 @@ class ECCPUScheduler:
429 return True
430
431 def _poll_step(self) -> None:
322 - import time
323 -
432 now = time.monotonic()
433 all_messages = self._transport.poll()
434 for addr, session in list(self._sessions.items()):
@@ -352,8 +460,7 @@ class ECCPUScheduler:
460 mm_hash,
461 )
462 for mm_hash in r.retryable:
355 - # No tombstone: dropping every trace of the attempt is what lets
356 - # the next admit pass re-issue the read as if it were the first.
463 + # Release this attempt's blocks, retaining the request wait budget.
464 self._in_flight.discard(mm_hash)
465 self._cache.discard(mm_hash)
466 logger.debug(
@@ -472,6 +579,8 @@ class ECCPUScheduler:
579 def request_finished(
580 self, request: "Request"
581 ) -> tuple[bool, "dict[str, Any] | None"]:
582 + for feature in request.mm_features:
583 + self._deferred_since.pop((request.request_id, feature.identifier), None)
584 if not (self._nixl_enabled and self._is_producer):
585 return False, None
586 if not request.mm_features:
@@ -520,6 +629,7 @@ class ECCPUScheduler:
629 self._pending_saves.clear()
630 self._pending_loads.clear()
631 self._load_acks.clear()
632 + self._deferred_since.clear()
633
634 self._is_producer = False
635 self._is_consumer = False
vllm/distributed/ec_transfer/ec_connector/cpu/session.py
+13 -2
@@ -71,7 +71,7 @@ class XferState(enum.Enum):
71 WAITING_ACK = "waiting_ack" # XferReq sent, awaiting XferAck
72 READING = "reading" # NIXL READ in flight
73 DONE = "done" # READ completed; promote to local cache
74 - ACK_TIMEOUT = "ack_timeout" # XferAck never arrived; free blocks + tombstone
74 + ACK_TIMEOUT = "ack_timeout" # XferAck never arrived; free blocks + retry
75 READ_FAILED = "read_failed" # Unexpected NIXL state; free blocks + tombstone
76 QUARANTINED = "quarantined" # Read timed out, NIXL unabortable; keep blocks
77 SETTLED = "settled" # Quarantined DMA done; safe to free blocks
@@ -500,7 +500,18 @@ class ConsumerSession:
500 if state == XferState.DONE:
501 self._completed.add(mm_hash)
502 del self._xfers[mm_hash]
503 - elif state in (XferState.ACK_TIMEOUT, XferState.READ_FAILED):
503 + elif state == XferState.ACK_TIMEOUT:
504 + # The scheduler bounds retries with a per-request deadline.
505 + logger.warning(
506 + "EC consumer: no XferAck for mm_hash=%s from %s:%d in time; "
507 + "will re-request",
508 + mm_hash,
509 + self._addr[0],
510 + self._addr[1],
511 + )
512 + self._retryable.add(mm_hash)
513 + del self._xfers[mm_hash]
514 + elif state == XferState.READ_FAILED:
515 self._tombstoned.add(mm_hash)
516 del self._xfers[mm_hash]
517 elif state == XferState.QUARANTINED: