From 4ccfc2b85c3d25c406c0c452086a08a3c3b4866c Mon Sep 17 00:00:00 2001 From: Liran Schour Date: Thu, 17 Sep 2026 17:31:32 +0300 Subject: [PATCH] [KVConnector][Tiering] Drive the tiering control plane from a dedicated thread MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The secondary tiers' control plane only advanced from the per-step scheduler hooks: on_schedule_end polled every tier for finished jobs and then let each one serve inbound peer requests. The p2p tier's transports are non-blocking by design, so nothing moved between steps, and a rank busy with long prefills made every peer control-plane round trip wait for a step boundary — several round trips per transfer. TieringOffloadingManager now runs both from its own daemon thread, and serializes it against the scheduler with one lock held for a whole scheduler step rather than per call. Step granularity is required, not just tidier: the connector learns a chunk is a HIT in lookup() and only reads it in a later prepare_load() hook, and the chunk is evictable in between, so a promotion started while serving a peer lookup in that window could evict it out from under the pending load. The lock is released at the end of on_schedule_end, which puts the control plane's window over model execution — the time a rank mid-prefill has to spare. The read-mostly hooks (has_pending_work, get_stats, take_events) instead hold the lock only for their own duration. The engine ticks them every iteration, including iterations that schedule nothing and run no model, so folding them into the step-wide hold would leave an idle engine with no gap between steps at all — the state a producer waiting for a consumer to connect is in, and exactly when the control plane matters. Tiers re-enter the manager through _SecondaryTierFacingParent, whose caller already holds the lock, so the four ParentManager methods gained *_unlocked variants; exclude_tier_idx moved onto those, restoring the public signatures to match OffloadingManager. A sweep that raises is logged and re-raised on the next scheduler-side call, matching the engine crash these failures produced when they ran on the scheduler thread, instead of leaving the control plane silently dead. No interface changes, and no changes under tiering/p2p. The thread is on whenever a secondary tier is configured; control_plane_thread=False keeps everything on the calling thread, which is what the step-driven tests need. Signed-off-by: Liran Schour --- tests/v1/kv_offload/tiering/test_fs_tier.py | 4 +- tests/v1/kv_offload/tiering/test_obj_tier.py | 4 +- .../tiering/test_tiering_offloading.py | 226 +++++++++ vllm/v1/kv_offload/tiering/manager.py | 454 +++++++++++++++--- 4 files changed, 625 insertions(+), 63 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 9c3a64d22a2a..ef1825b20962 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -902,7 +902,9 @@ def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path): root_dir=str(tmp_path), enable_kv_events=True, ) - manager = TieringOffloadingManager(primary_tier=primary, secondary_tiers=[tier]) + manager = TieringOffloadingManager( + primary_tier=primary, secondary_tiers=[tier], control_plane_thread=False + ) try: keys = [key(1), key(2)] manager.on_new_request(_CTX) diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index bf8afca42428..4bbec6ba7fdc 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -546,7 +546,9 @@ def test_xfer_cleanup_retry_finalizes_parent_job_and_primary_pin(self, monkeypat num_blocks=num_blocks, primary_kv_view=primary_kv_view ) manager = TieringOffloadingManager( - primary_tier=primary_tier, secondary_tiers=[obj_tier] + primary_tier=primary_tier, + secondary_tiers=[obj_tier], + control_plane_thread=False, ) keys = [key(1)] diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 63b4c2c4d0b5..954ff798e129 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -10,6 +10,8 @@ 5. Eviction coordination between tiers """ +import threading +import time from collections.abc import Iterable from unittest.mock import MagicMock @@ -199,6 +201,7 @@ def test_tiering_manager_aggregates_secondary_stats(): manager = TieringOffloadingManager( primary_tier=primary_tier, secondary_tiers=[secondary_tier], + control_plane_thread=False, ) stats = manager.get_stats() @@ -306,6 +309,7 @@ def manager_setup(self): self.manager = TieringOffloadingManager( primary_tier=self.primary_tier, secondary_tiers=[self.secondary_tier1, self.secondary_tier2], + control_plane_thread=False, ) def _simulate_on_schedule_end(self, new_req_ids: list[str] | None = None): @@ -1506,5 +1510,227 @@ def test_parse_tier_filter_skips_bad_entries(): ) +class _ControlPlaneTier(SecondaryTierManager): + """Test-only tier that records how the control plane drives it.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Bumped once per get_finished_jobs() / serve_external_requests() call. + # Read from the test thread; int writes are atomic under the GIL. + self.polls = 0 + self.serves = 0 + self.swept = threading.Event() + self.serve_idents: set[int] = set() + self.poll_error: Exception | None = None + self.finished_jobs: list[JobResult] = [] + + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + return LookupResult.MISS + + def submit_store(self, job_metadata: TransferJob) -> None: + return + + def submit_load(self, job_metadata: TransferJob) -> None: + return + + def get_finished_jobs(self) -> Iterable[JobResult]: + if self.poll_error is not None: + raise self.poll_error + self.polls += 1 + results = self.finished_jobs + self.finished_jobs = [] + return results + + def serve_external_requests(self, parent) -> None: + self.serves += 1 + self.serve_idents.add(threading.get_ident()) + self.swept.set() + + def drain_jobs(self) -> None: + return + + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + return RequestOffloadingContext() + + def has_pending_work(self) -> bool: + return True + + +def _make_threaded_manager(**kwargs): + """Build a manager with one _ControlPlaneTier and the control thread on.""" + primary_tier = CPUPrimaryTierOffloadingManager( + num_chunks=4, mmap_region=_mock_mmap_region(4) + ) + tier = _ControlPlaneTier( + _MOCK_OFFLOADING_SPEC, primary_tier.get_kv_memoryview(), "control_plane_test" + ) + manager = TieringOffloadingManager( + primary_tier=primary_tier, + secondary_tiers=[tier], + control_poll_interval_s=0.001, + **kwargs, + ) + return manager, tier + + +def _step(manager) -> None: + manager.on_schedule_end(ScheduleEndContext(new_req_ids=[], preempted_req_ids=())) + + +class TestControlPlaneThread: + """The control plane must advance between steps, and only between steps.""" + + def test_sweeps_without_further_scheduler_steps(self): + """One step starts the thread; it then sweeps on its own.""" + manager, tier = _make_threaded_manager() + try: + _step(manager) + assert tier.swept.wait(timeout=5.0), "control thread never swept" + assert tier.polls > 0 + assert tier.serves > 0 + finally: + manager.shutdown() + + def test_sweeps_run_off_the_scheduler_thread(self): + """Serving happens on the control thread, not the caller's.""" + manager, tier = _make_threaded_manager() + try: + _step(manager) + assert tier.swept.wait(timeout=5.0) + assert threading.get_ident() not in tier.serve_idents + assert len(tier.serve_idents) == 1 + finally: + manager.shutdown() + + def test_step_locks_out_the_control_thread(self): + """A sweep cannot interleave with a step in progress. + + The window between lookup() and prepare_load() is exactly where a + sweep-initiated promotion could evict a chunk the scheduler was just + told is a HIT, so the lock has to span the whole step. + """ + manager, tier = _make_threaded_manager() + try: + _step(manager) + assert tier.swept.wait(timeout=5.0) + + # Any scheduler-side call takes the step lock and holds it until + # on_schedule_end. + manager.lookup(to_keys([1])[0], _CTX) + before = tier.serves + time.sleep(0.05) + assert tier.serves == before, "control thread ran mid-step" + + tier.swept.clear() + _step(manager) + assert tier.swept.wait(timeout=5.0), "control thread stayed blocked" + finally: + manager.shutdown() + + def test_idle_engine_hooks_do_not_block_the_control_thread(self): + """The read-mostly hooks must not hold the lock until the next step. + + An engine with nothing scheduled still ticks has_pending_work(), + get_stats() and take_events() every iteration and runs no model, so if + those joined the step-wide hold the control plane would never get in — + which is the state a producer waiting for a consumer sits in. + """ + manager, tier = _make_threaded_manager() + try: + _step(manager) + assert tier.swept.wait(timeout=5.0) + + for _ in range(20): + manager.has_pending_work() + manager.get_stats() + list(manager.take_events()) + before = tier.serves + tier.swept.clear() + assert tier.swept.wait(timeout=5.0), "control thread starved when idle" + assert tier.serves > before + finally: + manager.shutdown() + + def test_on_schedule_end_does_not_serve_when_threaded(self): + """The per-step hook stops serving once the thread owns it.""" + manager, tier = _make_threaded_manager() + try: + _step(manager) # starts the thread; must not serve inline + assert threading.get_ident() not in tier.serve_idents + finally: + manager.shutdown() + + def test_sweep_failure_reraises_on_scheduler_thread(self): + """A failed sweep surfaces on the next scheduler call, not silently.""" + manager, tier = _make_threaded_manager() + try: + tier.poll_error = RuntimeError("tier exploded") + _step(manager) + thread = manager._control_thread + assert thread is not None + thread.join(timeout=5.0) + assert not thread.is_alive(), "thread should stop after a failure" + + with pytest.raises(RuntimeError, match="tier exploded"): + manager.lookup(to_keys([1])[0], _CTX) + finally: + tier.poll_error = None + manager.shutdown() + + def test_shutdown_joins_the_control_thread(self): + manager, tier = _make_threaded_manager() + _step(manager) + assert tier.swept.wait(timeout=5.0) + thread = manager._control_thread + assert thread is not None + + manager.shutdown() + + assert not thread.is_alive() + assert manager._control_thread is None + # Nothing may touch a torn-down tier afterwards. + quiesced = tier.polls + time.sleep(0.02) + assert tier.polls == quiesced + + def test_reset_cache_runs_with_the_control_thread_live(self): + manager, tier = _make_threaded_manager() + try: + _step(manager) + assert tier.swept.wait(timeout=5.0) + manager.reset_cache() + assert not manager._jobs + # reset_cache releases the lock, so sweeps resume. + tier.swept.clear() + assert tier.swept.wait(timeout=5.0) + finally: + manager.shutdown() + + def test_disabled_by_request(self): + """control_plane_thread=False keeps everything on the caller.""" + manager, tier = _make_threaded_manager(control_plane_thread=False) + try: + _step(manager) + assert manager._control_thread is None + # The per-step hook is still the driver. + assert tier.serve_idents == {threading.get_ident()} + assert tier.polls > 0 + finally: + manager.shutdown() + + def test_no_thread_without_secondary_tiers(self): + primary_tier = CPUPrimaryTierOffloadingManager( + num_chunks=4, mmap_region=_mock_mmap_region(4) + ) + manager = TieringOffloadingManager( + primary_tier=primary_tier, secondary_tiers=[] + ) + try: + _step(manager) + assert manager._control_thread is None + finally: + manager.shutdown() + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 6e0bf90a332a..0f1ac4cc8e75 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -19,8 +19,10 @@ protecting chunks from eviction until complete_read() is called """ +import contextlib +import threading import time -from collections.abc import Collection, Iterable, Sequence +from collections.abc import Collection, Iterable, Iterator, Sequence from dataclasses import dataclass, field from typing import NamedTuple @@ -57,6 +59,23 @@ logger = init_logger(__name__) +# Sleep between control-plane sweeps. A tier's poll is non-blocking (the p2p +# tier's ZMQ transport returns immediately when there is no traffic), so the +# thread has to pace itself. Small enough to be irrelevant next to a step, +# large enough not to fight the scheduler thread for the GIL. +_CONTROL_POLL_INTERVAL_S = 0.001 + +# How long a sweep waits for the step lock before looping. Bounded so the stop +# event is always observed promptly, and so a step that never reached +# on_schedule_end is reported instead of silently wedging the control plane. +_CONTROL_LOCK_TIMEOUT_S = 1.0 + +# Warn once the control plane has been unable to run for this long. +_CONTROL_STARVATION_WARN_S = 10.0 + +# How long shutdown waits for the control thread to finish its current sweep. +_CONTROL_THREAD_JOIN_TIMEOUT_S = 2.0 + @dataclass class PendingPromotion: @@ -143,7 +162,14 @@ def shutdown(self) -> None: class _SecondaryTierFacingParent(ParentManager): """Wrapper that implements ParentManager by delegating to the - TieringOffloadingManager with exclude_tier_idx set to the origin tier.""" + TieringOffloadingManager with exclude_tier_idx set to the origin tier. + + Handed to a tier only for the duration of serve_external_requests(), whose + caller — the control-plane thread, or on_schedule_end when that thread is + disabled — already holds the manager lock. So every method here delegates + to the *_unlocked variant: taking the lock again would deadlock, since it + is not reentrant. + """ __slots__ = ("_m", "_origin_idx") @@ -156,20 +182,18 @@ def __init__( self._origin_idx = tier_idx def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: - return self._m.on_new_request(req_context, exclude_tier_idx=self._origin_idx) + return self._m._on_new_request_unlocked(req_context, self._origin_idx) def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: - return self._m.lookup(key, req_context, exclude_tier_idx=self._origin_idx) + return self._m._lookup_unlocked(key, req_context, self._origin_idx) def create_store_job( self, keys: Collection[OffloadKey], req_context: ReqContext ) -> TransferJob: - return self._m.create_store_job(keys, req_context, self._origin_idx) + return self._m._create_store_job_unlocked(keys, req_context, self._origin_idx) def on_request_finished(self, req_context: ReqContext) -> None: - return self._m.on_request_finished( - req_context, exclude_tier_idx=self._origin_idx - ) + return self._m._on_request_finished_unlocked(req_context, self._origin_idx) class TieringOffloadingManager(OffloadingManager): @@ -185,12 +209,43 @@ class TieringOffloadingManager(OffloadingManager): - Secondary tiers return JobResult objects containing all necessary information - job_id_counter: monotonically increasing counter for job IDs + + Threading + --------- + Two threads touch this manager, serialized by ``_lock``: + + * The scheduler thread runs every public method. It takes ``_lock`` on its + first call of a scheduler step and drops it at the end of + ``on_schedule_end`` — step granularity, not per call. That matters: the + connector learns a chunk is a HIT from ``lookup()`` and only reads it in + a later ``prepare_load()`` hook, and in between the chunk is still + evictable. A promotion started while serving a peer lookup in that window + could evict it out from under the pending load. + * The control-plane thread (``_control_plane_loop``) polls tiers for + finished jobs and lets them serve inbound peer requests. It holds + ``_lock`` for a whole sweep, so it runs only between steps — in practice + during model execution, which is exactly the time a rank busy with a long + prefill has to spare. Without it, a peer's control-plane round trip costs + a full step boundary. + + The read-mostly hooks (``has_pending_work``, ``get_stats``, + ``take_events``) are the exception: they hold the lock only for their own + duration, via ``_observation_lock``. The engine calls them on every tick, + so folding them into the step-wide hold would leave an idle engine with no + gap between steps at all. + + Tiers re-enter the manager from that thread through + ``_SecondaryTierFacingParent``, which calls the ``*_unlocked`` variants + because its caller already holds the lock. """ def __init__( self, primary_tier: CPUPrimaryTierOffloadingManager, secondary_tiers: list[SecondaryTierManager] | None = None, + *, + control_plane_thread: bool = True, + control_poll_interval_s: float = _CONTROL_POLL_INTERVAL_S, ): """Initialize the TieringOffloadingManager. @@ -198,6 +253,12 @@ 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_thread: Whether to poll tiers and serve their inbound + peer requests from a dedicated thread rather than + from the per-step hooks. Pass False to keep all of + it on the calling thread, which is what + step-driven tests need. + control_poll_interval_s: Sleep between control-plane sweeps. """ self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier @@ -238,6 +299,84 @@ def __init__( for tier_idx, tier in enumerate(self.secondary_tiers) } + # Serializes the scheduler thread against the control-plane thread. See + # the class docstring for why it is held for a whole step. + self._lock = threading.Lock() + # True while the scheduler side holds _lock for the current step. Only + # scheduler-side callers read or write it, and that side is + # single-threaded, so it needs no protection of its own; the + # control-plane thread never touches it. + self._step_locked: bool = False + + self._control_thread_enabled: bool = control_plane_thread and bool( + self.secondary_tiers + ) + self._control_poll_interval_s: float = control_poll_interval_s + self._control_thread: threading.Thread | None = None + self._control_thread_stop = threading.Event() + # Exception raised inside the control-plane thread, re-raised on the + # next scheduler-side call. These failures used to happen on the + # scheduler thread and take the engine down with a real traceback; + # swallowing them here would instead leave the control plane silently + # dead, which looks like a hang. + self._control_thread_exc: BaseException | None = None + + # ------------------------------------------------------------------ + # Step lock + # ------------------------------------------------------------------ + + def _enter_step(self, *, raise_control_error: bool = True) -> None: + """Take the manager lock for the current scheduler step. + + Idempotent within a step: the first scheduler-side call acquires, later + ones are no-ops, and on_schedule_end releases. Never called from the + control-plane thread, which holds _lock around its whole sweep and + re-enters the manager through the *_unlocked helpers. + + Args: + raise_control_error: Whether to re-raise a control-plane thread + failure. Only shutdown passes False, so a failed sweep cannot + mask teardown. + + Raises: + BaseException: Whatever the control-plane thread failed with. + + """ + if raise_control_error and self._control_thread_exc is not None: + raise self._control_thread_exc + if self._step_locked: + return + self._lock.acquire() + self._step_locked = True + + def _exit_step(self) -> None: + """Release the manager lock, letting the control plane run.""" + if not self._step_locked: + return + self._step_locked = False + self._lock.release() + + @contextlib.contextmanager + def _observation_lock(self) -> Iterator[None]: + """Hold the manager lock only for the duration of one observation. + + For the read-mostly hooks (has_pending_work, get_stats, take_events). + They take no part in the lookup()->prepare_load() invariant that makes + the step-wide hold necessary, and the engine calls them on every tick, + including ticks that schedule nothing and run no model. Letting them + extend a step's hold would leave almost no gap between one step and the + next on an idle engine — exactly the state a producer waiting for a + consumer to connect is in, and exactly when the control plane matters. + """ + if self._control_thread_exc is not None: + raise self._control_thread_exc + if self._step_locked: + # Already held for this step; the step's release point owns it. + yield + return + with self._lock: + yield + @property def _transfer_jobs(self) -> dict[JobId, JobMetadata]: return self._jobs @@ -262,7 +401,13 @@ def _maybe_process_finished_jobs(self): Guarded by _processed_jobs_this_step: the first call in an engine step does the actual polling; subsequent calls are no-ops. The flag is reset in on_schedule_end() at the end of each step. + + A no-op once the control-plane thread is running: that thread owns tier + polling, and repeating it here would put the very transport polls we + moved off the scheduler thread straight back onto it. """ + if self._control_thread is not None: + return if self._processed_jobs_this_step: return self._processed_jobs_this_step = True @@ -339,15 +484,24 @@ def _process_finished_jobs(self): ) @override - def lookup( + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + """Check whether a single chunk is offloaded and ready. + + See _lookup_unlocked for the algorithm and the return values. + """ + self._enter_step() + return self._lookup_unlocked(key, req_context) + + def _lookup_unlocked( self, key: OffloadKey, req_context: ReqContext, - *, exclude_tier_idx: int | None = None, ) -> LookupResult: """Check whether a single chunk is offloaded and ready. + The caller must hold the manager lock. + Algorithm: 1. Process any completed async jobs first. 2. Query primary tier — short-circuit on hit or in-flight. @@ -518,6 +672,7 @@ def prepare_load( LoadStoreSpec for reading from primary tier. """ + self._enter_step() return self.primary_tier.prepare_load(keys, req_context) @override @@ -529,6 +684,7 @@ def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): req_context: Per-request context. """ + self._enter_step() self.primary_tier.touch(keys, req_context) for tier in self.secondary_tiers: tier.touch(keys, req_context) @@ -545,6 +701,7 @@ def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext): req_context: Per-request context. """ + self._enter_step() self.primary_tier.complete_load(keys, req_context) @override @@ -579,6 +736,7 @@ 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._enter_step() self._maybe_process_finished_jobs() # Step 2: Store to primary tier (new chunks only). @@ -643,7 +801,9 @@ def _cascade_existing_chunks_to_request_level_tiers( return for tier_idx in request_level_tiers: - job_metadata = self.create_store_job(ready_keys, req_context, tier_idx) + job_metadata = self._create_store_job_unlocked( + ready_keys, req_context, tier_idx + ) tier = self.secondary_tiers[tier_idx] tier.submit_store(job_metadata) @@ -688,6 +848,7 @@ def complete_store( req_context: Per-request context forwarded to primary.prepare_read(). """ + self._enter_step() # Step 1: Complete store in primary tier (makes chunks loadable) self.primary_tier.complete_store(keys, req_context, success) @@ -698,7 +859,9 @@ def complete_store( # eviction during the async transfer). One prepare_read() call per # secondary tier. for tier_idx, tier in enumerate(self.secondary_tiers): - job_metadata = self.create_store_job(keys, req_context, tier_idx) + job_metadata = self._create_store_job_unlocked( + keys, req_context, tier_idx + ) tier.submit_store(job_metadata) # Note: The async transfers are now in flight. Their completion is @@ -717,9 +880,22 @@ def create_store_job( ) -> TransferJob: """Pin chunks in the primary tier and create a tracked store job. + See _create_store_job_unlocked. + """ + self._enter_step() + return self._create_store_job_unlocked(keys, req_context, tier_idx) + + def _create_store_job_unlocked( + self, + keys: Collection[OffloadKey], + req_context: ReqContext, + tier_idx: int = 0, + ) -> TransferJob: + """Pin chunks in the primary tier and create a tracked store job. + Calls prepare_read() to increment ref_cnt (protecting chunks from eviction during the async transfer), allocates a job ID, - and registers the job in _jobs. + and registers the job in _jobs. The caller must hold the manager lock. The caller is responsible for the actual data transfer and reporting completion via get_finished_jobs(). @@ -738,16 +914,24 @@ def create_store_job( return job_metadata @override - def on_new_request( + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + """Query each secondary tier for its offload policy preference. + + See _on_new_request_unlocked. + """ + self._enter_step() + return self._on_new_request_unlocked(req_context) + + def _on_new_request_unlocked( self, req_context: ReqContext, - *, exclude_tier_idx: int | None = None, ) -> RequestOffloadingContext: """Query each secondary tier for its offload policy preference. Returns REQUEST_LEVEL if ANY secondary tier wants request-level. Only stores REQUEST_LEVEL tier decisions for use in prepare_store. + The caller must hold the manager lock. """ state = RequestState(req_context=req_context) self._metrics.on_new_request(req_context) @@ -769,12 +953,16 @@ def on_new_request( return RequestOffloadingContext(policy=policy) @override - def on_request_finished( + def on_request_finished(self, req_context: ReqContext) -> None: + self._enter_step() + self._on_request_finished_unlocked(req_context) + + def _on_request_finished_unlocked( self, req_context: ReqContext, - *, exclude_tier_idx: int | None = None, ) -> None: + """Finalize a finished request. The caller must hold the manager lock.""" self.primary_tier.on_request_finished(req_context) state = self._req_state[req_context.req_id] state.is_finished = True @@ -812,52 +1000,170 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None: Called once per scheduler step from OffloadingConnectorScheduler.build_connector_meta(). + + Also where the step's hold on the manager lock is released, handing the + control-plane thread the model-execution window. """ - # Catch-all poll: guarantees jobs are processed even on steps where - # lookup()/prepare_store() were never called (e.g. no requests - # scheduled but a tier still has_pending_work()). - self._maybe_process_finished_jobs() + self._enter_step() + try: + self._start_control_thread() + + if self._control_thread is None: + # Catch-all poll: guarantees jobs are processed even on steps + # where lookup()/prepare_store() were never called (e.g. no + # requests scheduled but a tier still has_pending_work()). + # Once the control thread runs, both of these are its job. + self._maybe_process_finished_jobs() + + for tier in self.secondary_tiers: + tier.serve_external_requests(self._tier_parents[tier]) + + # Reset the per-step gate AFTER serve_external_requests so that + # lookup() calls within it skip redundant _process_finished_jobs(). + self._processed_jobs_this_step = False + + self._flush_pending_promotions() + self._flush_pending_cascades() + for tier in self.secondary_tiers: + tier.on_schedule_end(context) + + for req_id in context.new_req_ids: + state = self._req_state.get(req_id) + if state is None: + continue + self._metrics.on_request_allocated(state.req_context) + finally: + self._exit_step() + + # ------------------------------------------------------------------ + # Control plane + # ------------------------------------------------------------------ + + def _start_control_thread(self) -> None: + """Start the control-plane thread, on the first scheduler step. + + Deferred out of __init__ so no sweep runs before the engine is + stepping: accepting a peer builds a session, and a tier may need state + that is only set up after this manager is constructed (the p2p tier + resolves the block-hash seed, which init_none_hash() sets). It also + keeps the thread out of the spec's partial-construction cleanup path. + """ + if self._control_thread is not None or not self._control_thread_enabled: + return + self._control_thread = threading.Thread( + target=self._control_plane_loop, + name="vllm_tiering_control_plane", + daemon=True, + ) + self._control_thread.start() + logger.info( + "Tiering control-plane thread started for %d secondary tier(s), " + "polling every %.3fs", + len(self.secondary_tiers), + self._control_poll_interval_s, + ) - for tier in self.secondary_tiers: - tier.serve_external_requests(self._tier_parents[tier]) + def _stop_control_thread(self) -> None: + """Signal the control-plane thread and wait out its current sweep. - # Reset the per-step gate AFTER serve_external_requests so that - # lookup() calls within it skip redundant _process_finished_jobs(). - self._processed_jobs_this_step = False + Releases the step lock first: the thread may be blocked on it, and it + only observes the stop event once it is no longer waiting. Because the + loop re-checks that event after acquiring, a straggler outliving the + join can no longer start a sweep against a torn-down tier. + """ + thread = self._control_thread + self._control_thread_stop.set() + self._exit_step() + if thread is None: + return + thread.join(timeout=_CONTROL_THREAD_JOIN_TIMEOUT_S) + if thread.is_alive(): + logger.error( + "Tiering control-plane thread did not stop within %.1fs; " + "continuing with tier shutdown under the manager lock.", + _CONTROL_THREAD_JOIN_TIMEOUT_S, + ) + self._control_thread = None - self._flush_pending_promotions() - self._flush_pending_cascades() - for tier in self.secondary_tiers: - tier.on_schedule_end(context) + def _control_plane_loop(self) -> None: + """Sweep the secondary tiers' control plane until shutdown. - for req_id in context.new_req_ids: - state = self._req_state.get(req_id) - if state is None: + Runs between scheduler steps, holding the manager lock for each sweep. + """ + stop = self._control_thread_stop + starved_since: float | None = None + while not stop.is_set(): + if not self._lock.acquire(timeout=_CONTROL_LOCK_TIMEOUT_S): + now = time.monotonic() + if starved_since is None: + starved_since = now + elif now - starved_since >= _CONTROL_STARVATION_WARN_S: + logger.warning( + "Tiering control plane has not run for %.0fs: the " + "scheduler still holds the manager lock, so a step " + "never reached on_schedule_end.", + now - starved_since, + ) + starved_since = now continue - self._metrics.on_request_allocated(state.req_context) + starved_since = None + try: + if stop.is_set(): + # Shutdown began while we waited for the lock. + return + self._sweep_control_plane() + except BaseException as exc: + # This work used to run on the scheduler thread, where a + # failure took the engine down with a real traceback. Hand it + # back there rather than leaving the control plane dead, which + # from the outside is indistinguishable from a hang. + logger.exception("Tiering control-plane sweep failed") + self._control_thread_exc = exc + return + finally: + self._lock.release() + stop.wait(self._control_poll_interval_s) + + def _sweep_control_plane(self) -> None: + """Poll tiers for finished jobs, then let them serve inbound peers. + + Same order as the on_schedule_end path this replaces: polling binds a + fetch that just arrived and enqueues inbound lookups, so serving + resolves them within the same sweep. + """ + self._process_finished_jobs() + for tier in self.secondary_tiers: + tier.serve_external_requests(self._tier_parents[tier]) @override def has_pending_work(self) -> bool: # In-flight primary<->secondary transfers (pending promotions are # translated to transfer jobs in on_schedule_end), plus any work the # secondary tiers themselves still have outstanding. - return ( - bool(self._jobs) - or any(state.pending_cascade_keys for state in self._req_state.values()) - or any(tier.has_pending_work() for tier in self.secondary_tiers) - ) + with self._observation_lock(): + return ( + bool(self._jobs) + or any(state.pending_cascade_keys for state in self._req_state.values()) + or any(tier.has_pending_work() for tier in self.secondary_tiers) + ) @override def take_events(self) -> Iterable[OffloadingEvent]: - """Yield events owned by the primary and secondary tiers. + """Collect events owned by the primary and secondary tiers. + + Materialized rather than generated: a generator body would not run + until the caller iterates, so the step lock would be taken (and the + tiers drained) at some later, unrelated point. - Yields: + Returns: New OffloadingEvents collected by each tier since the last call. """ - yield from self.primary_tier.take_events() - for tier in self.secondary_tiers: - yield from tier.take_events() + with self._observation_lock(): + events = list(self.primary_tier.take_events()) + for tier in self.secondary_tiers: + events.extend(tier.take_events()) + return events @override def reset_cache(self) -> None: @@ -873,7 +1179,19 @@ def reset_cache(self) -> None: (FS, network) keep their data across resets. Active request state is retained so those requests can continue after the reset; finished requests are finalized and removed. + + Runs outside a scheduler step, so it takes and drops the manager lock + itself. The control-plane thread is locked out for the whole reset; + drain_jobs() polls the tiers itself and does not need it. """ + self._enter_step() + try: + self._reset_cache_unlocked() + finally: + self._exit_step() + + def _reset_cache_unlocked(self) -> None: + """Body of reset_cache(). The caller must hold the manager lock.""" for tier in self.secondary_tiers: tier.drain_jobs() # All tier I/O has stopped; consume their completion notifications @@ -906,26 +1224,27 @@ def reset_cache(self) -> None: @override def get_stats(self) -> OffloadingConnectorStats | None: - stats = self.primary_tier.get_stats() + with self._observation_lock(): + stats = self.primary_tier.get_stats() - if stats is not None and stats.is_empty(): - stats = None + if stats is not None and stats.is_empty(): + stats = None - metrics_stats = self._metrics.take_stats() - if metrics_stats is not None: - if stats is None: - stats = metrics_stats - else: - stats.aggregate(metrics_stats) + metrics_stats = self._metrics.take_stats() + if metrics_stats is not None: + if stats is None: + stats = metrics_stats + else: + stats.aggregate(metrics_stats) - for tier in self.secondary_tiers: - tier_stats = tier.get_stats() - if tier_stats is None or tier_stats.is_empty(): - continue - if stats is None: - stats = tier_stats - else: - stats.aggregate(tier_stats) + for tier in self.secondary_tiers: + tier_stats = tier.get_stats() + if tier_stats is None or tier_stats.is_empty(): + continue + if stats is None: + stats = tier_stats + else: + stats.aggregate(tier_stats) return stats @@ -935,7 +1254,20 @@ def shutdown(self) -> None: Every secondary tier is given a shutdown attempt. If any shutdown fails, preserve the primary mmap because a failed tier may still use it. + + The control-plane thread is stopped first, so nothing sweeps a tier + while it is being torn down. A sweep that failed earlier is not + re-raised here: it must not mask the teardown. """ + self._stop_control_thread() + self._enter_step(raise_control_error=False) + try: + self._shutdown_unlocked() + finally: + self._exit_step() + + def _shutdown_unlocked(self) -> None: + """Body of shutdown(). The caller must hold the manager lock.""" shutdown_error: Exception | None = None for tier_idx, tier in enumerate(self.secondary_tiers): try: