From a0bd6a20a5d6eb3e37ea859afba33f026736bd0f Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Thu, 3 Sep 2026 04:47:53 +0000 Subject: [PATCH 1/8] [EPD] Fail a request whose remote encoding cannot arrive, instead of the engine Signed-off-by: Tianyu Guo --- .../unit/test_scheduler_nixl_consumer.py | 58 ++++++--- tests/v1/ec_connector/unit/test_session.py | 19 ++- .../ec_transfer/ec_connector/base.py | 16 +++ .../ec_transfer/ec_connector/cpu/connector.py | 4 + .../ec_connector/cpu/scheduler/__init__.py | 115 ++++++++++++++---- .../ec_transfer/ec_connector/cpu/session.py | 18 ++- vllm/v1/core/sched/scheduler.py | 6 + 7 files changed, 193 insertions(+), 43 deletions(-) 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 0b7d0070b42ca..90dc80d724365 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py @@ -156,7 +156,13 @@ def test_no_params_is_ready(monkeypatch): s.shutdown() -def test_tombstoned_read_discards_and_blocks_retry(monkeypatch): +def test_tombstoned_read_fails_the_request(monkeypatch): + """A failed remote read must fail the request, not admit it. + + The media reached the producer only, so an admitted request has no + embedding and no way to compute one: the model would receive an empty + multimodal batch and assert. + """ s = _consumer_sched(monkeypatch) fake = _FakeSession() @@ -177,16 +183,13 @@ def _fake_start(mm_hash, info, size): 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 + # Step 2: the entry is discarded and the request is reported unschedulable. + assert s.ensure_cache_available(req, 0) is False assert s._cache.get("h1") is None assert "h1" not in s._in_flight - s.build_connector_meta(scheduler_output=None) - - # 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 + assert s.get_unrecoverable_requests() == {"r1"} + # Drained by the read, so the scheduler cannot abort it twice. + assert s.get_unrecoverable_requests() == set() s.shutdown() @@ -225,14 +228,16 @@ def _fake_start(mm_hash, info, size): s.shutdown() -def test_size_mismatch_skips_transfer(monkeypatch): +def test_size_mismatch_fails_the_request(monkeypatch): + """A bad announcement is terminal: there is no second source to try.""" 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 + assert s.ensure_cache_available(req, 0) is False assert "h1" not in s._in_flight assert s._cache.get("h1") is None + assert s.get_unrecoverable_requests() == {"r1"} s.shutdown() @@ -250,7 +255,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 +279,42 @@ def _spy_alloc(key, n): 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 + assert s.get_unrecoverable_requests() == set() s.shutdown() -def test_alloc_failure_falls_back_to_local(monkeypatch): +def test_full_pool_defers_then_fails_the_request(monkeypatch): + """A full pool is transient, so defer — but not forever. + + Deferring without a budget would hold the request until the client gave + up, with no record of why. + """ + import time as _time + 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 + t0 = _time.monotonic() + monkeypatch.setattr(_time, "monotonic", lambda: t0) + assert s.ensure_cache_available(req, 0) is False assert "h1" not in s._in_flight + assert s.get_unrecoverable_requests() == set() + + # Still full once the budget has elapsed: fail rather than defer again. + monkeypatch.setattr( + _time, "monotonic", lambda: t0 + sched_mod._ADMIT_DEFER_TIMEOUT_S + 1 + ) + assert s.ensure_cache_available(req, 0) is False + assert s.get_unrecoverable_requests() == {"r1"} s.shutdown() diff --git a/tests/v1/ec_connector/unit/test_session.py b/tests/v1/ec_connector/unit/test_session.py index d86e6263a6f7c..056ed60fa8ef2 100644 --- a/tests/v1/ec_connector/unit/test_session.py +++ b/tests/v1/ec_connector/unit/test_session.py @@ -451,14 +451,29 @@ def test_consumer_session_poll_not_ready_nack_goes_to_retryable(): assert not results.completed +def test_consumer_session_ack_timeout_goes_to_retryable(): + """A late XferAck means a busy producer, not a broken one. + + The producer answers XferReqs from its scheduler step, so the reply can + lag a whole encoder forward pass. Tombstoning on that admits the request + with no embedding, which is fatal. + """ + s = _make_consumer_session() + s.start_xfer("h1", [0, 1], deadline=time.monotonic() - 1) + s.poll([], time.monotonic()) + results = s.take_results() + assert "h1" in results.retryable + assert not results.tombstoned + + def test_consumer_session_take_results_clears_state(): 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 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/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index 9138203bf4ddc..5c0c00f5ae822 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -261,6 +261,22 @@ def ensure_cache_available( """ return True + def get_unrecoverable_requests(self) -> set[str]: + """Requests that can never be scheduled and must be failed. + + A remote encoding has no local fallback: unlike KV, which is always + recomputable from the prompt, an encoding announced by a producer + exists only there, because the media never reached this instance. + A connector reports a request here once it is certain the encoding + will not arrive, and the scheduler finishes it with an error reason — + the same treatment an unrecoverable KV load gets — rather than + running the model without its encoder outputs. + + Returns: + Request IDs to fail. The set is drained by the call. + """ + return set() + @abstractmethod def update_state_after_alloc(self, request: "Request", index: int): """ diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py b/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py index a0996d9035158..3ec90991345e9 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 @@ def ensure_cache_available( request, num_computed_tokens ) + def get_unrecoverable_requests(self) -> set[str]: + assert self.connector_scheduler is not None + return self.connector_scheduler.get_unrecoverable_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 bc52e1fb627bf..5f196e55866d7 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -34,6 +34,10 @@ logger = init_logger(__name__) +# How long a transient condition (a settling DMA, a full pool) may hold a +# request back before it is failed rather than deferred indefinitely. +_ADMIT_DEFER_TIMEOUT_S = 60.0 + class ECCPUScheduler: """Scheduler delegate for the ECCPUConnector.""" @@ -91,6 +95,13 @@ def __init__(self, vllm_config: "VllmConfig") -> None: self._in_flight: set[str] = set() self._tombstones: set[str] = set() self._step_completed: set[str] = set() + # mm_hash -> when a transient condition first held a request back, so + # one that never clears fails the request instead of deferring it + # forever. Cleared as soon as the encoding lands. + self._deferred_since: dict[str, float] = {} + # Requests needing a remote encoding that will not arrive. Drained by + # get_unrecoverable_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 @@ -201,67 +212,129 @@ def _nixl_consumer_admit( ) -> bool: """Admit a request only once all its remote encodings are cached. - 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. + An item announced with a producer address lives only on that producer. + Unlike KV, it cannot be recomputed from the request: the media never + reached this instance. Scheduling a request whose remote item is + missing therefore hands the model an empty multimodal batch, so such + an item defers the request while it may still arrive, and fails the + request once it cannot. """ + import time + params: dict[str, dict[str, Any]] = ( getattr(request, "ec_transfer_params", None) or {} ) if not 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 + announced = params.get(mm_hash) + # Without a producer address there is nothing to fetch: the request + # carries what the model needs and the encoder runs locally. + remote: dict[str, Any] | None = ( + announced + if announced is not None 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(mm_hash, None) continue - if mm_hash in self._in_flight: - pending = True - continue - if mm_hash in self._step_completed: + if mm_hash in self._in_flight or mm_hash in self._step_completed: pending = True continue + if mm_hash in self._tombstones: self._tombstones.discard(mm_hash) - continue + if remote is None: + continue + self._fail(request, mm_hash, "the remote read failed") + return False + 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. + if remote is None: + continue + if not self._defer(request, mm_hash, now, "a DMA was settling"): + return False + pending = True 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: + + if remote is None: continue + expected = pos.length * self._hidden_dim * self._element_size - if int(info.get("size_bytes", -1)) != expected: + size = int(remote.get("size_bytes", -1)) + if size != expected: logger.warning( - "EC consumer: size mismatch mm_hash=%s; local encode", mm_hash + "EC consumer: size mismatch mm_hash=%s announced=%d expected=%d", + mm_hash, + size, + expected, ) - continue + self._fail(request, mm_hash, "the announced size was wrong") + return False + try: - started = self._start_xfer(mm_hash, info, expected) + started = self._start_xfer(mm_hash, remote, expected) except Exception: logger.exception( "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash ) - continue + self._fail(request, mm_hash, "the read could not be started") + return False if not started: + # The local pool is full; it drains as other reads complete. + if not self._defer(request, mm_hash, now, "the pool was full"): + return False + pending = True continue self._in_flight.add(mm_hash) + self._deferred_since.pop(mm_hash, None) pending = True return not pending + def _defer(self, request: "Request", mm_hash: str, now: float, why: str) -> bool: + """Hold a request back while a transient condition clears. + + Returns False once the condition has outlived its budget, having + marked the request unschedulable. + """ + since = self._deferred_since.setdefault(mm_hash, now) + waited = now - since + if waited <= _ADMIT_DEFER_TIMEOUT_S: + return True + self._fail(request, mm_hash, f"{why} for {waited:.0f}s") + return False + + def _fail(self, request: "Request", mm_hash: str, why: str) -> None: + self._deferred_since.pop(mm_hash, None) + self._unrecoverable.add(request.request_id) + logger.error( + "EC consumer: request %s needs remote encoding mm_hash=%s but %s. " + "The media never reached this instance, so it cannot be encoded " + "locally; failing the request.", + request.request_id, + mm_hash, + why, + ) + + def get_unrecoverable_requests(self) -> set[str]: + 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: diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/session.py b/vllm/distributed/ec_transfer/ec_connector/cpu/session.py index 54c69a05f32d8..73153499c3995 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,21 @@ def poll(self, messages: list[bytes], now: float) -> None: 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: + # A busy producer is a transient condition, not a broken one: + # drop every trace of the attempt so a later admit pass can + # re-issue the read as if it were the first. This is also the + # only path to a tombstone that would otherwise log nothing. + 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: diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index d9d86668a62ee..3d414181befe1 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2136,6 +2136,12 @@ def update_from_output( self.grammar_compile_error_reqs.clear() if failed_kv_load_req_ids and not self.recompute_kv_load_failures: error_req_ids.update(failed_kv_load_req_ids) + if self.ec_connector is not None: + # A remote encoding has no local fallback: unlike KV, it cannot be + # recomputed from the request, because the media only ever reached + # the producer. Error the request rather than schedule it without + # its encoder outputs. + error_req_ids.update(self.ec_connector.get_unrecoverable_requests()) if error_req_ids: error_reqs = self.finish_requests( From f01314c07f5b5c62bcac2b531ca58d24396d2585 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Thu, 3 Sep 2026 13:26:23 +0000 Subject: [PATCH 2/8] [EPD] Hold an announced encoding until its consumer has read it Signed-off-by: Tianyu Guo --- .../unit/test_scheduler_nixl_producer.py | 120 ++++++++++++++++++ tests/v1/ec_connector/unit/test_session.py | 2 +- .../ec_connector/cpu/scheduler/__init__.py | 85 ++++++++++++- .../ec_transfer/ec_connector/cpu/session.py | 10 ++ 4 files changed, 215 insertions(+), 2 deletions(-) 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 e55ac23769bed..383bdd1f194ba 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -111,3 +111,123 @@ def test_request_finished_skips_unallocated_entry(monkeypatch): # no cache entry to transfer (empty "metadata": no fields, no transfer). assert params == {"h1": {"metadata": {}}} s.shutdown() + + +# ── announced encodings are held until read ────────────────────────────────── + + +class _FakeProducerSession: + """Stands in for ProducerSession so build_connector_meta can run.""" + + def __init__(self): + self.served = set() + + def poll_step(self): + pass + + def take_served(self): + served, self.served = self.served, set() + return served + + +def _announcing_sched(monkeypatch, lease=30.0): + s = _sched_gate_off(monkeypatch) + s._nixl_enabled = True + s._peer_host, s._peer_port = "1.2.3.4", 5601 + s._hidden_dim, s._element_size = 32, 2 + s._announce_lease_s = lease + s._producer_session = _FakeProducerSession() + return s + + +def _announce(s, mm_hash, length=2): + return s.request_finished(_Request([_Feature(mm_hash, length=length)])) + + +def test_announced_encoding_is_not_evictable(monkeypatch): + """Announcing publishes an address the consumer uses on a later step. + + By then the orchestrator has rewritten the media off the request, so an + eviction inside that window leaves the consumer nothing to fall back on. + """ + s = _announcing_sched(monkeypatch) + entry = s._cache.alloc("h1", 2) + assert entry is not None + s._cache.mark_ready("h1") + assert entry.evictable + + _delay, params = _announce(s, "h1") + assert "peer_host" in params["h1"] + assert not entry.evictable + s.shutdown() + + +def test_hold_is_released_once_the_read_lands(monkeypatch): + s = _announcing_sched(monkeypatch) + entry = s._cache.alloc("h1", 2) + s._cache.mark_ready("h1") + _announce(s, "h1") + assert not entry.evictable + + s._producer_session.served.add("h1") + s.build_connector_meta(scheduler_output=None) + assert entry.evictable + s.shutdown() + + +def test_hold_lapses_when_no_consumer_ever_reads(monkeypatch): + """A consumer that never asks must not pin the pool forever.""" + s = _announcing_sched(monkeypatch, lease=0.0) + entry = s._cache.alloc("h1", 2) + s._cache.mark_ready("h1") + _announce(s, "h1") + assert not entry.evictable + + import time as _time + + _time.sleep(0.01) + s.build_connector_meta(scheduler_output=None) + assert entry.evictable + s.shutdown() + + +def test_hold_is_taken_when_a_late_save_lands(monkeypatch): + """An entry announced mid-save is pinned when it becomes evictable. + + A not-ready entry cannot be evicted, so the hold has to start at + mark_ready rather than at the announcement. + """ + from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( + ECCPUWorkerMetadata, + ) + from vllm.v1.outputs import ECConnectorOutput + + s = _announcing_sched(monkeypatch) + entry = s._cache.alloc("h1", 2) + assert entry is not None and not entry.ready + _announce(s, "h1") + assert "h1" in s._announce_pending + + s.update_connector_output( + ECConnectorOutput( + ec_connector_worker_meta=ECCPUWorkerMetadata(completed_saves=["h1"]) + ) + ) + assert entry.ready + assert not entry.evictable + assert "h1" not in s._announce_pending + s.shutdown() + + +def test_reannouncing_extends_the_hold_without_a_second_pin(monkeypatch): + """One release must be enough however many requests announced it.""" + s = _announcing_sched(monkeypatch) + entry = s._cache.alloc("h1", 2) + s._cache.mark_ready("h1") + _announce(s, "h1") + _announce(s, "h1") + + s._producer_session.served.add("h1") + s.build_connector_meta(scheduler_output=None) + assert entry.evictable + s.shutdown() diff --git a/tests/v1/ec_connector/unit/test_session.py b/tests/v1/ec_connector/unit/test_session.py index 056ed60fa8ef2..96ffe7926ec5b 100644 --- a/tests/v1/ec_connector/unit/test_session.py +++ b/tests/v1/ec_connector/unit/test_session.py @@ -456,7 +456,7 @@ def test_consumer_session_ack_timeout_goes_to_retryable(): The producer answers XferReqs from its scheduler step, so the reply can lag a whole encoder forward pass. Tombstoning on that admits the request - with no embedding, which is fatal. + with no embedding, which is fatal once the media has been rewritten away. """ s = _make_consumer_session() s.start_xfer("h1", [0, 1], deadline=time.monotonic() - 1) 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 5f196e55866d7..2b18e0ee1d13c 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -102,6 +102,17 @@ def __init__(self, vllm_config: "VllmConfig") -> None: # Requests needing a remote encoding that will not arrive. Drained by # get_unrecoverable_requests(); the scheduler aborts them. self._unrecoverable: set[str] = set() + # Announcing an encoding publishes an address a consumer will use on a + # later step, by which time the orchestrator has rewritten the media + # off the request. These hold the encoding across that window: FIFO + # eviction inside it leaves the consumer nothing to fall back on. + # mm_hash -> lease deadline of a pin already taken. + self._announce_pins: dict[str, float] = {} + # mm_hash -> lease deadline, for entries announced before their save + # landed. A not-ready entry is already non-evictable, so the pin is + # taken when the worker reports the save. + self._announce_pending: dict[str, float] = {} + self._announce_lease_s: float = 0.0 self._peer_host: str | None = None self._peer_port: int | None = None # Model shape for size checks + compat hash; only set by @@ -127,6 +138,7 @@ def _setup_nixl(self, vllm_config: "VllmConfig") -> None: 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 @@ -145,8 +157,22 @@ def _setup_nixl(self, vllm_config: "VllmConfig") -> None: # How long a consumer waits for an XferAck. The producer answers from # its scheduler step, so its reply latency scales with the encoder's # batch size: a deployment whose steps run longer must raise this. + # The default is the producer's own pin lease, the only value that + # cannot strand a grant: giving up earlier leaves the producer + # pinning blocks for a consumer that has already left. 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 + ) + ) + # How long a producer holds an announced encoding for a consumer that + # has not asked for it yet. It must cover the orchestrator's forward + # plus the consumer's queueing delay; the pin is released as soon as + # the read completes, so the steady-state cost is the work in flight. + self._announce_lease_s = float( + self._ec_config.get_from_extra_config( + "producer_announce_lease_s", PRODUCER_PIN_LEASE_S + ) ) self._compat_hash = compute_ec_compatibility_hash( vllm_version=VLLM_VERSION, @@ -387,6 +413,55 @@ def _start_xfer( ) return True + def _hold_announced(self, mm_hash: str, entry: Any) -> None: + """Hold an announced encoding until its consumer has read it.""" + import time + + deadline = time.monotonic() + self._announce_lease_s + if mm_hash in self._announce_pins: + # Announced again by a later request: extend the hold rather than + # taking a second pin, so one release is enough. + self._announce_pins[mm_hash] = deadline + return + if mm_hash in self._announce_pending: + self._announce_pending[mm_hash] = deadline + return + if not entry.ready: + self._announce_pending[mm_hash] = deadline + return + self._cache.pin(mm_hash) + self._announce_pins[mm_hash] = deadline + + def _release_announce_pins(self) -> None: + """Release holds whose read has landed, and sweep lapsed leases.""" + import time + + assert self._producer_session is not None + for mm_hash in self._producer_session.take_served(): + if self._announce_pins.pop(mm_hash, None) is not None: + self._unpin_announced(mm_hash) + else: + self._announce_pending.pop(mm_hash, None) + now = time.monotonic() + for mm_hash, deadline in list(self._announce_pins.items()): + if now > deadline: + del self._announce_pins[mm_hash] + self._unpin_announced(mm_hash) + logger.debug( + "EC producer: hold on mm_hash=%s lapsed; releasing", mm_hash + ) + for mm_hash, deadline in list(self._announce_pending.items()): + if now > deadline: + del self._announce_pending[mm_hash] + + def _unpin_announced(self, mm_hash: str) -> None: + # A pinned entry cannot be evicted or discarded, so it is normally + # still here; tolerate its absence rather than assert on a shutdown + # race. + if self._cache.get(mm_hash) is None: + return + self._cache.unpin(mm_hash) + def _poll_step(self) -> None: import time @@ -480,6 +555,7 @@ def build_connector_meta( if self._is_producer: if self._nixl_enabled and self._producer_session is not None: self._producer_session.poll_step() + self._release_announce_pins() meta.saves = self._pending_saves self._pending_saves = {} if self._is_consumer: @@ -512,6 +588,12 @@ def update_connector_output(self, connector_output: "ECConnectorOutput") -> None ) elif not entry.ready: self._cache.mark_ready(mm_hash) + # It has just become evictable, so an announcement made while + # the save was in flight takes its pin now. + deadline = self._announce_pending.pop(mm_hash, None) + if deadline is not None: + self._cache.pin(mm_hash) + self._announce_pins[mm_hash] = deadline logger.debug("EC producer: mm_hash=%s marked ready", mm_hash) for transfer_id in meta.completed_loads: pending = self._load_acks.get(transfer_id) @@ -578,6 +660,7 @@ def request_finished( peer_port=self._peer_port, size_bytes=size_bytes, ) + self._hold_announced(mm_hash, entry) logger.debug( "EC producer: announcing NIXL-readable encodings req_id=%s items=%s", request.request_id, diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/session.py b/vllm/distributed/ec_transfer/ec_connector/cpu/session.py index 73153499c3995..22f2f1fc43511 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/session.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/session.py @@ -260,6 +260,9 @@ def __init__( # Key: "{consumer_session_id}:{mm_hash}" — matches the NIXL notif_msg exactly, # so completion notification lookup is a direct dict pop with no parsing. self._active_xfers: dict[str, ProducerXfer] = {} + # mm_hash values whose read finished or whose grant lapsed, for the + # scheduler to release the pin it took when announcing them. + self._served: set[str] = set() self._req_decoder = msgspec.msgpack.Decoder(XferReq) self._req_encoder = msgspec.msgpack.Encoder() @@ -365,12 +368,18 @@ def _drain_notifs(self) -> None: xfer = self._active_xfers.pop(key, None) if xfer is not None: self._cache.unpin(xfer.mm_hash) + self._served.add(xfer.mm_hash) logger.debug( "EC producer: NIXL READ completed mm_hash=%s key=%s", xfer.mm_hash, key, ) + def take_served(self) -> set[str]: + """Return and clear the encodings whose read finished or lapsed.""" + served, self._served = self._served, set() + return served + def _sweep_timeouts(self) -> None: for key, xfer in list(self._active_xfers.items()): if xfer.is_expired(): @@ -382,6 +391,7 @@ def _sweep_timeouts(self) -> None: key, ) self._cache.unpin(xfer.mm_hash) + self._served.add(xfer.mm_hash) del self._active_xfers[key] From 5f72947ae375268c4ecd5a7d6b034cf300b856f3 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Fri, 4 Sep 2026 07:02:43 +0000 Subject: [PATCH 3/8] [EPD] Harden the EC admit path against a hostile announcement Signed-off-by: Tianyu Guo --- .../unit/test_scheduler_nixl_consumer.py | 64 ++++++++ .../unit/test_scheduler_nixl_producer.py | 25 +++- .../ec_connector/cpu/scheduler/__init__.py | 138 +++++++++++------- .../ec_transfer/ec_connector/cpu/session.py | 18 ++- 4 files changed, 179 insertions(+), 66 deletions(-) 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 90dc80d724365..204198bd7a831 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time import uuid +import pytest + 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 ( @@ -241,6 +244,67 @@ def test_size_mismatch_fails_the_request(monkeypatch): s.shutdown() +@pytest.mark.parametrize( + "announced", + [ + {"peer_host": "h", "peer_port": 1, "size_bytes": None}, + {"peer_host": "h", "peer_port": 1, "size_bytes": "big"}, + {"peer_host": "h", "peer_port": 1}, + 123, + ], + ids=["null-size", "text-size", "no-size", "not-a-mapping"], +) +def test_unusable_announcement_fails_the_request(monkeypatch, announced): + """`ec_transfer_params` arrives on the request, so it may be anything. + + Raising here would propagate out of Scheduler.schedule() and take the + engine down, which is the failure this whole path exists to prevent. + """ + s = _consumer_sched(monkeypatch) + req = _Request([_Feature("h1", 1)], params={"h1": announced}) + admitted = s.ensure_cache_available(req, 0) + if announced == 123: + # Not a mapping at all: indistinguishable from "never remote", so the + # encoder runs locally rather than the request failing. + assert admitted is True + assert s.get_unrecoverable_requests() == set() + else: + assert admitted is False + assert s.get_unrecoverable_requests() == {"r1"} + assert "h1" not in s._in_flight + s.shutdown() + + +def test_deferral_budget_is_per_request(monkeypatch): + """A request must not inherit an earlier one's clock for a shared item. + + Two requests needing the same encoding hit the same transient condition; + the newcomer gets its own budget, so the one that has actually waited too + long is the one that fails. + """ + s = _consumer_sched(monkeypatch) + # A not-ready orphan entry defers every request that needs it. + assert s._cache.alloc("h1", 1) is not 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.get_unrecoverable_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.get_unrecoverable_requests() == set() + assert s.ensure_cache_available(old, 0) is False + assert s.get_unrecoverable_requests() == {"old"} + 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 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 383bdd1f194ba..915d9eaa032b9 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -120,13 +120,13 @@ class _FakeProducerSession: """Stands in for ProducerSession so build_connector_meta can run.""" def __init__(self): - self.served = set() + self.served: list[str] = [] def poll_step(self): pass def take_served(self): - served, self.served = self.served, set() + served, self.served = self.served, [] return served @@ -169,7 +169,7 @@ def test_hold_is_released_once_the_read_lands(monkeypatch): _announce(s, "h1") assert not entry.evictable - s._producer_session.served.add("h1") + s._producer_session.served.append("h1") s.build_connector_meta(scheduler_output=None) assert entry.evictable s.shutdown() @@ -206,7 +206,7 @@ def test_hold_is_taken_when_a_late_save_lands(monkeypatch): entry = s._cache.alloc("h1", 2) assert entry is not None and not entry.ready _announce(s, "h1") - assert "h1" in s._announce_pending + assert s._announce["h1"].pending == 1 s.update_connector_output( ECConnectorOutput( @@ -215,19 +215,28 @@ def test_hold_is_taken_when_a_late_save_lands(monkeypatch): ) assert entry.ready assert not entry.evictable - assert "h1" not in s._announce_pending + assert not s._announce["h1"].pending s.shutdown() -def test_reannouncing_extends_the_hold_without_a_second_pin(monkeypatch): - """One release must be enough however many requests announced it.""" +def test_each_announcement_keeps_its_own_hold(monkeypatch): + """Two consumers told to read the same encoding each need it to survive. + + Releasing on the first read would let the entry be evicted while the + second consumer, whose media has already been rewritten to a remote + reference, has not started its read. + """ s = _announcing_sched(monkeypatch) entry = s._cache.alloc("h1", 2) s._cache.mark_ready("h1") _announce(s, "h1") _announce(s, "h1") - s._producer_session.served.add("h1") + s._producer_session.served.append("h1") + s.build_connector_meta(scheduler_output=None) + assert not entry.evictable + + s._producer_session.served.append("h1") s.build_connector_meta(scheduler_output=None) assert entry.evictable s.shutdown() 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 2b18e0ee1d13c..76bad0c599a63 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 @@ for the ECCPUConnector. """ +from dataclasses import dataclass from typing import TYPE_CHECKING, Any import torch @@ -39,6 +40,22 @@ _ADMIT_DEFER_TIMEOUT_S = 60.0 +@dataclass +class _AnnounceHold: + """Pins kept for consumers that were told where to read an encoding. + + One per announcement, not one per encoding: two requests routed to two + consumers each need the entry to survive until their own read lands. + """ + + deadline: float + # Pins taken on a ready entry. + holds: int = 0 + # Announcements made before the save landed. A not-ready entry cannot be + # evicted, so these take their pins at mark_ready. + pending: int = 0 + + class ECCPUScheduler: """Scheduler delegate for the ECCPUConnector.""" @@ -95,10 +112,12 @@ def __init__(self, vllm_config: "VllmConfig") -> None: self._in_flight: set[str] = set() self._tombstones: set[str] = set() self._step_completed: set[str] = set() - # mm_hash -> when a transient condition first held a request back, so - # one that never clears fails the request instead of deferring it - # forever. Cleared as soon as the encoding lands. - self._deferred_since: dict[str, float] = {} + # (request_id, mm_hash) -> when a transient condition first held this + # request back, so one that never clears fails the request instead of + # deferring it forever. Per request: a shared mm_hash must not make a + # newly arrived request inherit an older one's clock. Cleared as soon + # as the encoding lands. + self._deferred_since: dict[tuple[str, str], float] = {} # Requests needing a remote encoding that will not arrive. Drained by # get_unrecoverable_requests(); the scheduler aborts them. self._unrecoverable: set[str] = set() @@ -106,12 +125,7 @@ def __init__(self, vllm_config: "VllmConfig") -> None: # later step, by which time the orchestrator has rewritten the media # off the request. These hold the encoding across that window: FIFO # eviction inside it leaves the consumer nothing to fall back on. - # mm_hash -> lease deadline of a pin already taken. - self._announce_pins: dict[str, float] = {} - # mm_hash -> lease deadline, for entries announced before their save - # landed. A not-ready entry is already non-evictable, so the pin is - # taken when the worker reports the save. - self._announce_pending: dict[str, float] = {} + self._announce: dict[str, _AnnounceHold] = {} self._announce_lease_s: float = 0.0 self._peer_host: str | None = None self._peer_port: int | None = None @@ -262,9 +276,11 @@ def _nixl_consumer_admit( announced = params.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 announced is not None and "peer_host" in announced + if isinstance(announced, dict) and "peer_host" in announced else None ) @@ -272,7 +288,7 @@ def _nixl_consumer_admit( 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(mm_hash, None) + self._deferred_since.pop((request.request_id, mm_hash), None) continue if mm_hash in self._in_flight or mm_hash in self._step_completed: pending = True @@ -299,7 +315,11 @@ def _nixl_consumer_admit( continue expected = pos.length * self._hidden_dim * self._element_size - size = int(remote.get("size_bytes", -1)) + try: + size = int(remote["size_bytes"]) + except (KeyError, TypeError, ValueError): + self._fail(request, mm_hash, "the announced size was unusable") + return False if size != expected: logger.warning( "EC consumer: size mismatch mm_hash=%s announced=%d expected=%d", @@ -325,7 +345,7 @@ def _nixl_consumer_admit( pending = True continue self._in_flight.add(mm_hash) - self._deferred_since.pop(mm_hash, None) + self._deferred_since.pop((request.request_id, mm_hash), None) pending = True return not pending @@ -335,7 +355,7 @@ def _defer(self, request: "Request", mm_hash: str, now: float, why: str) -> bool Returns False once the condition has outlived its budget, having marked the request unschedulable. """ - since = self._deferred_since.setdefault(mm_hash, now) + since = self._deferred_since.setdefault((request.request_id, mm_hash), now) waited = now - since if waited <= _ADMIT_DEFER_TIMEOUT_S: return True @@ -343,7 +363,7 @@ def _defer(self, request: "Request", mm_hash: str, now: float, why: str) -> bool return False def _fail(self, request: "Request", mm_hash: str, why: str) -> None: - self._deferred_since.pop(mm_hash, None) + self._deferred_since.pop((request.request_id, mm_hash), None) self._unrecoverable.add(request.request_id) logger.error( "EC consumer: request %s needs remote encoding mm_hash=%s but %s. " @@ -414,45 +434,49 @@ def _start_xfer( return True def _hold_announced(self, mm_hash: str, entry: Any) -> None: - """Hold an announced encoding until its consumer has read it.""" + """Hold an announced encoding until its consumer has read it. + + One hold per announcement: the same encoding announced to two + consumers must survive until both reads land, so re-announcing takes + a second pin rather than only extending the window. + """ import time - deadline = time.monotonic() + self._announce_lease_s - if mm_hash in self._announce_pins: - # Announced again by a later request: extend the hold rather than - # taking a second pin, so one release is enough. - self._announce_pins[mm_hash] = deadline - return - if mm_hash in self._announce_pending: - self._announce_pending[mm_hash] = deadline - return - if not entry.ready: - self._announce_pending[mm_hash] = deadline - return - self._cache.pin(mm_hash) - self._announce_pins[mm_hash] = deadline + hold = self._announce.get(mm_hash) + if hold is None: + hold = self._announce[mm_hash] = _AnnounceHold(0.0) + hold.deadline = time.monotonic() + self._announce_lease_s + if entry.ready: + self._cache.pin(mm_hash) + hold.holds += 1 + else: + hold.pending += 1 def _release_announce_pins(self) -> None: - """Release holds whose read has landed, and sweep lapsed leases.""" + """Release one hold per landed read, and sweep lapsed leases.""" import time assert self._producer_session is not None for mm_hash in self._producer_session.take_served(): - if self._announce_pins.pop(mm_hash, None) is not None: + hold = self._announce.get(mm_hash) + if hold is None: + # Already swept by a lapsed lease; nothing left to release. + continue + if hold.holds: + hold.holds -= 1 self._unpin_announced(mm_hash) - else: - self._announce_pending.pop(mm_hash, None) + elif hold.pending: + hold.pending -= 1 + if not hold.holds and not hold.pending: + del self._announce[mm_hash] now = time.monotonic() - for mm_hash, deadline in list(self._announce_pins.items()): - if now > deadline: - del self._announce_pins[mm_hash] + for mm_hash, hold in list(self._announce.items()): + if now <= hold.deadline: + continue + for _ in range(hold.holds): self._unpin_announced(mm_hash) - logger.debug( - "EC producer: hold on mm_hash=%s lapsed; releasing", mm_hash - ) - for mm_hash, deadline in list(self._announce_pending.items()): - if now > deadline: - del self._announce_pending[mm_hash] + del self._announce[mm_hash] + logger.debug("EC producer: hold on mm_hash=%s lapsed; releasing", mm_hash) def _unpin_announced(self, mm_hash: str) -> None: # A pinned entry cannot be evicted or discarded, so it is normally @@ -473,6 +497,18 @@ def _poll_step(self) -> None: self._on_peer_down(addr) for session in self._sessions.values(): self._process_session_results(session) + self._sweep_deferrals(now) + + def _sweep_deferrals(self, now: float) -> None: + """Drop deferral clocks left behind by requests that went away. + + A request still being scheduled is failed, and its clock removed, as + soon as it passes the budget; anything older belongs to a request the + scheduler has since dropped. + """ + for key, since in list(self._deferred_since.items()): + if now - since > 2 * _ADMIT_DEFER_TIMEOUT_S: + del self._deferred_since[key] def _process_session_results(self, session: Any) -> None: r = session.take_results() @@ -588,12 +624,14 @@ def update_connector_output(self, connector_output: "ECConnectorOutput") -> None ) elif not entry.ready: self._cache.mark_ready(mm_hash) - # It has just become evictable, so an announcement made while - # the save was in flight takes its pin now. - deadline = self._announce_pending.pop(mm_hash, None) - if deadline is not None: - self._cache.pin(mm_hash) - self._announce_pins[mm_hash] = deadline + # It has just become evictable, so announcements made while + # the save was in flight take their pins now. + hold = self._announce.get(mm_hash) + if hold is not None and hold.pending: + for _ in range(hold.pending): + self._cache.pin(mm_hash) + hold.holds += hold.pending + hold.pending = 0 logger.debug("EC producer: mm_hash=%s marked ready", mm_hash) for transfer_id in meta.completed_loads: pending = self._load_acks.get(transfer_id) diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/session.py b/vllm/distributed/ec_transfer/ec_connector/cpu/session.py index 22f2f1fc43511..1e9fd7c7573a7 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/session.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/session.py @@ -260,9 +260,11 @@ def __init__( # Key: "{consumer_session_id}:{mm_hash}" — matches the NIXL notif_msg exactly, # so completion notification lookup is a direct dict pop with no parsing. self._active_xfers: dict[str, ProducerXfer] = {} - # mm_hash values whose read finished or whose grant lapsed, for the - # scheduler to release the pin it took when announcing them. - self._served: set[str] = set() + # One entry per grant whose read finished or whose lease lapsed, for + # the scheduler to release the matching pin it took when announcing. + # A list, not a set: two consumers reading the same encoding hold two + # pins, and collapsing their events would release only one. + self._served: list[str] = [] self._req_decoder = msgspec.msgpack.Decoder(XferReq) self._req_encoder = msgspec.msgpack.Encoder() @@ -368,16 +370,16 @@ def _drain_notifs(self) -> None: xfer = self._active_xfers.pop(key, None) if xfer is not None: self._cache.unpin(xfer.mm_hash) - self._served.add(xfer.mm_hash) + self._served.append(xfer.mm_hash) logger.debug( "EC producer: NIXL READ completed mm_hash=%s key=%s", xfer.mm_hash, key, ) - def take_served(self) -> set[str]: - """Return and clear the encodings whose read finished or lapsed.""" - served, self._served = self._served, set() + def take_served(self) -> list[str]: + """Return and clear one entry per grant that finished or lapsed.""" + served, self._served = self._served, [] return served def _sweep_timeouts(self) -> None: @@ -391,7 +393,7 @@ def _sweep_timeouts(self) -> None: key, ) self._cache.unpin(xfer.mm_hash) - self._served.add(xfer.mm_hash) + self._served.append(xfer.mm_hash) del self._active_xfers[key] From e5ef04d1122ad24c1fa507b715d7fc8987f6b28a Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Fri, 4 Sep 2026 07:48:33 +0000 Subject: [PATCH 4/8] [EPD] Give each announced encoding its own lease Signed-off-by: Tianyu Guo --- .../unit/test_scheduler_nixl_producer.py | 36 +++++++++- .../ec_connector/cpu/scheduler/__init__.py | 72 +++++++++++-------- 2 files changed, 78 insertions(+), 30 deletions(-) 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 915d9eaa032b9..593be69543a9a 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -206,7 +206,7 @@ def test_hold_is_taken_when_a_late_save_lands(monkeypatch): entry = s._cache.alloc("h1", 2) assert entry is not None and not entry.ready _announce(s, "h1") - assert s._announce["h1"].pending == 1 + assert len(s._announce["h1"].pending) == 1 s.update_connector_output( ECConnectorOutput( @@ -240,3 +240,37 @@ def test_each_announcement_keeps_its_own_hold(monkeypatch): s.build_connector_meta(scheduler_output=None) assert entry.evictable s.shutdown() + + +def test_a_reannouncement_does_not_extend_the_earlier_hold(monkeypatch): + """Each announcement waits out its own lease, not the newest one. + + An encoding announced again before its lease expires, to consumers that + never read, would otherwise keep a pin that no read and no expiry ever + releases, leaving the entry non-evictable for as long as it stays hot. + """ + import time + + clock = [1000.0] + monkeypatch.setattr(time, "monotonic", lambda: clock[0]) + + s = _announcing_sched(monkeypatch, lease=30.0) + entry = s._cache.alloc("h1", 2) + s._cache.mark_ready("h1") + + _announce(s, "h1") + clock[0] += 29.0 + _announce(s, "h1") + assert len(s._announce["h1"].holds) == 2 + + # Past the first lease, inside the second. + clock[0] += 2.0 + s.build_connector_meta(scheduler_output=None) + assert len(s._announce["h1"].holds) == 1 + assert not entry.evictable + + clock[0] += 30.0 + s.build_connector_meta(scheduler_output=None) + assert "h1" not in s._announce + assert entry.evictable + s.shutdown() 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 76bad0c599a63..9287793be4557 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -7,7 +7,8 @@ for the ECCPUConnector. """ -from dataclasses import dataclass +from collections import deque +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any import torch @@ -42,18 +43,20 @@ @dataclass class _AnnounceHold: - """Pins kept for consumers that were told where to read an encoding. + """Leases kept for consumers that were told where to read an encoding. - One per announcement, not one per encoding: two requests routed to two - consumers each need the entry to survive until their own read lands. + One lease per announcement, not one per encoding: two requests routed to + two consumers each need the entry to survive until their own read lands, + and each waits out its own window. A later announcement must not extend + an earlier one, or an encoding announced faster than the lease to + consumers that never read would stay pinned indefinitely. """ - deadline: float - # Pins taken on a ready entry. - holds: int = 0 - # Announcements made before the save landed. A not-ready entry cannot be - # evicted, so these take their pins at mark_ready. - pending: int = 0 + # Deadlines of the announcements holding a pin, oldest first. + holds: deque[float] = field(default_factory=deque) + # Deadlines of announcements made before the save landed. A not-ready + # entry cannot be evicted, so these take their pins at mark_ready. + pending: deque[float] = field(default_factory=deque) class ECCPUScheduler: @@ -436,24 +439,22 @@ def _start_xfer( def _hold_announced(self, mm_hash: str, entry: Any) -> None: """Hold an announced encoding until its consumer has read it. - One hold per announcement: the same encoding announced to two + One lease per announcement: the same encoding announced to two consumers must survive until both reads land, so re-announcing takes - a second pin rather than only extending the window. + its own pin and its own window rather than extending an earlier one. """ import time - hold = self._announce.get(mm_hash) - if hold is None: - hold = self._announce[mm_hash] = _AnnounceHold(0.0) - hold.deadline = time.monotonic() + self._announce_lease_s + hold = self._announce.setdefault(mm_hash, _AnnounceHold()) + deadline = time.monotonic() + self._announce_lease_s if entry.ready: self._cache.pin(mm_hash) - hold.holds += 1 + hold.holds.append(deadline) else: - hold.pending += 1 + hold.pending.append(deadline) def _release_announce_pins(self) -> None: - """Release one hold per landed read, and sweep lapsed leases.""" + """Release one lease per landed read, and sweep lapsed leases.""" import time assert self._producer_session is not None @@ -462,21 +463,34 @@ def _release_announce_pins(self) -> None: if hold is None: # Already swept by a lapsed lease; nothing left to release. continue + # A read names only the encoding, so it retires the oldest lease + # on it; every lease is the same length, so which one it was does + # not change when the others expire. if hold.holds: - hold.holds -= 1 + hold.holds.popleft() self._unpin_announced(mm_hash) elif hold.pending: - hold.pending -= 1 + hold.pending.popleft() if not hold.holds and not hold.pending: del self._announce[mm_hash] now = time.monotonic() for mm_hash, hold in list(self._announce.items()): - if now <= hold.deadline: - continue - for _ in range(hold.holds): + lapsed = 0 + while hold.holds and hold.holds[0] < now: + hold.holds.popleft() self._unpin_announced(mm_hash) - del self._announce[mm_hash] - logger.debug("EC producer: hold on mm_hash=%s lapsed; releasing", mm_hash) + lapsed += 1 + while hold.pending and hold.pending[0] < now: + hold.pending.popleft() + lapsed += 1 + if not hold.holds and not hold.pending: + del self._announce[mm_hash] + if lapsed: + logger.debug( + "EC producer: %d hold(s) on mm_hash=%s lapsed; releasing", + lapsed, + mm_hash, + ) def _unpin_announced(self, mm_hash: str) -> None: # A pinned entry cannot be evicted or discarded, so it is normally @@ -628,10 +642,10 @@ def update_connector_output(self, connector_output: "ECConnectorOutput") -> None # the save was in flight take their pins now. hold = self._announce.get(mm_hash) if hold is not None and hold.pending: - for _ in range(hold.pending): + for _ in range(len(hold.pending)): self._cache.pin(mm_hash) - hold.holds += hold.pending - hold.pending = 0 + hold.holds.extend(hold.pending) + hold.pending.clear() logger.debug("EC producer: mm_hash=%s marked ready", mm_hash) for transfer_id in meta.completed_loads: pending = self._load_acks.get(transfer_id) From 282be5d7b46ec57a907a264666b44d089f1d0c32 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Mon, 7 Sep 2026 08:00:14 +0000 Subject: [PATCH 5/8] [EPD] Simplify NIXL request failure regression tests Co-authored-by: OpenAI Codex Signed-off-by: Tianyu Guo --- .../unit/test_scheduler_nixl_consumer.py | 104 ++++-------------- .../unit/test_scheduler_nixl_producer.py | 73 ++---------- tests/v1/ec_connector/unit/test_session.py | 19 +--- 3 files changed, 32 insertions(+), 164 deletions(-) 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 92b12ecb373f5..88b9f49575542 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py @@ -153,19 +153,16 @@ def _fake_start(mm_hash, info, size): 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_fails_the_request(monkeypatch): - """A failed remote read must fail the request, not admit it. - - The media reached the producer only, so an admitted request has no - embedding and no way to compute one: the model would receive an empty - multimodal batch and assert. - """ + """A failed remote read must not admit a request with no local media.""" s = _consumer_sched(monkeypatch) fake = _FakeSession() @@ -231,60 +228,31 @@ def _fake_start(mm_hash, info, size): s.shutdown() -def test_size_mismatch_fails_the_request(monkeypatch): - """A bad announcement is terminal: there is no second source to try.""" - 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 False - assert "h1" not in s._in_flight - assert s._cache.get("h1") is None - assert s.take_unavailable_requests() == {"r1"} - s.shutdown() - - @pytest.mark.parametrize( - "announced", - [ - {"peer_host": "h", "peer_port": 1, "size_bytes": None}, - {"peer_host": "h", "peer_port": 1, "size_bytes": "big"}, - {"peer_host": "h", "peer_port": 1}, - 123, - ], - ids=["null-size", "text-size", "no-size", "not-a-mapping"], + "size_fields", + [{"size_bytes": 999}, {"size_bytes": None}, {"size_bytes": "big"}, {}], + ids=["mismatch", "null-size", "text-size", "no-size"], ) -def test_unusable_announcement_fails_the_request(monkeypatch, announced): - """`ec_transfer_params` arrives on the request, so it may be anything. - - Raising here would propagate out of Scheduler.schedule() and take the - engine down, which is the failure this whole path exists to prevent. - """ +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) + announced = {"peer_host": "h", "peer_port": 1, **size_fields} req = _Request([_Feature("h1", 1)], params={"h1": announced}) - admitted = s.ensure_cache_available(req, 0) - if announced == 123: - # Not a mapping at all: indistinguishable from "never remote", so the - # encoder runs locally rather than the request failing. - assert admitted is True - assert s.take_unavailable_requests() == set() - else: - assert admitted is False - assert s.take_unavailable_requests() == {"r1"} + 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): - """A request must not inherit an earlier one's clock for a shared item. - - Two requests needing the same encoding hit the same transient condition; - the newcomer gets its own budget, so the one that has actually waited too - long is the one that fails. - """ +@pytest.mark.parametrize("pool_full", [False, True], ids=["not-ready", "full-pool"]) +def test_deferral_budget_is_per_request(monkeypatch, pool_full): + """Transient failures defer until each request exhausts its own budget.""" s = _consumer_sched(monkeypatch) - # A not-ready orphan entry defers every request that needs it. - assert s._cache.alloc("h1", 1) is not None + if pool_full: + monkeypatch.setattr(s._cache, "alloc", lambda key, n: None) + else: + assert s._cache.alloc("h1", 1) is not 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") @@ -302,6 +270,7 @@ def test_deferral_budget_is_per_request(monkeypatch): assert s.take_unavailable_requests() == set() assert s.ensure_cache_available(old, 0) is False assert s.take_unavailable_requests() == {"old"} + assert "h1" not in s._in_flight s.shutdown() @@ -353,35 +322,6 @@ def _spy_alloc(key, n): s.shutdown() -def test_full_pool_defers_then_fails_the_request(monkeypatch): - """A full pool is transient, so defer — but not forever. - - Deferring without a budget would hold the request until the client gave - up, with no record of why. - """ - import time as _time - - 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)) - - t0 = _time.monotonic() - monkeypatch.setattr(_time, "monotonic", lambda: t0) - assert s.ensure_cache_available(req, 0) is False - assert "h1" not in s._in_flight - assert s.take_unavailable_requests() == set() - - # Still full once the budget has elapsed: fail rather than defer again. - monkeypatch.setattr( - _time, "monotonic", lambda: t0 + sched_mod._ADMIT_DEFER_TIMEOUT_S + 1 - ) - assert s.ensure_cache_available(req, 0) is False - assert s.take_unavailable_requests() == {"r1"} - s.shutdown() - - def test_already_computed_feature_admits_without_transfer(monkeypatch): # A feature entirely within num_computed_tokens is skipped by admit; no # cache lookup/alloc or transfer should be triggered for it. 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 593be69543a9a..c223f0b108f6f 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time import uuid import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod @@ -144,59 +145,8 @@ def _announce(s, mm_hash, length=2): return s.request_finished(_Request([_Feature(mm_hash, length=length)])) -def test_announced_encoding_is_not_evictable(monkeypatch): - """Announcing publishes an address the consumer uses on a later step. - - By then the orchestrator has rewritten the media off the request, so an - eviction inside that window leaves the consumer nothing to fall back on. - """ - s = _announcing_sched(monkeypatch) - entry = s._cache.alloc("h1", 2) - assert entry is not None - s._cache.mark_ready("h1") - assert entry.evictable - - _delay, params = _announce(s, "h1") - assert "peer_host" in params["h1"] - assert not entry.evictable - s.shutdown() - - -def test_hold_is_released_once_the_read_lands(monkeypatch): - s = _announcing_sched(monkeypatch) - entry = s._cache.alloc("h1", 2) - s._cache.mark_ready("h1") - _announce(s, "h1") - assert not entry.evictable - - s._producer_session.served.append("h1") - s.build_connector_meta(scheduler_output=None) - assert entry.evictable - s.shutdown() - - -def test_hold_lapses_when_no_consumer_ever_reads(monkeypatch): - """A consumer that never asks must not pin the pool forever.""" - s = _announcing_sched(monkeypatch, lease=0.0) - entry = s._cache.alloc("h1", 2) - s._cache.mark_ready("h1") - _announce(s, "h1") - assert not entry.evictable - - import time as _time - - _time.sleep(0.01) - s.build_connector_meta(scheduler_output=None) - assert entry.evictable - s.shutdown() - - def test_hold_is_taken_when_a_late_save_lands(monkeypatch): - """An entry announced mid-save is pinned when it becomes evictable. - - A not-ready entry cannot be evicted, so the hold has to start at - mark_ready rather than at the announcement. - """ + """A save completing after its announcement must acquire the pending pin.""" from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( ECCPUWorkerMetadata, ) @@ -220,16 +170,14 @@ def test_hold_is_taken_when_a_late_save_lands(monkeypatch): def test_each_announcement_keeps_its_own_hold(monkeypatch): - """Two consumers told to read the same encoding each need it to survive. - - Releasing on the first read would let the entry be evicted while the - second consumer, whose media has already been rewritten to a remote - reference, has not started its read. - """ + """An announced encoding stays pinned until both consumers have read it.""" s = _announcing_sched(monkeypatch) entry = s._cache.alloc("h1", 2) + assert entry is not None s._cache.mark_ready("h1") + assert entry.evictable _announce(s, "h1") + assert not entry.evictable _announce(s, "h1") s._producer_session.served.append("h1") @@ -243,14 +191,7 @@ def test_each_announcement_keeps_its_own_hold(monkeypatch): def test_a_reannouncement_does_not_extend_the_earlier_hold(monkeypatch): - """Each announcement waits out its own lease, not the newest one. - - An encoding announced again before its lease expires, to consumers that - never read, would otherwise keep a pin that no read and no expiry ever - releases, leaving the entry non-evictable for as long as it stays hot. - """ - import time - + """Unread announcements expire independently, eventually freeing the entry.""" clock = [1000.0] monkeypatch.setattr(time, "monotonic", lambda: clock[0]) diff --git a/tests/v1/ec_connector/unit/test_session.py b/tests/v1/ec_connector/unit/test_session.py index 96ffe7926ec5b..4d78b0426afca 100644 --- a/tests/v1/ec_connector/unit/test_session.py +++ b/tests/v1/ec_connector/unit/test_session.py @@ -451,27 +451,14 @@ def test_consumer_session_poll_not_ready_nack_goes_to_retryable(): assert not results.completed -def test_consumer_session_ack_timeout_goes_to_retryable(): - """A late XferAck means a busy producer, not a broken one. - - The producer answers XferReqs from its scheduler step, so the reply can - lag a whole encoder forward pass. Tombstoning on that admits the request - with no embedding, which is fatal once the media has been rewritten away. - """ - s = _make_consumer_session() - s.start_xfer("h1", [0, 1], deadline=time.monotonic() - 1) - s.poll([], time.monotonic()) - results = s.take_results() - assert "h1" in results.retryable - assert not results.tombstoned - - -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.retryable + assert not r1.tombstoned r2 = s.take_results() assert not r2.retryable # cleared after first take From 37ce970dbd43cd2d24af6612eb9057d2e5e2e042 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Tue, 8 Sep 2026 00:20:07 +0000 Subject: [PATCH 6/8] [EC Connector] Bound remote retries and preserve local fallback Remove announcement leases from this fix. Keep per-request wait budgets across retries and allow local fallback when model input is available. Signed-off-by: Tianyu Guo --- .../unit/test_scheduler_nixl_consumer.py | 96 ++++++- .../unit/test_scheduler_nixl_producer.py | 105 +------ .../ec_connector/cpu/scheduler/__init__.py | 264 +++++------------- .../ec_transfer/ec_connector/cpu/session.py | 17 +- 4 files changed, 158 insertions(+), 324 deletions(-) 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 88b9f49575542..03bfbd8e9c9f8 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py @@ -4,6 +4,7 @@ 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 @@ -11,6 +12,11 @@ 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 @@ -21,10 +27,12 @@ def __init__(self, offset, length): 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: @@ -97,6 +105,7 @@ def _region(cfg): # 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 @@ -161,8 +170,10 @@ def test_no_remote_announcement_is_ready(monkeypatch, params): s.shutdown() -def test_tombstoned_read_fails_the_request(monkeypatch): - """A failed remote read must not admit a request with no local media.""" +@pytest.mark.parametrize("payload", [None, "metadata", "pixels", "embeds", "shm"]) +@pytest.mark.parametrize("failure", ["tombstone", "timeout"]) +def test_failed_read_falls_back_only_with_local_input(monkeypatch, payload, failure): + """A remote miss is fatal only when local model input is unavailable.""" s = _consumer_sched(monkeypatch) fake = _FakeSession() @@ -173,23 +184,90 @@ def _fake_start(mm_hash, info, size): 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)) + now = [1000.0] + monkeypatch.setattr(time, "monotonic", lambda: now[0]) # 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") + if failure == "tombstone": + fake._results.tombstoned.add("h1") + else: + now[0] += sched_mod._ADMIT_DEFER_TIMEOUT_S + 1 + fake._results.retryable.add("h1") s._sessions[("h", 1)] = fake - # Step 2: the entry is discarded and the request is reported unschedulable. - assert s.ensure_cache_available(req, 0) is False + 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 - assert s.take_unavailable_requests() == {"r1"} + 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() + + +def test_request_finished_clears_only_its_wait_budget(monkeypatch): + """Cancelled requests leave no clocks behind, even for a shared hash.""" + s = _consumer_sched(monkeypatch) + monkeypatch.setattr(s, "_start_xfer", lambda *args: False) + old = _Request([_Feature("h1")], params=_params("h1", 1), req_id="old") + new = _Request([_Feature("h1")], params=_params("h1", 1), req_id="new") + assert not s.ensure_cache_available(old, 0) + assert not s.ensure_cache_available(new, 0) + assert s.request_finished(old) == (False, None) + assert set(s._deferred_since) == {("new", "h1")} + s.shutdown() + + +@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() 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 c223f0b108f6f..156d98ae794ce 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import time import uuid import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod @@ -64,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": {}, @@ -112,106 +112,3 @@ def test_request_finished_skips_unallocated_entry(monkeypatch): # no cache entry to transfer (empty "metadata": no fields, no transfer). assert params == {"h1": {"metadata": {}}} s.shutdown() - - -# ── announced encodings are held until read ────────────────────────────────── - - -class _FakeProducerSession: - """Stands in for ProducerSession so build_connector_meta can run.""" - - def __init__(self): - self.served: list[str] = [] - - def poll_step(self): - pass - - def take_served(self): - served, self.served = self.served, [] - return served - - -def _announcing_sched(monkeypatch, lease=30.0): - s = _sched_gate_off(monkeypatch) - s._nixl_enabled = True - s._peer_host, s._peer_port = "1.2.3.4", 5601 - s._hidden_dim, s._element_size = 32, 2 - s._announce_lease_s = lease - s._producer_session = _FakeProducerSession() - return s - - -def _announce(s, mm_hash, length=2): - return s.request_finished(_Request([_Feature(mm_hash, length=length)])) - - -def test_hold_is_taken_when_a_late_save_lands(monkeypatch): - """A save completing after its announcement must acquire the pending pin.""" - from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( - ECCPUWorkerMetadata, - ) - from vllm.v1.outputs import ECConnectorOutput - - s = _announcing_sched(monkeypatch) - entry = s._cache.alloc("h1", 2) - assert entry is not None and not entry.ready - _announce(s, "h1") - assert len(s._announce["h1"].pending) == 1 - - s.update_connector_output( - ECConnectorOutput( - ec_connector_worker_meta=ECCPUWorkerMetadata(completed_saves=["h1"]) - ) - ) - assert entry.ready - assert not entry.evictable - assert not s._announce["h1"].pending - s.shutdown() - - -def test_each_announcement_keeps_its_own_hold(monkeypatch): - """An announced encoding stays pinned until both consumers have read it.""" - s = _announcing_sched(monkeypatch) - entry = s._cache.alloc("h1", 2) - assert entry is not None - s._cache.mark_ready("h1") - assert entry.evictable - _announce(s, "h1") - assert not entry.evictable - _announce(s, "h1") - - s._producer_session.served.append("h1") - s.build_connector_meta(scheduler_output=None) - assert not entry.evictable - - s._producer_session.served.append("h1") - s.build_connector_meta(scheduler_output=None) - assert entry.evictable - s.shutdown() - - -def test_a_reannouncement_does_not_extend_the_earlier_hold(monkeypatch): - """Unread announcements expire independently, eventually freeing the entry.""" - clock = [1000.0] - monkeypatch.setattr(time, "monotonic", lambda: clock[0]) - - s = _announcing_sched(monkeypatch, lease=30.0) - entry = s._cache.alloc("h1", 2) - s._cache.mark_ready("h1") - - _announce(s, "h1") - clock[0] += 29.0 - _announce(s, "h1") - assert len(s._announce["h1"].holds) == 2 - - # Past the first lease, inside the second. - clock[0] += 2.0 - s.build_connector_meta(scheduler_output=None) - assert len(s._announce["h1"].holds) == 1 - assert not entry.evictable - - clock[0] += 30.0 - s.build_connector_meta(scheduler_output=None) - assert "h1" not in s._announce - assert entry.evictable - s.shutdown() 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 07ba68f2ad558..1d72c55298db1 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -7,9 +7,8 @@ for the ECCPUConnector. """ -from collections import deque +import time from collections.abc import Collection -from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any import torch @@ -31,35 +30,17 @@ 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__) -# How long a transient condition (a settling DMA, a full pool) may hold a -# request back before it is failed rather than deferred indefinitely. +# Total remote wait budget per request item, including all retries. _ADMIT_DEFER_TIMEOUT_S = 60.0 -@dataclass -class _AnnounceHold: - """Leases kept for consumers that were told where to read an encoding. - - One lease per announcement, not one per encoding: two requests routed to - two consumers each need the entry to survive until their own read lands, - and each waits out its own window. A later announcement must not extend - an earlier one, or an encoding announced faster than the lease to - consumers that never read would stay pinned indefinitely. - """ - - # Deadlines of the announcements holding a pin, oldest first. - holds: deque[float] = field(default_factory=deque) - # Deadlines of announcements made before the save landed. A not-ready - # entry cannot be evicted, so these take their pins at mark_ready. - pending: deque[float] = field(default_factory=deque) - - class ECCPUScheduler: """Scheduler delegate for the ECCPUConnector.""" @@ -116,21 +97,11 @@ def __init__(self, vllm_config: "VllmConfig") -> None: self._in_flight: set[str] = set() self._tombstones: set[str] = set() self._step_completed: set[str] = set() - # (request_id, mm_hash) -> when a transient condition first held this - # request back, so one that never clears fails the request instead of - # deferring it forever. Per request: a shared mm_hash must not make a - # newly arrived request inherit an older one's clock. Cleared as soon - # as the encoding lands. + # 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() - # Announcing an encoding publishes an address a consumer will use on a - # later step, by which time the orchestrator has rewritten the media - # off the request. These hold the encoding across that window: FIFO - # eviction inside it leaves the consumer nothing to fall back on. - self._announce: dict[str, _AnnounceHold] = {} - self._announce_lease_s: float = 0.0 self._peer_host: str | None = None self._peer_port: int | None = None # Model shape for size checks + compat hash; only set by @@ -175,23 +146,11 @@ def _setup_nixl(self, vllm_config: "VllmConfig") -> None: # How long a consumer waits for an XferAck. The producer answers from # its scheduler step, so its reply latency scales with the encoder's # batch size: a deployment whose steps run longer must raise this. - # The default is the producer's own pin lease, the only value that - # cannot strand a grant: giving up earlier leaves the producer - # pinning blocks for a consumer that has already left. self._ack_timeout_s = float( self._ec_config.get_from_extra_config( "consumer_ack_timeout_s", PRODUCER_PIN_LEASE_S ) ) - # How long a producer holds an announced encoding for a consumer that - # has not asked for it yet. It must cover the orchestrator's forward - # plus the consumer's queueing delay; the pin is released as soon as - # the read completes, so the steady-state cost is the work in flight. - self._announce_lease_s = float( - self._ec_config.get_from_extra_config( - "producer_announce_lease_s", PRODUCER_PIN_LEASE_S - ) - ) self._compat_hash = compute_ec_compatibility_hash( vllm_version=VLLM_VERSION, model=str(vllm_config.model_config.model), @@ -257,30 +216,18 @@ def ensure_cache_available( def _nixl_consumer_admit( self, request: "Request", num_computed_tokens: int ) -> bool: - """Admit a request only once all its remote encodings are cached. - - An item announced with a producer address lives only on that producer. - Unlike KV, it cannot be recomputed from the request: the media never - reached this instance. Scheduling a request whose remote item is - missing therefore hands the model an empty multimodal batch, so such - an item defers the request while it may still arrive, and fails the - request once it cannot. - """ - import time - - params: dict[str, dict[str, Any]] = ( - getattr(request, "ec_transfer_params", None) or {} - ) - if not params: + """Wait for remote encodings, falling back only with local input.""" + if not request.ec_transfer_params: return True now = time.monotonic() pending = False for feature in request.mm_features: 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) continue - mm_hash = feature.identifier - announced = params.get(mm_hash) + 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 @@ -297,36 +244,42 @@ def _nixl_consumer_admit( # loads it through the same path as a natively cached entry. self._deferred_since.pop((request.request_id, mm_hash), None) continue + if remote is None: + continue + since = self._deferred_since.setdefault((request.request_id, mm_hash), now) + if now - since > _ADMIT_DEFER_TIMEOUT_S: + if not self._fail_or_fallback( + request, feature, f"the remote wait exceeded {now - since:.0f}s" + ): + return False + continue if mm_hash in self._in_flight or mm_hash in self._step_completed: pending = True continue if mm_hash in self._tombstones: self._tombstones.discard(mm_hash) - if remote is None: - continue - self._fail(request, mm_hash, "the remote read failed") - return False + if not self._fail_or_fallback( + request, feature, "the remote read failed" + ): + return False + 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. - if remote is None: - continue - if not self._defer(request, mm_hash, now, "a DMA was settling"): - return False pending = True continue - if remote is None: - continue - expected = pos.length * self._hidden_dim * self._element_size try: size = int(remote["size_bytes"]) except (KeyError, TypeError, ValueError): - self._fail(request, mm_hash, "the announced size was unusable") - return False + if not self._fail_or_fallback( + request, feature, "the announced size was unusable" + ): + return False + continue if size != expected: logger.warning( "EC consumer: size mismatch mm_hash=%s announced=%d expected=%d", @@ -334,8 +287,11 @@ def _nixl_consumer_admit( size, expected, ) - self._fail(request, mm_hash, "the announced size was wrong") - return False + if not self._fail_or_fallback( + request, feature, "the announced size was wrong" + ): + return False + continue try: started = self._start_xfer(mm_hash, remote, expected) @@ -343,43 +299,49 @@ def _nixl_consumer_admit( logger.exception( "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash ) - self._fail(request, mm_hash, "the read could not be started") - return False - if not started: - # The local pool is full; it drains as other reads complete. - if not self._defer(request, mm_hash, now, "the pool was full"): + if not self._fail_or_fallback( + request, feature, "the read could not be started" + ): return False - pending = True continue - self._in_flight.add(mm_hash) - self._deferred_since.pop((request.request_id, mm_hash), None) + if started: + self._in_flight.add(mm_hash) pending = True return not pending - def _defer(self, request: "Request", mm_hash: str, now: float, why: str) -> bool: - """Hold a request back while a transient condition clears. - - Returns False once the condition has outlived its budget, having - marked the request unschedulable. - """ - since = self._deferred_since.setdefault((request.request_id, mm_hash), now) - waited = now - since - if waited <= _ADMIT_DEFER_TIMEOUT_S: - return True - self._fail(request, mm_hash, f"{why} for {waited:.0f}s") - return False - - def _fail(self, request: "Request", mm_hash: str, why: str) -> None: + def _fail_or_fallback( + self, request: "Request", feature: "MultiModalFeatureSpec", why: str + ) -> bool: + """Allow local execution only with data beyond placeholder metadata.""" + 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. " - "The media never reached this instance, so it cannot be encoded " - "locally; failing the request.", + "No local model input is available; failing the request.", request.request_id, mm_hash, why, ) + return False def take_unavailable_requests(self) -> set[str]: if not self._unrecoverable: @@ -394,10 +356,8 @@ def _start_xfer( """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 ( @@ -408,7 +368,7 @@ def _start_xfer( 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, ) @@ -440,73 +400,7 @@ def _start_xfer( ) return True - def _hold_announced(self, mm_hash: str, entry: Any) -> None: - """Hold an announced encoding until its consumer has read it. - - One lease per announcement: the same encoding announced to two - consumers must survive until both reads land, so re-announcing takes - its own pin and its own window rather than extending an earlier one. - """ - import time - - hold = self._announce.setdefault(mm_hash, _AnnounceHold()) - deadline = time.monotonic() + self._announce_lease_s - if entry.ready: - self._cache.pin(mm_hash) - hold.holds.append(deadline) - else: - hold.pending.append(deadline) - - def _release_announce_pins(self) -> None: - """Release one lease per landed read, and sweep lapsed leases.""" - import time - - assert self._producer_session is not None - for mm_hash in self._producer_session.take_served(): - hold = self._announce.get(mm_hash) - if hold is None: - # Already swept by a lapsed lease; nothing left to release. - continue - # A read names only the encoding, so it retires the oldest lease - # on it; every lease is the same length, so which one it was does - # not change when the others expire. - if hold.holds: - hold.holds.popleft() - self._unpin_announced(mm_hash) - elif hold.pending: - hold.pending.popleft() - if not hold.holds and not hold.pending: - del self._announce[mm_hash] - now = time.monotonic() - for mm_hash, hold in list(self._announce.items()): - lapsed = 0 - while hold.holds and hold.holds[0] < now: - hold.holds.popleft() - self._unpin_announced(mm_hash) - lapsed += 1 - while hold.pending and hold.pending[0] < now: - hold.pending.popleft() - lapsed += 1 - if not hold.holds and not hold.pending: - del self._announce[mm_hash] - if lapsed: - logger.debug( - "EC producer: %d hold(s) on mm_hash=%s lapsed; releasing", - lapsed, - mm_hash, - ) - - def _unpin_announced(self, mm_hash: str) -> None: - # A pinned entry cannot be evicted or discarded, so it is normally - # still here; tolerate its absence rather than assert on a shutdown - # race. - if self._cache.get(mm_hash) is None: - return - self._cache.unpin(mm_hash) - def _poll_step(self) -> None: - import time - now = time.monotonic() all_messages = self._transport.poll() for addr, session in list(self._sessions.items()): @@ -515,18 +409,6 @@ def _poll_step(self) -> None: self._on_peer_down(addr) for session in self._sessions.values(): self._process_session_results(session) - self._sweep_deferrals(now) - - def _sweep_deferrals(self, now: float) -> None: - """Drop deferral clocks left behind by requests that went away. - - A request still being scheduled is failed, and its clock removed, as - soon as it passes the budget; anything older belongs to a request the - scheduler has since dropped. - """ - for key, since in list(self._deferred_since.items()): - if now - since > 2 * _ADMIT_DEFER_TIMEOUT_S: - del self._deferred_since[key] def _process_session_results(self, session: Any) -> None: r = session.take_results() @@ -550,8 +432,7 @@ def _process_session_results(self, session: Any) -> None: 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( @@ -609,7 +490,6 @@ def build_connector_meta( if self._is_producer: if self._nixl_enabled and self._producer_session is not None: self._producer_session.poll_step() - self._release_announce_pins() meta.saves = self._pending_saves self._pending_saves = {} if self._is_consumer: @@ -642,14 +522,6 @@ def update_connector_output(self, connector_output: "ECConnectorOutput") -> None ) elif not entry.ready: self._cache.mark_ready(mm_hash) - # It has just become evictable, so announcements made while - # the save was in flight take their pins now. - hold = self._announce.get(mm_hash) - if hold is not None and hold.pending: - for _ in range(len(hold.pending)): - self._cache.pin(mm_hash) - hold.holds.extend(hold.pending) - hold.pending.clear() logger.debug("EC producer: mm_hash=%s marked ready", mm_hash) for transfer_id in meta.completed_loads: pending = self._load_acks.get(transfer_id) @@ -679,6 +551,8 @@ def has_pending_push_work(self) -> bool: 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: @@ -716,7 +590,6 @@ def request_finished( peer_port=self._peer_port, size_bytes=size_bytes, ) - self._hold_announced(mm_hash, entry) logger.debug( "EC producer: announcing NIXL-readable encodings req_id=%s items=%s", request.request_id, @@ -728,6 +601,7 @@ def shutdown(self) -> None: 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 1e9fd7c7573a7..80bbd1fd40516 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/session.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/session.py @@ -260,11 +260,6 @@ def __init__( # Key: "{consumer_session_id}:{mm_hash}" — matches the NIXL notif_msg exactly, # so completion notification lookup is a direct dict pop with no parsing. self._active_xfers: dict[str, ProducerXfer] = {} - # One entry per grant whose read finished or whose lease lapsed, for - # the scheduler to release the matching pin it took when announcing. - # A list, not a set: two consumers reading the same encoding hold two - # pins, and collapsing their events would release only one. - self._served: list[str] = [] self._req_decoder = msgspec.msgpack.Decoder(XferReq) self._req_encoder = msgspec.msgpack.Encoder() @@ -370,18 +365,12 @@ def _drain_notifs(self) -> None: xfer = self._active_xfers.pop(key, None) if xfer is not None: self._cache.unpin(xfer.mm_hash) - self._served.append(xfer.mm_hash) logger.debug( "EC producer: NIXL READ completed mm_hash=%s key=%s", xfer.mm_hash, key, ) - def take_served(self) -> list[str]: - """Return and clear one entry per grant that finished or lapsed.""" - served, self._served = self._served, [] - return served - def _sweep_timeouts(self) -> None: for key, xfer in list(self._active_xfers.items()): if xfer.is_expired(): @@ -393,7 +382,6 @@ def _sweep_timeouts(self) -> None: key, ) self._cache.unpin(xfer.mm_hash) - self._served.append(xfer.mm_hash) del self._active_xfers[key] @@ -513,10 +501,7 @@ def poll(self, messages: list[bytes], now: float) -> None: self._completed.add(mm_hash) del self._xfers[mm_hash] elif state == XferState.ACK_TIMEOUT: - # A busy producer is a transient condition, not a broken one: - # drop every trace of the attempt so a later admit pass can - # re-issue the read as if it were the first. This is also the - # only path to a tombstone that would otherwise log nothing. + # 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", From 89031d6ca339bd9583cbafa2fa64178e9c3f21b6 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Tue, 8 Sep 2026 01:15:27 +0000 Subject: [PATCH 7/8] [Test][EC Connector] Add NIXL remote failure isolation E2E Assisted-by: OpenAI Codex Signed-off-by: Tianyu Guo --- tests/v1/ec_connector/integration/README.md | 20 ++ .../integration/test_nixl_failure.py | 202 ++++++++++++++++++ 2 files changed, 222 insertions(+) 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 bf2dcd76d6327..a89bbfd3479bf 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 0000000000000..ec5cc8518e11f --- /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 From 10cc51d7280b221b6c5c63fec0fddd549d02fe0a Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Wed, 9 Sep 2026 00:15:59 +0000 Subject: [PATCH 8/8] [EC Connector] Simplify NIXL admission failure handling Centralize per-item remote failures, clarify admission contracts, and trim redundant unit test coverage. Assisted-by: OpenAI Codex Signed-off-by: Tianyu Guo --- .../unit/test_scheduler_nixl_consumer.py | 36 +--- .../ec_connector/cpu/scheduler/__init__.py | 194 ++++++++++-------- 2 files changed, 118 insertions(+), 112 deletions(-) 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 03bfbd8e9c9f8..a3e87f9b78883 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py @@ -171,8 +171,7 @@ def test_no_remote_announcement_is_ready(monkeypatch, params): @pytest.mark.parametrize("payload", [None, "metadata", "pixels", "embeds", "shm"]) -@pytest.mark.parametrize("failure", ["tombstone", "timeout"]) -def test_failed_read_falls_back_only_with_local_input(monkeypatch, payload, failure): +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() @@ -201,18 +200,12 @@ def _fake_start(mm_hash, info, size): else None ) req = _Request([_Feature("h1", 1, data)], params=_params("h1", 1)) - now = [1000.0] - monkeypatch.setattr(time, "monotonic", lambda: now[0]) # Step 1: start the read. assert s.ensure_cache_available(req, 0) is False s.build_connector_meta(scheduler_output=None) - if failure == "tombstone": - fake._results.tombstoned.add("h1") - else: - now[0] += sched_mod._ADMIT_DEFER_TIMEOUT_S + 1 - fake._results.retryable.add("h1") + fake._results.tombstoned.add("h1") s._sessions[("h", 1)] = fake can_fallback = payload in ("pixels", "embeds", "shm") @@ -230,19 +223,6 @@ def _fake_start(mm_hash, info, size): s.shutdown() -def test_request_finished_clears_only_its_wait_budget(monkeypatch): - """Cancelled requests leave no clocks behind, even for a shared hash.""" - s = _consumer_sched(monkeypatch) - monkeypatch.setattr(s, "_start_xfer", lambda *args: False) - old = _Request([_Feature("h1")], params=_params("h1", 1), req_id="old") - new = _Request([_Feature("h1")], params=_params("h1", 1), req_id="new") - assert not s.ensure_cache_available(old, 0) - assert not s.ensure_cache_available(new, 0) - assert s.request_finished(old) == (False, None) - assert set(s._deferred_since) == {("new", "h1")} - s.shutdown() - - @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.""" @@ -323,14 +303,10 @@ def test_invalid_remote_size_fails_the_request(monkeypatch, size_fields): s.shutdown() -@pytest.mark.parametrize("pool_full", [False, True], ids=["not-ready", "full-pool"]) -def test_deferral_budget_is_per_request(monkeypatch, pool_full): - """Transient failures defer until each request exhausts its own budget.""" +def test_deferral_budget_is_per_request(monkeypatch): + """Shared hashes have independent wait budgets and cancellation cleanup.""" s = _consumer_sched(monkeypatch) - if pool_full: - monkeypatch.setattr(s._cache, "alloc", lambda key, n: None) - else: - assert s._cache.alloc("h1", 1) is not None + 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") @@ -346,6 +322,8 @@ def test_deferral_budget_is_per_request(monkeypatch, pool_full): # `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 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 1d72c55298db1..cf0a92dcc26cc 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -41,6 +41,10 @@ _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.""" @@ -216,103 +220,121 @@ def ensure_cache_available( def _nixl_consumer_admit( self, request: "Request", num_computed_tokens: int ) -> bool: - """Wait for remote encodings, falling back only with local input.""" + """Admit a request once its required remote encodings are ready. + + 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. + """ if not request.ec_transfer_params: return True now = time.monotonic() pending = False for feature in request.mm_features: - 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) - continue - 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) - continue - if remote is None: - continue - since = self._deferred_since.setdefault((request.request_id, mm_hash), now) - if now - since > _ADMIT_DEFER_TIMEOUT_S: - if not self._fail_or_fallback( - request, feature, f"the remote wait exceeded {now - since:.0f}s" - ): + try: + 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 - continue - if mm_hash in self._in_flight or mm_hash in self._step_completed: - pending = True - continue + return not pending - if mm_hash in self._tombstones: - self._tombstones.discard(mm_hash) - if not self._fail_or_fallback( - request, feature, "the remote read failed" - ): - return False - continue + def _admit_item( + self, + request: "Request", + feature: "MultiModalFeatureSpec", + num_computed_tokens: int, + now: float, + ) -> bool: + """Check one encoding, starting a remote read if needed. - 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. - pending = True - continue + Returns: + True if already computed, cached, or not remotely sourced. False + while waiting for a read or cache space, including across retries. - expected = pos.length * self._hidden_dim * self._element_size - try: - size = int(remote["size_bytes"]) - except (KeyError, TypeError, ValueError): - if not self._fail_or_fallback( - request, feature, "the announced size was unusable" - ): - return False - continue - if size != expected: - logger.warning( - "EC consumer: size mismatch mm_hash=%s announced=%d expected=%d", - mm_hash, - size, - expected, - ) - if not self._fail_or_fallback( - request, feature, "the announced size was wrong" - ): - return False - continue + 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 + ) - try: - started = self._start_xfer(mm_hash, remote, expected) - except Exception: - logger.exception( - "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash - ) - if not self._fail_or_fallback( - request, feature, "the read could not be started" - ): - return False - continue - if started: - self._in_flight.add(mm_hash) - pending = True - return not pending + 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: - """Allow local execution only with data beyond placeholder metadata.""" + """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 @@ -344,6 +366,12 @@ def _fail_or_fallback( 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