[Bugfix][KV Offload] Let idle P2P-only engines park

docs/features/kv_offloading_usage.md

@@ -224,9 +224,9 @@ Most tiers only ever see the scheduler thread, which holds the lock for the leng
 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.
+- Work the thread deliberately leaves alone is batched per step: a promotion it initiates is submitted at the next `on_schedule_end()`, not immediately. Because only an engine step submits that promotion and collects its completion, a multi-tier setup that includes an opted-in tier (for example `p2p` alongside `fs`) keeps the engine stepping even when idle, so pause and graceful shutdown wait there too. A setup whose only secondary tier is the opted-in one still goes idle, since a peer's lookup cannot promote out of it. Work the thread can finish by itself, such as serving a peer's lookup and fetch, does not count as pending, so a producer that peers keep fetching from still goes idle. Pause does not stop the P2P server either: it stops scheduling and quiesces the device, and a paused engine keeps serving peers from its CPU tier.
 
-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.
+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. With the thread disabled the engine keeps stepping for such a tier's lifetime, so an idle engine does not report itself idle and pause and graceful shutdown do not complete on their own. A round that keeps failing does not stop the thread: after 30 s of consecutive failures it backs off to a slower interval — the larger of 1 s and `tier_poll_interval_s` — and keeps retrying, so peers are still served while the engine is idle. Failing rounds log at most one message every 30 s: a full traceback for the first failure of a run, a one-line summary for the rest.
 
 ### Filesystem (FS)
 

tests/v1/kv_offload/tiering/p2p/test_manager.py

@@ -146,6 +146,38 @@ def _make_manager() -> P2PSecondaryTierManager:
     return mgr
 
 
+def _attach_fake_transports(mgr: P2PSecondaryTierManager) -> None:
+    """Give a stubbed manager inert control/data transports.
+
+    The shipped configuration runs the tier-poll thread, which drives
+    _poll_once() -> _control.poll()/session.poll() and shutdown()'s
+    _control.close()/_data.close(). A stub built with __new__ has none.
+    """
+
+    class _Control:
+        def poll(self):
+            return []
+
+        def close(self) -> None:
+            pass
+
+    class _Data:
+        def remove_remote_peer(self, pid):
+            pass
+
+        def cancel(self, ids, mode="immediate"):
+            return []
+
+        def poll(self, peer_id=None):
+            return SimpleNamespace(done=[], failed=[])
+
+        def close(self) -> None:
+            pass
+
+    mgr._control = _Control()
+    mgr._data = _Data()
+
+
 def _init_offloading_spec() -> SimpleNamespace:
     """Minimal offloading_spec for driving the real __init__."""
     return SimpleNamespace(
@@ -1003,25 +1035,175 @@ def test_submit_store_parks_unbound_batch(self):
 
 
 class TestHasPendingWork:
-    """has_pending_work() must always return True so the engine keeps
-    ticking the offload pipeline — that's the only thread driving
-    _control.poll() (incoming peer connects) and session.poll()
-    (incoming fetch messages on existing sessions)."""
-
-    def test_returns_true_unconditionally_to_keep_engine_ticking(self):
-        """Even with no sessions and no jobs, has_pending_work() returns
-        True so the engine keeps calling get_finished_jobs(), which is
-        what drives _control.poll() for inbound peer connects."""
+    """has_pending_work() reports peer-originated work that still needs a step.
+
+    Idle peer polling is left to the tier-poll thread, so a quiescent tier must
+    report False and let pause and graceful shutdown observe an idle engine.
+    The tiering manager adds the keep-alive itself when no thread services the
+    tier.
+    """
+
+    def test_idle_manager_reports_no_pending_work(self):
+        """With no sessions and no jobs, the tier is idle."""
+        mgr = _make_manager()
+        assert mgr.has_pending_work() is False
+
+    def test_connected_idle_session_is_not_pending_work(self):
+        """An idle connected session is not work: polling it is the
+        tier-poll thread's job, not a reason to keep the engine stepping."""
+        mgr = _make_manager()
+        mgr._sessions["peer:1"] = _FakeSession(peer_id="peer:1", connected=True)
+        assert mgr.has_pending_work() is False
+
+    def test_session_with_outstanding_transfer_is_pending_work(self):
+        """A session with an in-flight load keeps the engine stepping until
+        the transfer is drained."""
         mgr = _make_manager()
+        session = _FakeSession(peer_id="peer:1", connected=True)
+        session._client._inbound[1] = object()
+        mgr._sessions["peer:1"] = session
         assert mgr.has_pending_work() is True
 
-    def test_returns_true_even_when_sessions_present(self):
-        """The result is the same regardless of session state — there is
-        no 'idle' branch."""
+    def test_finished_jobs_awaiting_handoff_is_pending_work(self):
+        """Results not yet returned by get_finished_jobs() need a step."""
         mgr = _make_manager()
-        mgr._sessions["peer:1"] = _FakeSession(peer_id="peer:1")
+        mgr._finished_jobs.append(JobResult(job_id=1, success=True))
         assert mgr.has_pending_work() is True
 
+    def test_parked_unbound_store_is_pending_work(self):
+        """A parked store still pins primary slots until it is fetched or
+        reaped, so it keeps the engine stepping."""
+        from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch
+
+        mgr = _make_manager()
+        mgr._unbound_stores["req-1"] = [
+            _UnboundStoreBatch(job_id=1, keys=[OffloadKey(b"k")], block_ids=[0])
+        ]
+        assert mgr.has_pending_work() is True
+
+    def test_failed_serve_ctx_awaiting_release_is_pending_work(self):
+        """A synthetic lookup ctx still owing parent.on_request_finished must
+        be flushed on the next serve_external_requests()."""
+        mgr = _make_manager()
+        mgr._failed_serve_ctxs.append(_req_context(kv_params=None))
+        assert mgr.has_pending_work() is True
+
+    @pytest.mark.cpu_test
+    def test_B2_idle_manager_does_not_keep_engine_busy(self):
+        """A quiescent P2P tier must let pause and shutdown observe idle.
+
+        The shipped configuration runs the tier-poll thread, which services
+        idle peers between steps; the tier must not keep EngineCore.has_work()
+        True just for that.
+        """
+        from unittest.mock import MagicMock
+
+        from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
+            OffloadingConnectorScheduler,
+        )
+        from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
+            OffloadingConnector,
+        )
+        from vllm.v1.core.sched.interface import PauseState
+        from vllm.v1.core.sched.scheduler import Scheduler
+        from vllm.v1.engine.core import EngineCoreProc
+        from vllm.v1.kv_offload.tiering.manager import (
+            CPUPrimaryTierOffloadingManager,
+            TieringOffloadingManager,
+        )
+
+        p2p = _make_manager()
+        _attach_fake_transports(p2p)
+        mmap_region = MagicMock()
+        view = memoryview(np.zeros((1, 16), dtype=np.int8))
+        mmap_region.create_kv_memoryview.return_value = view
+        manager = TieringOffloadingManager(
+            CPUPrimaryTierOffloadingManager(1, mmap_region),
+            [p2p],
+        )
+        try:
+            connector_scheduler = object.__new__(OffloadingConnectorScheduler)
+            connector_scheduler.manager = manager
+            connector_scheduler._jobs = {}
+            connector_scheduler._reset_fence = {}
+            connector = object.__new__(OffloadingConnector)
+            connector.connector_scheduler = connector_scheduler
+            scheduler = object.__new__(Scheduler)
+            scheduler._pause_state = PauseState.PAUSED_NEW
+            scheduler.running = []
+            scheduler.waiting = []
+            scheduler.kv_holding_waiting = []
+            scheduler.num_waiting_for_streaming_input = 0
+            scheduler.finished_req_ids = set()
+            scheduler.requests = {}
+            scheduler.connector = connector
+            scheduler.ec_connector = None
+            core = object.__new__(EngineCoreProc)
+            core.scheduler = scheduler
+            core.engines_running = False
+            core.batch_queue = None
+
+            assert core._pause_complete() is True
+        finally:
+            manager.shutdown()
+
+    @pytest.mark.cpu_test
+    def test_B2_no_poll_thread_keeps_engine_busy(self):
+        """With the tier-poll thread disabled the engine tick is the only
+        driver, so the tiering manager must keep the engine stepping."""
+        from unittest.mock import MagicMock
+
+        from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
+            OffloadingConnectorScheduler,
+        )
+        from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
+            OffloadingConnector,
+        )
+        from vllm.v1.core.sched.interface import PauseState
+        from vllm.v1.core.sched.scheduler import Scheduler
+        from vllm.v1.engine.core import EngineCoreProc
+        from vllm.v1.kv_offload.tiering.manager import (
+            CPUPrimaryTierOffloadingManager,
+            TieringOffloadingManager,
+        )
+
+        p2p = _make_manager()
+        _attach_fake_transports(p2p)
+        mmap_region = MagicMock()
+        view = memoryview(np.zeros((1, 16), dtype=np.int8))
+        mmap_region.create_kv_memoryview.return_value = view
+        manager = TieringOffloadingManager(
+            CPUPrimaryTierOffloadingManager(1, mmap_region),
+            [p2p],
+            tier_poll_interval_s=0,
+        )
+        try:
+            assert manager._tier_poll_thread is None
+            connector_scheduler = object.__new__(OffloadingConnectorScheduler)
+            connector_scheduler.manager = manager
+            connector_scheduler._jobs = {}
+            connector_scheduler._reset_fence = {}
+            connector = object.__new__(OffloadingConnector)
+            connector.connector_scheduler = connector_scheduler
+            scheduler = object.__new__(Scheduler)
+            scheduler._pause_state = PauseState.PAUSED_NEW
+            scheduler.running = []
+            scheduler.waiting = []
+            scheduler.kv_holding_waiting = []
+            scheduler.num_waiting_for_streaming_input = 0
+            scheduler.finished_req_ids = set()
+            scheduler.requests = {}
+            scheduler.connector = connector
+            scheduler.ec_connector = None
+            core = object.__new__(EngineCoreProc)
+            core.scheduler = scheduler
+            core.engines_running = False
+            core.batch_queue = None
+
+            assert core._pause_complete() is False
+        finally:
+            manager.shutdown()
+
 
 # ---------------------------------------------------------------------------
 # Shutdown drain
@@ -2771,3 +2953,112 @@ def peer_lookup(kv_request_id: str) -> LookupResult:
         assert [job.keys for job in promotions] == [[key]]
     finally:
         tiering.shutdown()
+
+
+# ---------------------------------------------------------------------------
+# Idleness while serving peers
+# ---------------------------------------------------------------------------
+
+
+def _producer_with_cpu_hit(
+    key: OffloadKey, tier_poll_interval_s: float | None = None
+) -> tuple[P2PSecondaryTierManager, TieringOffloadingManager]:
+    """A consumer paired with a lone-P2P producer whose CPU tier holds key."""
+    consumer, server = _build_paired_managers()
+    kwargs = {}
+    if tier_poll_interval_s is not None:
+        kwargs["tier_poll_interval_s"] = tier_poll_interval_s
+    tiering = TieringOffloadingManager(
+        primary_tier=CPUPrimaryTierOffloadingManager(
+            num_chunks=4, mmap_region=_mock_mmap_region(4)
+        ),
+        secondary_tiers=[server],
+        **kwargs,
+    )
+    seed_ctx = ReqContext(req_id="seed")
+    with tiering.lock:
+        assert tiering.primary_tier.prepare_store([key], seed_ctx) is not None
+        tiering.primary_tier.complete_store([key], seed_ctx)
+    return consumer, tiering
+
+
+def _peer_hit(
+    consumer: P2PSecondaryTierManager,
+    tiering: TieringOffloadingManager,
+    key: OffloadKey,
+    step_producer: bool,
+) -> ReqContext:
+    """Drive a consumer lookup until the producer pins key as a HIT."""
+    step_end = ScheduleEndContext(new_req_ids=[], preempted_req_ids=[])
+    ctx = ReqContext(
+        req_id="req-peer",
+        kv_transfer_params=_remote_kv_source_kv_params("B", 2, "req-peer"),
+    )
+    _annotate_req_context(ctx)
+    consumer.on_new_request(ctx)
+    deadline = time.monotonic() + 5.0
+    while time.monotonic() < deadline:
+        list(consumer.get_finished_jobs())
+        if step_producer:
+            with tiering.lock:
+                tiering.on_schedule_end(step_end)
+        result = consumer.lookup(key, ctx)
+        if result is LookupResult.HIT:
+            return ctx
+        assert result is LookupResult.RETRY
+        consumer.on_schedule_end(step_end)
+        time.sleep(0.001)
+    raise AssertionError("peer lookup did not reach HIT")
+
+
+class TestIdleWhileServingPeers:
+    @pytest.mark.cpu_test
+    def test_pinned_peer_hit_does_not_block_idleness(self):
+        """A lone-P2P producer with a live tier-poll thread serves a whole peer
+        request without an engine step, so its pin must not count as pending
+        work, and the fetch must still complete and release the pin."""
+        key = OffloadKey(b"cpu-resident")
+        consumer, tiering = _producer_with_cpu_hit(key)
+        try:
+            assert tiering._tier_poll_thread is not None
+            ctx = _peer_hit(consumer, tiering, key, step_producer=False)
+            with tiering.lock:
+                assert tiering._jobs
+                assert tiering.has_pending_work() is False
+
+            consumer.submit_load(
+                TransferJob(
+                    job_id=7,
+                    keys=[key],
+                    chunk_ids=np.array([0]),
+                    is_promotion=True,
+                    req_context=ctx,
+                )
+            )
+            results: list[JobResult] = []
+            deadline = time.monotonic() + 5.0
+            while time.monotonic() < deadline:
+                results.extend(consumer.get_finished_jobs())
+                with tiering.lock:
+                    if not tiering._jobs and results:
+                        break
+                time.sleep(0.001)
+            assert JobResult(job_id=7, success=True) in results
+            with tiering.lock:
+                assert not tiering._jobs
+        finally:
+            tiering.shutdown()
+
+    @pytest.mark.cpu_test
+    def test_pinned_peer_hit_without_poll_thread_is_pending(self):
+        """With the tier-poll thread disabled only a step completes the
+        fetch, so the engine must keep stepping."""
+        key = OffloadKey(b"cpu-resident")
+        consumer, tiering = _producer_with_cpu_hit(key, tier_poll_interval_s=0)
+        try:
+            assert tiering._tier_poll_thread is None
+            _peer_hit(consumer, tiering, key, step_producer=True)
+            assert tiering._jobs
+            assert tiering.has_pending_work() is True
+        finally:
+            tiering.shutdown()