From 0fc251209454db524e15487bbca3b0bb5451ae8b Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 30 Jun 2026 17:07:12 +0300 Subject: [PATCH] [KV Offload] Pass `ScheduleEndContext` to `on_schedule_end` hook (#46450) Signed-off-by: Ronen Schaffer Co-authored-by: Or Ozeri --- tests/v1/kv_offload/tiering/p2p/test_manager.py | 9 +++++++-- tests/v1/kv_offload/tiering/test_fs_tier.py | 3 ++- tests/v1/kv_offload/tiering/test_obj_tier.py | 3 ++- .../kv_offload/tiering/test_tiering_offloading.py | 4 +++- .../kv_connector/v1/offloading/scheduler.py | 7 ++++++- vllm/v1/kv_offload/base.py | 13 +++++++++++-- vllm/v1/kv_offload/tiering/base.py | 3 ++- vllm/v1/kv_offload/tiering/fs/manager.py | 3 ++- vllm/v1/kv_offload/tiering/manager.py | 5 +++-- vllm/v1/kv_offload/tiering/obj/manager.py | 3 ++- vllm/v1/kv_offload/tiering/p2p/manager.py | 3 ++- 11 files changed, 42 insertions(+), 14 deletions(-) diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index 868b62587f..73e3655ecc 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/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 diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 0300fb5d4d..680a958478 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/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(): diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 28570926db..3df9d30691 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/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(): diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index f06b91aa20..fe9283a811 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/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): 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 7aa2b563ba..284098ce94 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/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 (): diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 507e457ac5..48838599d6 100644 --- a/vllm/v1/kv_offload/base.py +++ b/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 diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 662a826a06..b022e1f2f8 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/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 diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 816e88c522..a38ac8eb40 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/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 diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index ee5b0b5274..81151a7c0c 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/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: diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index 857c3c758a..4c3fa754bf 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/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: diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index c6a64268ad..3456f95b7e 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/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 # ------------------------------------------------------------------