[KV Offload] Pass ScheduleEndContext to on_schedule_end hook (#46450)

Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>

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

@@ -13,7 +13,7 @@ from types import SimpleNamespace
 
 import numpy as np
 
-from vllm.v1.kv_offload.base import LookupResult, ReqContext
+from vllm.v1.kv_offload.base import LookupResult, ReqContext, ScheduleEndContext
 from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
 from vllm.v1.kv_offload.tiering.p2p import manager as manager_module
 from vllm.v1.kv_offload.tiering.p2p.manager import (
@@ -1279,7 +1279,12 @@ class TestOnScheduleEnd:
         mgr = _make_manager()
         before_sessions = dict(mgr._sessions)
         before_jobs = list(mgr._finished_jobs)
-        assert mgr.on_schedule_end() is None
+        assert (
+            mgr.on_schedule_end(
+                ScheduleEndContext(new_req_ids=[], preempted_req_ids=())
+            )
+            is None
+        )
         assert mgr._sessions == before_sessions
         assert mgr._finished_jobs == before_jobs
 

tests/v1/kv_offload/tiering/test_fs_tier.py

@@ -22,6 +22,7 @@ from vllm.v1.kv_offload.base import (
     LookupResult,
     OffloadKey,
     ReqContext,
+    ScheduleEndContext,
     make_offload_key,
 )
 from vllm.v1.kv_offload.tiering.base import JobMetadata
@@ -107,7 +108,7 @@ def lookup_and_wait(
     """Perform a full async lookup cycle and return resolved results."""
     for k in keys:
         tier.lookup(k, ctx)
-    tier.on_schedule_end()
+    tier.on_schedule_end(ScheduleEndContext(new_req_ids=[], preempted_req_ids=()))
     deadline = time.monotonic() + timeout
     while time.monotonic() < deadline:
         if not tier._lookup_manager._pending_results.empty():

tests/v1/kv_offload/tiering/test_obj_tier.py

@@ -21,6 +21,7 @@ from vllm.v1.kv_offload.base import (
     LookupResult,
     OffloadKey,
     ReqContext,
+    ScheduleEndContext,
     make_offload_key,
 )
 from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
@@ -222,7 +223,7 @@ def lookup_and_wait(
     """Perform a full async lookup cycle and return resolved results."""
     for k in keys:
         tier.lookup(k, ctx)
-    tier.on_schedule_end()
+    tier.on_schedule_end(ScheduleEndContext(new_req_ids=[], preempted_req_ids=()))
     deadline = time.monotonic() + timeout
     while time.monotonic() < deadline:
         if not tier._lookup_manager._pending_results.empty():

tests/v1/kv_offload/tiering/test_tiering_offloading.py

@@ -27,6 +27,7 @@ from vllm.v1.kv_offload.base import (
     OffloadPolicy,
     ReqContext,
     RequestOffloadingContext,
+    ScheduleEndContext,
     make_offload_key,
 )
 from vllm.v1.kv_offload.tiering.base import (
@@ -233,7 +234,8 @@ class TestTieringOffloadingManager:
 
     def _simulate_on_schedule_end(self):
         """Simulate end of scheduler step: lifecycle flush + drain events."""
-        self.manager.on_schedule_end()
+        ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=())
+        self.manager.on_schedule_end(ctx)
         list(self.manager.take_events())
 
     def _start_request(self, req_context: ReqContext = _CTX):

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

@@ -42,6 +42,7 @@ from vllm.v1.kv_offload.base import (
     OffloadPolicy,
     ReqContext,
     RequestOffloadingContext,
+    ScheduleEndContext,
     make_offload_key,
 )
 from vllm.v1.outputs import KVConnectorOutput
@@ -1027,7 +1028,11 @@ class OffloadingConnectorScheduler:
         self, scheduler_output: SchedulerOutput
     ) -> KVConnectorMetadata:
         self._update_req_states(scheduler_output)
-        self.manager.on_schedule_end()
+        schedule_end_context = ScheduleEndContext(
+            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 ():

vllm/v1/kv_offload/base.py

@@ -8,7 +8,7 @@ from abc import ABC, abstractmethod
 from collections.abc import Collection, Iterable, Sequence
 from dataclasses import dataclass
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Any, NewType
+from typing import TYPE_CHECKING, Any, NamedTuple, NewType
 
 import numpy as np
 import torch
@@ -76,6 +76,15 @@ class RequestOffloadingContext:
     policy: OffloadPolicy = OffloadPolicy.BLOCK_LEVEL
 
 
+class ScheduleEndContext(NamedTuple):
+    """Per-step scheduling info passed to on_schedule_end()."""
+
+    # Request IDs scheduled for the first time this step.
+    new_req_ids: Collection[str]
+    # Request IDs preempted this step.
+    preempted_req_ids: Collection[str]
+
+
 class LoadStoreSpec(ABC):
     """
     Abstract metadata that encapsulates information allowing a worker
@@ -309,7 +318,7 @@ class OffloadingManager(ABC):
         """
         return ()
 
-    def on_schedule_end(self) -> None:
+    def on_schedule_end(self, context: ScheduleEndContext) -> None:
         """Called once at the end of each scheduler step.
 
         Managers may override this to flush deferred work accumulated

vllm/v1/kv_offload/tiering/base.py

@@ -17,6 +17,7 @@ from vllm.v1.kv_offload.base import (
     OffloadKey,
     ReqContext,
     RequestOffloadingContext,
+    ScheduleEndContext,
 )
 
 if TYPE_CHECKING:
@@ -209,7 +210,7 @@ class SecondaryTierManager(ABC):
         """
         return
 
-    def on_schedule_end(self) -> None:
+    def on_schedule_end(self, context: ScheduleEndContext) -> None:
         """Called once at the end of each scheduler step.
 
         Secondary tiers may override this for per-step cleanup or

vllm/v1/kv_offload/tiering/fs/manager.py

@@ -38,6 +38,7 @@ from vllm.v1.kv_offload.tiering.base import (
     JobMetadata,
     JobResult,
     RequestOffloadingContext,
+    ScheduleEndContext,
     SecondaryTierManager,
 )
 from vllm.v1.kv_offload.tiering.fs.io import load_block, store_block
@@ -201,7 +202,7 @@ class FileSystemTierManager(SecondaryTierManager):
         self._lookup_manager.cleanup(req_context.req_id)
 
     @override
-    def on_schedule_end(self) -> None:
+    def on_schedule_end(self, context: ScheduleEndContext) -> None:
         self._lookup_manager.flush()
 
     @override

vllm/v1/kv_offload/tiering/manager.py

@@ -40,6 +40,7 @@ from vllm.v1.kv_offload.base import (
     PrepareStoreOutput,
     ReqContext,
     RequestOffloadingContext,
+    ScheduleEndContext,
 )
 from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
@@ -603,7 +604,7 @@ class TieringOffloadingManager(OffloadingManager):
         del self._req_state[req_id]
 
     @override
-    def on_schedule_end(self) -> None:
+    def on_schedule_end(self, context: ScheduleEndContext) -> None:
         """End-of-schedule hook: process finished jobs, flush deferred
         promotions, and reset the per-step gate.
 
@@ -614,7 +615,7 @@ class TieringOffloadingManager(OffloadingManager):
         self._processed_jobs_this_step = False
         self._flush_pending_promotions()
         for tier in self.secondary_tiers:
-            tier.on_schedule_end()
+            tier.on_schedule_end(context)
 
     @override
     def has_pending_work(self) -> bool:

vllm/v1/kv_offload/tiering/obj/manager.py

@@ -17,6 +17,7 @@ from vllm.v1.kv_offload.tiering.base import (
     JobMetadata,
     JobResult,
     RequestOffloadingContext,
+    ScheduleEndContext,
     SecondaryTierManager,
 )
 from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig
@@ -242,7 +243,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
     def on_request_finished(self, req_context: ReqContext) -> None:
         self._lookup_manager.cleanup(req_context.req_id)
 
-    def on_schedule_end(self) -> None:
+    def on_schedule_end(self, context: ScheduleEndContext) -> None:
         self._lookup_manager.flush()
 
     def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:

vllm/v1/kv_offload/tiering/p2p/manager.py

@@ -26,6 +26,7 @@ from vllm.v1.kv_offload.file_mapper import FileMapper
 from vllm.v1.kv_offload.tiering.base import (
     JobMetadata,
     JobResult,
+    ScheduleEndContext,
     SecondaryTierManager,
 )
 from vllm.v1.kv_offload.tiering.p2p.control import ControlTransport, ZmqTransport
@@ -426,7 +427,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
             time.sleep(_DRAIN_SLEEP_S)
 
     @override
-    def on_schedule_end(self) -> None:
+    def on_schedule_end(self, context: ScheduleEndContext) -> None:
         return
 
     # ------------------------------------------------------------------