diff --git a/tests/v1/ec_connector/__init__.py b/tests/v1/ec_connector/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py index 0fe465c3e8..3a4cf86a67 100644 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py +++ b/tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py @@ -21,13 +21,15 @@ def _cache(num_blocks: int = 8) -> EmbeddingCache: # ── alloc ──────────────────────────────────────────────────────────────────── -def test_alloc_returns_unique_block_ids(): +def test_alloc_returns_unique_ascending_block_ids(): cache = _cache() entry = cache.alloc("a", 3) assert entry is not None assert len(entry.block_ids) == 3 assert len(set(entry.block_ids)) == 3 assert all(0 <= i < 8 for i in entry.block_ids) + # Ascending order lets consumers detect which blocks are adjacent. + assert list(entry.block_ids) == sorted(entry.block_ids) def test_alloc_returns_tuple(): @@ -161,6 +163,33 @@ def test_unpin_asserts_on_unpinned(): cache.unpin("a") +# ── has_held_entries ───────────────────────────────────────────────────────── + + +def test_has_held_entries_empty_cache(): + cache = _cache() + assert cache.has_held_entries() is False + + +def test_has_held_entries_true_while_not_ready(): + cache = _cache() + cache.alloc("a", 2) # not ready → held + assert cache.has_held_entries() is True + cache.mark_ready("a") # ready + unpinned → settled + assert cache.has_held_entries() is False + + +def test_has_held_entries_true_while_pinned(): + cache = _cache() + cache.alloc("a", 2) + cache.mark_ready("a") + assert cache.has_held_entries() is False + cache.pin("a") # pinned → held + assert cache.has_held_entries() is True + cache.unpin("a") + assert cache.has_held_entries() is False + + # ── eviction ───────────────────────────────────────────────────────────────── diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py index 0e224c1af3..24ea717cab 100644 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py +++ b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py @@ -3,6 +3,11 @@ import uuid 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.config.ec_transfer import ECRole +from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( + ECCPUWorkerMetadata, +) from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( ECSharedRegion, ) @@ -37,13 +42,26 @@ class _Request: self.request_id = request_id +class _WorkerOutput: + """Stand-in for ECConnectorOutput carrying a completion report. + + `saves` are mm_hashes; `loads` are transfer ids. One entry per reporting + rank, so a list holding the same id twice is two ranks reporting it. + """ + + def __init__(self, *, saves=None, loads=None): + self.ec_connector_worker_meta = ECCPUWorkerMetadata( + completed_saves=list(saves or []), + completed_loads=list(loads or []), + ) + + def _make_scheduler( monkeypatch, num_blocks=_N, - max_concurrent_batches=1, *, - is_producer=True, - is_consumer=True, + ec_role: ECRole = "ec_both", + tensor_parallel_size=1, ) -> ECCPUScheduler: region = ECSharedRegion( engine_id=str(uuid.uuid4()), @@ -52,27 +70,22 @@ def _make_scheduler( ) monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region) - _is_prod = is_producer - _is_cons = is_consumer - _mcb = max_concurrent_batches - - class _EC: - is_ec_producer = _is_prod - is_ec_consumer = _is_cons - - class _Cfg: - ec_transfer_config = _EC() - - @property - def max_concurrent_batches(self): - return _mcb - - return ECCPUScheduler(_Cfg()) + return ECCPUScheduler( + create_ec_vllm_config( + ec_role=ec_role, + tensor_parallel_size=tensor_parallel_size, + ) + ) -class _SchedulerOutput: - def __init__(self, finished_req_ids=None): - self.finished_req_ids = finished_req_ids or set() +def _load_ids(meta) -> list[int]: + """Transfer ids the scheduler dispatched in this step's metadata.""" + return [transfer_id for transfer_id, _ in meta.loads.values()] + + +def _load_blocks(meta, mm_hash: str) -> list[int]: + """Block ids the scheduler dispatched for one load.""" + return meta.loads[mm_hash][1] def _seed_cached(s: ECCPUScheduler, mm_hash: str, n_blocks: int): @@ -83,7 +96,9 @@ def _seed_cached(s: ECCPUScheduler, mm_hash: str, n_blocks: int): def test_offload_reuse_cycle(monkeypatch): s = _make_scheduler(monkeypatch) - req = _Request([_Feature("h1", length=1)]) + # Two blocks, so the reload also pins down that every allocated block id + # comes back in meta.loads, not just the first. + req = _Request([_Feature("h1", length=2)]) # Step A: first sight — allocate and save. assert s.has_cache_item("h1") is False @@ -92,24 +107,24 @@ def test_offload_reuse_cycle(monkeypatch): meta_a = s.build_connector_meta(scheduler_output=None) assert "h1" in meta_a.saves assert meta_a.loads == {} - # Not ready yet (max_concurrent_batches=1 means ready next step). + # Not ready until the worker reports the save memcpy complete. assert s.has_cache_item("h1") is False - # Step B: _step_readiness marks h1 ready (no new allocs). - s.build_connector_meta(scheduler_output=None) + # Worker reports the save copy finished → entry becomes ready. + s.update_connector_output(_WorkerOutput(saves=["h1"])) assert s.has_cache_item("h1") is True # Step C: reload now works — entry is ready at update_state_after_alloc time. s.update_state_after_alloc(req, 0) meta_c = s.build_connector_meta(scheduler_output=None) assert "h1" in meta_c.loads - assert meta_c.loads["h1"] == meta_a.saves["h1"] + assert _load_blocks(meta_c, "h1") == meta_a.saves["h1"] s.shutdown() def test_has_cache_item_false_when_not_consumer(monkeypatch): - s = _make_scheduler(monkeypatch, is_consumer=False) + s = _make_scheduler(monkeypatch, ec_role="ec_producer") assert s.has_cache_item("anything") is False s.shutdown() @@ -123,10 +138,10 @@ def test_connector_keys_on_identifier_not_mm_hash(monkeypatch): meta = s.build_connector_meta(scheduler_output=None) assert "ENC_KEY" in meta.saves assert "PROC_KEY" not in meta.saves - # Not ready yet (1 step delay). + # Not ready until the completion report. assert s.has_cache_item("ENC_KEY") is False - # After another step it becomes ready — keyed on identifier. - s.build_connector_meta(scheduler_output=None) + # After the report it becomes ready — keyed on identifier. + s.update_connector_output(_WorkerOutput(saves=["ENC_KEY"])) assert s.has_cache_item("ENC_KEY") is True assert s.has_cache_item("PROC_KEY") is False s.shutdown() @@ -136,13 +151,11 @@ def test_cpu_region_fifo_eviction(monkeypatch): # Region holds exactly 2 one-block encodings. s = _make_scheduler(monkeypatch, num_blocks=2) - # Fill with A and B, marking each ready. + # Fill with A and B, marking each ready via a completion report. for h in ("A", "B"): s.update_state_after_alloc(_Request([_Feature(h, length=1)]), 0) s.build_connector_meta(scheduler_output=None) - # After the step above, A's step was called in B's build. A is ready. - # B was submitted in the last build, so one more step makes it ready. - s.build_connector_meta(scheduler_output=None) # marks B ready + s.update_connector_output(_WorkerOutput(saves=[h])) assert s.has_cache_item("A") is True assert s.has_cache_item("B") is True @@ -161,8 +174,8 @@ def test_multiple_mm_items_per_request(monkeypatch): meta = s.build_connector_meta(scheduler_output=None) assert set(meta.saves) == {"h1", "h2"} - # Step delay: mark both ready. - s.build_connector_meta(scheduler_output=None) + # Worker reports both saves complete. + s.update_connector_output(_WorkerOutput(saves=["h1", "h2"])) assert s.has_cache_item("h1") is True assert s.has_cache_item("h2") is True @@ -174,26 +187,6 @@ def test_multiple_mm_items_per_request(monkeypatch): s.shutdown() -def test_load_returns_correct_block_ids(monkeypatch): - """meta.loads must contain the same block IDs that meta.saves - allocated — verified through the public API only.""" - s = _make_scheduler(monkeypatch) - req = _Request([_Feature("a", length=2)]) - - s.update_state_after_alloc(req, 0) - meta_save = s.build_connector_meta(scheduler_output=None) - assert "a" in meta_save.saves - - # Readiness delay. - s.build_connector_meta(scheduler_output=None) - - # Reload. - s.update_state_after_alloc(req, 0) - meta_load = s.build_connector_meta(scheduler_output=None) - assert meta_load.loads["a"] == meta_save.saves["a"] - s.shutdown() - - def test_loads_only_serves_hashes_touched_this_step(monkeypatch): s = _make_scheduler(monkeypatch) _seed_cached(s, "a", n_blocks=2) @@ -228,10 +221,10 @@ def test_load_not_emitted_for_uncached_entry(monkeypatch): s.shutdown() -def test_delayed_unpin_protects_blocks_during_worker_read(monkeypatch): - """Blocks being loaded must not be evictable until the delayed unpin - fires. Observable: a save that needs eviction cannot reclaim a - still-pinned loaded entry.""" +def test_load_pin_protects_blocks_until_completion(monkeypatch): + """Loaded blocks stay pinned until the worker reports the load memcpy + complete. Observable: a save that needs eviction cannot reclaim a + still-pinned loaded entry, but can once the completion report arrives.""" # 2 blocks, both occupied by "a" (ready). s = _make_scheduler(monkeypatch, num_blocks=2) _seed_cached(s, "a", n_blocks=2) @@ -240,94 +233,79 @@ def test_delayed_unpin_protects_blocks_during_worker_read(monkeypatch): s.update_state_after_alloc(_Request([_Feature("a")]), 0) meta = s.build_connector_meta(scheduler_output=None) assert "a" in meta.loads + (transfer_id,) = _load_ids(meta) # Try to save "b" (needs 2 blocks, but "a" is pinned) — must fail. s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0) meta2 = s.build_connector_meta(scheduler_output=None) assert "b" not in meta2.saves # can't evict pinned "a" - # After the delay, "a" unpins. Now "b" can evict it. + # The one participant reports the load complete → "a" unpins. Now "b" can + # evict it. + s.update_connector_output(_WorkerOutput(loads=[transfer_id])) s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0) meta3 = s.build_connector_meta(scheduler_output=None) assert "b" in meta3.saves s.shutdown() -def test_delayed_unpin_depth_2(monkeypatch): - """With max_concurrent_batches=2, loaded blocks stay pinned for 2 steps - after build. Observable via eviction failure.""" - s = _make_scheduler(monkeypatch, num_blocks=2, max_concurrent_batches=2) - _seed_cached(s, "a", n_blocks=2) - - # Load "a" → pinned. - s.update_state_after_alloc(_Request([_Feature("a")]), 0) - s.build_connector_meta(scheduler_output=None) - - # Step 1 after load: still pinned — can't evict for "b". - s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0) - meta1 = s.build_connector_meta(scheduler_output=None) - assert "b" not in meta1.saves - - # Step 2 after load: still pinned. - s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0) - meta2 = s.build_connector_meta(scheduler_output=None) - assert "b" not in meta2.saves - - # Step 3: unpin fires, now eviction works. - s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0) - meta3 = s.build_connector_meta(scheduler_output=None) - assert "b" in meta3.saves - s.shutdown() +def _can_evict(s: ECCPUScheduler, mm_hash: str, n_blocks: int) -> bool: + """Whether a new save of `mm_hash` can claim `n_blocks`, i.e. whether the + blocks currently held by other entries are reclaimable.""" + s.update_state_after_alloc(_Request([_Feature(mm_hash, length=n_blocks)]), 0) + meta = s.build_connector_meta(scheduler_output=None) + return mm_hash in meta.saves -def test_shutdown_unpins_pending_loads(monkeypatch): - """Shutdown must unpin entries still in _pending_loads (not yet drained - to the unpin deque). Observable: after shutdown, the entry is evictable - (alloc can reclaim its blocks).""" - s = _make_scheduler(monkeypatch, num_blocks=2) +def test_load_needs_every_participant_before_unpin(monkeypatch): + """With two participating ranks, one report must not release the pin. + + This is the multi-rank case: each rank copies the same CPU blocks into its + own GPU memory, so the blocks stay in use until the last rank is done. + """ + s = _make_scheduler(monkeypatch, num_blocks=2, tensor_parallel_size=2) _seed_cached(s, "a", n_blocks=2) s.update_state_after_alloc(_Request([_Feature("a")]), 0) - # Don't call build_connector_meta — "a" stays in _pending_loads. + (transfer_id,) = _load_ids(s.build_connector_meta(scheduler_output=None)) + + # First rank reports: still in use by the second, so "a" must hold. + s.update_connector_output(_WorkerOutput(loads=[transfer_id])) + assert _can_evict(s, "b", 2) is False + assert s.has_pending_push_work() is True + + # Second rank reports: now it is genuinely free. + s.update_connector_output(_WorkerOutput(loads=[transfer_id])) + assert s.has_pending_push_work() is False + assert _can_evict(s, "c", 2) is True s.shutdown() - # After shutdown, scheduler reports no items (consumer role is off). - assert s.has_cache_item("a") is False +def test_late_report_does_not_release_a_later_load_of_same_hash(monkeypatch): + """A straggling report from one dispatch must not release the pin taken by + a subsequent dispatch of the same mm_hash. -def test_shutdown_unpins_deferred_unpin_queue(monkeypatch): - """Shutdown must unpin entries queued in the deferred unpin deque - (already drained from _pending_loads but not yet expired).""" - s = _make_scheduler(monkeypatch, num_blocks=2, max_concurrent_batches=2) + This is what the transfer id is for: mm_hash identifies a cache entry, not + a transfer, so it cannot tell the two dispatches apart. + """ + s = _make_scheduler(monkeypatch, num_blocks=2, tensor_parallel_size=2) _seed_cached(s, "a", n_blocks=2) - # Load "a" and build — moves to unpin deque. + # First dispatch, fully reported by both ranks → pin released. s.update_state_after_alloc(_Request([_Feature("a")]), 0) - s.build_connector_meta(scheduler_output=None) + (first_id,) = _load_ids(s.build_connector_meta(scheduler_output=None)) + s.update_connector_output(_WorkerOutput(loads=[first_id, first_id])) + assert s.has_pending_push_work() is False - # "a" is in the deferred unpin deque (depth=2, not yet expired). - # Shutdown must unpin it. - s.shutdown() - assert s.has_cache_item("a") is False + # Second dispatch of the same hash takes a fresh pin. + s.update_state_after_alloc(_Request([_Feature("a")]), 0) + (second_id,) = _load_ids(s.build_connector_meta(scheduler_output=None)) + assert second_id != first_id - -def test_eviction_skips_entry_pinned_by_pending_load(monkeypatch): - """A loaded entry still within the unpin delay window must survive - eviction attempts from new saves.""" - # 2 blocks: A occupies both, ready. - s = _make_scheduler(monkeypatch, num_blocks=2) - _seed_cached(s, "A", n_blocks=2) - - # Load A → pinned. - s.update_state_after_alloc(_Request([_Feature("A")]), 0) - meta_load = s.build_connector_meta(scheduler_output=None) - assert "A" in meta_load.loads - - # Try to save C (needs 2 blocks) — can't evict pinned A. - s.update_state_after_alloc(_Request([_Feature("C", length=2)]), 0) - meta_save = s.build_connector_meta(scheduler_output=None) - assert "C" not in meta_save.saves - assert s.has_cache_item("A") is True # survived + # A replayed report for the first dispatch must be ignored entirely. + s.update_connector_output(_WorkerOutput(loads=[first_id, first_id])) + assert s.has_pending_push_work() is True + assert _can_evict(s, "b", 2) is False s.shutdown() @@ -350,30 +328,10 @@ def test_region_full_skips_save_and_never_blocks(monkeypatch): s.shutdown() -def test_step_readiness_depth_2(monkeypatch): - """With max_concurrent_batches=2, entries need 2 steps to become ready.""" - s = _make_scheduler(monkeypatch, max_concurrent_batches=2) - req = _Request([_Feature("h1", length=1)]) - - s.update_state_after_alloc(req, 0) - meta_0 = s.build_connector_meta(scheduler_output=None) - assert "h1" in meta_0.saves - assert s.has_cache_item("h1") is False - - # One step later: still not ready (depth=2). - s.build_connector_meta(scheduler_output=None) - assert s.has_cache_item("h1") is False - - # Two steps later: now ready. - s.build_connector_meta(scheduler_output=None) - assert s.has_cache_item("h1") is True - s.shutdown() - - def test_producer_only_never_emits_loads(monkeypatch): """A producer-only scheduler must never populate meta.loads, even when entries are ready.""" - s = _make_scheduler(monkeypatch, is_consumer=False) + s = _make_scheduler(monkeypatch, ec_role="ec_producer") req = _Request([_Feature("h1", length=1)]) s.update_state_after_alloc(req, 0) @@ -381,8 +339,7 @@ def test_producer_only_never_emits_loads(monkeypatch): assert "h1" in meta.saves assert meta.loads == {} - # Mark ready. - s.build_connector_meta(scheduler_output=None) + s.update_connector_output(_WorkerOutput(saves=["h1"])) # Even if we re-encounter the same feature, no load. s.update_state_after_alloc(req, 0) @@ -393,7 +350,7 @@ def test_producer_only_never_emits_loads(monkeypatch): def test_consumer_only_never_emits_saves(monkeypatch): """A consumer-only scheduler must never populate meta.saves.""" - s = _make_scheduler(monkeypatch, is_producer=False) + s = _make_scheduler(monkeypatch, ec_role="ec_consumer") _seed_cached(s, "a", n_blocks=1) req = _Request([_Feature("a", length=1)]) @@ -406,7 +363,7 @@ def test_consumer_only_never_emits_saves(monkeypatch): def test_consumer_only_has_cache_item(monkeypatch): """Consumer-only: has_cache_item works normally for ready entries.""" - s = _make_scheduler(monkeypatch, is_producer=False) + s = _make_scheduler(monkeypatch, ec_role="ec_consumer") assert s.has_cache_item("a") is False _seed_cached(s, "a", n_blocks=1) assert s.has_cache_item("a") is True @@ -431,50 +388,75 @@ def test_save_not_emitted_for_already_cached_entry(monkeypatch): s.shutdown() -# ── first-finish fast-path (idle engine) ──────────────────────────────────── +# ── completion reporting (update_connector_output) ────────────────────────── -def test_first_finish_marks_ready_immediately(monkeypatch): - """When the originating request finishes, the save entry becomes ready - without waiting for the step-count delay.""" - s = _make_scheduler(monkeypatch, max_concurrent_batches=5) - req = _Request([_Feature("h1", length=1)], request_id="req_x") +def test_update_connector_output_ignores_foreign_meta(monkeypatch): + """A worker output whose payload is not an ECCPUWorkerMetadata is a + no-op — no cache mutation, no crash.""" + s = _make_scheduler(monkeypatch) + _seed_cached(s, "a", n_blocks=1) - s.update_state_after_alloc(req, 0) + class _Foreign: + ec_connector_worker_meta = object() + + s.update_connector_output(_Foreign()) # must not raise + assert s.has_cache_item("a") is True + s.shutdown() + + +def test_update_connector_output_ignores_unknown_report(monkeypatch): + """Reports for a never-seen save hash or a transfer id this scheduler never + dispatched are dropped rather than mutating unrelated state.""" + s = _make_scheduler(monkeypatch) + s.update_connector_output(_WorkerOutput(saves=["gone"], loads=[4242])) + assert s.has_cache_item("gone") is False + s.shutdown() + + +# ── has_pending_push_work ─────────────────────────────────────────────────── + + +def test_has_pending_push_work_tracks_inflight_save(monkeypatch): + """A dispatched-but-unconfirmed save keeps push work pending until the + completion report marks it ready.""" + s = _make_scheduler(monkeypatch) + assert s.has_pending_push_work() is False + + s.update_state_after_alloc(_Request([_Feature("h1", length=1)]), 0) s.build_connector_meta(scheduler_output=None) - assert s.has_cache_item("h1") is False + assert s.has_pending_push_work() is True # save not yet confirmed - # Request finishes — readiness should fire immediately. - s.build_connector_meta(_SchedulerOutput(finished_req_ids={"req_x"})) - assert s.has_cache_item("h1") is True + s.update_connector_output(_WorkerOutput(saves=["h1"])) + assert s.has_pending_push_work() is False s.shutdown() -def test_first_finish_unpins_immediately(monkeypatch): - """When the originating request finishes, the load entry is unpinned - without waiting for the step-count delay.""" - s = _make_scheduler(monkeypatch, num_blocks=2, max_concurrent_batches=5) - _seed_cached(s, "a", n_blocks=2) +def test_has_pending_push_work_tracks_inflight_load(monkeypatch): + """A dispatched-but-unconfirmed load keeps push work pending until the + unpin completion report arrives.""" + s = _make_scheduler(monkeypatch) + _seed_cached(s, "a", n_blocks=1) + assert s.has_pending_push_work() is False # ready + unpinned - req = _Request([_Feature("a")], request_id="req_y") - s.update_state_after_alloc(req, 0) + s.update_state_after_alloc(_Request([_Feature("a")]), 0) meta = s.build_connector_meta(scheduler_output=None) - assert "a" in meta.loads + assert s.has_pending_push_work() is True # pinned, awaiting unpin - # "a" is pinned. Try to evict — must fail. - s.update_state_after_alloc( - _Request([_Feature("b", length=2)], request_id="req_z"), 0 - ) - meta2 = s.build_connector_meta(scheduler_output=None) - assert "b" not in meta2.saves - - # req_y finishes — unpin fires in build_connector_meta (fast-path). - s.build_connector_meta(_SchedulerOutput(finished_req_ids={"req_y"})) - - # Next step: "a" is now evictable, so "b" can be allocated. - s.update_state_after_alloc( - _Request([_Feature("b", length=2)], request_id="req_w"), 0 - ) - meta3 = s.build_connector_meta(scheduler_output=None) - assert "b" in meta3.saves + s.update_connector_output(_WorkerOutput(loads=_load_ids(meta))) + assert s.has_pending_push_work() is False s.shutdown() + + +# ── shutdown ──────────────────────────────────────────────────────────────── + + +def test_shutdown_disables_roles_and_cleans_region(monkeypatch): + """After shutdown the scheduler serves no items and the region is + cleaned up.""" + s = _make_scheduler(monkeypatch, num_blocks=2) + _seed_cached(s, "a", n_blocks=2) + assert s.has_cache_item("a") is True + + s.shutdown() + assert s.has_cache_item("a") is False # consumer role disabled diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_step_tracker.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_step_tracker.py deleted file mode 100644 index 370b46412b..0000000000 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_step_tracker.py +++ /dev/null @@ -1,162 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for StepTracker.""" - -from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.step_tracker import ( - StepTracker, -) - - -def test_basic_step_count_expiry_depth_1(): - """With max_concurrent_batches=1, entry expires after 1 step.""" - t = StepTracker(max_concurrent_batches=1) - t.add("h1", "req_a") - - # Step 1: entry was added, committed to slot. Deque fills to maxlen=1, - # but no expiry yet (the entry is in the slot just committed). - result = t.step(set()) - assert result == [] - - # Step 2: deque is full, oldest slot (containing h1) expires. - result = t.step(set()) - assert result == ["h1"] - - -def test_basic_step_count_expiry_depth_2(): - """With max_concurrent_batches=2, entry expires after 2 steps.""" - t = StepTracker(max_concurrent_batches=2) - t.add("h1", "req_a") - - assert t.step(set()) == [] # step 1: committed - assert t.step(set()) == [] # step 2: deque fills but h1 in newer slot - assert t.step(set()) == ["h1"] # step 3: h1's slot is oldest, expires - - -def test_first_finish_fast_path(): - """Entry is returned immediately when request_id appears in finished.""" - t = StepTracker(max_concurrent_batches=3) - t.add("h1", "req_a") - - result = t.step({"req_a"}) - assert result == ["h1"] - - -def test_no_double_processing_fast_path_then_expiry(): - """Entry fast-pathed should NOT fire again on deque expiry.""" - t = StepTracker(max_concurrent_batches=1) - t.add("h1", "req_a") - - # Fast-path fires h1. - result = t.step({"req_a"}) - assert result == ["h1"] - - # Normal expiry would fire now, but h1 is already processed. - result = t.step(set()) - assert result == [] - - -def test_no_double_processing_expiry_then_fast_path(): - """Entry expired via step-count should NOT fire again via fast-path.""" - t = StepTracker(max_concurrent_batches=1) - t.add("h1", "req_a") - - t.step(set()) # commit - result = t.step(set()) # expiry fires h1 - assert result == ["h1"] - - # req_a finishes later — must not re-fire. - result = t.step({"req_a"}) - assert result == [] - - -def test_multiplicity_same_hash_different_requests(): - """Two adds of the same mm_hash (different requests) → returned twice.""" - t = StepTracker(max_concurrent_batches=3) - t.add("h1", "req_a") - t.add("h1", "req_b") - - # Both requests finish in the same step. - result = t.step({"req_a", "req_b"}) - assert sorted(result) == ["h1", "h1"] - - -def test_mixed_expiry_and_fast_path(): - """Some entries expire, others fast-path, in the same step.""" - t = StepTracker(max_concurrent_batches=1) - t.add("h1", "req_a") - - t.step(set()) # commit h1's slot - - # Add h2 in the next step. - t.add("h2", "req_b") - - # This step: h1's slot expires AND req_b finishes. - result = t.step({"req_b"}) - assert set(result) == {"h1", "h2"} - - -def test_drain_all_returns_unprocessed(): - """drain_all returns everything pending and clears state.""" - t = StepTracker(max_concurrent_batches=3) - t.add("h1", "req_a") - t.add("h2", "req_b") - t.step(set()) # commit both - t.add("h3", "req_c") # in _current, not yet committed - - result = t.drain_all() - assert set(result) == {"h1", "h2", "h3"} - - # After drain, everything is empty. - assert t.step(set()) == [] - assert t.drain_all() == [] - - -def test_drain_all_skips_already_processed(): - """drain_all should not re-return entries already fast-pathed.""" - t = StepTracker(max_concurrent_batches=3) - t.add("h1", "req_a") - t.add("h2", "req_b") - - # Fast-path h1. - result = t.step({"req_a"}) - assert result == ["h1"] - - # drain should only return h2. - result = t.drain_all() - assert result == ["h2"] - - -def test_empty_step_returns_nothing(): - """step() with nothing pending returns empty list.""" - t = StepTracker(max_concurrent_batches=2) - assert t.step(set()) == [] - assert t.step({"req_x"}) == [] - - -def test_add_after_step_goes_to_next_slot(): - """Entries added after a step go into the next slot, not the current.""" - t = StepTracker(max_concurrent_batches=1) - - t.step(set()) # empty step - t.add("h1", "req_a") - - # h1 is in _current. Next step commits it. - t.step(set()) - - # Now the deque is full (1 slot with h1). Next step expires it. - result = t.step(set()) - assert result == ["h1"] - - -def test_fast_path_on_current_entries(): - """Entries still in _current (not yet committed) can be fast-pathed.""" - t = StepTracker(max_concurrent_batches=3) - t.add("h1", "req_a") - - # req_a finishes in the same step h1 was added. - result = t.step({"req_a"}) - assert result == ["h1"] - - # Entry is committed to slot but marked processed — won't re-fire. - for _ in range(5): - assert t.step(set()) == [] diff --git a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py index 3daf4ee7de..8ae8f81f0a 100644 --- a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py +++ b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py @@ -22,14 +22,16 @@ Mocking policy import contextlib import logging +import time import uuid +from collections import deque from unittest.mock import MagicMock, Mock, patch +import numpy as np import pytest import torch -from vllm.config import VllmConfig -from vllm.config.parallel import ParallelConfig +from tests.v1.ec_connector.unit.utils import create_ec_vllm_config from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( ECCPUConnectorMetadata, ) @@ -55,6 +57,18 @@ _requires_accelerator = pytest.mark.skipif( reason="exercises real accelerator stream/event coordination in ECCPUWorker", ) +# The memory-lifetime tests hold a copy back with `torch.cuda._sleep`, which the +# CUDA-like platforms provide but XPU does not. +_requires_cuda_alike = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="stalling the compute stream requires a CUDA-like platform", +) + +# Cycles to stall the compute stream for. Must comfortably outlast the host-side +# work the lifetime tests do while a copy is deliberately held back — ~0.5 s on +# a current datacenter GPU. +_STALL_CYCLES = 1_000_000_000 + def _make_region() -> ECSharedRegion: """Fresh region backed by a real per-test mmap file.""" @@ -66,18 +80,28 @@ def _make_region() -> ECSharedRegion: def _vllm_config(rank: int = 0) -> Mock: - cfg = Mock(spec=VllmConfig) - cfg.parallel_config = Mock(spec=ParallelConfig) - cfg.parallel_config.rank = rank - cfg.model_config = Mock() - cfg.model_config.dtype = _DTYPE - return cfg + return create_ec_vllm_config(rank=rank, dtype=_DTYPE) def _meta( *, saves: dict | None = None, loads: dict | None = None ) -> ECCPUConnectorMetadata: - return ECCPUConnectorMetadata(saves=saves or {}, loads=loads or {}) + """Build step metadata from plain `{mm_hash: block_ids}` dicts. + + Load transfer ids are synthesized here so tests can keep naming loads by + mm_hash; `_load_id` recovers the id the worker will report. + """ + return ECCPUConnectorMetadata( + saves=saves or {}, + loads={ + mm_hash: (idx, block_ids) + for idx, (mm_hash, block_ids) in enumerate((loads or {}).items()) + }, + ) + + +def _load_id(meta: ECCPUConnectorMetadata, mm_hash: str) -> int: + return meta.loads[mm_hash][0] # ── fixtures ───────────────────────────────────────────────────────────────── @@ -136,19 +160,52 @@ def make_worker(): region.cleanup() +def _wait_for_completion( + worker: ECCPUWorker, expected, direction: str, timeout_s: float = 30.0 +) -> None: + """Poll ``build_connector_worker_meta`` until ``expected`` is reported done. + + Saves are reported by mm_hash, loads by transfer id. + """ + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + meta = worker.build_connector_worker_meta() + if meta is not None and expected in getattr(meta, f"completed_{direction}"): + return + time.sleep(0.001) + raise AssertionError(f"{direction} for {expected!r} never reported complete") + + +def _warm_up_stream_pools(worker: ECCPUWorker) -> None: + """Run one throwaway save so the stream and event pools are populated. + + The first CUDA stream creation in a process blocks the host for ~0.5 s on + driver initialization, which would swallow the stall the lifetime tests + below install. One flush puts a stream and its events in the worker's pools, + matching the steady state every step after the first sees. + """ + scratch = torch.zeros(_HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE) + worker.save_caches({"w": scratch}, "w", _meta(saves={"w": [0]})) + worker.flush_saves() + _wait_for_completion(worker, "w", "saves") + assert worker._stream_pool, "warm-up did not recycle a stream" + + # ── save_caches ────────────────────────────────────────────────────────────── @_requires_accelerator @pytest.mark.parametrize( - "n_elements,n_blocks", + "n_elements,block_ids", [ - (3 * _HIDDEN_DIM, 3), # exact fit: 3 full blocks, no remainder - (3 * _HIDDEN_DIM + 4, 4), # partial last block: 4 fp16 trail in slot 3 + (3 * _HIDDEN_DIM, [7, 2, 5]), # exact fit: 3 full blocks, no remainder + (3 * _HIDDEN_DIM + 4, [7, 2, 5, 0]), # partial last block: 4 fp16 trail + (3 * _HIDDEN_DIM, [2, 3, 4]), # adjacent: one coalesced copy + (3 * _HIDDEN_DIM + 4, [2, 3, 4, 6]), # coalesced run, then a gap ], - ids=["exact-fit", "partial-last-block"], + ids=["exact-fit", "partial-last-block", "adjacent-blocks", "run-then-gap"], ) -def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks): +def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, block_ids): """``save_caches`` + ``flush_saves`` copies the source GPU tensor's bytes into the block IDs named by ``meta.saves``, in the order they appear, via a single batched GPU→CPU copy. @@ -156,6 +213,8 @@ def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks When ``total_bytes`` is not a multiple of ``block_size_bytes`` the last block is partially written; the unwritten tail must remain whatever was there before. Blocks not named in ``meta.saves`` must be untouched. + Blocks adjacent in the region share one descriptor, which must land the + same bytes in the same blocks as one descriptor per block does. """ worker = make_worker() sentinel = 0x5A @@ -165,9 +224,9 @@ def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks expected_bytes = src.cpu().reshape(-1).view(torch.uint8) total_bytes = n_elements * _DTYPE.itemsize - block_ids = [7, 2, 5, 0][:n_blocks] worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids})) worker.flush_saves() + _wait_for_completion(worker, "h", "saves") for slot, block_idx in enumerate(block_ids): block_byte_start = slot * _BLOCK_SIZE_BYTES @@ -189,6 +248,42 @@ def test_save_caches_writes_to_assigned_blocks(make_worker, n_elements, n_blocks ) +@_requires_accelerator +@pytest.mark.parametrize( + "n_elements,block_ids,expected_sizes", + [ + (3 * _HIDDEN_DIM, [2, 3, 4], [3 * _BLOCK_SIZE_BYTES]), + (3 * _HIDDEN_DIM, [7, 2, 5], [_BLOCK_SIZE_BYTES] * 3), + (3 * _HIDDEN_DIM, [2, 3, 6], [2 * _BLOCK_SIZE_BYTES, _BLOCK_SIZE_BYTES]), + (2 * _HIDDEN_DIM + 4, [2, 3, 4], [(2 * _HIDDEN_DIM + 4) * _DTYPE.itemsize]), + (_HIDDEN_DIM, [2, 3, 4], [_BLOCK_SIZE_BYTES]), + ], + ids=[ + "one-run", + "no-run", + "run-then-gap", + "run-stops-at-output", + "blocks-past-output-skipped", + ], +) +def test_save_caches_coalesces_only_adjacent_blocks( + make_worker, n_elements, block_ids, expected_sizes +): + """One descriptor per run of region-adjacent blocks, not one per block. + + Per-descriptor overhead dominates these copies, so the byte tests above + would pass either way; this pins the descriptor layout that makes them + fast. A coalesced copy also stops at the encoder output, since entries are + sized by placeholder count, which can exceed the number of embeddings. + """ + worker = make_worker() + src = torch.arange(n_elements, dtype=_DTYPE, device=DEVICE_TYPE) + + worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids})) + + assert worker._save_bufs.sizes[: worker._save_count].tolist() == expected_sizes + + def test_save_caches_noop_when_mm_hash_not_in_saves(make_worker): """When the scheduler hasn't pre-allocated blocks for ``mm_hash``, ``save_caches`` + ``flush_saves`` is a pure no-op.""" @@ -245,6 +340,7 @@ def test_save_caches_batches_multiple_hashes(make_worker): worker.save_caches(cache, "a", _meta(saves={"a": [1], "b": [3]})) worker.save_caches(cache, "b", _meta(saves={"a": [1], "b": [3]})) worker.flush_saves() + _wait_for_completion(worker, "a", "saves") expected_a = src_a.cpu().view(torch.uint8) expected_b = src_b.cpu().view(torch.uint8) @@ -258,6 +354,50 @@ def test_save_caches_batches_multiple_hashes(make_worker): assert torch.all(worker._region.blocks[idx] == sentinel) +class _HighAddress: + """A device tensor reporting an address with the high bit set. + + XPU USM allocations live up there and no host allocator will hand out such + an address, so naming one is the only way to reach the arithmetic. Every + attribute other than the address comes from the real tensor. + """ + + def __init__(self, real: torch.Tensor, address: int): + self._real = real + self._address = address + + def __getattr__(self, name): + return getattr(self._real, name) + + def view(self, *args): + return self + + def data_ptr(self) -> int: + return self._address + + +@_requires_accelerator +def test_save_descriptors_hold_high_bit_source_addresses(make_worker): + """Source addresses >= 2**63 survive into the descriptor buffers.""" + worker = make_worker() + real = torch.zeros(2 * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE) + address = 2**64 - 8 * _BLOCK_SIZE_BYTES + + # Two runs, so the second descriptor's address is an offset off the base. + worker.save_caches( + {"a": _HighAddress(real, address)}, "a", _meta(saves={"a": [1, 5]}) + ) + + assert worker._save_count == 2 + stored = worker._save_bufs.src_np[:2].astype(np.uint64) + assert list(stored) == [address, address + _BLOCK_SIZE_BYTES] + + # The batch describes an address nothing may read from; drop it unflushed. + worker._save_bufs = None + worker._save_count = 0 + worker._save_mm_hashes = [] + + # ── start_load_caches ──────────────────────────────────────────────────────── @@ -288,19 +428,31 @@ def test_start_load_caches_copies_with_correct_shape_dtype_and_bytes(make_worker @_requires_accelerator -def test_start_load_caches_preserves_existing_encoder_cache_entry(make_worker): - """If ``encoder_cache`` already holds the ``mm_hash``, the worker must - not overwrite it.""" +@pytest.mark.parametrize( + "loads,expected_sizes", + [ + ({"a": [2, 3], "b": [4, 5]}, [4 * _BLOCK_SIZE_BYTES]), + ({"b": [4, 5], "a": [2, 3]}, [4 * _BLOCK_SIZE_BYTES]), + ({"a": [2, 3], "b": [6, 7]}, [2 * _BLOCK_SIZE_BYTES] * 2), + ], + ids=["adjacent-entries", "adjacent-entries-reversed", "gap-between-entries"], +) +def test_start_load_caches_coalesces_across_entries(make_worker, loads, expected_sizes): + """Blocks adjacent in the region share a descriptor across entries too. + + Entries are ordered by first block, so whether they coalesce follows from + the region layout rather than the order the scheduler dispatched them. + """ worker = make_worker() - worker._region.blocks[0].fill_(0x42) + captured: list[list[int]] = [] - sentinel = torch.full((_HIDDEN_DIM,), 7.0, dtype=_DTYPE, device=DEVICE_TYPE) - encoder_cache = {"h": sentinel} - worker.start_load_caches(encoder_cache, _meta(loads={"h": [0]})) + with patch( + "vllm.distributed.ec_transfer.ec_connector.cpu.worker.swap_blocks_batch", + side_effect=lambda src, dst, sizes, **kw: captured.append(sizes.tolist()), + ): + worker.start_load_caches({}, _meta(loads=loads)) - assert encoder_cache["h"] is sentinel, ( - "existing encoder_cache entry must not be replaced" - ) + assert captured == [expected_sizes] @_requires_accelerator @@ -315,30 +467,34 @@ def test_start_load_caches_noop_when_loads_is_empty(make_worker): @_requires_accelerator def test_start_load_caches_skips_cached_and_loads_new_in_same_step(make_worker): - """Cached entries are preserved while new ones are loaded.""" + """Every hash in ``meta.loads`` is copied from mmap, including one whose + key is already resident in ``encoder_cache`` — there is no worker-side + skip.""" worker = make_worker() - n_blocks = 2 + n_blocks = 3 src_orig = torch.arange(n_blocks * _HIDDEN_DIM, dtype=_DTYPE).reshape( n_blocks, _HIDDEN_DIM ) src_int8 = src_orig.view(torch.int8).reshape(n_blocks, _BLOCK_SIZE_BYTES) - new_block_ids = [4, 2] - for i, idx in enumerate(new_block_ids): + block_ids = [0, 4, 2] + for i, idx in enumerate(block_ids): worker._region.blocks[idx].copy_(src_int8[i]) cached_tensor = torch.full((1, _HIDDEN_DIM), 99.0, dtype=_DTYPE, device=DEVICE_TYPE) encoder_cache: dict[str, torch.Tensor] = {"cached_h": cached_tensor} worker.start_load_caches( encoder_cache, - _meta(loads={"cached_h": [0], "new_h": new_block_ids}), + _meta(loads={"resident_h": [block_ids[0]], "new_h": block_ids[1:]}), ) - assert encoder_cache["cached_h"] is cached_tensor - assert "new_h" in encoder_cache - out = encoder_cache["new_h"] - assert out.shape == (n_blocks, _HIDDEN_DIM) - assert out.dtype == _DTYPE - assert torch.equal(out.cpu(), src_orig) + resident = encoder_cache["resident_h"] + assert resident is not cached_tensor, "resident entry must be reloaded, not skipped" + assert torch.equal(resident.cpu(), src_orig[:1]) + + new = encoder_cache["new_h"] + assert new.shape == (2, _HIDDEN_DIM) + assert new.dtype == _DTYPE + assert torch.equal(new.cpu(), src_orig[1:]) @_requires_accelerator @@ -382,6 +538,7 @@ def test_save_then_load_round_trips_bytes(make_worker): ).reshape(n_blocks, _HIDDEN_DIM) worker.save_caches({"h": src}, "h", _meta(saves={"h": block_ids})) worker.flush_saves() + _wait_for_completion(worker, "h", "saves") encoder_cache: dict[str, torch.Tensor] = {} worker.start_load_caches(encoder_cache, _meta(loads={"h": block_ids})) @@ -392,18 +549,133 @@ def test_save_then_load_round_trips_bytes(make_worker): assert torch.equal(out.cpu(), src.cpu()) +# ── memory lifetime across in-flight copies ───────────────────────────────── + + +@_requires_cuda_alike +def test_save_survives_encoder_cache_free_before_copy_runs(make_worker): + """The bytes handed to ``save_caches`` must reach the mmap intact even when + the caller drops its ``encoder_cache`` reference before the copy has run. + + The model runner pops encoder cache entries on whatever step the scheduler + asks it to, with no dependency on the save copy having landed. Dropping the + last reference hands the GPU memory back to the caching allocator, which is + free to reuse it for the next same-sized allocation on the compute stream — + the copy reads that memory on a different stream, so the allocator has no + reason to hold it back. + + The stall makes the hazard deterministic rather than timing-dependent: the + save copy is gated behind the compute stream, so it provably has not started + while the overwrite (issued on an ungated stream and waited on) completes. + + The mmap must be pinned, as it is in production: a copy into pageable host + memory blocks the host until it lands, which would hide the hazard. + """ + worker = make_worker(pin_memory_available=True) + _warm_up_stream_pools(worker) + worker._region.blocks.fill_(0x5A) + + block_ids = [0, 1, 2] + saved = torch.full( + (len(block_ids), _HIDDEN_DIM), 1.0, dtype=_DTYPE, device=DEVICE_TYPE + ) + expected = saved.cpu().reshape(-1).view(torch.uint8) + + encoder_cache = {"h": saved} + worker.save_caches(encoder_cache, "h", _meta(saves={"h": block_ids})) + + torch.cuda._sleep(_STALL_CYCLES) + worker.flush_saves() + + saved_ptr = saved.data_ptr() + del saved + encoder_cache.clear() + + # Same size on the same stream, so the allocator hands back the block just + # freed. Overwriting it from an ungated stream — and waiting for that write + # — orders the overwrite strictly before the still-stalled save copy. + overwrite = torch.empty( + (len(block_ids), _HIDDEN_DIM), dtype=_DTYPE, device=DEVICE_TYPE + ) + reused = overwrite.data_ptr() == saved_ptr + other_stream = torch.cuda.Stream() + with torch.cuda.stream(other_stream): + overwrite.fill_(-1.0) + other_stream.synchronize() + + _wait_for_completion(worker, "h", "saves") + + actual = worker._region.blocks[block_ids].reshape(-1).view(torch.uint8) + assert torch.equal(actual, expected), ( + "mmap holds bytes written after save_caches was called; the copy read " + f"freed source memory (allocator reused the freed block: {reused})" + ) + + +@_requires_cuda_alike +def test_load_buffer_survives_eviction_while_consumer_read_is_queued(make_worker): + """A loaded encoder cache entry must keep its bytes for a consumer already + queued on the compute stream, even after the entry is dropped and another + load is dispatched. + + ``start_load_caches`` allocates its destination inside the load stream's + context, which ties that memory to the load stream in the caching allocator. + Dropping the encoder cache entry returns it to the load stream's pool, where + the next load can claim it — while the model's read of it is still queued on + the compute stream. + """ + worker = make_worker(pin_memory_available=True) + _warm_up_stream_pools(worker) + a_blocks, b_blocks = [4, 5], [6, 7] + for block_idx, fill in zip(a_blocks + b_blocks, (0x11, 0x12, 0x21, 0x22)): + worker._region.blocks[block_idx].fill_(fill) + expected_a = worker._region.blocks[a_blocks].reshape(-1).clone() + + encoder_cache: dict[str, torch.Tensor] = {} + meta_a = _meta(loads={"a": a_blocks}) + worker.start_load_caches(encoder_cache, meta_a) + loaded_a = encoder_cache["a"] + + # Queue a consumer of the loaded entry on the compute stream behind a stall, + # the way the model reads the encoder cache later in the step. + torch.cuda._sleep(_STALL_CYCLES) + consumed = loaded_a.clone() + + # Draining the completion report recycles the load stream into the pool, so + # the next load allocates from the same allocator pool the entry was freed to. + _wait_for_completion(worker, _load_id(meta_a, "a"), "loads") + + loaded_a_ptr = loaded_a.data_ptr() + del loaded_a + encoder_cache.clear() + + worker.start_load_caches(encoder_cache, _meta(loads={"b": b_blocks})) + reused = encoder_cache["b"].data_ptr() == loaded_a_ptr + torch.accelerator.synchronize() + + assert torch.equal(consumed.cpu().reshape(-1).view(torch.int8), expected_a), ( + "queued consumer read bytes from a later load; the load destination was " + f"recycled while still in use (allocator reused the buffer: {reused})" + ) + + # ── buffer recycling ──────────────────────────────────────────────────────── @_requires_accelerator def test_buffer_pool_is_reused_across_save_steps(make_worker): - """After flush_saves, descriptor buffers are returned to the pool and - reused on the next flush — no reallocation.""" + """Once a save copy completes its descriptor buffers return to the pool and + are reused by the next flush — no reallocation. + + The buffers hold the addresses the copy reads, so they stay out of the pool + until the transfer's end event fires. + """ worker = make_worker() src = torch.arange(_HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE) worker.save_caches({"h": src}, "h", _meta(saves={"h": [0]})) worker.flush_saves() + _wait_for_completion(worker, "h", "saves") assert len(worker._buf_pool._pool) == 1 buf_id = id(worker._buf_pool._pool[0].src_ptrs) @@ -411,6 +683,7 @@ def test_buffer_pool_is_reused_across_save_steps(make_worker): # Second step reuses the same buffer. worker.save_caches({"h": src}, "h", _meta(saves={"h": [1]})) worker.flush_saves() + _wait_for_completion(worker, "h", "saves") assert len(worker._buf_pool._pool) == 1 assert id(worker._buf_pool._pool[0].src_ptrs) == buf_id @@ -418,35 +691,29 @@ def test_buffer_pool_is_reused_across_save_steps(make_worker): @_requires_accelerator def test_buffer_pool_is_reused_across_load_steps(make_worker): - """After start_load_caches, descriptor buffers are returned to the pool - and reused on the next call.""" + """Once a load copy completes its descriptor buffers return to the pool and + are reused by the next call — no reallocation.""" worker = make_worker() worker._region.blocks[0].fill_(0x01) worker._region.blocks[1].fill_(0x02) encoder_cache: dict[str, torch.Tensor] = {} - worker.start_load_caches(encoder_cache, _meta(loads={"a": [0]})) + meta = _meta(loads={"a": [0]}) + worker.start_load_caches(encoder_cache, meta) + _wait_for_completion(worker, _load_id(meta, "a"), "loads") assert len(worker._buf_pool._pool) == 1 buf_id = id(worker._buf_pool._pool[0].src_ptrs) encoder_cache2: dict[str, torch.Tensor] = {} - worker.start_load_caches(encoder_cache2, _meta(loads={"b": [1]})) + meta2 = _meta(loads={"b": [1]}) + worker.start_load_caches(encoder_cache2, meta2) + _wait_for_completion(worker, _load_id(meta2, "b"), "loads") assert len(worker._buf_pool._pool) == 1 assert id(worker._buf_pool._pool[0].src_ptrs) == buf_id -# ── stream management ──────────────────────────────────────────────────────── - - -@_requires_accelerator -def test_stream_initialized_at_construction(make_worker): - """``_load_stream`` must be a fully initialized accelerator stream.""" - worker = make_worker() - assert isinstance(worker._load_stream, current_platform.Stream) - - # ── lifecycle ──────────────────────────────────────────────────────────────── @@ -500,12 +767,15 @@ def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm): worker = object.__new__(ECCPUWorker) mock_region = Mock(spec=ECSharedRegion) worker._region = mock_region - worker._load_stream = MagicMock() worker._save_bufs = None worker._save_count = 0 + worker._save_mm_hashes = [] + worker._inflight_saves = deque() + worker._inflight_loads = deque() + worker._stream_pool = [] + worker._event_pool = [] worker.shutdown() - worker._load_stream.synchronize.assert_called_once() mock_region.cleanup.assert_called_once() mock_region.cleanup.side_effect = RuntimeError("boom") @@ -525,12 +795,13 @@ def test_shutdown_calls_region_cleanup_and_swallows_errors(caplog_vllm): @_requires_accelerator def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch): - """Full pipeline: scheduler allocates blocks, worker saves GPU tensor to - mmap via flush_saves, scheduler marks ready after step delay, worker - loads from mmap back to GPU, and the result matches the original. + """Full pipeline: scheduler allocates blocks, worker saves a GPU tensor to + mmap via flush_saves, the worker's completion report marks the entry ready, + the worker loads from mmap back to GPU, and the result matches the original. Exercises the real scheduler + worker cooperation through a shared - ECSharedRegion, with real accelerator transfers and stream coordination. + ECSharedRegion, with real accelerator transfers and stream coordination, + and the event-driven mark_ready path. """ import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ( @@ -544,15 +815,7 @@ def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch): # Build scheduler sharing the same region. monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region) - class _EC: - is_ec_producer = True - is_ec_consumer = True - - class _Cfg: - ec_transfer_config = _EC() - max_concurrent_batches = 1 - - scheduler = ECCPUScheduler(_Cfg()) + scheduler = ECCPUScheduler(_vllm_config()) # -- Step 1: scheduler allocates, worker saves -- n_blocks = 3 @@ -580,16 +843,24 @@ def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch): worker.save_caches(encoder_cache, "img_001", meta_save) worker.flush_saves() - # -- Step 2: readiness delay (max_concurrent_batches=1) -- - meta_step2 = scheduler.build_connector_meta(scheduler_output=None) + # -- Step 2: worker reports the save memcpy complete → scheduler marks ready -- + torch.accelerator.synchronize() + worker_meta = worker.build_connector_worker_meta() + assert worker_meta is not None + assert "img_001" in worker_meta.completed_saves + + class _Output: + ec_connector_worker_meta = worker_meta + + scheduler.update_connector_output(_Output()) assert scheduler.has_cache_item("img_001") is True - assert meta_step2.loads == {} # no load requested yet # -- Step 3: scheduler emits load, worker loads -- scheduler.update_state_after_alloc(_Request(), 0) meta_load = scheduler.build_connector_meta(scheduler_output=None) assert "img_001" in meta_load.loads - assert meta_load.loads["img_001"] == meta_save.saves["img_001"] + _, load_blocks = meta_load.loads["img_001"] + assert load_blocks == meta_save.saves["img_001"] load_cache: dict[str, torch.Tensor] = {} worker.start_load_caches(load_cache, meta_load) diff --git a/tests/v1/ec_connector/unit/test_ec_cpu_connector.py b/tests/v1/ec_connector/unit/test_ec_cpu_connector.py index 6a1ccf8571..e2cb8b120a 100644 --- a/tests/v1/ec_connector/unit/test_ec_cpu_connector.py +++ b/tests/v1/ec_connector/unit/test_ec_cpu_connector.py @@ -48,9 +48,10 @@ def _make_unique_image(base_image, iteration: int): def _wait_for_ec_ready(llm: LLM) -> None: """Send a dummy request to force a scheduler step. - After generate() returns, the EC entry's data is in the mmap but - readiness may not have fired yet (first-finish fires in the NEXT - scheduler step). This forces that step. + When generate() returns, the save copy has been issued to the GPU but its + completion event may fire on a later step, so the entry is not yet marked + ready. This extra step lets the scheduler poll the worker's completion + report and mark the entry ready. """ llm.generate("hi", SamplingParams(max_tokens=1), use_tqdm=False) diff --git a/tests/v1/ec_connector/unit/test_metadata.py b/tests/v1/ec_connector/unit/test_metadata.py index 55579e128c..cca1df0171 100644 --- a/tests/v1/ec_connector/unit/test_metadata.py +++ b/tests/v1/ec_connector/unit/test_metadata.py @@ -38,19 +38,19 @@ def test_metadata_saves_dict_operations(): assert meta.saves["mm_hash_2"] == [10, 11] -def test_metadata_loads_dict_operations(): - """The loads dict supports standard dict operations.""" +def test_metadata_loads_carry_a_transfer_id_and_blocks(): + """Each load entry is a (transfer_id, block_ids) pair, so the worker can + report the transfer rather than the mm_hash.""" meta = ECCPUConnectorMetadata() - # Add multiple entries - meta.loads["mm_hash_1"] = [5, 6, 7] - meta.loads["mm_hash_2"] = [100] + meta.loads["mm_hash_1"] = (0, [5, 6, 7]) + meta.loads["mm_hash_2"] = (1, [100]) - assert "mm_hash_1" in meta.loads - assert "mm_hash_2" in meta.loads assert len(meta.loads) == 2 - assert meta.loads["mm_hash_1"] == [5, 6, 7] - assert meta.loads["mm_hash_2"] == [100] + assert meta.loads["mm_hash_1"] == (0, [5, 6, 7]) + transfer_id, block_ids = meta.loads["mm_hash_2"] + assert transfer_id == 1 + assert block_ids == [100] def test_metadata_saves_and_loads_are_independent(): @@ -58,11 +58,11 @@ def test_metadata_saves_and_loads_are_independent(): meta = ECCPUConnectorMetadata() meta.saves["key"] = [1, 2] - meta.loads["key"] = [3, 4] + meta.loads["key"] = (0, [3, 4]) # Same key but different dicts should not interfere assert meta.saves["key"] == [1, 2] - assert meta.loads["key"] == [3, 4] + assert meta.loads["key"] == (0, [3, 4]) # Modifying one should not affect the other meta.saves["other"] = [99] @@ -74,11 +74,11 @@ def test_metadata_block_indices_are_integers(): meta = ECCPUConnectorMetadata() meta.saves["hash"] = [0, 100, 999, 65535] - meta.loads["hash"] = [1, 2, 3] + meta.loads["hash"] = (0, [1, 2, 3]) for idx in meta.saves["hash"]: assert isinstance(idx, int) - for idx in meta.loads["hash"]: + for idx in meta.loads["hash"][1]: assert isinstance(idx, int) diff --git a/tests/v1/ec_connector/unit/utils.py b/tests/v1/ec_connector/unit/utils.py new file mode 100644 index 0000000000..93b97cb5ef --- /dev/null +++ b/tests/v1/ec_connector/unit/utils.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared config builders for EC connector unit tests.""" + +import uuid +from unittest.mock import Mock + +import torch + +from vllm.config import ModelConfig, VllmConfig +from vllm.config.ec_transfer import ECRole, ECTransferConfig +from vllm.config.parallel import ParallelConfig + +# The scheduler derives its per-load ack count from the PCP size only under the +# `mp` backend -- which also keeps the world size off the local GPU count, so +# TP > 1 stays expressible on a single-GPU host. +_EXECUTOR_BACKEND = "mp" +_PCP_SIZE = 1 + + +def create_ec_vllm_config( + *, + ec_role: ECRole = "ec_both", + tensor_parallel_size: int = 1, + rank: int = 0, + dtype: torch.dtype = torch.float16, +) -> Mock: + """Build a `VllmConfig` stand-in for EC connector unit tests. + + `ec_transfer_config` and `parallel_config` are the real config objects, so + role derivation and the rank arithmetic the connector depends on behave as + they do in production. + + `model_config` is a stub because nothing under test needs a real one: the + connector reads only `dtype` here, and the paths that inspect the model + itself are not exercised. Building a real `ModelConfig` would resolve an HF + model and inspect its architecture. + + Args: + ec_role: EC role this instance plays. `is_ec_producer` and + `is_ec_consumer` are derived from it, as in production. + tensor_parallel_size: TP degree. + rank: Global rank of this worker. + dtype: Encoder cache dtype. + + Returns: + A `VllmConfig`-specced mock carrying the configs above. + """ + parallel_config = ParallelConfig( + tensor_parallel_size=tensor_parallel_size, + prefill_context_parallel_size=_PCP_SIZE, + distributed_executor_backend=_EXECUTOR_BACKEND, + rank=rank, + ) + ec_transfer_config = ECTransferConfig( + # The role properties report False unless a connector is named. + ec_connector="ECCPUConnector", + ec_role=ec_role, + engine_id=str(uuid.uuid4()), + ) + + model_config = Mock(spec=ModelConfig) + model_config.dtype = dtype + + vllm_config = Mock(spec=VllmConfig) + vllm_config.ec_transfer_config = ec_transfer_config + vllm_config.parallel_config = parallel_config + vllm_config.model_config = model_config + vllm_config.instance_id = f"ec-test-{uuid.uuid4()}" + return vllm_config diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py index 771d4ca25a..e2c37b72a6 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py @@ -7,7 +7,10 @@ from typing import TYPE_CHECKING import torch -from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata +from vllm.distributed.ec_transfer.ec_connector.base import ( + ECConnectorMetadata, + ECConnectorWorkerMetadata, +) from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( ECSharedRegion, ) @@ -29,10 +32,38 @@ class ECCPUConnectorMetadata(ECConnectorMetadata): # encoder_cache[mm_hash] → mmap at these block IDs. saves: dict[str, list[int]] = field(default_factory=dict) - # Consumer role: mm_hashes whose bytes are available in the local mmap; - # the worker's start_load_caches copies mmap[block_ids] → GPU - # encoder_cache. - loads: dict[str, list[int]] = field(default_factory=dict) + # Consumer role: mm_hashes whose bytes are available in the local mmap, + # mapped to (transfer_id, block_ids); the worker's start_load_caches + # copies mmap[block_ids] → GPU encoder_cache and reports the transfer_id + # once its copy lands. The id is minted per dispatch, so two loads of the + # same mm_hash in different steps are distinct and a late report from the + # earlier one cannot release the later one's pin. + loads: dict[str, tuple[int, list[int]]] = field(default_factory=dict) + + +@dataclass +class ECCPUWorkerMetadata(ECConnectorWorkerMetadata): + """Per-step worker → scheduler payload for the ECCPUConnector. + + Reports the GPU copies that completed this step: saved mm_hashes become + safe to mark ready, and loaded transfers become safe to unpin once every + participating rank has reported them. Built by + `ECCPUWorker.build_connector_worker_meta`; consumed by + `ECCPUScheduler.update_connector_output`. + + Loads are reported by transfer id rather than mm_hash because every rank + copies the same blocks: `aggregate` concatenates, so an id appears once + per reporting rank and the scheduler can count participants off the list. + """ + + completed_saves: list[str] = field(default_factory=list) + completed_loads: list[int] = field(default_factory=list) + + def aggregate(self, other: ECConnectorWorkerMetadata) -> ECConnectorWorkerMetadata: + assert isinstance(other, ECCPUWorkerMetadata) + self.completed_saves.extend(other.completed_saves) + self.completed_loads.extend(other.completed_loads) + return self def _get_encoder_cache_hidden_dim(vllm_config: "VllmConfig") -> int: diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py b/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py index 7de5c7fae3..c606e98d74 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/connector.py @@ -15,6 +15,7 @@ import torch from vllm.distributed.ec_transfer.ec_connector.base import ( ECConnectorBase, ECConnectorRole, + ECConnectorWorkerMetadata, ) from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( ECCPUConnectorMetadata, @@ -24,6 +25,7 @@ from vllm.logger import init_logger if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.outputs import ECConnectorOutput from vllm.v1.request import Request logger = init_logger(__name__) @@ -113,6 +115,20 @@ class ECCPUConnector(ECConnectorBase): self.connector_worker.flush_saves() return None, None + def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None: + if self.connector_worker is not None: + return self.connector_worker.build_connector_worker_meta() + return None + + def update_connector_output(self, connector_output: "ECConnectorOutput") -> None: + if self.connector_scheduler is not None: + self.connector_scheduler.update_connector_output(connector_output) + + def has_pending_push_work(self) -> bool: + if self.connector_scheduler is not None: + return self.connector_scheduler.has_pending_push_work() + return False + # Shared. def shutdown(self) -> None: if self.connector_scheduler is not None: 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 132efc1df0..cee946fd26 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -11,19 +11,18 @@ from typing import TYPE_CHECKING from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( ECCPUConnectorMetadata, + ECCPUWorkerMetadata, create_ec_shared_region, ) from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import ( EmbeddingCache, ) -from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.step_tracker import ( - StepTracker, -) from vllm.logger import init_logger if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.outputs import ECConnectorOutput from vllm.v1.request import Request logger = init_logger(__name__) @@ -42,16 +41,27 @@ class ECCPUScheduler: # Block allocator + LRU eviction policy for the shared region. self._cache = EmbeddingCache(self._region.num_blocks) - max_batches = vllm_config.max_concurrent_batches - # Delays mark_ready until the GPU→mmap DMA is guaranteed complete. - self._ready_tracker = StepTracker(max_batches) - # Delays unpin until the mmap→GPU DMA is guaranteed complete. - self._unpin_tracker = StepTracker(max_batches) - # mm_hash → block IDs allocated this step for GPU→mmap saves. self._pending_saves: dict[str, list[int]] = {} - # mm_hash → block IDs to load from mmap→GPU this step. - self._pending_loads: dict[str, list[int]] = {} + # mm_hash → (transfer_id, block IDs) to load from mmap→GPU this step. + self._pending_loads: dict[str, tuple[int, list[int]]] = {} + + # Dispatched loads awaiting completion reports, keyed by transfer id: + # transfer_id → (mm_hash, reports still outstanding). The pin taken at + # dispatch is released only once the count reaches zero. + self._load_acks: dict[int, tuple[str, int]] = {} + self._next_transfer_id = 0 + + pc = vllm_config.parallel_config + # Reports to expect per load. Only pipeline stage 0 runs the EC + # connector, so the participants are the tp × pcp ranks of that stage. + # Other executors deliver a single rank's output, and expecting reports + # that cannot arrive would hold the pin forever. + self._expected_load_reports = ( + pc.tensor_parallel_size * pc.prefill_context_parallel_size + if pc.distributed_executor_backend == "mp" + else 1 + ) def has_cache_item(self, identifier: str) -> bool: if not self._is_consumer: @@ -72,30 +82,19 @@ class ECCPUScheduler: entry = self._cache.alloc(mm_hash, feature.mm_position.length) if entry is not None: self._pending_saves[mm_hash] = list(entry.block_ids) - self._ready_tracker.add(mm_hash, request.request_id) if self._is_consumer and mm_hash not in self._pending_loads: entry = self._cache.get(mm_hash) if entry is not None and entry.ready: self._cache.pin(mm_hash) - self._pending_loads[mm_hash] = list(entry.block_ids) - self._unpin_tracker.add(mm_hash, request.request_id) + transfer_id = self._next_transfer_id + self._next_transfer_id += 1 + self._pending_loads[mm_hash] = (transfer_id, list(entry.block_ids)) + self._load_acks[transfer_id] = (mm_hash, self._expected_load_reports) def build_connector_meta( self, scheduler_output: "SchedulerOutput" ) -> ECCPUConnectorMetadata: - finished = ( - scheduler_output.finished_req_ids if scheduler_output is not None else set() - ) - - for key in self._ready_tracker.step(finished): - entry = self._cache.get(key) - if entry is not None and not entry.ready: - self._cache.mark_ready(key) - - for key in self._unpin_tracker.step(finished): - self._cache.unpin(key) - meta = ECCPUConnectorMetadata() if self._is_producer: meta.saves = self._pending_saves @@ -105,13 +104,46 @@ class ECCPUScheduler: self._pending_loads = {} return meta - def shutdown(self) -> None: - # drain_all() covers both entries still in _current (never - # consumed by build_connector_meta) and entries in slots. - self._pending_loads.clear() - for mm_hash in self._unpin_tracker.drain_all(): + def update_connector_output(self, connector_output: "ECConnectorOutput") -> None: + """Apply the worker's memcpy-completion report to the cache. + + Completed saves become safe to mark ready. A load is unpinned once + every participating rank has reported its transfer id; reports for a + transfer that has already been released, or for one this scheduler + never dispatched, are ignored. + """ + meta = connector_output.ec_connector_worker_meta + if not isinstance(meta, ECCPUWorkerMetadata): + return + for mm_hash in meta.completed_saves: + entry = self._cache.get(mm_hash) + if entry is not None and not entry.ready: + self._cache.mark_ready(mm_hash) + for transfer_id in meta.completed_loads: + pending = self._load_acks.get(transfer_id) + if pending is None: + continue + mm_hash, outstanding = pending + if outstanding > 1: + self._load_acks[transfer_id] = (mm_hash, outstanding - 1) + continue + # Drop the entry before unpinning so a replayed report is treated + # as stale rather than releasing the pin a second time. + del self._load_acks[transfer_id] self._cache.unpin(mm_hash) - self._ready_tracker.drain_all() + + def has_pending_push_work(self) -> bool: + """True while any dispatched save or load has not been confirmed done. + + Keeps the engine stepping so the worker's completion reports are polled + even when no requests are otherwise runnable. + """ + return self._cache.has_held_entries() + + def shutdown(self) -> None: + self._pending_saves.clear() + self._pending_loads.clear() + self._load_acks.clear() self._is_producer = False self._is_consumer = False diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py index 8cd5f5c8e7..acc42a6ee7 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py @@ -67,6 +67,15 @@ class EmbeddingCache: with self._lock: return self._entries.get(key) + def has_held_entries(self) -> bool: + """True if any entry is not-ready or pinned (an in-flight transfer). + + Held entries are exactly those absent from the eviction free list, so + they represent saves awaiting mark_ready or loads awaiting unpin. + """ + with self._lock: + return len(self._entries) != len(self._entries_free_list) + def alloc(self, key: str, n_blocks: int) -> CacheEntry | None: """Allocate *n_blocks* for *key*, evicting as needed. @@ -84,7 +93,10 @@ class EmbeddingCache: if len(self._free_blocks) + self._evictable_block_count < n_blocks: return None self._evict_until(n_blocks) - block_ids = tuple(self._free_blocks.pop() for _ in range(n_blocks)) + # Canonical ascending order: keeps a freed entry's blocks contiguous + # in the free list when it is evicted, and lets consumers detect + # which of an entry's blocks are adjacent. + block_ids = tuple(sorted(self._free_blocks.pop() for _ in range(n_blocks))) entry = CacheEntry(block_ids) self._entries[key] = entry return entry diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/step_tracker.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/step_tracker.py deleted file mode 100644 index 9ef7e87a7b..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/step_tracker.py +++ /dev/null @@ -1,115 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""StepTracker — step-counting delay with first-finish fast-path. - -Manages a collection of (mm_hash, request_id) entries that need to be -processed after a configurable number of engine steps have elapsed. -Additionally, if the originating request finishes before the step-count -expires, the entry is processed immediately (the GPU transfer is -guaranteed complete by request finish time). - -Used by ECCPUScheduler: one instance for readiness marking, one for -unpins. -""" - -from collections import deque -from dataclasses import dataclass - - -@dataclass(slots=True) -class _PendingEntry: - mm_hash: str - request_id: str - processed: bool = False - - -class StepTracker: - """Track pending operations with step-count delay and first-finish. - - Each ``add(mm_hash, request_id)`` registers an operation that will be - returned from ``step()`` either when: - (a) ``max_concurrent_batches`` steps elapse (deque expiry), OR - (b) ``request_id`` appears in ``finished_req_ids`` (first-finish). - - Each add is processed exactly once — never both (a) and (b). - """ - - def __init__(self, max_concurrent_batches: int) -> None: - # One slot per past step; oldest at the right. When the deque is - # full the rightmost slot is popped and its entries expire. - self._slots: deque[list[_PendingEntry]] = deque(maxlen=max_concurrent_batches) - # Reverse index: request_id → entries across all slots + _current. - # Enables O(1) lookup for the first-finish fast-path. - self._req_index: dict[str, list[_PendingEntry]] = {} - # Entries added this step (via add()), committed to a new slot at - # the start of the next step() call. - self._current: list[_PendingEntry] = [] - - def add(self, mm_hash: str, request_id: str) -> None: - """Register an mm_hash needing processing after a delay.""" - entry = _PendingEntry(mm_hash=mm_hash, request_id=request_id) - self._current.append(entry) - self._req_index.setdefault(request_id, []).append(entry) - - def step(self, finished_req_ids: set[str]) -> list[str]: - """Advance one step. Return mm_hashes ready for processing. - - Returns: - List of mm_hashes to process. May contain duplicates if the - same mm_hash was added multiple times (from different requests). - """ - result: list[str] = [] - - # Phase 1: first-finish fast-path. - for req_id in finished_req_ids: - entries = self._req_index.pop(req_id, None) - if entries is None: - continue - for entry in entries: - if not entry.processed: - entry.processed = True - result.append(entry.mm_hash) - - # Phase 2: step-count expiry. - if len(self._slots) == self._slots.maxlen: - expired_slot = self._slots.pop() - for entry in expired_slot: - if not entry.processed: - entry.processed = True - result.append(entry.mm_hash) - self._cleanup_expired(expired_slot) - - # Phase 3: commit current entries to new slot. - self._slots.appendleft(self._current) - self._current = [] - - return result - - def drain_all(self) -> list[str]: - """Return all unprocessed entries and clear internal state.""" - result: list[str] = [] - for entry in self._current: - if not entry.processed: - entry.processed = True - result.append(entry.mm_hash) - self._current = [] - - for slot in self._slots: - for entry in slot: - if not entry.processed: - entry.processed = True - result.append(entry.mm_hash) - self._slots.clear() - self._req_index.clear() - return result - - def _cleanup_expired(self, slot: list[_PendingEntry]) -> None: - """Prune reverse index for entries in an expired slot.""" - for entry in slot: - entries_for_req = self._req_index.get(entry.request_id) - if entries_for_req is None: - # Already removed by fast-path (pop in step()). - continue - entries_for_req.remove(entry) - if not entries_for_req: - del self._req_index[entry.request_id] diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py index 5a957c4966..56bc09e709 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py @@ -7,13 +7,19 @@ per-step connector metadata (`ECCPUConnectorMetadata`) to decide which blocks to copy in each direction. """ +from collections import deque +from collections.abc import Sequence +from dataclasses import dataclass +from itertools import chain from typing import TYPE_CHECKING +import numpy as np import torch from vllm._custom_ops import swap_blocks_batch from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( ECCPUConnectorMetadata, + ECCPUWorkerMetadata, create_ec_shared_region, ) from vllm.distributed.ec_transfer.ec_connector.cpu.worker.descriptor_buffers import ( @@ -34,6 +40,50 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _coalesce_runs( + block_ids: Sequence[int], +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Split block ids into runs that are consecutive in the region. + + Returns `(slots, first_blocks, num_blocks)`, one element per run, where + `slots` is the run's offset in the flat block sequence. Per-descriptor + overhead dominates these copies, so describing a run with one descriptor + rather than one per block is worth an order of magnitude. + + Both callers derive the non-region side of a run from its slot, so that + side advances in lockstep with the region and only the region has to be + split. `EmbeddingCache` hands out ascending ids, which is what makes runs + findable at all. + """ + + ids = np.fromiter(block_ids, dtype=np.int64, count=len(block_ids)) + breaks = np.flatnonzero(np.diff(ids) != 1) + 1 + slots = np.concatenate(([0], breaks)) + ends = np.concatenate((breaks, [ids.size])) + return slots, ids[slots], ends - slots + + +@dataclass +class Transfer: + """A batched copy in flight on a dedicated GPU stream. + + Its descriptor buffers and stream are held until `end_event` fires; only + then are the buffers and stream safe to recycle and the `completions` safe + to report. `start_event`/`end_event` bracket the copy so its elapsed time + can be measured on completion. + + `completions` is what the scheduler is told about this copy: mm_hashes for + saves, transfer ids for loads. + """ + + start_event: torch.Event + end_event: torch.Event + completions: list[str] | list[int] + bufs: DescriptorBuffers + stream: torch.Stream + num_bytes: int + + class ECCPUWorker: """Worker-side delegate for the ECCPUConnector. @@ -43,8 +93,18 @@ class ECCPUWorker: batched call in `flush_saves`. - Consumer role: copies `mmap[block_ids]` → `encoder_cache[mm_hash]` for all entries in `metadata.loads` via a single `swap_blocks_batch` - call on the load stream. + call on a dedicated stream. - On `ec_both` nodes both paths run back-to-back in a single step. + + Every batched copy runs on a stream drawn from a pool so it overlaps the + compute stream, bracketed by a start/end event pair. The stream, events, + and descriptor buffers are recycled once the end event fires. + + Both paths move memory between the two streams, so each buffer is registered + against the stream that is not its own: save sources against the save stream, + load destinations against the compute stream. Without that, the caching + allocator would be free to reissue memory a copy or a queued read still + needs, since it only tracks the stream a buffer was allocated on. """ def __init__(self, vllm_config: "VllmConfig") -> None: @@ -62,15 +122,72 @@ class ECCPUWorker: get_tensor_model_parallel_rank() == 0 and get_pcp_group().rank_in_group == 0 ) - # Dedicated stream for async mmap→GPU loads (overlaps with compute). - self._load_stream = current_platform.Stream() - # Descriptor buffer pool (recycled across steps, shared by both paths). self._buf_pool = DescriptorBufferPool() # Active save buffer being filled during save_caches calls this step. self._save_bufs: DescriptorBuffers | None = None self._save_count: int = 0 + # Bytes accumulated into the active save buffer this step, for logging. + self._save_bytes: int = 0 + # mm_hashes accumulated into the active save buffer this step. + self._save_mm_hashes: list[str] = [] + # Stream the active save buffer's copy will run on, held from the moment + # the buffer opens so `save_caches` can register its sources against it. + self._save_stream: torch.Stream | None = None + + # Batched copies whose end event has not yet fired. Buffers and stream + # stay held until the event fires; the completions are reported to the + # scheduler only then. Separate per direction because saves and loads + # can be in flight concurrently. + self._inflight_saves: deque[Transfer] = deque() + self._inflight_loads: deque[Transfer] = deque() + + # Streams and timing events recycled across steps to avoid per-step + # create/destroy churn. Shared by both directions. + self._stream_pool: list[torch.Stream] = [] + self._event_pool: list[torch.Event] = [] + + def _acquire_stream(self) -> torch.Stream: + if self._stream_pool: + return self._stream_pool.pop() + return current_platform.Stream() + + def _acquire_event(self) -> torch.Event: + return ( + self._event_pool.pop() + if self._event_pool + else torch.Event(enable_timing=True) + ) + + def _collect_finished(self, inflight: deque[Transfer], direction: str) -> list: + """Pop transfers whose end event has fired, recycle their stream, + events, and buffers, and return their completions. + + The front check is conservative: a transfer is reported only once its + own end event fires, so completions are always genuinely done, and each + is reported exactly once because the transfer is popped as it is + reported. A later transfer that happens to finish first simply waits + behind the front and is reported on a subsequent poll. + """ + done: list = [] + while inflight and inflight[0].end_event.query(): + transfer = inflight.popleft() + done.extend(transfer.completions) + elapsed_ms = transfer.start_event.elapsed_time(transfer.end_event) + logger.debug( + "EC %s: %d entr%s (%d bytes) took %.3f ms", + direction, + len(transfer.completions), + "y" if len(transfer.completions) == 1 else "ies", + transfer.num_bytes, + elapsed_ms, + ) + self._buf_pool.release(transfer.bufs) + self._stream_pool.append(transfer.stream) + self._event_pool.append(transfer.start_event) + self._event_pool.append(transfer.end_event) + return done def save_caches( self, @@ -78,7 +195,14 @@ class ECCPUWorker: mm_hash: str, connector_metadata: ECCPUConnectorMetadata, ) -> None: - """Fill descriptor buffers directly for batched flush.""" + """Fill descriptor buffers directly for batched flush. + + Registers `encoder_cache[mm_hash]` against the stream the copy will run + on, so the caching allocator will not hand that memory to another + allocation before the copy has read it. The descriptor buffers hold raw + addresses, and the caller is free to evict the entry as soon as this + step's saves are dispatched. + """ if not self._is_save_rank: return block_ids = connector_metadata.saves.get(mm_hash) @@ -99,43 +223,95 @@ class ECCPUWorker: if self._save_bufs is None: total = sum(len(v) for v in connector_metadata.saves.values()) self._save_bufs = self._buf_pool.acquire(total) + self._save_stream = self._acquire_stream() + self._save_mm_hashes = [] assert self._save_count + len(block_ids) <= self._save_bufs.src_ptrs.numel() bufs = self._save_bufs - src_base = src.view(-1).view(torch.uint8).data_ptr() + assert self._save_stream is not None + src.record_stream(self._save_stream) + # Device addresses are unsigned and can have the high bit set (XPU USM), + # so they are added up in uint64; signed arithmetic rejects them. + src_base = np.uint64(src.view(-1).view(torch.uint8).data_ptr()) dst_base = self._region.blocks.data_ptr() - idx = self._save_count + # Entries are sized by placeholder count, which can exceed the number + # of embeddings, so blocks the output never reaches get no descriptor. + n_used = -(-total_bytes // block_size) + if n_used: + slots, first_blocks, num_blocks = _coalesce_runs(block_ids[:n_used]) + run_bytes = num_blocks * block_size + # The output can stop partway into the final run. + run_bytes[-1] = total_bytes - slots[-1] * block_size + bufs.add_copies( + self._save_count, + src_base + slots.astype(np.uint64) * np.uint64(block_size), + dst_base + first_blocks * block_size, + run_bytes, + ) + self._save_count += slots.size - for i, block_idx in enumerate(block_ids): - start = i * block_size - bufs.set_ptrs(idx, src_base + start, dst_base + block_idx * block_size) - bufs.sizes[idx] = min(block_size, total_bytes - start) - idx += 1 - - self._save_count = idx + self._save_bytes += total_bytes + self._save_mm_hashes.append(mm_hash) def flush_saves(self) -> None: - """Flush all accumulated saves in a single swap_blocks_batch call.""" + """Flush all accumulated saves in a single swap_blocks_batch call. + + Runs the copy on the stream the batch acquired when it opened, gated + behind the compute stream that produced the encoder outputs, brackets it + with start/end events, and enqueues the batch as in-flight. The stream, + events, and descriptor buffers are recycled once the end event fires (see + `_collect_finished`), which is also when the saved mm_hashes become + safe to mark ready. + """ if self._save_count == 0: return bufs = self._save_bufs - assert bufs is not None + stream = self._save_stream + assert bufs is not None and stream is not None src_ptrs, dst_ptrs, sizes = bufs.src_ptrs, bufs.dst_ptrs, bufs.sizes n = self._save_count - swap_blocks_batch(src_ptrs[:n], dst_ptrs[:n], sizes[:n]) + num_bytes = self._save_bytes + + # Gate the GPU→CPU copy behind the compute stream: it reads the encoder + # outputs the model just produced there. + stream.wait_stream(current_platform.current_stream()) + start_event = self._acquire_event() + end_event = self._acquire_event() + with current_platform.stream(stream): + start_event.record(stream) + swap_blocks_batch(src_ptrs[:n], dst_ptrs[:n], sizes[:n]) + end_event.record(stream) + + self._inflight_saves.append( + Transfer( + start_event=start_event, + end_event=end_event, + completions=self._save_mm_hashes, + bufs=bufs, + stream=stream, + num_bytes=num_bytes, + ) + ) - self._buf_pool.release(bufs) self._save_bufs = None + self._save_stream = None self._save_count = 0 + self._save_bytes = 0 + self._save_mm_hashes = [] def start_load_caches( self, encoder_cache: dict[str, torch.Tensor], connector_metadata: ECCPUConnectorMetadata, ) -> None: - """Consumer path: single batched copy of all loads from mmap→GPU.""" + """Consumer path: single batched copy of all loads from mmap→GPU. + + The destination is registered against the compute stream that consumes + it, so the caching allocator will not hand that memory to a later load + while reads of it are still queued there. + """ if not connector_metadata.loads: return @@ -145,59 +321,99 @@ class ECCPUWorker: device_type = current_platform.device_type src_base = blocks.data_ptr() - # Pre-filter: only hashes not already in encoder_cache. - load_items = { - h: idxs - for h, idxs in connector_metadata.loads.items() - if h not in encoder_cache - } - if not load_items: - return + # Copy every dispatched load. The scheduler routes an input here only + # after missing the GPU encoder cache, so a resident hit is not expected; + # copying unconditionally means every participating rank reports each + # transfer exactly once, which is what the scheduler counts to release + # the pin. + # Ordered by first block so entries that neighbour each other in the + # region land next to each other in the destination and coalesce into + # one descriptor. Every loop below consumes this same order. + load_items = sorted( + connector_metadata.loads.items(), key=lambda kv: kv[1][1][0] + ) + total_blocks = sum(len(block_ids) for _, (_, block_ids) in load_items) - total_blocks = sum(len(idxs) for idxs in load_items.values()) - - with current_platform.stream(self._load_stream): + stream = self._acquire_stream() + compute_stream = current_platform.current_stream() + start_event = self._acquire_event() + end_event = self._acquire_event() + with current_platform.stream(stream): # Single contiguous destination buffer for all loads. dst_buf = torch.empty( total_blocks, block_size, dtype=torch.int8, device=device_type ) - dst_buf_base = dst_buf.data_ptr() + dst_buf.record_stream(compute_stream) + # See save_caches: device addresses are added up in uint64. + dst_buf_base = np.uint64(dst_buf.data_ptr()) bufs = self._buf_pool.acquire(total_blocks) - src_ptrs = bufs.src_ptrs[:total_blocks] - dst_ptrs = bufs.dst_ptrs[:total_blocks] - sizes = bufs.sizes[:total_blocks] - sizes[:] = block_size - - op_idx = 0 - for block_ids in load_items.values(): - for block_idx in block_ids: - bufs.set_ptrs( - op_idx, - src_base + block_idx * block_size, - dst_buf_base + op_idx * block_size, - ) - op_idx += 1 + slots, first_blocks, num_blocks = _coalesce_runs( + list(chain.from_iterable(ids for _, (_, ids) in load_items)) + ) + bufs.add_copies( + 0, + src_base + first_blocks * block_size, + dst_buf_base + slots.astype(np.uint64) * np.uint64(block_size), + num_blocks * block_size, + ) + op_idx = slots.size + src_ptrs = bufs.src_ptrs[:op_idx] + dst_ptrs = bufs.dst_ptrs[:op_idx] + sizes = bufs.sizes[:op_idx] + start_event.record(stream) swap_blocks_batch(src_ptrs, dst_ptrs, sizes, is_src_access_order_any=True) - - self._buf_pool.release(bufs) + end_event.record(stream) + self._inflight_loads.append( + Transfer( + start_event=start_event, + end_event=end_event, + completions=[tid for _, (tid, _) in load_items], + bufs=bufs, + stream=stream, + num_bytes=total_blocks * block_size, + ) + ) # Slice contiguous buffer into per-hash views. offset = 0 - for mm_hash, block_ids in load_items.items(): + for mm_hash, (_, block_ids) in load_items: n = len(block_ids) encoder_cache[mm_hash] = ( dst_buf[offset : offset + n].view(dtype).reshape(n, -1) ) offset += n - current_platform.current_stream().wait_stream(self._load_stream) + current_platform.current_stream().wait_stream(stream) + + def build_connector_worker_meta(self) -> ECCPUWorkerMetadata | None: + """Report the GPU copies that completed on this rank this step: saved + mm_hashes and loaded transfer ids. + + Returns None when nothing finished, so the scheduler sees no payload. + """ + completed_saves = self._collect_finished(self._inflight_saves, "save") + completed_loads = self._collect_finished(self._inflight_loads, "load") + if not completed_saves and not completed_loads: + return None + return ECCPUWorkerMetadata( + completed_saves=completed_saves, + completed_loads=completed_loads, + ) def shutdown(self) -> None: - self._load_stream.synchronize() + for transfer in (*self._inflight_saves, *self._inflight_loads): + transfer.end_event.synchronize() self._save_bufs = None + self._save_stream = None self._save_count = 0 + self._save_bytes = 0 + self._save_mm_hashes = [] + self._inflight_saves.clear() + self._inflight_loads.clear() + self._stream_pool.clear() + self._event_pool.clear() try: self._region.cleanup() except Exception: diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py index a42d6ebda4..5b2109187f 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py @@ -22,24 +22,43 @@ class DescriptorBuffers(NamedTuple): src_ptrs: torch.Tensor dst_ptrs: torch.Tensor sizes: torch.Tensor - # Numpy aliases of src_ptrs/dst_ptrs, written via set_ptrs(). + # Numpy aliases of the three tensors, written via add_copies(). src_np: np.ndarray dst_np: np.ndarray + sizes_np: np.ndarray - def set_ptrs(self, idx: int, src: int, dst: int) -> None: - """Record the source and destination address of descriptor *idx*. + def add_copies( + self, + first: int, + sources: np.ndarray, + destinations: np.ndarray, + byte_counts: np.ndarray, + ) -> None: + """Store a batch of copies in these buffers, the first at *first*. - TODO(torch>=2.14): drop this indirection and assign the tensors - directly once the minimum supported torch is 2.14. The numpy detour - exists only because torch's setitem unpacks the value as a signed - long long before 2.14 (pytorch#191458), rejecting XPU USM pointers - >= 2**63, while the two's-complement rewrite those versions accept is - in turn rejected by uint64 from 2.14 on. Numpy casts against the - array dtype and so works on either. `sizes` holds byte counts and - needs no such care. + One copy is one source address, one destination address, and a number + of bytes. `swap_blocks_batch` performs every copy held in the buffers + in a single call, so a caller fills in all of a step's copies before + asking for any of them to run, and passes *first* to say where its own + copies begin. + + Whole batches go in at once because storing a single number into a + torch tensor costs microseconds. Adding copies one at a time can take + longer than performing them. + + TODO(torch>=2.14): write the tensors directly and drop these numpy + aliases once the minimum supported torch is 2.14. They exist only + because torch's setitem unpacks the value as a signed long long before + 2.14 (pytorch#191458), rejecting XPU USM addresses >= 2**63, while the + two's-complement rewrite those versions accept is in turn rejected by + uint64 from 2.14 on. Numpy casts against the array's own type and so + works on either. Byte counts are small enough to need none of this + care, and use their alias only to match the addresses. """ - self.src_np[idx] = src - self.dst_np[idx] = dst + last = first + len(sources) + self.src_np[first:last] = sources + self.dst_np[first:last] = destinations + self.sizes_np[first:last] = byte_counts class DescriptorBufferPool: @@ -62,7 +81,9 @@ class DescriptorBufferPool: if bufs.src_ptrs.numel() >= n: return bufs src, dst, sizes = (torch.empty(n, dtype=_PTR_DTYPE) for _ in range(3)) - return DescriptorBuffers(src, dst, sizes, src.numpy(), dst.numpy()) + return DescriptorBuffers( + src, dst, sizes, src.numpy(), dst.numpy(), sizes.numpy() + ) def release(self, bufs: DescriptorBuffers) -> None: """Return a buffer triple to the pool for reuse."""