From 0e81e711aeb12fc55c3d6492f1e783968fc3345d Mon Sep 17 00:00:00 2001 From: Liran Schour Date: Tue, 22 Sep 2026 17:02:52 +0300 Subject: [PATCH 1/8] [KV Offload] Run on_schedule_end as the last manager call of a step on_schedule_end is documented as running at the end of each scheduler step, but build_connector_meta called it ~40 lines early, before the preempted-job flush, before _build_store_jobs (which calls prepare_store, and so evicts), and before on_request_finished. Move the call to just before the return so the contract holds. This matters for callers that want to bound a step's exclusive access to the manager: releasing at on_schedule_end today leaves prepare_store and on_request_finished outside the window. prepare_store's poll becomes unconditional. Its fresh poll previously came from the per-step gate being reset mid-build_connector_meta; with on_schedule_end at the end, the gate would still be set by the step's first lookup() and the poll would be skipped, leaving eviction decisions on stale completion data. _flush_pending_cascades now runs after this step's prepare_store, so keys it parked are re-parked and retried next step. That is safe -- _maybe_finalize_request holds off while pending_cascade_keys is non-empty -- and is noted in the code. Test asserts on_schedule_end is the last manager call of every step; without the move it observes ['on_schedule_end', 'prepare_store']. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Liran Schour --- .../offloading_connector/test_scheduler.py | 43 +++++++++++++++++++ .../kv_connector/v1/offloading/scheduler.py | 7 ++- vllm/v1/kv_offload/base.py | 4 ++ vllm/v1/kv_offload/tiering/manager.py | 21 ++++++--- 4 files changed, 68 insertions(+), 7 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index a69129561620..516b293c9c9b 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -734,6 +734,49 @@ def test_scheduler_reports_lookup_sync_delay(request_runner): assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_sum"] > 0 +def test_on_schedule_end_is_the_last_manager_call_of_a_step(request_runner): + """on_schedule_end must run after every other manager call of the step. + + Callers bound a step's exclusive access to the manager by releasing at + on_schedule_end, so anything issued after it — notably prepare_store, which + evicts — would fall outside that window. + """ + block_size = 4 + runner = request_runner( + block_size=block_size, + num_gpu_blocks=8, + async_scheduling=False, + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output([]) + ) + + connector_scheduler = runner.connector_scheduler + build_meta = connector_scheduler.build_connector_meta + windows: list[list[str]] = [] + + def spy(scheduler_output): + first = len(runner.manager.mock_calls) + meta = build_meta(scheduler_output) + windows.append([name for name, _, _ in runner.manager.mock_calls[first:]]) + return meta + + connector_scheduler.build_connector_meta = spy + + runner.new_request(token_ids=[0] * (block_size * 2)) + runner.run(decoded_tokens=[EOS_TOKEN_ID]) + + assert windows, "build_connector_meta was never called" + assert any("prepare_store" in w for w in windows), ( + "step must exercise the calls that follow on_schedule_end today" + ) + assert any("on_request_finished" in w for w in windows) + for names in windows: + assert names[-1] == "on_schedule_end", ( + f"on_schedule_end must be the last manager call of a step, got {names}" + ) + + def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner): """A deferred lookup reports its async delay once it resolves.""" runner = request_runner( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 47aa8654feb8..77434589f08a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1782,7 +1782,6 @@ def build_connector_meta( new_req_ids=[req.req_id for req in scheduler_output.scheduled_new_reqs], preempted_req_ids=scheduler_output.preempted_req_ids or (), ) - self.manager.on_schedule_end(schedule_end_context) # Flush jobs for preempted requests. for req_id in scheduler_output.preempted_req_ids or (): @@ -1828,6 +1827,12 @@ def build_connector_meta( self._current_batch_load_jobs = {} self._current_batch_jobs_to_flush = set() self._current_batch_allocated_block_ids = set() + + # Last manager call of the step, after every prepare_store and + # on_request_finished above. on_schedule_end is documented as running at + # the end of the step, and callers rely on that to bound the step's + # exclusive access to the manager. + self.manager.on_schedule_end(schedule_end_context) return meta def has_pending_push_work(self) -> bool: diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index f8e92abed604..bebf396ce08e 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -382,6 +382,10 @@ def take_events(self) -> Iterable[OffloadingEvent]: def on_schedule_end(self, context: ScheduleEndContext) -> None: """Called once at the end of each scheduler step. + This is the last manager call of a step: every lookup(), prepare_load(), + prepare_store() and on_request_finished() for the step has already been + issued. + Managers may override this to flush deferred work accumulated during the step (e.g., batched promotions). """ diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index ce518042b656..5177d3f2204f 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -594,9 +594,9 @@ def prepare_store( ) -> PrepareStoreOutput | None: """Prepare chunks to be stored from GPU to primary tier. - CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure - that any completed async transfers have their ref_cnt decremented - before the primary tier makes eviction decisions. + CRITICAL: This method polls for finished jobs FIRST to ensure that any + completed async transfers have their ref_cnt decremented before the + primary tier makes eviction decisions. For request-level tiers, chunks already present in the primary tier are immediately cascaded via submit_store(). @@ -620,7 +620,12 @@ def prepare_store( # not-yet-ready chunk's ref_cnt from -1 to 0 via complete_write(), # making it evictable for the first time. # Both must be accounted for before the eviction decision below. - self._maybe_process_finished_jobs() + # Unconditional, not _maybe_process_finished_jobs(): eviction must see + # the freshest completions, and on_schedule_end now runs at the very end + # of the step, so the once-per-step gate would otherwise still be set by + # this step's first lookup() and skip the poll. + self._processed_jobs_this_step = True + self._process_finished_jobs() # Step 2: Store to primary tier (new chunks only). # Cascading of these newly-stored chunks to ALL secondary tiers @@ -855,8 +860,8 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None: """End-of-schedule hook: process finished jobs, flush deferred promotions, and reset the per-step gate. - Called once per scheduler step from - OffloadingConnectorScheduler.build_connector_meta(). + Called once per scheduler step, as the last manager call of the step, + from OffloadingConnectorScheduler.build_connector_meta(). """ # Catch-all poll: guarantees jobs are processed even on steps where # lookup()/prepare_store() were never called (e.g. no requests @@ -871,6 +876,10 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None: self._processed_jobs_this_step = False self._flush_pending_promotions() + # Keys parked by this step's prepare_store cannot have completed yet, so + # they are simply re-parked here and retried on the next step. That is + # safe: _maybe_finalize_request() holds off while pending_cascade_keys is + # non-empty, so nothing finalizes early and the list still drains. self._flush_pending_cascades() for tier in self.secondary_tiers: tier.on_schedule_end(context) From d35c9320c94ffa41472b72852a41c673c5b3a6e4 Mon Sep 17 00:00:00 2001 From: Liran Schour Date: Tue, 22 Sep 2026 18:04:31 +0300 Subject: [PATCH 2/8] [KV Offload] Add OffloadingManager.lock()/unlock() and a tiering control thread Expose exclusive access to an OffloadingManager as a public API, and use it for the first case that needs it. The API. lock()/unlock(), plus a locked() helper, guard a manager's whole state subtree: the manager, the tiers it composes, and tier front-end state. Manager methods never acquire it themselves, so a caller holds it across as many operations as its invariants need. That matters because several manager operations only make sense in sequence -- prepare_load() documents that its keys are offloaded, a fact lookup() established one hook earlier -- and a caller has to be able to span both. OffloadingConnectorScheduler is the first caller. It holds one region per engine step, opened by the first manager-touching hook and closed at the end of build_connector_meta, so the manager is free while the engine runs the model. Hooks outside schedule() take a short region instead. has_pending_push_work must use the short one: the engine asks it before schedule() and returns early when the answer is False, so a step-scoped acquire there would never reach the release and would hold the manager for good -- the state an idle engine settles into. It does need the lock, since has_pending_work() iterates _req_state while another thread may insert or delete. take_events now drains the manager inside the region instead of lazily. manager.take_events() is a generator whose backing list is cleared once exhausted, so iterating it lazily held the manager for as long as the consumer took and discarded events appended in the meantime, unpublished. The consumer. The p2p tier answers a remote peer, which the engine does not drive, so its control plane only advanced from on_schedule_end and a peer's lookup or fetch waited for a step boundary -- seconds on a saturated rank, against a transfer of well under one. Tiers now opt in with needs_control_plane_thread, and the tiering manager runs one thread that polls them and lets them serve, once per control_plane_interval_s (default 1 ms, 0 disables). _process_finished_jobs takes a tier filter so a round does not drag tiers that expect the scheduler thread along with it. The thread yields unconditionally after unlock(), which is what bounds the scheduler's wait to one round: Python locks are not fair, so a release-then-reacquire loop could starve it. A failed round is logged rather than fatal, and on_schedule_end still serves every tier, so the fallback is per-step servicing. shutdown() stops and joins the thread before any tier teardown, since closing a transport under a live round can end the process rather than raise. Docstrings that claimed single-threaded tier access are corrected, including the eviction-race argument in _pin_and_register_hits, which still holds -- now because serve_external_requests runs inside one hold of the lock. One assumption is recorded rather than resolved, in nixl.py: strict serialization is sufficient only if libnixl and UCX keep no state tied to the creating thread. The agent runs with NIXL_THREAD_SYNC_NONE and the config offers no way to ask for STRICT without also starting NIXL's listen thread. This needs a real two-node soak before it is trusted. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Liran Schour --- docs/features/kv_offloading_usage.md | 14 + .../offloading_connector/test_scheduler.py | 78 ++++- .../tiering/test_control_plane_thread.py | 329 ++++++++++++++++++ .../kv_connector/v1/offloading/scheduler.py | 129 ++++++- vllm/v1/kv_offload/base.py | 60 +++- vllm/v1/kv_offload/cpu/manager.py | 1 + vllm/v1/kv_offload/tiering/async_lookup.py | 8 +- vllm/v1/kv_offload/tiering/base.py | 26 +- vllm/v1/kv_offload/tiering/manager.py | 177 +++++++++- .../v1/kv_offload/tiering/p2p/control/base.py | 7 +- vllm/v1/kv_offload/tiering/p2p/data/base.py | 6 +- vllm/v1/kv_offload/tiering/p2p/data/nixl.py | 19 +- vllm/v1/kv_offload/tiering/p2p/manager.py | 22 +- .../kv_offload/tiering/p2p/session/server.py | 11 +- vllm/v1/kv_offload/tiering/spec.py | 9 + 15 files changed, 865 insertions(+), 31 deletions(-) create mode 100644 tests/v1/kv_offload/tiering/test_control_plane_thread.py diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 22a2350d8095..e2fabf4c5f7c 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -108,6 +108,7 @@ vllm serve \ | `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. | | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. | | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). | +| `control_plane_interval_s` | no | `0.001` | multi-tier | Pause between control-plane rounds for tiers that must be serviced between engine steps (currently `p2p`). See [Threading model](#threading-model). Set to `0` to disable the thread and service those tiers once per engine step instead. | | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. | | `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Partial recurrent tails emit the hash-aligned portion from the physical block start through the tail boundary. Other sliding-window/SSM chunks keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). | @@ -150,6 +151,19 @@ The filesystem and object-store tiers can publish hash-only `BlockStored` KV eve Set the optional `locality` tier field to `LOCAL` or `REMOTE` to describe the tier's storage location relative to the publishing vLLM instance. `LOCAL` marks storage local to that instance, while `REMOTE` marks storage that is not local to it. When the setting is omitted, locality is unspecified. vLLM does not infer it from the tier type, so an `obj` tier is not implicitly `REMOTE`. A KV event includes `locality` only when the tier explicitly configures it. This metadata describes the tier property without implying that a consumer can already route requests to its blocks. +### Threading model + +Tier methods run in the scheduler process, under the tiering manager's lock, so they are never entered concurrently. They are not always called from the same thread, though. + +Most tiers only ever see the scheduler thread, which holds the lock for the length of one engine step. A tier whose counterpart is not driven by the engine — the `p2p` tier, answering a remote peer — needs servicing while the engine is busy running the model, because otherwise a peer's lookup or fetch waits for the next step boundary. On a saturated rank that wait can be seconds, dwarfing the transfer it gates. Such tiers set `needs_control_plane_thread`, and the tiering manager runs one thread that polls them and lets them serve, once per `control_plane_interval_s`. + +Two consequences worth knowing: + +- Any tier can be reached from that thread indirectly, because serving a peer looks keys up through the tiering manager, which fans out to the other tiers. Tier state must therefore not assume a particular thread, only that the manager lock is held. +- Work the thread deliberately leaves alone is batched per step: a promotion it initiates is submitted at the next `on_schedule_end()`, not immediately. + +Setting `control_plane_interval_s` to `0` disables the thread; opted-in tiers then get serviced once per engine step, as they were before it existed. + ### Filesystem (FS) The filesystem tier (`type: "fs"`) writes blocks to a filesystem directory. diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 516b293c9c9b..f593f24a570c 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -758,7 +758,15 @@ def test_on_schedule_end_is_the_last_manager_call_of_a_step(request_runner): def spy(scheduler_output): first = len(runner.manager.mock_calls) meta = build_meta(scheduler_output) - windows.append([name for name, _, _ in runner.manager.mock_calls[first:]]) + # lock()/unlock() bracket the step; only the state operations they guard + # are ordered relative to on_schedule_end. + windows.append( + [ + name + for name, _, _ in runner.manager.mock_calls[first:] + if name not in ("lock", "unlock") + ] + ) return meta connector_scheduler.build_connector_meta = spy @@ -777,6 +785,74 @@ def spy(scheduler_output): ) +def test_step_releases_the_manager_before_the_model_runs(request_runner): + """The step-scoped region must close when build_connector_meta returns. + + Everything after it -- the model future above all -- runs with the manager + free, which is the only window another thread has to use it. + """ + runner = request_runner( + block_size=4, + num_gpu_blocks=10, + async_scheduling=False, + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output([]) + ) + + runner.new_request(token_ids=[1] * 4) + runner.run(decoded_tokens=[EOS_TOKEN_ID]) + + scheduler = runner.connector_scheduler + assert scheduler._manager_held is False + assert runner.manager.lock.call_count > 0 + assert runner.manager.lock.call_count == runner.manager.unlock.call_count + + +def test_has_pending_push_work_does_not_open_the_step_region(request_runner): + """It must take a short region, never the step-scoped one. + + The engine asks this before schedule() and returns early when the answer is + False, so a step-scoped acquire here would never reach the release in + build_connector_meta and would hold the manager for good. + """ + runner = request_runner( + block_size=4, + num_gpu_blocks=10, + async_scheduling=False, + ) + scheduler = runner.connector_scheduler + runner.manager.has_pending_work.return_value = False + + assert scheduler.has_pending_push_work() is False + assert scheduler._manager_held is False + assert runner.manager.lock.call_count == runner.manager.unlock.call_count + + +def test_take_events_drains_the_manager_inside_the_region(request_runner): + """take_events must not hold the manager across its own iteration. + + manager.take_events() is a generator whose backing list is cleared once + exhausted, so draining it lazily would both pin the manager for as long as + the consumer takes and lose events appended in the meantime. + """ + runner = request_runner( + block_size=4, + num_gpu_blocks=10, + async_scheduling=False, + ) + scheduler = runner.connector_scheduler + runner.manager.take_events.return_value = iter(()) + runner.manager.reset_mock() + + events = list(scheduler.take_events()) + + assert events == [] + assert runner.manager.take_events.call_count == 1 + assert scheduler._manager_held is False + assert runner.manager.lock.call_count == runner.manager.unlock.call_count + + def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner): """A deferred lookup reports its async delay once it resolves.""" runner = request_runner( diff --git a/tests/v1/kv_offload/tiering/test_control_plane_thread.py b/tests/v1/kv_offload/tiering/test_control_plane_thread.py new file mode 100644 index 000000000000..308a09fc7644 --- /dev/null +++ b/tests/v1/kv_offload/tiering/test_control_plane_thread.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for OffloadingManager.lock() and the tiering control-plane thread. + +The control plane of a tier that answers something the engine does not drive -- +a remote peer, say -- used to advance only from on_schedule_end(), so a peer's +request waited for a model-step boundary. These tests cover the thread that +services such tiers between steps and the exclusion that makes it safe. +""" + +import threading +import time +from collections.abc import Iterable +from typing import ClassVar +from unittest.mock import MagicMock + +import pytest + +from vllm.v1.kv_offload.base import ( + LookupResult, + Medium, + OffloadKey, + ReqContext, + RequestOffloadingContext, + ScheduleEndContext, +) +from vllm.v1.kv_offload.tiering.base import ( + JobResult, + ParentManager, + SecondaryTierManager, + TransferJob, +) +from vllm.v1.kv_offload.tiering.manager import ( + CPUPrimaryTierOffloadingManager, + TieringOffloadingManager, +) + +from .test_tiering_offloading import _mock_mmap_region + +_CTX = ReqContext(req_id="test") +_EMPTY_SCHEDULE_END = ScheduleEndContext(new_req_ids=(), preempted_req_ids=()) + +# Long enough that a 1ms-interval thread gets many rounds, short enough to keep +# the suite quick. +_SETTLE_S = 0.5 + + +class _FakeTier(SecondaryTierManager): + """Tier that records which thread serviced it, and when.""" + + medium: ClassVar[Medium] = Medium.CPU + needs_control_plane_thread: ClassVar[bool] = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Counters are only ever incremented, and int writes are atomic under + # the GIL, so the test thread may read them without the manager lock. + self.serves = 0 + self.polls = 0 + self.serve_threads: set[int] = set() + self.served = threading.Event() + self.shutdown_saw_live_thread: bool | None = None + self.on_serve = None + + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + return LookupResult.MISS + + def submit_store(self, job_metadata: TransferJob) -> None: + pass + + def submit_load(self, job_metadata: TransferJob) -> None: + pass + + def get_finished_jobs(self) -> Iterable[JobResult]: + self.polls += 1 + return () + + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + return RequestOffloadingContext() + + def drain_jobs(self) -> None: + pass + + def serve_external_requests(self, parent: ParentManager) -> None: + self.serves += 1 + self.serve_threads.add(threading.get_ident()) + if self.on_serve is not None: + self.on_serve(parent) + self.served.set() + + def shutdown(self) -> None: + self.shutdown_saw_live_thread = any( + t.name == "vllm_offload_control_plane" and t.is_alive() + for t in threading.enumerate() + ) + + +class _QuietTier(_FakeTier): + """Tier that does not want servicing between steps.""" + + needs_control_plane_thread: ClassVar[bool] = False + + +def _make_manager(tier_cls=_FakeTier, num_chunks: int = 8, **kwargs): + primary = CPUPrimaryTierOffloadingManager( + num_chunks=num_chunks, + mmap_region=_mock_mmap_region(num_chunks), + ) + tier = tier_cls( + offloading_spec=MagicMock(), + primary_kv_view=primary.get_kv_memoryview(), + tier_type="fake", + ) + manager = TieringOffloadingManager( + primary_tier=primary, + secondary_tiers=[tier], + **kwargs, + ) + return manager, tier + + +@pytest.fixture +def manager_and_tier(request): + kwargs = getattr(request, "param", {}) + manager, tier = _make_manager(**kwargs) + try: + yield manager, tier + finally: + manager.shutdown() + + +def _live_control_threads() -> list[threading.Thread]: + return [ + t + for t in threading.enumerate() + if t.name == "vllm_offload_control_plane" and t.is_alive() + ] + + +def test_control_thread_and_scheduler_never_overlap(manager_and_tier): + """lock() must keep the two threads out of the manager at the same time. + + Both sides must also make progress: a thread that never runs would pass a + mutual-exclusion check trivially, and so would one that starves the + scheduler. + """ + manager, tier = manager_and_tier + state = {"inside": False, "overlaps": 0} + + def critical_section(): + if state["inside"]: + state["overlaps"] += 1 + state["inside"] = True + time.sleep(0.002) + state["inside"] = False + + tier.on_serve = lambda parent: critical_section() + + steps = 0 + deadline = time.monotonic() + _SETTLE_S + while time.monotonic() < deadline: + manager.lock() + try: + critical_section() + finally: + manager.unlock() + steps += 1 + + assert state["overlaps"] == 0 + assert steps > 0, "scheduler thread was starved by the control thread" + assert tier.serves > 0, "control thread never ran" + + +def test_serves_while_the_scheduler_is_outside_the_lock(manager_and_tier): + """The point of the thread: serve during the model-execution window. + + The engine releases the manager at the end of a step and then blocks on the + model future. Here the "scheduler" does the same -- one step, then a sleep + with the lock free -- and the tier must be serviced during that sleep. + """ + manager, tier = manager_and_tier + + manager.lock() + try: + manager.on_schedule_end(_EMPTY_SCHEDULE_END) + finally: + manager.unlock() + + tier.served.clear() + serves_before = tier.serves + + # Stand in for future.result(): the engine thread is blocked, holding + # nothing. + assert tier.served.wait(_SETTLE_S), ( + "tier was not serviced while the scheduler held no lock" + ) + assert tier.serves > serves_before + assert tier.polls > 0 + + scheduler_thread_id = threading.get_ident() + assert scheduler_thread_id not in tier.serve_threads or len(tier.serve_threads) > 1 + + +@pytest.mark.parametrize( + "manager_and_tier", [{"control_plane_interval_s": 0.0}], indirect=True +) +def test_zero_interval_disables_the_thread(manager_and_tier): + """The escape hatch back to per-step servicing. + + Without the thread the tier is serviced only by on_schedule_end, which is + the behaviour the control plane had before it existed. + """ + manager, tier = manager_and_tier + + assert not _live_control_threads() + assert not tier.served.wait(0.05) + assert tier.serves == 0 + + manager.lock() + try: + manager.on_schedule_end(_EMPTY_SCHEDULE_END) + finally: + manager.unlock() + + assert tier.serves == 1 + + +def test_no_thread_when_no_tier_asks_for_one(): + """A tier that does not opt in must not cost a thread.""" + manager, tier = _make_manager(tier_cls=_QuietTier) + try: + assert not _live_control_threads() + finally: + manager.shutdown() + + +def test_tier_that_did_not_opt_in_is_not_serviced_off_thread(): + """Only opted-in tiers are touched by the control thread. + + A tier that expects the scheduler thread must not be dragged onto another + one just because it shares a manager with a tier that opted in. + """ + num_chunks = 8 + primary = CPUPrimaryTierOffloadingManager( + num_chunks=num_chunks, + mmap_region=_mock_mmap_region(num_chunks), + ) + common = { + "offloading_spec": MagicMock(), + "primary_kv_view": primary.get_kv_memoryview(), + } + opted_in = _FakeTier(tier_type="fake", **common) + quiet = _QuietTier(tier_type="quiet", **common) + manager = TieringOffloadingManager( + primary_tier=primary, + secondary_tiers=[opted_in, quiet], + ) + try: + assert opted_in.served.wait(_SETTLE_S) + assert quiet.serves == 0 + assert quiet.polls == 0 + + manager.lock() + try: + manager.on_schedule_end(_EMPTY_SCHEDULE_END) + finally: + manager.unlock() + + # on_schedule_end still serves every tier, so the thread dying degrades + # to per-step servicing rather than to none. + assert quiet.serves == 1 + finally: + manager.shutdown() + + +def test_per_step_gate_is_not_touched_by_the_control_thread(manager_and_tier): + """_processed_jobs_this_step stays owned by the scheduler thread. + + Sharing it would let a fast step skip a poll it would otherwise do, pushing + a completion out by a whole step. + """ + manager, tier = manager_and_tier + + assert tier.served.wait(_SETTLE_S) + polls_after_rounds = tier.polls + assert polls_after_rounds > 0 + assert manager._processed_jobs_this_step is False + + +def test_shutdown_joins_the_control_thread_before_tier_teardown(manager_and_tier): + """The thread drives tier transports, so it must be stopped first. + + Closing a transport under a running round can take the process down rather + than raise. + """ + manager, tier = manager_and_tier + assert tier.served.wait(_SETTLE_S) + assert _live_control_threads() + + manager.shutdown() + + assert tier.shutdown_saw_live_thread is False + assert not _live_control_threads() + + +def test_locked_helper_releases_on_error(manager_and_tier): + """locked() must not leak the lock when its body raises.""" + manager, _ = manager_and_tier + + with pytest.raises(RuntimeError), manager.locked(): + raise RuntimeError("boom") + + assert manager.lock(timeout=0.5), "locked() leaked the lock" + manager.unlock() + + +def test_lock_timeout_reports_failure(manager_and_tier): + """A timed lock() returns False rather than blocking forever.""" + manager, _ = manager_and_tier + + manager.lock() + try: + acquired: list[bool] = [] + t = threading.Thread(target=lambda: acquired.append(manager.lock(timeout=0.01))) + t.start() + t.join(timeout=5.0) + assert acquired == [False] + finally: + manager.unlock() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 77434589f08a..2cd0d1b1552c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time -from collections.abc import Iterable, Sequence +from collections.abc import Iterable, Iterator, Sequence +from contextlib import contextmanager from dataclasses import dataclass, field, replace from itertools import chain, islice from typing import Any, NamedTuple @@ -633,6 +634,68 @@ def _sliding_window_sort_key(i: int) -> int: self._events_tracker = OffloadingEventsTracker(spec.kv_events_config) + # True while this thread holds the manager lock for the current schedule(). + # Read and written only by the scheduler thread, so it needs no + # synchronization of its own. Declared on the class so an instance built + # without __init__ still reads False rather than raising. + _manager_held: bool = False + + # ------------------------------------------------------------------ + # Manager exclusion + # + # The manager is entered from the scheduler thread and, for managers that + # run one, from a control-plane thread. Manager methods never take their + # own lock, so the exclusion boundary lives here. + # + # schedule() holds one region for the whole step: opened by the first + # manager-touching hook and closed at the end of build_connector_meta. + # That span is required, not tidiness -- get_num_new_matched_tokens + # establishes lookup() HITs that update_state_after_alloc later pins with + # prepare_load(), and prepare_store() evicts after both. Releasing between + # them would let another thread evict a block the step already treated as + # present. + # + # Hooks outside schedule() use a short region instead, so the manager is + # free while the engine runs the model. + # ------------------------------------------------------------------ + + def _acquire_step(self) -> None: + """Open the step-scoped region. Idempotent within a step.""" + if self._manager_held: + return + self.manager.lock() + self._manager_held = True + + def _release_step(self) -> None: + """Close the step-scoped region, letting other threads in.""" + if not self._manager_held: + return + self._manager_held = False + self.manager.unlock() + + @contextmanager + def _manager_locked(self) -> Iterator[None]: + """Hold the manager for the duration of one hook outside schedule(). + + None of these hooks can legitimately run inside the step-scoped region: + on_new_request comes from add_request, has_pending_push_work from the + engine's liveness check before schedule(), the rest from output + processing or a client RPC. So finding the region open here means a + previous schedule() raised between _acquire_step() and the release in + build_connector_meta. Recover instead of deadlocking. + """ + if self._manager_held: + logger.error( + "Offloading manager lock was left held by a previous " + "schedule(); releasing it before continuing." + ) + self._release_step() + self.manager.lock() + try: + yield + finally: + self.manager.unlock() + def _maybe_observe_lookup_async_delay( self, req_status: RequestOffloadState ) -> None: @@ -1024,7 +1087,8 @@ def _lookup( def on_new_request(self, request: Request) -> None: """Called when a new request is added to the scheduler.""" req_context = _create_req_context(request) - offloading_context = self.manager.on_new_request(req_context) + with self._manager_locked(): + offloading_context = self.manager.on_new_request(req_context) req_status = RequestOffloadState( config=self.config, req=request, @@ -1060,6 +1124,10 @@ def get_num_new_matched_tokens( (between scheduler steps). """ + # First manager-touching hook of a step: opens the step-scoped region, + # closed at the end of build_connector_meta. + self._acquire_step() + req_status = self._req_status[request.request_id] for group_state in req_status.group_states: group_state.block_ids.clear() @@ -1099,6 +1167,10 @@ def update_state_after_alloc( if num_external_tokens == 0: return + # prepare_load() below pins keys that get_num_new_matched_tokens() + # already saw as HIT, so both must sit inside the same region. + self._acquire_step() + req_status = self._req_status[request.request_id] num_locally_computed_tokens = req_status.num_locally_computed_tokens @@ -1776,6 +1848,18 @@ def _build_store_jobs( def build_connector_meta( self, scheduler_output: SchedulerOutput + ) -> KVConnectorMetadata: + # A step that scheduled nothing reaches build_connector_meta without + # any earlier manager touch, so open the region here too. + self._acquire_step() + try: + return self._build_connector_meta(scheduler_output) + finally: + # The step is over: hand the manager back. + self._release_step() + + def _build_connector_meta( + self, scheduler_output: SchedulerOutput ) -> KVConnectorMetadata: self._update_req_states(scheduler_output) schedule_end_context = ScheduleEndContext( @@ -1840,8 +1924,17 @@ def has_pending_push_work(self) -> bool: While True, build_connector_meta() and update_connector_output() continue to be called even when no requests are scheduled. + + Must use a short region, never the step-scoped one. The engine asks this + before schedule() and returns early when the answer is False, so a + step-scoped acquire here would never reach build_connector_meta's + release and would hold the manager for good -- which is precisely the + state an idle engine settles into. """ - return bool(self._jobs) or self.manager.has_pending_work() + if self._jobs: + return True + with self._manager_locked(): + return self.manager.has_pending_work() def update_connector_output(self, connector_output: KVConnectorOutput): """Update KVConnector state from worker-side connectors output. @@ -1851,6 +1944,10 @@ def update_connector_output(self, connector_output: KVConnectorOutput): connectors output. """ + with self._manager_locked(): + self._update_connector_output(connector_output) + + def _update_connector_output(self, connector_output: KVConnectorOutput): meta = connector_output.kv_connector_worker_meta if not isinstance(meta, OffloadingWorkerMetadata): assert meta is None @@ -1929,7 +2026,8 @@ def get_stats(self) -> OffloadingConnectorStats | None: stats = self._connector_stats self._connector_stats = OffloadingConnectorStats() - manager_stats = self.manager.get_stats() + with self._manager_locked(): + manager_stats = self.manager.get_stats() if manager_stats is not None: if stats is None: stats = manager_stats @@ -1958,8 +2056,9 @@ def request_finished( # Untracked request (offloading never started): no in-flight jobs, # nothing was deferred, so finalize immediately. req_context = _create_req_context(request) - self.manager.on_new_request(req_context) - self.manager.on_request_finished(req_context) + with self._manager_locked(): + self.manager.on_new_request(req_context) + self.manager.on_request_finished(req_context) return False, None self._maybe_observe_lookup_async_delay(req_status) @@ -1991,7 +2090,16 @@ def take_events(self) -> Iterable[KVCacheEvent]: the underlying :class:`OffloadingEvent` stream. """ - yield from self._events_tracker.take_events(self.manager.take_events()) + # Drain the manager eagerly, inside the region. manager.take_events() is + # itself a generator whose backing list is cleared once exhausted, so + # iterating it lazily would both hold the manager for as long as the + # consumer takes and let a concurrent writer append events that the + # clear then discards unpublished. + with self._manager_locked(): + manager_events = list(self.manager.take_events()) + # Still a generator, so callers testing the result for truthiness keep + # seeing a non-empty object the way they do today. + yield from self._events_tracker.take_events(manager_events) def reset_cache(self) -> None: """Reset the offloading manager cache, evicting all stored chunks.""" @@ -2000,6 +2108,13 @@ def reset_cache(self) -> None: assert not self._current_batch_jobs_to_flush assert not self._current_batch_allocated_block_ids + # Held for the whole reset. manager.reset_cache() drains tiers by + # polling them itself, so a control-plane thread locked out here loses + # nothing; it simply skips rounds until the reset completes. + with self._manager_locked(): + self._reset_cache() + + def _reset_cache(self) -> None: # Flush all in-flight jobs self._current_batch_jobs_to_flush.update(self._jobs.keys()) diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index bebf396ce08e..95c15e60a354 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Core abstractions for KV cache offloading in vLLM v1.""" +import threading from abc import ABC, abstractmethod -from collections.abc import Collection, Iterable, Sequence +from collections.abc import Collection, Iterable, Iterator, Sequence +from contextlib import contextmanager from dataclasses import dataclass, field from enum import Enum, auto from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, NewType, TypeVar @@ -197,6 +199,14 @@ class OffloadingEvent: as well as a list of blocks that were evicted as a result. complete_store() - marks a previous store as completed. Following this call, the given blocks will become loadable. + +Exclusion: + lock() / unlock() guard a manager's whole state subtree: the manager + itself, every tier it composes, and any front-end state those tiers own. + Manager methods never take the lock themselves, so a caller holds it + across as many operations as its invariants need -- notably a lookup() + HIT and the prepare_load() that pins the result, which must not be + separated by another thread's eviction. """ @@ -232,6 +242,54 @@ class OffloadingKVEventsConfig: class OffloadingManager(ABC): + def __init__(self) -> None: + self._manager_lock = threading.Lock() + + def lock(self, timeout: float | None = None) -> bool: + """Acquire exclusive access to this manager's state subtree. + + Scope is the manager, every tier it composes, and any front-end state + those tiers own (e.g. a tier's async lookup front end). Every caller of + an OffloadingManager method must hold this lock. Manager methods never + acquire it themselves, so a caller may hold it across an arbitrary + sequence of operations -- which is required wherever one call + establishes a fact a later call relies on, such as a lookup() HIT + followed by the prepare_load() that pins it. + + Not reentrant. A thread must not call lock() twice without an + intervening unlock(); use the held state the caller already tracks, or + a *_unlocked-style helper, rather than re-acquiring. + + Args: + timeout: Seconds to wait. None blocks until acquired. + + Returns: + True if the lock was acquired, False only on timeout. + + """ + if timeout is None: + return self._manager_lock.acquire() + return self._manager_lock.acquire(timeout=timeout) + + def unlock(self) -> None: + """Release a lock() acquired by the calling thread.""" + self._manager_lock.release() + + @contextmanager + def locked(self, timeout: float | None = None) -> Iterator[bool]: + """Scoped lock()/unlock() for callers that need no step granularity. + + Yields whether the lock was acquired; the body still runs on timeout, + so a caller passing a timeout must check the value before touching the + manager. + """ + acquired = self.lock(timeout) + try: + yield acquired + finally: + if acquired: + self.unlock() + @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """Checks whether a single block is offloaded and ready to be read. diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index e42cb47f2530..e8b3011f5f02 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -63,6 +63,7 @@ def __init__( store_threshold: int = 1, max_tracker_size: int = 64_000, ): + super().__init__() self.medium: Medium = Medium.CPU self._num_chunks: int = num_chunks self._num_allocated_chunks: int = 0 diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py index a6879d8c44e5..e57fb36bab5e 100644 --- a/vllm/v1/kv_offload/tiering/async_lookup.py +++ b/vllm/v1/kv_offload/tiering/async_lookup.py @@ -11,8 +11,12 @@ -------------- There is no explicit lock. Thread safety is achieved by ownership: -* _lookup_state and _lookup_batch are owned exclusively by the scheduler - thread. lookup(), flush(), and cleanup() read and write them directly. +* _lookup_state and _lookup_batch are owned by whichever thread holds the + OffloadingManager lock. lookup(), flush(), and cleanup() read and write + them directly. That is usually the scheduler thread, but a tier reached + through ParentManager -- a peer lookup fanning out from another tier, say -- + can arrive on the tiering manager's control-plane thread instead. Both hold + the manager lock, so the accesses stay serialized. * _lookup_queue is written by the scheduler (flush → put_nowait, one item per step) and read by the background thread (get). queue.Queue is diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index a90d33eb1136..1e0bee499abc 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -138,10 +138,23 @@ class SecondaryTierManager(ABC): IMPORTANT: All methods run in the Scheduler process and must be lightweight and non-blocking. submit_load() and submit_store() submit async jobs; get_finished_jobs() polls for completion. + + Methods are called under the tiering manager's lock, so they are never + entered concurrently -- but not always from the same thread. A tier that + sets needs_control_plane_thread is also serviced from the manager's + control-plane thread, and any tier can be reached from there via + ParentManager fan-out. Keep per-thread assumptions out of tier state. """ medium: ClassVar[Medium | None] = None + # Whether this tier's control plane must be serviced between engine steps + # and not only from on_schedule_end(). Set it on tiers that answer a + # counterpart the engine does not drive -- a remote peer, say -- whose + # requests would otherwise wait for a step boundary. The tiering manager + # runs a thread on behalf of such tiers; see serve_control_plane(). + needs_control_plane_thread: ClassVar[bool] = False + def __init__( self, offloading_spec: "OffloadingSpec", @@ -312,8 +325,17 @@ def on_request_finished(self, req_context: ReqContext) -> None: def serve_external_requests(self, parent: ParentManager) -> None: """Process remotely-originated requests using the parent manager. - Called once per scheduler step, BEFORE _flush_pending_promotions(). - The parent handle is valid only for the duration of this call. + Called once per scheduler step, BEFORE _flush_pending_promotions(), and + additionally once per control-plane round for tiers that set + needs_control_plane_thread. The _flush_pending_promotions() ordering + holds only for the per-step call: a promotion this method initiates from + a control-plane round is submitted at the next on_schedule_end(). + + The parent handle is valid only for the duration of this call, and the + caller holds the manager lock throughout it -- so a lookup() HIT taken + here stays valid until this call pins it. An implementation must not + release that lock partway through. + Tiers that don't serve external requests leave this as a no-op. """ return diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 5177d3f2204f..795b4d8d6026 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -19,6 +19,7 @@ protecting chunks from eviction until complete_read() is called """ +import threading import time from collections.abc import Collection, Iterable, Sequence from dataclasses import dataclass, field @@ -58,6 +59,25 @@ logger = init_logger(__name__) +# Pause between control-plane rounds. A round is a non-blocking sweep, so this +# sets how long a peer's request can sit unserved; 1 ms is three orders of +# magnitude below the model step it hides, and matches the pause the p2p tier's +# own drain loops and EngineCoreProc's GIL-yield sleep already use. +_CONTROL_PLANE_INTERVAL_S = 0.001 +# How long a round waits for the manager lock before giving up and retrying. +# Bounded so the thread stays responsive to shutdown even if a caller holds the +# lock for a long time, or leaked it. +_CONTROL_PLANE_LOCK_TIMEOUT_S = 1.0 +# Warn after this much continuous failure to acquire the lock. Above the 5 s +# warning drain_jobs() emits, so a slow reset_cache() does not warn twice. +_CONTROL_PLANE_STALL_WARN_S = 30.0 +# Give up after this many consecutive failed rounds. on_schedule_end() still +# serves the tiers, so the fallback is per-step servicing rather than none. +_CONTROL_PLANE_MAX_CONSECUTIVE_ERRORS = 100 +# How long shutdown() waits for the thread to finish its round. +_CONTROL_PLANE_JOIN_TIMEOUT_S = 5.0 + + @dataclass class PendingPromotion: """Accumulator for chunks awaiting submit_load() for one (tier, request).""" @@ -191,6 +211,7 @@ def __init__( self, primary_tier: CPUPrimaryTierOffloadingManager, secondary_tiers: list[SecondaryTierManager] | None = None, + control_plane_interval_s: float = _CONTROL_PLANE_INTERVAL_S, ): """Initialize the TieringOffloadingManager. @@ -198,8 +219,13 @@ def __init__( primary_tier: The primary tier manager (CPU-based). secondary_tiers: List of secondary tier managers (e.g., Storage, Network). Can be None or empty list. + control_plane_interval_s: Pause between control-plane rounds for + tiers that set needs_control_plane_thread. Zero or negative + disables the thread, leaving those tiers serviced once per + engine step as before. """ + super().__init__() self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier self.secondary_tiers = secondary_tiers or [] @@ -242,6 +268,118 @@ def __init__( tier: i for i, tier in enumerate(self.secondary_tiers) } + # Tiers whose control plane must be serviced between engine steps. + self._control_plane_tiers: list[tuple[int, SecondaryTierManager]] = [ + (i, tier) + for i, tier in enumerate(self.secondary_tiers) + if tier.needs_control_plane_thread + ] + self._control_plane_interval_s = control_plane_interval_s + self._control_plane_stop = threading.Event() + self._control_plane_thread: threading.Thread | None = None + self._start_control_plane() + + # ------------------------------------------------------------------ + # Control plane + # ------------------------------------------------------------------ + + def _start_control_plane(self) -> None: + """Start the control-plane thread, if any tier asked for one.""" + if not self._control_plane_tiers or self._control_plane_interval_s <= 0: + return + self._control_plane_thread = threading.Thread( + target=self._control_plane_loop, + name="vllm_offload_control_plane", + daemon=True, + ) + self._control_plane_thread.start() + logger.info( + "KV offload control-plane thread started for tier(s) %s, interval %.3fs", + [tier.tier_type for _, tier in self._control_plane_tiers], + self._control_plane_interval_s, + ) + + def _stop_control_plane(self) -> None: + """Stop and join the control-plane thread. + + Must run before any tier teardown: the thread drives tier transports, + and closing one underneath it can crash the process rather than raise. + """ + self._control_plane_stop.set() + thread = self._control_plane_thread + if thread is None: + return + thread.join(timeout=_CONTROL_PLANE_JOIN_TIMEOUT_S) + if thread.is_alive(): + logger.error( + "KV offload control-plane thread did not exit within %.1fs", + _CONTROL_PLANE_JOIN_TIMEOUT_S, + ) + else: + self._control_plane_thread = None + + def _control_plane_loop(self) -> None: + stalled_since: float | None = None + errors = 0 + while not self._control_plane_stop.is_set(): + if not self.lock(timeout=_CONTROL_PLANE_LOCK_TIMEOUT_S): + now = time.monotonic() + if stalled_since is None: + stalled_since = now + elif now - stalled_since > _CONTROL_PLANE_STALL_WARN_S: + logger.warning( + "KV offload control-plane thread has not acquired the " + "manager lock for %.0fs; peer requests are only being " + "served at engine-step boundaries.", + now - stalled_since, + ) + stalled_since = None + continue + stalled_since = None + try: + self.serve_control_plane() + errors = 0 + except Exception: + # Never let a round kill the thread silently: on_schedule_end() + # still serves these tiers, so persistent failure degrades to + # per-step servicing instead of stopping the control plane. + errors += 1 + logger.exception("KV offload control-plane round failed") + finally: + self.unlock() + if errors >= _CONTROL_PLANE_MAX_CONSECUTIVE_ERRORS: + logger.error( + "KV offload control-plane thread stopping after %d " + "consecutive failures; falling back to per-step servicing.", + errors, + ) + return + # Yield unconditionally, and only after unlock(). Python locks are + # not fair, so a release-then-reacquire loop could starve the + # scheduler thread; this pause is what bounds its wait to one round. + self._control_plane_stop.wait(self._control_plane_interval_s) + + def serve_control_plane(self) -> None: + """Advance the control plane of tiers that cannot wait for a step. + + Polls those tiers for finished jobs and lets them serve whatever their + counterpart has asked for -- the same two steps on_schedule_end() runs, + narrowed to the tiers that opted in. + + Excludes _flush_pending_promotions() and _flush_pending_cascades() by + design: both are per-step batching points, and running them here would + fragment batches the scheduler thread accumulated during schedule(). A + promotion that a lookup() from this path initiates therefore waits for + the next on_schedule_end() to be submitted. + + The caller must hold lock() for the whole call. serve_external_requests() + relies on that: it establishes lookup() HITs and then pins them, and the + two must not be separated by an eviction. + """ + self._process_finished_jobs(self._control_plane_tiers) + for _, tier in self._control_plane_tiers: + tier.serve_external_requests(self._tier_parents[tier]) + @property def _transfer_jobs(self) -> dict[JobId, JobMetadata]: return self._jobs @@ -306,8 +444,11 @@ def _complete_promotion( False, ) - def _process_finished_jobs(self): - """Unconditionally poll all secondary tiers for completed jobs. + def _process_finished_jobs( + self, + tiers: Sequence[tuple[int, SecondaryTierManager]] | None = None, + ): + """Unconditionally poll secondary tiers for completed jobs. This method: 1. Calls get_finished_jobs() on each secondary tier @@ -315,8 +456,15 @@ def _process_finished_jobs(self): to decrement ref_cnt 3. For completed loads (secondary→primary): calls primary.complete_write() to make chunks available + + Args: + tiers: (index, tier) pairs to poll. Defaults to every secondary + tier. Pass a subset to poll only the tiers a caller is + responsible for, so a tier that expects to be polled on the + scheduler thread is not dragged elsewhere. + """ - for i, tier in enumerate(self.secondary_tiers): + for i, tier in tiers if tiers is not None else enumerate(self.secondary_tiers): for completed_job in tier.get_finished_jobs(): job_id = completed_job.job_id job_metadata = self._pop_job(job_id) @@ -992,9 +1140,32 @@ def get_stats(self) -> OffloadingConnectorStats | None: def shutdown(self) -> None: """Shut down secondary tiers before releasing primary resources. + Stops the control-plane thread first: it drives tier transports, and + tearing one down underneath it (closing a socket, destroying a ZMQ + context) can take the process down rather than raise. If the thread + will not exit, hold the manager lock across tier teardown so it cannot + be mid-round, and skip teardown entirely if even that is unavailable -- + a lingering daemon thread in an exiting process beats a crash. + Every secondary tier is given a shutdown attempt. If any shutdown fails, preserve the primary mmap because a failed tier may still use it. """ + self._stop_control_plane() + if self._control_plane_thread is not None: + if not self.lock(timeout=_CONTROL_PLANE_JOIN_TIMEOUT_S): + logger.error( + "KV offload control-plane thread is still running and the " + "manager lock is unavailable; skipping tier shutdown." + ) + return + try: + self._shutdown_tiers() + finally: + self.unlock() + return + self._shutdown_tiers() + + def _shutdown_tiers(self) -> None: shutdown_error: Exception | None = None for tier_idx, tier in enumerate(self.secondary_tiers): try: diff --git a/vllm/v1/kv_offload/tiering/p2p/control/base.py b/vllm/v1/kv_offload/tiering/p2p/control/base.py index 5a7030737bdb..890bc16b2441 100644 --- a/vllm/v1/kv_offload/tiering/p2p/control/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/control/base.py @@ -20,8 +20,11 @@ ├── mark_dead() — signal that the peer is gone └── close() — tear down the connection -Threading model: all I/O is driven by the caller invoking poll(). -No background threads. poll() must be called periodically to: +Threading model: all I/O is driven by the caller invoking poll(). No background +threads of its own. More than one caller thread may drive poll() over a +transport's lifetime, but never concurrently -- callers serialize themselves +(the tiering manager does so with its lock), because sockets here are not +thread-safe. poll() must be called periodically to: - receive messages (buffered per-connection) - accept new inbound peers - detect disconnections diff --git a/vllm/v1/kv_offload/tiering/p2p/data/base.py b/vllm/v1/kv_offload/tiering/p2p/data/base.py index 62304b7544a3..866450bf04a2 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/base.py @@ -66,7 +66,11 @@ - close() releases all resources (memory registrations, handles). After close(), no other methods may be called. -Threading model: no background threads. All I/O driven by poll(). +Threading model: no background Python threads. All I/O driven by poll(). More +than one caller thread may drive a transport over its lifetime, but never +concurrently -- callers serialize themselves (the tiering manager does so with +its lock). Note that serialization alone is not obviously sufficient for a +backend that binds state to the thread that created it; see NixlTransport. """ from __future__ import annotations diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py index e8fb66dd6e30..e9fa61e3d780 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -1,6 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""NixlTransport: Data-plane transport for RDMA-based KV block transfers via NIXL.""" +"""NixlTransport: Data-plane transport for RDMA-based KV block transfers via NIXL. + +Threading: the agent is reached from the scheduler thread and, when the tiering +manager runs one, from its control-plane thread. The manager's lock serializes +them strictly -- one Python thread inside any agent call at a time, with a full +memory barrier between -- which is sufficient only if libnixl and UCX keep no +state tied to the creating thread. + +That is an assumption, not a verified fact, and worth knowing when debugging +transfer corruption or hangs. nixl_agent picks NIXL_THREAD_SYNC_STRICT only when +``enable_listen`` is set, which this transport does not set, so the agent runs +with NIXL_THREAD_SYNC_NONE and does no internal locking; the config exposes no +way to request STRICT without also starting NIXL's listen thread. If a thread +affinity problem does surface, the contained fix is to funnel every ``_agent`` +call in this class through a single-worker executor, as the NIXL connector's +worker already does for its handshakes, so one OS thread owns the agent +regardless of which caller initiated the work. +""" from __future__ import annotations diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 59b9c320695f..251692f6faef 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -11,7 +11,7 @@ import uuid from collections.abc import Iterable, Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, ClassVar from typing_extensions import override @@ -200,12 +200,18 @@ class P2PSecondaryTierManager(SecondaryTierManager): blocks from the peer) and server-role (serving blocks to the peer) over the same control connection. - Single-threaded: every public method runs on the scheduler thread, and - the engine drives polling via ``get_finished_jobs()`` once per step. + Never entered concurrently: every public method runs under the tiering + manager's ``lock()``, taken either by the scheduler thread for the length of + a step or by the control-plane thread for the length of one round. Peers are + not driven by the engine, so this tier opts into that thread + (``needs_control_plane_thread``); without it a peer's lookup or fetch waits + for a step boundary, which on a saturated rank can be seconds. ``has_pending_work()`` keeps the engine ticking so the control transport and existing sessions are polled even when no requests are scheduled. """ + needs_control_plane_thread: ClassVar[bool] = True + def __init__( self, offloading_spec: OffloadingSpec, @@ -604,9 +610,10 @@ def submit_load(self, job_metadata: TransferJob) -> None: @override def get_finished_jobs(self) -> Iterable[JobResult]: - # Drive one polling sweep on the scheduler thread, then hand off - # whatever has accumulated. The engine calls this once per step - # (and keeps stepping while has_pending_work() is True). + # Drive one polling sweep, then hand off whatever has accumulated. The + # engine calls this once per step (and keeps stepping while + # has_pending_work() is True); the control-plane thread calls it between + # steps, so a peer is not left waiting for a step boundary. self._poll_once() result = self._finished_jobs self._finished_jobs = [] @@ -833,7 +840,8 @@ def _poll_once(self) -> None: Drains the control transport, polls every session, accumulates their results into ``_finished_jobs``, and reaps any dead sessions. - Runs on the scheduler thread. + Runs under the tiering manager's lock, on the scheduler thread or on + the control-plane thread. """ new_connections = self._control.poll() if new_connections: diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index c78918323266..d0a193c56879 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -584,10 +584,13 @@ def _pin_and_register_hits( """Pin primary slots for HIT keys and park them as the lookup's round supply via ``add_stored_blocks``. - Caller has already confirmed every key is HIT (single-threaded - scheduler ⇒ no eviction race), so the JobMetadata returned by - ``parent.create_store_job`` carries parallel ``keys``/``block_ids`` - of length ``len(keys)``. + Caller has already confirmed every key is HIT, so the JobMetadata + returned by ``parent.create_store_job`` carries parallel + ``keys``/``block_ids`` of length ``len(keys)``. Nothing can evict + between that confirmation and the pin because the whole + ``serve_external_requests`` call runs inside one hold of the tiering + manager's lock -- which is why that call must never release it + partway through. """ meta = parent.create_store_job(keys, lookup.ctx) self.add_stored_blocks( diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 93c6cbabcb85..a4b5e37f3191 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -73,6 +73,7 @@ from vllm.v1.kv_offload.tiering.base import TieringOffloadingMetrics from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory from vllm.v1.kv_offload.tiering.manager import ( + _CONTROL_PLANE_INTERVAL_S, CPUPrimaryTierOffloadingManager, TieringOffloadingManager, ) @@ -287,6 +288,13 @@ def __init__(self, config: OffloadingConfig): if not isinstance(self.secondary_tier_configs, list): raise ValueError("secondary_tiers must be a list of tier configurations") + # Pause between control-plane rounds for tiers that need servicing + # between engine steps. Zero or negative disables the thread, which + # restores per-step servicing. + self.control_plane_interval_s = float( + self.extra_config.get("control_plane_interval_s", _CONTROL_PLANE_INTERVAL_S) + ) + # Backpressure config is merged field-by-field in priority order # (highest first): # 1. Per-tier ``backpressure`` dict in the tier config @@ -387,6 +395,7 @@ def get_manager(self) -> OffloadingManager: tiering_manager = TieringOffloadingManager( primary_tier=primary_tier, secondary_tiers=secondary_tiers, + control_plane_interval_s=self.control_plane_interval_s, ) self._manager = tiering_manager except Exception: From 98ac81140e1e134d7d0c4efc73fa2034746e9cc2 Mon Sep 17 00:00:00 2001 From: Liran Schour Date: Wed, 23 Sep 2026 10:29:05 +0300 Subject: [PATCH 3/8] [KV Offload] Correct the NIXL thread-safety note for 1.4.1 The note was written against the locally installed NIXL 0.6.0 and generalized wrongly. On 1.4.1, which is what deployments run, nixl_agent_config does accept sync_mode, and NIXL_THREAD_SYNC_RW does exist. More importantly the open question is now answered rather than flagged. makeMtType never produces mt_mode_t::CONTEXT, the only value mapping to UCS_THREAD_MODE_SINGLE and so the only UCX mode with creator-thread affinity; the weakest mode NIXL requests is UCS_THREAD_MODE_SERIALIZED, documented as "multiple threads can access, but only one at a time", and this transport's default prog thread gets it MULTI with a hard library check. Serialized use from a second thread is supported. What does need stating is that the manager lock is the only protection, since syncMode NONE makes NIXL's own agent lock a no-op, and that NIXL binds a shared UCX worker per calling thread -- harmless only while there is one shared worker, as there is here. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Liran Schour --- vllm/v1/kv_offload/tiering/p2p/data/nixl.py | 36 ++++++++++++++------- 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py index e9fa61e3d780..0c7383bae828 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -5,18 +5,30 @@ Threading: the agent is reached from the scheduler thread and, when the tiering manager runs one, from its control-plane thread. The manager's lock serializes them strictly -- one Python thread inside any agent call at a time, with a full -memory barrier between -- which is sufficient only if libnixl and UCX keep no -state tied to the creating thread. - -That is an assumption, not a verified fact, and worth knowing when debugging -transfer corruption or hangs. nixl_agent picks NIXL_THREAD_SYNC_STRICT only when -``enable_listen`` is set, which this transport does not set, so the agent runs -with NIXL_THREAD_SYNC_NONE and does no internal locking; the config exposes no -way to request STRICT without also starting NIXL's listen thread. If a thread -affinity problem does surface, the contained fix is to funnel every ``_agent`` -call in this class through a single-worker executor, as the NIXL connector's -worker already does for its handshakes, so one OS thread owns the agent -regardless of which caller initiated the work. +memory barrier between. + +That lock is the only protection, not a belt over NIXL's own. nixl_agent picks +NIXL_THREAD_SYNC_STRICT only when ``enable_listen`` is set, which this transport +does not set, so the agent runs with NIXL_THREAD_SYNC_NONE -- under which NIXL's +internal agent lock compiles down to no-ops, leaving every agent entry point +unguarded. So the manager lock must cover all of them, and the lifetime of every +handle passed between calls. + +Serialized use from a second thread is supported (checked against NIXL 1.4.1 +sources): NIXL never creates its UCX worker in UCS_THREAD_MODE_SINGLE, the only +creator-thread-affine mode. The weakest mode it ever requests is +UCS_THREAD_MODE_SERIALIZED, documented as "multiple threads can access, but only +one at a time", and because ``enable_prog_thread`` defaults on here it actually +requests UCS_THREAD_MODE_MULTI and fails construction if the UCX build cannot +provide it. + +Two caveats. NIXL binds a shared UCX worker per calling thread +(``tlsSharedWorkerMap``), which is harmless only while there is exactly one +shared worker -- true here, but it would stop being true if a ``num_workers`` +UCX backend param were ever set. And NIXL 1.4.1 does accept an explicit +``sync_mode`` (including NIXL_THREAD_SYNC_RW), which would make its own locking +real; that parameter does not exist before 1.x, so passing it would need a +version gate, the way ``capture_telemetry`` does. """ from __future__ import annotations From 46d150c0fc0c94621660f3b3b69d091d4af9f382 Mon Sep 17 00:00:00 2001 From: Liran Schour Date: Thu, 1 Oct 2026 10:29:20 +0300 Subject: [PATCH 4/8] [KV Offload] Make OffloadingManager.lock a context-manager property Address review feedback on the lock API: - Replace lock()/unlock()/locked() with a single `lock` property returning a context manager. Callers write `with manager.lock:`. - Drop the lock from the base class. The default is a no-op nullcontext() for managers only ever entered from one thread; CPUOffloadingManager and TieringOffloadingManager each own a threading.Lock and expose it. The tiering manager's control-plane loop and shutdown() still need a bounded wait, so they call acquire(timeout=)/release() on the lock they own directly; the timeout is no longer part of the public API. The scheduler's step-scoped region spans schedule() to the end of build_connector_meta(), which no single with-block can cover, so it is now carried by an ExitStack that holds the entered lock between _acquire_step() and _release_step(). Signed-off-by: Liran Schour --- .../offloading_connector/test_scheduler.py | 24 +++--- .../tiering/test_control_plane_thread.py | 66 ++++++++-------- .../kv_connector/v1/offloading/scheduler.py | 35 +++++---- vllm/v1/kv_offload/base.py | 77 ++++++------------- vllm/v1/kv_offload/cpu/manager.py | 8 +- vllm/v1/kv_offload/tiering/manager.py | 22 ++++-- vllm/v1/kv_offload/tiering/p2p/manager.py | 2 +- 7 files changed, 117 insertions(+), 117 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 0908d50ef75a..36c27c17b5b1 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -758,13 +758,13 @@ def test_on_schedule_end_is_the_last_manager_call_of_a_step(request_runner): def spy(scheduler_output): first = len(runner.manager.mock_calls) meta = build_meta(scheduler_output) - # lock()/unlock() bracket the step; only the state operations they guard - # are ordered relative to on_schedule_end. + # Entering and exiting manager.lock brackets the step; only the state + # operations it guards are ordered relative to on_schedule_end. windows.append( [ name for name, _, _ in runner.manager.mock_calls[first:] - if name not in ("lock", "unlock") + if not name.startswith("lock.") ] ) return meta @@ -804,9 +804,10 @@ def test_step_releases_the_manager_before_the_model_runs(request_runner): runner.run(decoded_tokens=[EOS_TOKEN_ID]) scheduler = runner.connector_scheduler - assert scheduler._manager_held is False - assert runner.manager.lock.call_count > 0 - assert runner.manager.lock.call_count == runner.manager.unlock.call_count + assert scheduler._step_region is None + lock = runner.manager.lock + assert lock.__enter__.call_count > 0 + assert lock.__enter__.call_count == lock.__exit__.call_count def test_has_pending_push_work_does_not_open_the_step_region(request_runner): @@ -825,8 +826,9 @@ def test_has_pending_push_work_does_not_open_the_step_region(request_runner): runner.manager.has_pending_work.return_value = False assert scheduler.has_pending_push_work() is False - assert scheduler._manager_held is False - assert runner.manager.lock.call_count == runner.manager.unlock.call_count + assert scheduler._step_region is None + lock = runner.manager.lock + assert lock.__enter__.call_count == lock.__exit__.call_count def test_take_events_drains_the_manager_inside_the_region(request_runner): @@ -849,8 +851,10 @@ def test_take_events_drains_the_manager_inside_the_region(request_runner): assert events == [] assert runner.manager.take_events.call_count == 1 - assert scheduler._manager_held is False - assert runner.manager.lock.call_count == runner.manager.unlock.call_count + assert scheduler._step_region is None + lock = runner.manager.lock + assert lock.__enter__.call_count == 1 + assert lock.__exit__.call_count == 1 def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner): diff --git a/tests/v1/kv_offload/tiering/test_control_plane_thread.py b/tests/v1/kv_offload/tiering/test_control_plane_thread.py index 308a09fc7644..089b2765eab4 100644 --- a/tests/v1/kv_offload/tiering/test_control_plane_thread.py +++ b/tests/v1/kv_offload/tiering/test_control_plane_thread.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for OffloadingManager.lock() and the tiering control-plane thread. +"""Tests for OffloadingManager.lock and the tiering control-plane thread. The control plane of a tier that answers something the engine does not drive -- a remote peer, say -- used to advance only from on_schedule_end(), so a peer's @@ -11,6 +11,7 @@ import threading import time from collections.abc import Iterable +from contextlib import nullcontext from typing import ClassVar from unittest.mock import MagicMock @@ -19,11 +20,13 @@ from vllm.v1.kv_offload.base import ( LookupResult, Medium, + OffloadingManager, OffloadKey, ReqContext, RequestOffloadingContext, ScheduleEndContext, ) +from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.tiering.base import ( JobResult, ParentManager, @@ -138,7 +141,7 @@ def _live_control_threads() -> list[threading.Thread]: def test_control_thread_and_scheduler_never_overlap(manager_and_tier): - """lock() must keep the two threads out of the manager at the same time. + """The lock must keep the two threads out of the manager at the same time. Both sides must also make progress: a thread that never runs would pass a mutual-exclusion check trivially, and so would one that starves the @@ -159,11 +162,8 @@ def critical_section(): steps = 0 deadline = time.monotonic() + _SETTLE_S while time.monotonic() < deadline: - manager.lock() - try: + with manager.lock: critical_section() - finally: - manager.unlock() steps += 1 assert state["overlaps"] == 0 @@ -180,11 +180,8 @@ def test_serves_while_the_scheduler_is_outside_the_lock(manager_and_tier): """ manager, tier = manager_and_tier - manager.lock() - try: + with manager.lock: manager.on_schedule_end(_EMPTY_SCHEDULE_END) - finally: - manager.unlock() tier.served.clear() serves_before = tier.serves @@ -216,11 +213,8 @@ def test_zero_interval_disables_the_thread(manager_and_tier): assert not tier.served.wait(0.05) assert tier.serves == 0 - manager.lock() - try: + with manager.lock: manager.on_schedule_end(_EMPTY_SCHEDULE_END) - finally: - manager.unlock() assert tier.serves == 1 @@ -260,11 +254,8 @@ def test_tier_that_did_not_opt_in_is_not_serviced_off_thread(): assert quiet.serves == 0 assert quiet.polls == 0 - manager.lock() - try: + with manager.lock: manager.on_schedule_end(_EMPTY_SCHEDULE_END) - finally: - manager.unlock() # on_schedule_end still serves every tier, so the thread dying degrades # to per-step servicing rather than to none. @@ -303,27 +294,42 @@ def test_shutdown_joins_the_control_thread_before_tier_teardown(manager_and_tier assert not _live_control_threads() -def test_locked_helper_releases_on_error(manager_and_tier): - """locked() must not leak the lock when its body raises.""" +def test_lock_releases_on_error(manager_and_tier): + """Leaving a with-block on an exception must not leak the lock.""" manager, _ = manager_and_tier - with pytest.raises(RuntimeError), manager.locked(): + with pytest.raises(RuntimeError), manager.lock: raise RuntimeError("boom") - assert manager.lock(timeout=0.5), "locked() leaked the lock" - manager.unlock() + assert manager.lock.acquire(timeout=0.5), "lock leaked after an error" + manager.lock.release() -def test_lock_timeout_reports_failure(manager_and_tier): - """A timed lock() returns False rather than blocking forever.""" +def test_lock_excludes_other_threads(manager_and_tier): + """While one thread holds the lock, another cannot take it.""" manager, _ = manager_and_tier - manager.lock() - try: + with manager.lock: acquired: list[bool] = [] - t = threading.Thread(target=lambda: acquired.append(manager.lock(timeout=0.01))) + t = threading.Thread( + target=lambda: acquired.append(manager.lock.acquire(timeout=0.01)) + ) t.start() t.join(timeout=5.0) assert acquired == [False] - finally: - manager.unlock() + + +def test_lock_defaults_to_no_op(): + """Managers entered from one thread only pay nothing for the lock. + + The base default must be a no-op context manager, while the CPU manager, + which runs under the plain offloading connector, keeps a real one. + """ + base_lock = OffloadingManager.lock.fget(MagicMock()) + assert isinstance(base_lock, nullcontext) + with base_lock, base_lock: + pass + + cpu = CPUOffloadingManager(num_chunks=4) + with cpu.lock: + assert not cpu.lock.acquire(blocking=False) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 2cd0d1b1552c..da13f98c768d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time from collections.abc import Iterable, Iterator, Sequence -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager from dataclasses import dataclass, field, replace from itertools import chain, islice from typing import Any, NamedTuple @@ -634,18 +634,19 @@ def _sliding_window_sort_key(i: int) -> int: self._events_tracker = OffloadingEventsTracker(spec.kv_events_config) - # True while this thread holds the manager lock for the current schedule(). + # Holds the manager lock for the current schedule(), or None between steps. # Read and written only by the scheduler thread, so it needs no # synchronization of its own. Declared on the class so an instance built - # without __init__ still reads False rather than raising. - _manager_held: bool = False + # without __init__ still reads None rather than raising. + _step_region: ExitStack | None = None # ------------------------------------------------------------------ # Manager exclusion # # The manager is entered from the scheduler thread and, for managers that # run one, from a control-plane thread. Manager methods never take their - # own lock, so the exclusion boundary lives here. + # own lock, so the exclusion boundary lives here. manager.lock may be a + # no-op for a manager that needs no exclusion; it is entered regardless. # # schedule() holds one region for the whole step: opened by the first # manager-touching hook and closed at the end of build_connector_meta. @@ -657,21 +658,26 @@ def _sliding_window_sort_key(i: int) -> int: # # Hooks outside schedule() use a short region instead, so the manager is # free while the engine runs the model. + # + # The step region outlives any single call, so it cannot be a with-block; + # an ExitStack carries the entered lock from open to close instead. # ------------------------------------------------------------------ def _acquire_step(self) -> None: """Open the step-scoped region. Idempotent within a step.""" - if self._manager_held: + if self._step_region is not None: return - self.manager.lock() - self._manager_held = True + region = ExitStack() + region.enter_context(self.manager.lock) + self._step_region = region def _release_step(self) -> None: """Close the step-scoped region, letting other threads in.""" - if not self._manager_held: + region = self._step_region + if region is None: return - self._manager_held = False - self.manager.unlock() + self._step_region = None + region.close() @contextmanager def _manager_locked(self) -> Iterator[None]: @@ -684,17 +690,14 @@ def _manager_locked(self) -> Iterator[None]: previous schedule() raised between _acquire_step() and the release in build_connector_meta. Recover instead of deadlocking. """ - if self._manager_held: + if self._step_region is not None: logger.error( "Offloading manager lock was left held by a previous " "schedule(); releasing it before continuing." ) self._release_step() - self.manager.lock() - try: + with self.manager.lock: yield - finally: - self.manager.unlock() def _maybe_observe_lookup_async_delay( self, req_status: RequestOffloadState diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 95c15e60a354..9f77db4fda31 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -2,10 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Core abstractions for KV cache offloading in vLLM v1.""" -import threading from abc import ABC, abstractmethod -from collections.abc import Collection, Iterable, Iterator, Sequence -from contextlib import contextmanager +from collections.abc import Collection, Iterable, Sequence +from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, field from enum import Enum, auto from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, NewType, TypeVar @@ -201,12 +200,14 @@ class OffloadingEvent: Following this call, the given blocks will become loadable. Exclusion: - lock() / unlock() guard a manager's whole state subtree: the manager - itself, every tier it composes, and any front-end state those tiers own. - Manager methods never take the lock themselves, so a caller holds it - across as many operations as its invariants need -- notably a lookup() - HIT and the prepare_load() that pins the result, which must not be - separated by another thread's eviction. + lock is a context manager guarding a manager's whole state subtree: the + manager itself, every tier it composes, and any front-end state those + tiers own. Manager methods never take the lock themselves, so a caller + holds it across as many operations as its invariants need -- notably a + lookup() HIT and the prepare_load() that pins the result, which must not + be separated by another thread's eviction. The default is a no-op; + managers that can be entered from more than one thread supply a real + lock. """ @@ -242,53 +243,25 @@ class OffloadingKVEventsConfig: class OffloadingManager(ABC): - def __init__(self) -> None: - self._manager_lock = threading.Lock() - - def lock(self, timeout: float | None = None) -> bool: - """Acquire exclusive access to this manager's state subtree. + @property + def lock(self) -> AbstractContextManager[Any]: + """Context manager granting exclusive access to this manager's state. Scope is the manager, every tier it composes, and any front-end state those tiers own (e.g. a tier's async lookup front end). Every caller of - an OffloadingManager method must hold this lock. Manager methods never - acquire it themselves, so a caller may hold it across an arbitrary - sequence of operations -- which is required wherever one call - establishes a fact a later call relies on, such as a lookup() HIT - followed by the prepare_load() that pins it. - - Not reentrant. A thread must not call lock() twice without an - intervening unlock(); use the held state the caller already tracks, or - a *_unlocked-style helper, rather than re-acquiring. - - Args: - timeout: Seconds to wait. None blocks until acquired. - - Returns: - True if the lock was acquired, False only on timeout. - - """ - if timeout is None: - return self._manager_lock.acquire() - return self._manager_lock.acquire(timeout=timeout) - - def unlock(self) -> None: - """Release a lock() acquired by the calling thread.""" - self._manager_lock.release() - - @contextmanager - def locked(self, timeout: float | None = None) -> Iterator[bool]: - """Scoped lock()/unlock() for callers that need no step granularity. - - Yields whether the lock was acquired; the body still runs on timeout, - so a caller passing a timeout must check the value before touching the - manager. + an OffloadingManager method must hold it. Manager methods never enter + it themselves, so a caller may hold it across an arbitrary sequence of + operations -- which is required wherever one call establishes a fact a + later call relies on, such as a lookup() HIT followed by the + prepare_load() that pins it. + + The default is a no-op, for managers that are only ever entered from + one thread. A manager that may be entered from another thread -- its + own background thread, say -- returns a real lock instead. Callers must + not assume either: always enter it, and never re-enter it while held, + since a real lock need not be reentrant. """ - acquired = self.lock(timeout) - try: - yield acquired - finally: - if acquired: - self.unlock() + return nullcontext() @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index e8b3011f5f02..dc19bbb97667 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading from collections import OrderedDict from collections.abc import Collection, Iterable from dataclasses import dataclass, field @@ -63,7 +64,7 @@ def __init__( store_threshold: int = 1, max_tracker_size: int = 64_000, ): - super().__init__() + self._lock = threading.Lock() self.medium: Medium = Medium.CPU self._num_chunks: int = num_chunks self._num_allocated_chunks: int = 0 @@ -89,6 +90,11 @@ def __init__( OrderedDict() if store_threshold >= 2 else None ) + @property + @override + def lock(self) -> threading.Lock: + return self._lock + # --- chunk pool --- def _get_num_free_chunks(self) -> int: diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 795b4d8d6026..5767fec7355b 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -225,7 +225,10 @@ def __init__( engine step as before. """ - super().__init__() + # Guards this manager and every tier it composes, the primary tier + # included; the primary's own lock is not used underneath it. A real + # lock because the control-plane thread enters the manager too. + self._lock = threading.Lock() self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier self.secondary_tiers = secondary_tiers or [] @@ -279,6 +282,11 @@ def __init__( self._control_plane_thread: threading.Thread | None = None self._start_control_plane() + @property + @override + def lock(self) -> threading.Lock: + return self._lock + # ------------------------------------------------------------------ # Control plane # ------------------------------------------------------------------ @@ -322,7 +330,7 @@ def _control_plane_loop(self) -> None: stalled_since: float | None = None errors = 0 while not self._control_plane_stop.is_set(): - if not self.lock(timeout=_CONTROL_PLANE_LOCK_TIMEOUT_S): + if not self._lock.acquire(timeout=_CONTROL_PLANE_LOCK_TIMEOUT_S): now = time.monotonic() if stalled_since is None: stalled_since = now @@ -346,7 +354,7 @@ def _control_plane_loop(self) -> None: errors += 1 logger.exception("KV offload control-plane round failed") finally: - self.unlock() + self._lock.release() if errors >= _CONTROL_PLANE_MAX_CONSECUTIVE_ERRORS: logger.error( "KV offload control-plane thread stopping after %d " @@ -354,7 +362,7 @@ def _control_plane_loop(self) -> None: errors, ) return - # Yield unconditionally, and only after unlock(). Python locks are + # Yield unconditionally, and only after release(). Python locks are # not fair, so a release-then-reacquire loop could starve the # scheduler thread; this pause is what bounds its wait to one round. self._control_plane_stop.wait(self._control_plane_interval_s) @@ -372,7 +380,7 @@ def serve_control_plane(self) -> None: promotion that a lookup() from this path initiates therefore waits for the next on_schedule_end() to be submitted. - The caller must hold lock() for the whole call. serve_external_requests() + The caller must hold self.lock for the whole call. serve_external_requests() relies on that: it establishes lookup() HITs and then pins them, and the two must not be separated by an eviction. """ @@ -1152,7 +1160,7 @@ def shutdown(self) -> None: """ self._stop_control_plane() if self._control_plane_thread is not None: - if not self.lock(timeout=_CONTROL_PLANE_JOIN_TIMEOUT_S): + if not self._lock.acquire(timeout=_CONTROL_PLANE_JOIN_TIMEOUT_S): logger.error( "KV offload control-plane thread is still running and the " "manager lock is unavailable; skipping tier shutdown." @@ -1161,7 +1169,7 @@ def shutdown(self) -> None: try: self._shutdown_tiers() finally: - self.unlock() + self._lock.release() return self._shutdown_tiers() diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 251692f6faef..19aa16a95fc9 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -201,7 +201,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): over the same control connection. Never entered concurrently: every public method runs under the tiering - manager's ``lock()``, taken either by the scheduler thread for the length of + manager's ``lock``, held either by the scheduler thread for the length of a step or by the control-plane thread for the length of one round. Peers are not driven by the engine, so this tier opts into that thread (``needs_control_plane_thread``); without it a peer's lookup or fetch waits From 633efb74c1bd41334627b046adef316addfc1e8f Mon Sep 17 00:00:00 2001 From: Liran Schour Date: Thu, 1 Oct 2026 10:38:29 +0300 Subject: [PATCH 5/8] [KV Offload] Rename control-plane polling knobs; make default public Address review feedback: - serve_control_plane() -> poll_tiers() - control_plane_interval_s -> tier_poll_interval_s, for the constructor argument, the attribute, and the user-facing kv_connector_extra_config key (docs updated). - _CONTROL_PLANE_INTERVAL_S -> DEFAULT_TIER_POLL_INTERVAL_S. spec.py imports it and it is the documented default of a user-facing key, so it is public. It stays in tiering/manager.py next to the thread it paces: spec.py already imports manager.py, so moving it to spec.py would leave the manager's constructor without a default to reference. The remaining control-plane constants stay private; nothing outside the module reads them. Signed-off-by: Liran Schour --- docs/features/kv_offloading_usage.md | 6 ++--- .../tiering/test_control_plane_thread.py | 2 +- vllm/v1/kv_offload/tiering/base.py | 2 +- vllm/v1/kv_offload/tiering/manager.py | 23 ++++++++++--------- vllm/v1/kv_offload/tiering/spec.py | 8 +++---- 5 files changed, 21 insertions(+), 20 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index e2fabf4c5f7c..364395f38ccb 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -108,7 +108,7 @@ vllm serve \ | `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. | | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. | | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). | -| `control_plane_interval_s` | no | `0.001` | multi-tier | Pause between control-plane rounds for tiers that must be serviced between engine steps (currently `p2p`). See [Threading model](#threading-model). Set to `0` to disable the thread and service those tiers once per engine step instead. | +| `tier_poll_interval_s` | no | `0.001` | multi-tier | Pause between control-plane rounds for tiers that must be serviced between engine steps (currently `p2p`). See [Threading model](#threading-model). Set to `0` to disable the thread and service those tiers once per engine step instead. | | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. | | `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Partial recurrent tails emit the hash-aligned portion from the physical block start through the tail boundary. Other sliding-window/SSM chunks keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). | @@ -155,14 +155,14 @@ Set the optional `locality` tier field to `LOCAL` or `REMOTE` to describe the ti Tier methods run in the scheduler process, under the tiering manager's lock, so they are never entered concurrently. They are not always called from the same thread, though. -Most tiers only ever see the scheduler thread, which holds the lock for the length of one engine step. A tier whose counterpart is not driven by the engine — the `p2p` tier, answering a remote peer — needs servicing while the engine is busy running the model, because otherwise a peer's lookup or fetch waits for the next step boundary. On a saturated rank that wait can be seconds, dwarfing the transfer it gates. Such tiers set `needs_control_plane_thread`, and the tiering manager runs one thread that polls them and lets them serve, once per `control_plane_interval_s`. +Most tiers only ever see the scheduler thread, which holds the lock for the length of one engine step. A tier whose counterpart is not driven by the engine — the `p2p` tier, answering a remote peer — needs servicing while the engine is busy running the model, because otherwise a peer's lookup or fetch waits for the next step boundary. On a saturated rank that wait can be seconds, dwarfing the transfer it gates. Such tiers set `needs_control_plane_thread`, and the tiering manager runs one thread that polls them and lets them serve, once per `tier_poll_interval_s`. Two consequences worth knowing: - Any tier can be reached from that thread indirectly, because serving a peer looks keys up through the tiering manager, which fans out to the other tiers. Tier state must therefore not assume a particular thread, only that the manager lock is held. - Work the thread deliberately leaves alone is batched per step: a promotion it initiates is submitted at the next `on_schedule_end()`, not immediately. -Setting `control_plane_interval_s` to `0` disables the thread; opted-in tiers then get serviced once per engine step, as they were before it existed. +Setting `tier_poll_interval_s` to `0` disables the thread; opted-in tiers then get serviced once per engine step, as they were before it existed. ### Filesystem (FS) diff --git a/tests/v1/kv_offload/tiering/test_control_plane_thread.py b/tests/v1/kv_offload/tiering/test_control_plane_thread.py index 089b2765eab4..434451ea45f4 100644 --- a/tests/v1/kv_offload/tiering/test_control_plane_thread.py +++ b/tests/v1/kv_offload/tiering/test_control_plane_thread.py @@ -199,7 +199,7 @@ def test_serves_while_the_scheduler_is_outside_the_lock(manager_and_tier): @pytest.mark.parametrize( - "manager_and_tier", [{"control_plane_interval_s": 0.0}], indirect=True + "manager_and_tier", [{"tier_poll_interval_s": 0.0}], indirect=True ) def test_zero_interval_disables_the_thread(manager_and_tier): """The escape hatch back to per-step servicing. diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 1e0bee499abc..673209824e1f 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -152,7 +152,7 @@ class SecondaryTierManager(ABC): # and not only from on_schedule_end(). Set it on tiers that answer a # counterpart the engine does not drive -- a remote peer, say -- whose # requests would otherwise wait for a step boundary. The tiering manager - # runs a thread on behalf of such tiers; see serve_control_plane(). + # runs a thread on behalf of such tiers; see poll_tiers(). needs_control_plane_thread: ClassVar[bool] = False def __init__( diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 5767fec7355b..007873d9c32a 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -59,11 +59,12 @@ logger = init_logger(__name__) -# Pause between control-plane rounds. A round is a non-blocking sweep, so this -# sets how long a peer's request can sit unserved; 1 ms is three orders of +# Default pause between control-plane rounds; the user-facing default of the +# ``tier_poll_interval_s`` extra-config key. A round is a non-blocking sweep, so +# this sets how long a peer's request can sit unserved; 1 ms is three orders of # magnitude below the model step it hides, and matches the pause the p2p tier's # own drain loops and EngineCoreProc's GIL-yield sleep already use. -_CONTROL_PLANE_INTERVAL_S = 0.001 +DEFAULT_TIER_POLL_INTERVAL_S = 0.001 # How long a round waits for the manager lock before giving up and retrying. # Bounded so the thread stays responsive to shutdown even if a caller holds the # lock for a long time, or leaked it. @@ -211,7 +212,7 @@ def __init__( self, primary_tier: CPUPrimaryTierOffloadingManager, secondary_tiers: list[SecondaryTierManager] | None = None, - control_plane_interval_s: float = _CONTROL_PLANE_INTERVAL_S, + tier_poll_interval_s: float = DEFAULT_TIER_POLL_INTERVAL_S, ): """Initialize the TieringOffloadingManager. @@ -219,7 +220,7 @@ def __init__( primary_tier: The primary tier manager (CPU-based). secondary_tiers: List of secondary tier managers (e.g., Storage, Network). Can be None or empty list. - control_plane_interval_s: Pause between control-plane rounds for + tier_poll_interval_s: Pause between control-plane rounds for tiers that set needs_control_plane_thread. Zero or negative disables the thread, leaving those tiers serviced once per engine step as before. @@ -277,7 +278,7 @@ def __init__( for i, tier in enumerate(self.secondary_tiers) if tier.needs_control_plane_thread ] - self._control_plane_interval_s = control_plane_interval_s + self._tier_poll_interval_s = tier_poll_interval_s self._control_plane_stop = threading.Event() self._control_plane_thread: threading.Thread | None = None self._start_control_plane() @@ -293,7 +294,7 @@ def lock(self) -> threading.Lock: def _start_control_plane(self) -> None: """Start the control-plane thread, if any tier asked for one.""" - if not self._control_plane_tiers or self._control_plane_interval_s <= 0: + if not self._control_plane_tiers or self._tier_poll_interval_s <= 0: return self._control_plane_thread = threading.Thread( target=self._control_plane_loop, @@ -304,7 +305,7 @@ def _start_control_plane(self) -> None: logger.info( "KV offload control-plane thread started for tier(s) %s, interval %.3fs", [tier.tier_type for _, tier in self._control_plane_tiers], - self._control_plane_interval_s, + self._tier_poll_interval_s, ) def _stop_control_plane(self) -> None: @@ -345,7 +346,7 @@ def _control_plane_loop(self) -> None: continue stalled_since = None try: - self.serve_control_plane() + self.poll_tiers() errors = 0 except Exception: # Never let a round kill the thread silently: on_schedule_end() @@ -365,9 +366,9 @@ def _control_plane_loop(self) -> None: # Yield unconditionally, and only after release(). Python locks are # not fair, so a release-then-reacquire loop could starve the # scheduler thread; this pause is what bounds its wait to one round. - self._control_plane_stop.wait(self._control_plane_interval_s) + self._control_plane_stop.wait(self._tier_poll_interval_s) - def serve_control_plane(self) -> None: + def poll_tiers(self) -> None: """Advance the control plane of tiers that cannot wait for a step. Polls those tiers for finished jobs and lets them serve whatever their diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index a4b5e37f3191..73921f447571 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -73,7 +73,7 @@ from vllm.v1.kv_offload.tiering.base import TieringOffloadingMetrics from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory from vllm.v1.kv_offload.tiering.manager import ( - _CONTROL_PLANE_INTERVAL_S, + DEFAULT_TIER_POLL_INTERVAL_S, CPUPrimaryTierOffloadingManager, TieringOffloadingManager, ) @@ -291,8 +291,8 @@ def __init__(self, config: OffloadingConfig): # Pause between control-plane rounds for tiers that need servicing # between engine steps. Zero or negative disables the thread, which # restores per-step servicing. - self.control_plane_interval_s = float( - self.extra_config.get("control_plane_interval_s", _CONTROL_PLANE_INTERVAL_S) + self.tier_poll_interval_s = float( + self.extra_config.get("tier_poll_interval_s", DEFAULT_TIER_POLL_INTERVAL_S) ) # Backpressure config is merged field-by-field in priority order @@ -395,7 +395,7 @@ def get_manager(self) -> OffloadingManager: tiering_manager = TieringOffloadingManager( primary_tier=primary_tier, secondary_tiers=secondary_tiers, - control_plane_interval_s=self.control_plane_interval_s, + tier_poll_interval_s=self.tier_poll_interval_s, ) self._manager = tiering_manager except Exception: From 4955042eee1dab16ae1480af7708f1bf3e7be999 Mon Sep 17 00:00:00 2001 From: Liran Schour Date: Thu, 1 Oct 2026 10:55:51 +0300 Subject: [PATCH 6/8] [KV Offload] Rename needs_control_plane_thread to serves_external_requests Address review feedback: the opt-in flag now states what the tier does -- answers a counterpart the engine does not drive -- rather than what the manager should do for it, and matches the serve_external_requests() method it gates. The flag's docstring now says that opting in moves both get_finished_jobs() and serve_external_requests() onto the polling thread, and that overriding serve_external_requests() alone does not opt in. _control_plane_tiers is renamed _external_serving_tiers to match. KVCR implements serve_external_requests() but stays opted out, now explicitly: opting in would also poll the KVCR library from the polling thread, which it is not yet verified to tolerate. Signed-off-by: Liran Schour --- docs/features/kv_offloading_usage.md | 2 +- .../tiering/test_control_plane_thread.py | 4 ++-- vllm/v1/kv_offload/tiering/base.py | 23 +++++++++++-------- vllm/v1/kv_offload/tiering/kvcr/manager.py | 8 ++++++- vllm/v1/kv_offload/tiering/manager.py | 17 +++++++------- vllm/v1/kv_offload/tiering/p2p/manager.py | 6 ++--- 6 files changed, 36 insertions(+), 24 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 364395f38ccb..5b72abf762c1 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -155,7 +155,7 @@ Set the optional `locality` tier field to `LOCAL` or `REMOTE` to describe the ti Tier methods run in the scheduler process, under the tiering manager's lock, so they are never entered concurrently. They are not always called from the same thread, though. -Most tiers only ever see the scheduler thread, which holds the lock for the length of one engine step. A tier whose counterpart is not driven by the engine — the `p2p` tier, answering a remote peer — needs servicing while the engine is busy running the model, because otherwise a peer's lookup or fetch waits for the next step boundary. On a saturated rank that wait can be seconds, dwarfing the transfer it gates. Such tiers set `needs_control_plane_thread`, and the tiering manager runs one thread that polls them and lets them serve, once per `tier_poll_interval_s`. +Most tiers only ever see the scheduler thread, which holds the lock for the length of one engine step. A tier whose counterpart is not driven by the engine — the `p2p` tier, answering a remote peer — needs servicing while the engine is busy running the model, because otherwise a peer's lookup or fetch waits for the next step boundary. On a saturated rank that wait can be seconds, dwarfing the transfer it gates. Such tiers set `serves_external_requests`, and the tiering manager runs one thread that, once per `tier_poll_interval_s`, polls them for finished jobs and lets them serve external requests. A tier that implements `serve_external_requests()` without setting the flag is served once per engine step only. Two consequences worth knowing: diff --git a/tests/v1/kv_offload/tiering/test_control_plane_thread.py b/tests/v1/kv_offload/tiering/test_control_plane_thread.py index 434451ea45f4..68ac31edfd9d 100644 --- a/tests/v1/kv_offload/tiering/test_control_plane_thread.py +++ b/tests/v1/kv_offload/tiering/test_control_plane_thread.py @@ -52,7 +52,7 @@ class _FakeTier(SecondaryTierManager): """Tier that records which thread serviced it, and when.""" medium: ClassVar[Medium] = Medium.CPU - needs_control_plane_thread: ClassVar[bool] = True + serves_external_requests: ClassVar[bool] = True def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -101,7 +101,7 @@ def shutdown(self) -> None: class _QuietTier(_FakeTier): """Tier that does not want servicing between steps.""" - needs_control_plane_thread: ClassVar[bool] = False + serves_external_requests: ClassVar[bool] = False def _make_manager(tier_cls=_FakeTier, num_chunks: int = 8, **kwargs): diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 673209824e1f..8b32e2ad879b 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -141,19 +141,22 @@ class SecondaryTierManager(ABC): Methods are called under the tiering manager's lock, so they are never entered concurrently -- but not always from the same thread. A tier that - sets needs_control_plane_thread is also serviced from the manager's + sets serves_external_requests is also serviced from the manager's control-plane thread, and any tier can be reached from there via ParentManager fan-out. Keep per-thread assumptions out of tier state. """ medium: ClassVar[Medium | None] = None - # Whether this tier's control plane must be serviced between engine steps - # and not only from on_schedule_end(). Set it on tiers that answer a - # counterpart the engine does not drive -- a remote peer, say -- whose - # requests would otherwise wait for a step boundary. The tiering manager - # runs a thread on behalf of such tiers; see poll_tiers(). - needs_control_plane_thread: ClassVar[bool] = False + # Whether this tier answers a counterpart the engine does not drive -- a + # remote peer, say -- whose requests would otherwise wait for a step + # boundary. The tiering manager then services the tier between engine + # steps too, not only from on_schedule_end(): on its polling thread, it + # calls both get_finished_jobs() and serve_external_requests(); see + # poll_tiers(). Overriding serve_external_requests() alone does not opt + # in: setting this asserts the tier tolerates those calls off the + # scheduler thread. + serves_external_requests: ClassVar[bool] = False def __init__( self, @@ -327,7 +330,7 @@ def serve_external_requests(self, parent: ParentManager) -> None: Called once per scheduler step, BEFORE _flush_pending_promotions(), and additionally once per control-plane round for tiers that set - needs_control_plane_thread. The _flush_pending_promotions() ordering + serves_external_requests. The _flush_pending_promotions() ordering holds only for the per-step call: a promotion this method initiates from a control-plane round is submitted at the next on_schedule_end(). @@ -336,7 +339,9 @@ def serve_external_requests(self, parent: ParentManager) -> None: here stays valid until this call pins it. An implementation must not release that lock partway through. - Tiers that don't serve external requests leave this as a no-op. + Tiers that don't serve external requests leave this as a no-op. A tier + that overrides it without setting serves_external_requests is served + once per step only. """ return diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py index 18da2a6c2764..07f429d9ed70 100644 --- a/vllm/v1/kv_offload/tiering/kvcr/manager.py +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -9,7 +9,7 @@ import uuid from collections.abc import Collection, Iterable from pathlib import Path -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, ClassVar, cast import msgspec from kvcr import ( @@ -264,6 +264,12 @@ def decode(self, key: BlockKey) -> int | bytes: class KVCRSecondaryTierManager(SecondaryTierManager): """Secondary tier wrapper around the KVCR KV P2P API.""" + # Serves framework pin requests from serve_external_requests(), but only + # once per step for now: opting in would also poll the KVCR library from + # the tiering manager's polling thread, which it is not yet verified to + # tolerate. + serves_external_requests: ClassVar[bool] = False + @classmethod @override def build_metric_definitions( diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 007873d9c32a..e2b2d41c607a 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -221,7 +221,7 @@ def __init__( secondary_tiers: List of secondary tier managers (e.g., Storage, Network). Can be None or empty list. tier_poll_interval_s: Pause between control-plane rounds for - tiers that set needs_control_plane_thread. Zero or negative + tiers that set serves_external_requests. Zero or negative disables the thread, leaving those tiers serviced once per engine step as before. @@ -272,11 +272,12 @@ def __init__( tier: i for i, tier in enumerate(self.secondary_tiers) } - # Tiers whose control plane must be serviced between engine steps. - self._control_plane_tiers: list[tuple[int, SecondaryTierManager]] = [ + # Tiers that serve external requests, and so must also be serviced + # between engine steps. + self._external_serving_tiers: list[tuple[int, SecondaryTierManager]] = [ (i, tier) for i, tier in enumerate(self.secondary_tiers) - if tier.needs_control_plane_thread + if tier.serves_external_requests ] self._tier_poll_interval_s = tier_poll_interval_s self._control_plane_stop = threading.Event() @@ -294,7 +295,7 @@ def lock(self) -> threading.Lock: def _start_control_plane(self) -> None: """Start the control-plane thread, if any tier asked for one.""" - if not self._control_plane_tiers or self._tier_poll_interval_s <= 0: + if not self._external_serving_tiers or self._tier_poll_interval_s <= 0: return self._control_plane_thread = threading.Thread( target=self._control_plane_loop, @@ -304,7 +305,7 @@ def _start_control_plane(self) -> None: self._control_plane_thread.start() logger.info( "KV offload control-plane thread started for tier(s) %s, interval %.3fs", - [tier.tier_type for _, tier in self._control_plane_tiers], + [tier.tier_type for _, tier in self._external_serving_tiers], self._tier_poll_interval_s, ) @@ -385,8 +386,8 @@ def poll_tiers(self) -> None: relies on that: it establishes lookup() HITs and then pins them, and the two must not be separated by an eviction. """ - self._process_finished_jobs(self._control_plane_tiers) - for _, tier in self._control_plane_tiers: + self._process_finished_jobs(self._external_serving_tiers) + for _, tier in self._external_serving_tiers: tier.serve_external_requests(self._tier_parents[tier]) @property diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 19aa16a95fc9..3bf8d66cdc90 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -203,14 +203,14 @@ class P2PSecondaryTierManager(SecondaryTierManager): Never entered concurrently: every public method runs under the tiering manager's ``lock``, held either by the scheduler thread for the length of a step or by the control-plane thread for the length of one round. Peers are - not driven by the engine, so this tier opts into that thread - (``needs_control_plane_thread``); without it a peer's lookup or fetch waits + not driven by the engine, so this tier sets ``serves_external_requests`` + to opt into that thread; without it a peer's lookup or fetch waits for a step boundary, which on a saturated rank can be seconds. ``has_pending_work()`` keeps the engine ticking so the control transport and existing sessions are polled even when no requests are scheduled. """ - needs_control_plane_thread: ClassVar[bool] = True + serves_external_requests: ClassVar[bool] = True def __init__( self, From fc7b2ef246434a4e5a0432926c325dfbba70739f Mon Sep 17 00:00:00 2001 From: Liran Schour Date: Thu, 1 Oct 2026 11:08:17 +0300 Subject: [PATCH 7/8] [KV Offload] Keep a polling round's parent lookups out of the step gate A tier served by the polling thread can look keys up through its parent -- the p2p server role does so for every peer lookup -- and TieringOffloadingManager.lookup() runs the once-per-step _maybe_process_finished_jobs(). Between steps the gate is clear, so the first such lookup in a round: - polled every secondary tier from the polling thread, including tiers that did not set serves_external_requests; - re-entered the serving tier's get_finished_jobs() from inside its own serve_external_requests(), which for p2p means _poll_once() adding or reaping sessions while serve iterates over them; - left the gate set, so the next step's lookups and on_schedule_end()'s catch-all poll skipped their poll. poll_tiers() now holds the gate shut for the duration of the round and restores it afterwards. The round has just polled its own tiers, so the skipped poll loses nothing, and the scheduler's per-step gate is left as it found it. Add a regression test with a tier that does a parent lookup while serving, and narrow the opt-in test's docstring: ParentManager fan-out can still reach tiers that did not opt in, as tiering/base.py documents. Signed-off-by: Liran Schour --- .../tiering/test_control_plane_thread.py | 75 ++++++++++++++++++- vllm/v1/kv_offload/tiering/manager.py | 18 ++++- 2 files changed, 88 insertions(+), 5 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_control_plane_thread.py b/tests/v1/kv_offload/tiering/test_control_plane_thread.py index 68ac31edfd9d..7358565f53e9 100644 --- a/tests/v1/kv_offload/tiering/test_control_plane_thread.py +++ b/tests/v1/kv_offload/tiering/test_control_plane_thread.py @@ -229,10 +229,12 @@ def test_no_thread_when_no_tier_asks_for_one(): def test_tier_that_did_not_opt_in_is_not_serviced_off_thread(): - """Only opted-in tiers are touched by the control thread. + """Only opted-in tiers are polled and served by the control thread. A tier that expects the scheduler thread must not be dragged onto another - one just because it shares a manager with a tier that opted in. + one just because it shares a manager with a tier that opted in. Fan-out + through ParentManager (lookup, on_new_request, on_request_finished) can + still reach it from that thread; that is the documented exception. """ num_chunks = 8 primary = CPUPrimaryTierOffloadingManager( @@ -278,6 +280,75 @@ def test_per_step_gate_is_not_touched_by_the_control_thread(manager_and_tier): assert manager._processed_jobs_this_step is False +class _LookingUpTier(_FakeTier): + """Opted-in tier that looks keys up through its parent while serving. + + That is what the p2p server role does with a peer's lookup, and the parent + lookup lands in TieringOffloadingManager.lookup(), which polls tiers for + finished jobs. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.in_serve = False + self.nested_polls = 0 + self.on_serve = self._look_up + + def _look_up(self, parent: ParentManager) -> None: + self.in_serve = True + try: + parent.lookup(OffloadKey(b"\x01" * 8), ReqContext(req_id="peer")) + finally: + self.in_serve = False + + def get_finished_jobs(self) -> Iterable[JobResult]: + if self.in_serve: + self.nested_polls += 1 + return super().get_finished_jobs() + + +def test_parent_lookup_while_serving_keeps_the_round_contained(): + """A lookup issued while serving must not escape the control round. + + Before the fix it ran the step's once-per-step poll from the control + thread: every tier got polled off-thread, the opted-in tier re-entrantly + from inside its own serve, and the gate was left set, so the next step + skipped its poll. + """ + num_chunks = 8 + primary = CPUPrimaryTierOffloadingManager( + num_chunks=num_chunks, + mmap_region=_mock_mmap_region(num_chunks), + ) + common = { + "offloading_spec": MagicMock(), + "primary_kv_view": primary.get_kv_memoryview(), + } + opted_in = _LookingUpTier(tier_type="fake", **common) + quiet = _QuietTier(tier_type="quiet", **common) + manager = TieringOffloadingManager( + primary_tier=primary, + secondary_tiers=[opted_in, quiet], + ) + try: + # Several rounds, each issuing a parent lookup between steps. + deadline = time.monotonic() + _SETTLE_S + while opted_in.serves < 3 and time.monotonic() < deadline: + time.sleep(0.01) + assert opted_in.serves >= 3, "control thread never served" + + with manager.lock: + assert manager._processed_jobs_this_step is False + assert opted_in.nested_polls == 0 + assert quiet.polls == 0, "quiet tier was polled off-thread" + + # The next step's first lookup still does its poll. + manager.lookup(OffloadKey(b"\x02" * 8), _CTX) + assert quiet.polls == 1 + finally: + manager.shutdown() + + def test_shutdown_joins_the_control_thread_before_tier_teardown(manager_and_tier): """The thread drives tier transports, so it must be stopped first. diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index e2b2d41c607a..5d3c320dae5b 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -385,10 +385,22 @@ def poll_tiers(self) -> None: The caller must hold self.lock for the whole call. serve_external_requests() relies on that: it establishes lookup() HITs and then pins them, and the two must not be separated by an eviction. + + Serving can look keys up through the parent, and lookup() runs the + step's once-per-step poll. From here that poll would reach every tier + off-thread, re-enter the serving tier's get_finished_jobs() from inside + its own serve, and leave the gate set so the next step skips its poll. + The round has just polled its tiers itself, so it holds the gate shut + for its duration and then puts it back untouched. """ - self._process_finished_jobs(self._external_serving_tiers) - for _, tier in self._external_serving_tiers: - tier.serve_external_requests(self._tier_parents[tier]) + gate = self._processed_jobs_this_step + self._processed_jobs_this_step = True + try: + self._process_finished_jobs(self._external_serving_tiers) + for _, tier in self._external_serving_tiers: + tier.serve_external_requests(self._tier_parents[tier]) + finally: + self._processed_jobs_this_step = gate @property def _transfer_jobs(self) -> dict[JobId, JobMetadata]: From 2fdfd08acd28cd23ca038cfdeef249b0ebfdcc9a Mon Sep 17 00:00:00 2001 From: Liran Schour Date: Thu, 1 Oct 2026 11:36:25 +0300 Subject: [PATCH 8/8] [KV Offload] Note the registration-worker exception to the NIXL lock rule The NixlTransport threading note says the manager lock must cover every agent entry point. Peer registration moved onto a worker thread, as NixlConnector's handshake executor already does and as #57139 proposes for this transport, deliberately runs outside that lock. Say so, so the note stays accurate when that lands. Signed-off-by: Liran Schour --- vllm/v1/kv_offload/tiering/p2p/data/nixl.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py index 0c7383bae828..2c79acd52929 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -12,7 +12,9 @@ does not set, so the agent runs with NIXL_THREAD_SYNC_NONE -- under which NIXL's internal agent lock compiles down to no-ops, leaving every agent entry point unguarded. So the manager lock must cover all of them, and the lifetime of every -handle passed between calls. +handle passed between calls. The exception is peer registration moved onto a +worker thread, as NixlConnector's handshake executor does: it runs outside the +lock and relies, like that connector, on NIXL tolerating it beside transfers. Serialized use from a second thread is supported (checked against NIXL 1.4.1 sources): NIXL never creates its UCX worker in UCS_THREAD_MODE_SINGLE, the only