From f1aeff60073dda82a423d7d081f132c8a458e16b Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Wed, 9 Sep 2026 11:13:33 +0800 Subject: [PATCH] [Bugfix][EC Connector] Fail the request, not the engine, when a remote encoding cannot arrive (#55290) Signed-off-by: Tianyu Guo Co-authored-by: OpenAI Codex --- tests/v1/ec_connector/integration/README.md | 20 ++ .../integration/test_nixl_failure.py | 202 +++++++++++++++ .../unit/test_scheduler_nixl_consumer.py | 152 ++++++++--- .../unit/test_scheduler_nixl_producer.py | 1 + tests/v1/ec_connector/unit/test_session.py | 8 +- .../ec_transfer/ec_connector/cpu/connector.py | 4 + .../ec_connector/cpu/scheduler/__init__.py | 238 +++++++++++++----- .../ec_transfer/ec_connector/cpu/session.py | 15 +- 8 files changed, 536 insertions(+), 104 deletions(-) create mode 100644 tests/v1/ec_connector/integration/test_nixl_failure.py diff --git a/tests/v1/ec_connector/integration/README.md b/tests/v1/ec_connector/integration/README.md index bf2dcd76d6..a89bbfd347 100644 --- a/tests/v1/ec_connector/integration/README.md +++ b/tests/v1/ec_connector/integration/README.md @@ -58,6 +58,26 @@ EC_SHARED_STORAGE_PATH="/tmp/my_ec_cache" bash ./tests/v1/ec_connector/integrati ## How It Works +### NIXL EC failure isolation (1E + 1PD) + +```bash +PATH="$PWD/.venv/bin:$PATH" CUDA_VISIBLE_DEVICES=0,1 \ + .venv/bin/python -m pytest \ + tests/v1/ec_connector/integration/test_nixl_failure.py -v -s +``` + +Requires two CUDA GPUs and NIXL. `MODEL` overrides the default +`Qwen/Qwen2.5-VL-3B-Instruct`. The test starts real E/PD servers with CUDA +graphs enabled and uses the proxy's rewrite helper. Only the failing request's +control endpoint is replaced: a ZMQ peer returns `NACK_MISSING` and verifies +that the consumer actually requested the encoding. + +Metadata-only requests must return a request-level error; requests retaining +the image must succeed through local fallback. After each case, a fresh +metadata-only image request must succeed through the real NIXL transfer path. +No server restart, proxy retry, cache hit, or timeout race can mask the result. +This is a correctness test, not a throughput benchmark. + ### Mooncake EC (1E + 1PD) ```bash diff --git a/tests/v1/ec_connector/integration/test_nixl_failure.py b/tests/v1/ec_connector/integration/test_nixl_failure.py new file mode 100644 index 0000000000..ec5cc8518e --- /dev/null +++ b/tests/v1/ec_connector/integration/test_nixl_failure.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NIXL request failure isolation with real E/PD servers and an injected NACK. + +Requires two CUDA GPUs and NIXL. Set MODEL to override the small default model. +""" + +import asyncio +import copy +import json +import os +import uuid +from pathlib import Path + +import aiohttp +import msgspec +import pytest +import zmq +import zmq.asyncio +from PIL import Image + +from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import ( + rewrite_for_decode, +) +from tests.utils import RemoteOpenAIServer +from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import ( + XferAck, + XferReq, + XferStatus, +) +from vllm.multimodal.utils import encode_image_url +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port + +MODEL = os.environ.get("MODEL", "Qwen/Qwen2.5-VL-3B-Instruct") + + +@pytest.fixture(scope="module") +def servers(): + pytest.importorskip("nixl") + if not current_platform.is_cuda() or current_platform.device_count() < 2: + pytest.skip("Requires two CUDA GPUs") + devices = os.environ.get("CUDA_VISIBLE_DEVICES", "0,1").split(",") + common = [ + "--host", + "127.0.0.1", + "--dtype", + "bfloat16", + "--max-model-len", + "4096", + "--max-num-seqs", + "4", + "--gpu-memory-utilization", + "0.7", + "--no-enable-prefix-caching", + ] + + def config(role): + return json.dumps( + { + "ec_connector": "ECCPUConnector", + "ec_role": role, + "ec_connector_extra_config": { + "ec_enable_nixl": True, + "ec_cpu_bytes": 64 << 20, + }, + } + ) + + with ( + RemoteOpenAIServer( + MODEL, + common + + [ + "--mm-encoder-only", + "--disable-hybrid-kv-cache-manager", + "--ec-transfer-config", + config("ec_producer"), + ], + env_dict={ + "CUDA_VISIBLE_DEVICES": devices[0], + "VLLM_EC_SIDE_CHANNEL_HOST": "127.0.0.1", + "VLLM_EC_SIDE_CHANNEL_PORT": str(get_open_port()), + }, + ) as encoder, + RemoteOpenAIServer( + MODEL, + common + + ["--enable-mm-embeds", "--ec-transfer-config", config("ec_consumer")], + env_dict={"CUDA_VISIBLE_DEVICES": devices[1]}, + ) as consumer, + ): + yield encoder, consumer + + +async def _post(session, server, body): + async with session.post(server.url_for("v1/chat/completions"), json=body) as r: + return r.status, await r.json() + + +async def _encode(session, encoder): + image = Image.open(Path(__file__).with_name("hato.jpg")).convert("RGB") + image.thumbnail((224, 224)) + item_uuid = uuid.uuid4().hex # Avoid both EC and processor cache hits. + body = { + "model": MODEL, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": encode_image_url(image)}, + "uuid": item_uuid, + }, + {"type": "text", "text": "Describe the image briefly."}, + ], + } + ], + "max_tokens": 16, + "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}, + } + status, result = await _post(session, encoder, body) + assert status == 200, result + params = result["ec_transfer_params"] + assert len(params) == 1, params + ((mm_hash, handle),) = params.items() + assert handle["metadata"] and handle["peer_host"] and handle["size_bytes"] > 0 + rewritten = rewrite_for_decode( + body, + { + 0: { + **handle["metadata"], + "mm_hash": item_uuid, + "ec_mm_hash": mm_hash, + } + }, + ) + assert rewritten["messages"][0]["content"][0]["type"] == "image_embeds" + for request in (body, rewritten): + request["ec_transfer_params"] = {item_uuid: copy.deepcopy(handle)} + return body, rewritten, item_uuid + + +@pytest.mark.asyncio +@pytest.mark.parametrize("local_image", [False, True], ids=["metadata-only", "pixels"]) +async def test_missing_remote_encoding_is_request_scoped(servers, local_image): + """A terminal NACK must preserve local fallback and subsequent serving.""" + encoder, consumer = servers + timeout = aiohttp.ClientTimeout(total=120) + async with aiohttp.ClientSession(timeout=timeout) as session: + original, rewritten, mm_hash = await _encode(session, encoder) + request = original if local_image else rewritten + + # Keep the real encoder's shape/metadata; replace only the control peer. + # The socket stays bound until completion, so no port or timing race is + # needed to make the consumer receive NACK_MISSING. + with zmq.asyncio.Context() as ctx, ctx.socket(zmq.ROUTER) as router: + router.setsockopt(zmq.LINGER, 0) + port = router.bind_to_random_port("tcp://127.0.0.1") + request["ec_transfer_params"][mm_hash].update( + peer_host="127.0.0.1", peer_port=port + ) + + async def reject_read(): + identity, delimiter, payload = await router.recv_multipart() + req = msgspec.msgpack.decode(payload, type=XferReq) + assert req.mm_hash == mm_hash + ack = XferAck( + mm_hash=req.mm_hash, + status=XferStatus.NACK_MISSING, + session_id=req.session_id, + ) + await router.send_multipart( + [identity, delimiter, msgspec.msgpack.encode(ack)] + ) + + rejection = asyncio.create_task(asyncio.wait_for(reject_read(), 30)) + try: + status, result = await _post(session, consumer, request) + # A frontend rejection must not count as successful injection. + await rejection + finally: + rejection.cancel() + await asyncio.gather(rejection, return_exceptions=True) + + if local_image: + assert status == 200, result + assert result["choices"][0]["message"]["content"], result + else: + assert status == 500, result + assert result["error"]["type"] == "InternalServerError", result + + async with session.get(consumer.url_for("health")) as health: + assert health.status == 200, "The failed read killed EngineCore" + # A fresh metadata-only image must cross the real NIXL data path after + # the fault, not just pass a frontend health check or hit a local cache. + _, healthy, _ = await _encode(session, encoder) + status, result = await _post(session, consumer, healthy) + assert status == 200, result + assert result["choices"][0]["message"]["content"], result diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py index 0b7d0070b4..a3e87f9b78 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py @@ -1,13 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time import uuid +import pytest +import torch + import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod from tests.v1.ec_connector.unit.utils import create_ec_vllm_config from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( ECSharedRegion, ) from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler +from vllm.multimodal.inputs import ( + MultiModalFieldElem, + MultiModalKwargsItem, + MultiModalSharedField, +) _N, _BS, _HID, _ES = 16, 64, 32, 2 @@ -18,10 +27,12 @@ class _Pos: class _Feature: - def __init__(self, mm_hash, length=1): + def __init__(self, mm_hash, length=1, data=None): self.mm_hash = mm_hash self.identifier = mm_hash self.mm_position = _Pos(0, length) + self.data = data + self.modality = "image" class _Request: @@ -94,6 +105,7 @@ def _consumer_sched(monkeypatch): # directly since this helper builds gate-off then flips fields on. s._hidden_dim = _HID s._element_size = _ES + s._metadata_resolver._cache["image"] = {"image_grid_thw"} return s @@ -150,13 +162,17 @@ def test_new_remote_read_defers_then_completes(monkeypatch): s.shutdown() -def test_no_params_is_ready(monkeypatch): +@pytest.mark.parametrize("params", [None, {"x": 123}]) +def test_no_remote_announcement_is_ready(monkeypatch, params): s = _consumer_sched(monkeypatch) - assert s.ensure_cache_available(_Request([_Feature("x")], params=None), 0) is True + assert s.ensure_cache_available(_Request([_Feature("x")], params=params), 0) + assert s.take_unavailable_requests() == set() s.shutdown() -def test_tombstoned_read_discards_and_blocks_retry(monkeypatch): +@pytest.mark.parametrize("payload", [None, "metadata", "pixels", "embeds", "shm"]) +def test_failed_read_falls_back_only_with_local_input(monkeypatch, payload): + """A remote miss is fatal only when local model input is unavailable.""" s = _consumer_sched(monkeypatch) fake = _FakeSession() @@ -167,26 +183,71 @@ def test_tombstoned_read_discards_and_blocks_retry(monkeypatch): return True monkeypatch.setattr(s, "_start_xfer", _fake_start) - req = _Request([_Feature("h1", 1)], params=_params("h1", 1)) + fields = { + "metadata": {"image_grid_thw": torch.tensor([1, 2, 2])}, + "pixels": {"pixel_values": torch.ones(1, 3)}, + "embeds": {"image_embeds": torch.ones(1, _HID)}, + "shm": {"address": 123, "monotonic_id": 1}, + } + data = ( + MultiModalKwargsItem( + { + key: MultiModalFieldElem(value, MultiModalSharedField(batch_size=1)) + for key, value in fields[payload].items() + } + ) + if payload is not None + else None + ) + req = _Request([_Feature("h1", 1, data)], params=_params("h1", 1)) # Step 1: start the read. assert s.ensure_cache_available(req, 0) is False s.build_connector_meta(scheduler_output=None) - # Producer rejected the read -> tombstoned. fake._results.tombstoned.add("h1") s._sessions[("h", 1)] = fake - # Step 2: poll discards the entry and records a tombstone; the tombstone - # is consumed by admit so the request proceeds to local compute. - assert s.ensure_cache_available(req, 0) is True + can_fallback = payload in ("pixels", "embeds", "shm") + assert s.ensure_cache_available(req, 0) is can_fallback assert s._cache.get("h1") is None assert "h1" not in s._in_flight - s.build_connector_meta(scheduler_output=None) + assert s.take_unavailable_requests() == (set() if can_fallback else {"r1"}) + # Drained by the read, so the scheduler cannot abort it twice. + assert s.take_unavailable_requests() == set() + if can_fallback: + # A later scheduling step must not retry the failed remote source. + s.build_connector_meta(scheduler_output=None) + assert s.ensure_cache_available(req, 0) + assert fake.started == ["h1"] + s.shutdown() - # Step 3: tombstone was consumed, so a fresh transfer starts again. - assert s.ensure_cache_available(req, 0) is False - assert "h1" in s._in_flight + +@pytest.mark.parametrize("retry", [False, True], ids=["in-flight", "retryable"]) +def test_remote_wait_budget_survives_retries_and_long_steps(monkeypatch, retry): + """Neither retries nor delayed scheduling may restart the request budget.""" + s = _consumer_sched(monkeypatch) + fake = _FakeSession() + s._sessions[("h", 1)] = fake + + def start(mm_hash, info, size): + assert s._cache.alloc(mm_hash, 1) is not None + fake.started.append(mm_hash) + return True + + monkeypatch.setattr(s, "_start_xfer", start) + now = [1000.0] + monkeypatch.setattr(time, "monotonic", lambda: now[0]) + req = _Request([_Feature("h1")], params=_params("h1", 1)) + assert not s.ensure_cache_available(req, 0) + for elapsed in (20, 40, 121): + now[0] = 1000.0 + elapsed + if retry: + fake._results.retryable.add("h1") + s.build_connector_meta(scheduler_output=None) + assert not s.ensure_cache_available(req, 0) + assert s.take_unavailable_requests() == ({"r1"} if elapsed > 60 else set()) + assert len(fake.started) == (3 if retry else 1) s.shutdown() @@ -225,17 +286,50 @@ def test_retryable_read_re_requests_without_admitting(monkeypatch): s.shutdown() -def test_size_mismatch_skips_transfer(monkeypatch): +@pytest.mark.parametrize( + "size_fields", + [{"size_bytes": 999}, {"size_bytes": None}, {"size_bytes": "big"}, {}], + ids=["mismatch", "null-size", "text-size", "no-size"], +) +def test_invalid_remote_size_fails_the_request(monkeypatch, size_fields): + """Invalid remote sizes fail the request without raising in the scheduler.""" s = _consumer_sched(monkeypatch) - # Advertised size disagrees with pos.length * hidden_dim * element_size. - bad = {"h1": {"peer_host": "h", "peer_port": 1, "size_bytes": 999}} - req = _Request([_Feature("h1", 1)], params=bad) - assert s.ensure_cache_available(req, 0) is True + announced = {"peer_host": "h", "peer_port": 1, **size_fields} + req = _Request([_Feature("h1", 1)], params={"h1": announced}) + assert s.ensure_cache_available(req, 0) is False + assert s.take_unavailable_requests() == {"r1"} assert "h1" not in s._in_flight assert s._cache.get("h1") is None s.shutdown() +def test_deferral_budget_is_per_request(monkeypatch): + """Shared hashes have independent wait budgets and cancellation cleanup.""" + s = _consumer_sched(monkeypatch) + monkeypatch.setattr(s._cache, "alloc", lambda key, n: None) + params = _params("h1", 1) + old = _Request([_Feature("h1", 1)], params=params, req_id="old") + new = _Request([_Feature("h1", 1)], params=params, req_id="new") + + now = [1000.0] + monkeypatch.setattr(time, "monotonic", lambda: now[0]) + assert s.ensure_cache_available(old, 0) is False + now[0] += sched_mod._ADMIT_DEFER_TIMEOUT_S - 1 + assert s.ensure_cache_available(new, 0) is False + assert s.take_unavailable_requests() == set() + + now[0] += 2 + # `old` is now past its budget; `new` is not. + assert s.ensure_cache_available(new, 0) is False + assert s.take_unavailable_requests() == set() + assert s.request_finished(new) == (False, None) + assert set(s._deferred_since) == {("old", "h1")} + assert s.ensure_cache_available(old, 0) is False + assert s.take_unavailable_requests() == {"old"} + assert "h1" not in s._in_flight + s.shutdown() + + def test_metadata_only_entry_admits_without_transfer(monkeypatch): # A producer that reported placeholder metadata but had no cache entry # 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): s.shutdown() -def test_orphan_not_ready_entry_falls_back_no_realloc(monkeypatch): +def test_orphan_not_ready_entry_defers_no_realloc(monkeypatch): # Post-quarantine, post-tombstone-consumed orphan: a not-ready cache # entry exists for the mm_hash, but the hash is tracked in none of # _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): monkeypatch.setattr(s._cache, "alloc", _spy_alloc) req = _Request([_Feature("h1", 1)], params=_params("h1", 1)) - # Must not raise; request falls back to local compute (admitted True), - # the orphan entry is untouched, and alloc is never called again. - assert s.ensure_cache_available(req, 0) is True + # Must not raise; the request is deferred while the DMA settles, the + # orphan entry is untouched, and alloc is never called again. + assert s.ensure_cache_available(req, 0) is False assert calls == [] assert s._cache.get("h1") is entry assert not entry.ready - s.shutdown() - - -def test_alloc_failure_falls_back_to_local(monkeypatch): - s = _consumer_sched(monkeypatch) - - # Force the cache to reject the allocation. - monkeypatch.setattr(s._cache, "alloc", lambda key, n: None) - req = _Request([_Feature("h1", 1)], params=_params("h1", 1)) - - # _start_xfer returns False -> request admitted for local recompute. - assert s.ensure_cache_available(req, 0) is True - assert "h1" not in s._in_flight + assert s.take_unavailable_requests() == set() s.shutdown() diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py index e55ac23769..156d98ae79 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -63,6 +63,7 @@ def test_request_finished_producer_emits_params(monkeypatch): delay, params = s.request_finished(_Request([_Feature("h1", length=2)])) assert delay is False + assert entry.evictable assert params == { "h1": { "metadata": {}, diff --git a/tests/v1/ec_connector/unit/test_session.py b/tests/v1/ec_connector/unit/test_session.py index d86e6263a6..4d78b0426a 100644 --- a/tests/v1/ec_connector/unit/test_session.py +++ b/tests/v1/ec_connector/unit/test_session.py @@ -451,14 +451,16 @@ def test_consumer_session_poll_not_ready_nack_goes_to_retryable(): assert not results.completed -def test_consumer_session_take_results_clears_state(): +def test_consumer_session_ack_timeout_reports_retry_once(): + """A delayed producer ACK is retryable, and its result is drained once.""" s = _make_consumer_session() s.start_xfer("h1", [0, 1], deadline=time.monotonic() - 1) s.poll([], time.monotonic()) r1 = s.take_results() - assert "h1" in r1.tombstoned + assert "h1" in r1.retryable + assert not r1.tombstoned r2 = s.take_results() - assert not r2.tombstoned # cleared after first take + assert not r2.retryable # cleared after first take def test_consumer_session_on_peer_down_cancels_waiting_ack(): diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py b/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py index a0996d9035..a47688f64b 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py @@ -104,6 +104,10 @@ class ECCPUConnector(ECConnectorBase): request, num_computed_tokens ) + def take_unavailable_requests(self) -> set[str]: + assert self.connector_scheduler is not None + return self.connector_scheduler.take_unavailable_requests() + def update_state_after_alloc(self, request: "Request", index: int) -> None: assert self.connector_scheduler is not None self.connector_scheduler.update_state_after_alloc(request, index) diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py index cd65740108..cf0a92dcc2 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -7,6 +7,7 @@ Owns the mmap region and the embedding cache, and handles the producer for the ECCPUConnector. """ +import time from collections.abc import Collection from typing import TYPE_CHECKING, Any @@ -29,12 +30,20 @@ from vllm.logger import init_logger if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import ECConnectorOutput from vllm.v1.request import Request logger = init_logger(__name__) +# Total remote wait budget per request item, including all retries. +_ADMIT_DEFER_TIMEOUT_S = 60.0 + + +class _RemoteUnavailable(Exception): + """An item cannot be fetched within its request's remote wait budget.""" + class ECCPUScheduler: """Scheduler delegate for the ECCPUConnector.""" @@ -92,6 +101,11 @@ class ECCPUScheduler: self._in_flight: set[str] = set() self._tombstones: set[str] = set() self._step_completed: set[str] = set() + # Per-request clocks survive retries and are cleared on completion. + self._deferred_since: dict[tuple[str, str], float] = {} + # Requests needing a remote encoding that will not arrive. Drained by + # take_unavailable_requests(); the scheduler aborts them. + self._unrecoverable: set[str] = set() self._peer_host: str | None = None self._peer_port: int | None = None # Model shape for size checks + compat hash; only set by @@ -117,6 +131,7 @@ class ECCPUScheduler: compute_ec_compatibility_hash, ) from vllm.distributed.ec_transfer.ec_connector.cpu.session import ( + PRODUCER_PIN_LEASE_S, ProducerSession, ) from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config @@ -136,7 +151,9 @@ class ECCPUScheduler: # its scheduler step, so its reply latency scales with the encoder's # batch size: a deployment whose steps run longer must raise this. self._ack_timeout_s = float( - self._ec_config.get_from_extra_config("consumer_ack_timeout_s", 2.0) + self._ec_config.get_from_extra_config( + "consumer_ack_timeout_s", PRODUCER_PIN_LEASE_S + ) ) self._compat_hash = compute_ec_compatibility_hash( vllm_version=VLLM_VERSION, @@ -203,79 +220,172 @@ class ECCPUScheduler: def _nixl_consumer_admit( self, request: "Request", num_computed_tokens: int ) -> bool: - """Admit a request only once all its remote encodings are cached. + """Admit a request once its required remote encodings are ready. - Returns True when every encoding the request needs is either already - ready in the local cache or has no remote source. Any encoding whose - NIXL READ is still in flight (or was just started this step) defers - the request, which the scheduler re-presents on a later step. + Each item has a per-request wait budget that survives transfer retries. + Unavailable items fall back to local input; without it, the request is + reported by take_unavailable_requests() for the scheduler to abort. + + Returns: + True if no item needs to wait for a remote encoding. False if any + item is pending or the request must be aborted. """ - params: dict[str, dict[str, Any]] = ( - getattr(request, "ec_transfer_params", None) or {} - ) - if not params: + if not request.ec_transfer_params: return True + now = time.monotonic() pending = False for feature in request.mm_features: - pos = feature.mm_position - if pos.offset + pos.length <= num_computed_tokens: - continue - mm_hash = feature.identifier - entry = self._cache.get(mm_hash) - if entry is not None and entry.ready: - # Local hit: upstream's update_state_after_alloc pins and - # loads it through the same path as a natively cached entry. - continue - if mm_hash in self._in_flight: - pending = True - continue - if mm_hash in self._step_completed: - pending = True - continue - if mm_hash in self._tombstones: - self._tombstones.discard(mm_hash) - continue - if entry is not None: - # Present but not ready and not being fetched: its blocks are - # held by a quarantined/settling DMA and cannot be reused. - # Fall back to local compute this step; never re-alloc a - # mm_hash already in the cache. - continue - info = params.get(mm_hash) - # Absent entirely, or present with only placeholder metadata (no - # cache entry existed on the producer side at request_finished - # time): nothing to fetch, fall back to local compute. - if not info or "peer_host" not in info: - continue - expected = pos.length * self._hidden_dim * self._element_size - if int(info.get("size_bytes", -1)) != expected: - logger.warning( - "EC consumer: size mismatch mm_hash=%s; local encode", mm_hash - ) - continue try: - started = self._start_xfer(mm_hash, info, expected) - except Exception: - logger.exception( - "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash - ) - continue - if not started: - continue - self._in_flight.add(mm_hash) - pending = True + if not self._admit_item(request, feature, num_computed_tokens, now): + pending = True + except _RemoteUnavailable as error: + if not self._fail_or_fallback(request, feature, str(error)): + return False return not pending + def _admit_item( + self, + request: "Request", + feature: "MultiModalFeatureSpec", + num_computed_tokens: int, + now: float, + ) -> bool: + """Check one encoding, starting a remote read if needed. + + Returns: + True if already computed, cached, or not remotely sourced. False + while waiting for a read or cache space, including across retries. + + Raises: + _RemoteUnavailable: The remote source is invalid, the read failed, + or the item's total wait budget expired. + """ + pos = feature.mm_position + mm_hash = feature.identifier + if pos.offset + pos.length <= num_computed_tokens: + self._deferred_since.pop((request.request_id, mm_hash), None) + return True + announced = (request.ec_transfer_params or {}).get(mm_hash) + # Without a producer address there is nothing to fetch: the request + # carries what the model needs and the encoder runs locally. + # `ec_transfer_params` reaches us from the request, so its shape is + # checked rather than assumed. + remote: dict[str, Any] | None = ( + announced + if isinstance(announced, dict) and "peer_host" in announced + else None + ) + + entry = self._cache.get(mm_hash) + if entry is not None and entry.ready: + # Local hit: upstream's update_state_after_alloc pins and + # loads it through the same path as a natively cached entry. + self._deferred_since.pop((request.request_id, mm_hash), None) + return True + if remote is None: + return True + since = self._deferred_since.setdefault((request.request_id, mm_hash), now) + if now - since > _ADMIT_DEFER_TIMEOUT_S: + raise _RemoteUnavailable(f"the remote wait exceeded {now - since:.0f}s") + if mm_hash in self._in_flight or mm_hash in self._step_completed: + return False + + if mm_hash in self._tombstones: + self._tombstones.discard(mm_hash) + raise _RemoteUnavailable("the remote read failed") + + if entry is not None: + # Present but not ready and not being fetched: its blocks are + # held by a quarantined/settling DMA and cannot be reused. + return False + + expected = pos.length * self._hidden_dim * self._element_size + try: + size = int(remote["size_bytes"]) + except (KeyError, TypeError, ValueError) as error: + raise _RemoteUnavailable("the announced size was unusable") from error + if size != expected: + logger.warning( + "EC consumer: size mismatch mm_hash=%s announced=%d expected=%d", + mm_hash, + size, + expected, + ) + raise _RemoteUnavailable("the announced size was wrong") + + try: + started = self._start_xfer(mm_hash, remote, expected) + except Exception as error: + logger.exception( + "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash + ) + raise _RemoteUnavailable("the read could not be started") from error + if started: + self._in_flight.add(mm_hash) + return False + + def _fail_or_fallback( + self, request: "Request", feature: "MultiModalFeatureSpec", why: str + ) -> bool: + """Resolve an unavailable remote encoding and clear its wait budget. + + Returns: + True if local media or embeddings are available, removing this + item's remote announcement so subsequent steps do not retry it. + False if only placeholder metadata (or no input) is available, + recording the request for take_unavailable_requests() to drain. + """ + mm_hash = feature.identifier + self._deferred_since.pop((request.request_id, mm_hash), None) + data = feature.data + metadata = ( + self._metadata_resolver.fields_for(feature.modality) if data else set() + ) + if data and any( + value.data is not None for key, value in data.items() if key not in metadata + ): + # Passthrough embeddings bypass the processor cache. SHM address + # items therefore refer to media that the worker can encode. + request.ec_transfer_params = dict(request.ec_transfer_params or {}) + request.ec_transfer_params.pop(mm_hash, None) + logger.warning( + "EC consumer: request %s mm_hash=%s: %s; using local input", + request.request_id, + mm_hash, + why, + ) + return True + self._unrecoverable.add(request.request_id) + logger.error( + "EC consumer: request %s needs remote encoding mm_hash=%s but %s. " + "No local model input is available; failing the request.", + request.request_id, + mm_hash, + why, + ) + return False + + def take_unavailable_requests(self) -> set[str]: + """Return and clear IDs of requests that cannot obtain their encodings. + + The scheduler must abort these requests rather than leave them waiting. + Requests that can fall back to local input are not included, and each + recorded failure is returned only once. + """ + if not self._unrecoverable: + return set() + failed = self._unrecoverable + self._unrecoverable = set() + return failed + def _start_xfer( self, mm_hash: str, info: "dict[str, Any]", size_bytes: int ) -> bool: """Allocate a not-ready cache entry and start a NIXL READ into it. Returns True if the transfer was started. Returns False when the - cache cannot accommodate the encoding, in which case the request - falls back to local recomputation. + cache cannot accommodate the encoding; admission will retry it. """ - import time from math import ceil from vllm.distributed.ec_transfer.ec_connector.cpu.session import ( @@ -286,7 +396,7 @@ class ECCPUScheduler: entry = self._cache.alloc(mm_hash, n_blocks) if entry is None: logger.debug( - "EC consumer: cache full for mm_hash=%s (%d blocks); local encode", + "EC consumer: cache full for mm_hash=%s (%d blocks); deferring", mm_hash, n_blocks, ) @@ -319,8 +429,6 @@ class ECCPUScheduler: return True def _poll_step(self) -> None: - import time - now = time.monotonic() all_messages = self._transport.poll() for addr, session in list(self._sessions.items()): @@ -352,8 +460,7 @@ class ECCPUScheduler: mm_hash, ) for mm_hash in r.retryable: - # No tombstone: dropping every trace of the attempt is what lets - # the next admit pass re-issue the read as if it were the first. + # Release this attempt's blocks, retaining the request wait budget. self._in_flight.discard(mm_hash) self._cache.discard(mm_hash) logger.debug( @@ -472,6 +579,8 @@ class ECCPUScheduler: def request_finished( self, request: "Request" ) -> tuple[bool, "dict[str, Any] | None"]: + for feature in request.mm_features: + self._deferred_since.pop((request.request_id, feature.identifier), None) if not (self._nixl_enabled and self._is_producer): return False, None if not request.mm_features: @@ -520,6 +629,7 @@ class ECCPUScheduler: self._pending_saves.clear() self._pending_loads.clear() self._load_acks.clear() + self._deferred_since.clear() self._is_producer = False self._is_consumer = False diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/session.py b/vllm/distributed/ec_transfer/ec_connector/cpu/session.py index 54c69a05f3..80bbd1fd40 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/session.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/session.py @@ -71,7 +71,7 @@ class XferState(enum.Enum): WAITING_ACK = "waiting_ack" # XferReq sent, awaiting XferAck READING = "reading" # NIXL READ in flight DONE = "done" # READ completed; promote to local cache - ACK_TIMEOUT = "ack_timeout" # XferAck never arrived; free blocks + tombstone + ACK_TIMEOUT = "ack_timeout" # XferAck never arrived; free blocks + retry READ_FAILED = "read_failed" # Unexpected NIXL state; free blocks + tombstone QUARANTINED = "quarantined" # Read timed out, NIXL unabortable; keep blocks SETTLED = "settled" # Quarantined DMA done; safe to free blocks @@ -500,7 +500,18 @@ class ConsumerSession: if state == XferState.DONE: self._completed.add(mm_hash) del self._xfers[mm_hash] - elif state in (XferState.ACK_TIMEOUT, XferState.READ_FAILED): + elif state == XferState.ACK_TIMEOUT: + # The scheduler bounds retries with a per-request deadline. + logger.warning( + "EC consumer: no XferAck for mm_hash=%s from %s:%d in time; " + "will re-request", + mm_hash, + self._addr[0], + self._addr[1], + ) + self._retryable.add(mm_hash) + del self._xfers[mm_hash] + elif state == XferState.READ_FAILED: self._tombstoned.add(mm_hash) del self._xfers[mm_hash] elif state == XferState.QUARANTINED: