0fc2512094 · tduncan/Nvidia-Optimization-Stack-VLLM

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

+7 -2

@@ -13,7 +13,7 @@ from types import SimpleNamespace

13

14 import numpy as np

15

16 -from vllm.v1.kv_offload.base import LookupResult, ReqContext

16 +from vllm.v1.kv_offload.base import LookupResult, ReqContext, ScheduleEndContext

17 from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult

18 from vllm.v1.kv_offload.tiering.p2p import manager as manager_module

19 from vllm.v1.kv_offload.tiering.p2p.manager import (

@@ -1279,7 +1279,12 @@ class TestOnScheduleEnd:

1279 mgr = _make_manager()

1280 before_sessions = dict(mgr._sessions)

1281 before_jobs = list(mgr._finished_jobs)

1282 - assert mgr.on_schedule_end() is None

1282 + assert (

1283 + mgr.on_schedule_end(

1284 + ScheduleEndContext(new_req_ids=[], preempted_req_ids=())

1285 + )

1286 + is None

1287 + )

1288 assert mgr._sessions == before_sessions

1289 assert mgr._finished_jobs == before_jobs

1290

tests/v1/kv_offload/tiering/test_fs_tier.py

+2 -1

@@ -22,6 +22,7 @@ from vllm.v1.kv_offload.base import (

22 LookupResult,

23 OffloadKey,

24 ReqContext,

25 + ScheduleEndContext,

26 make_offload_key,

27 )

28 from vllm.v1.kv_offload.tiering.base import JobMetadata

@@ -107,7 +108,7 @@ def lookup_and_wait(

108 """Perform a full async lookup cycle and return resolved results."""

109 for k in keys:

110 tier.lookup(k, ctx)

110 - tier.on_schedule_end()

111 + tier.on_schedule_end(ScheduleEndContext(new_req_ids=[], preempted_req_ids=()))

112 deadline = time.monotonic() + timeout

113 while time.monotonic() < deadline:

114 if not tier._lookup_manager._pending_results.empty():

tests/v1/kv_offload/tiering/test_obj_tier.py

+2 -1

@@ -21,6 +21,7 @@ from vllm.v1.kv_offload.base import (

21 LookupResult,

22 OffloadKey,

23 ReqContext,

24 + ScheduleEndContext,

25 make_offload_key,

26 )

27 from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult

@@ -222,7 +223,7 @@ def lookup_and_wait(

223 """Perform a full async lookup cycle and return resolved results."""

224 for k in keys:

225 tier.lookup(k, ctx)

225 - tier.on_schedule_end()

226 + tier.on_schedule_end(ScheduleEndContext(new_req_ids=[], preempted_req_ids=()))

227 deadline = time.monotonic() + timeout

228 while time.monotonic() < deadline:

229 if not tier._lookup_manager._pending_results.empty():

tests/v1/kv_offload/tiering/test_tiering_offloading.py

+3 -1

@@ -27,6 +27,7 @@ from vllm.v1.kv_offload.base import (

27 OffloadPolicy,

28 ReqContext,

29 RequestOffloadingContext,

30 + ScheduleEndContext,

31 make_offload_key,

32 )

33 from vllm.v1.kv_offload.tiering.base import (

@@ -233,7 +234,8 @@ class TestTieringOffloadingManager:

234

235 def _simulate_on_schedule_end(self):

236 """Simulate end of scheduler step: lifecycle flush + drain events."""

236 - self.manager.on_schedule_end()

237 + ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=())

238 + self.manager.on_schedule_end(ctx)

239 list(self.manager.take_events())

240

241 def _start_request(self, req_context: ReqContext = _CTX):

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

+6 -1

@@ -42,6 +42,7 @@ from vllm.v1.kv_offload.base import (

42 OffloadPolicy,

43 ReqContext,

44 RequestOffloadingContext,

45 + ScheduleEndContext,

46 make_offload_key,

47 )

48 from vllm.v1.outputs import KVConnectorOutput

@@ -1027,7 +1028,11 @@ class OffloadingConnectorScheduler:

1028 self, scheduler_output: SchedulerOutput

1029 ) -> KVConnectorMetadata:

1030 self._update_req_states(scheduler_output)

1030 - self.manager.on_schedule_end()

1031 + schedule_end_context = ScheduleEndContext(

1032 + new_req_ids=[req.req_id for req in scheduler_output.scheduled_new_reqs],

1033 + preempted_req_ids=scheduler_output.preempted_req_ids or (),

1034 + )

1035 + self.manager.on_schedule_end(schedule_end_context)

1036

1037 # Flush jobs for preempted requests.

1038 for req_id in scheduler_output.preempted_req_ids or ():

vllm/v1/kv_offload/base.py

+11 -2

@@ -8,7 +8,7 @@ from abc import ABC, abstractmethod

8 from collections.abc import Collection, Iterable, Sequence

9 from dataclasses import dataclass

10 from enum import Enum, auto

11 -from typing import TYPE_CHECKING, Any, NewType

11 +from typing import TYPE_CHECKING, Any, NamedTuple, NewType

12

13 import numpy as np

14 import torch

@@ -76,6 +76,15 @@ class RequestOffloadingContext:

76 policy: OffloadPolicy = OffloadPolicy.BLOCK_LEVEL

77

78

79 +class ScheduleEndContext(NamedTuple):

80 + """Per-step scheduling info passed to on_schedule_end()."""

81 +

82 + # Request IDs scheduled for the first time this step.

83 + new_req_ids: Collection[str]

84 + # Request IDs preempted this step.

85 + preempted_req_ids: Collection[str]

86 +

87 +

88 class LoadStoreSpec(ABC):

89 """

90 Abstract metadata that encapsulates information allowing a worker

@@ -309,7 +318,7 @@ class OffloadingManager(ABC):

318 """

319 return ()

320

312 - def on_schedule_end(self) -> None:

321 + def on_schedule_end(self, context: ScheduleEndContext) -> None:

322 """Called once at the end of each scheduler step.

323

324 Managers may override this to flush deferred work accumulated

vllm/v1/kv_offload/tiering/base.py

+2 -1

@@ -17,6 +17,7 @@ from vllm.v1.kv_offload.base import (

17 OffloadKey,

18 ReqContext,

19 RequestOffloadingContext,

20 + ScheduleEndContext,

21 )

22

23 if TYPE_CHECKING:

@@ -209,7 +210,7 @@ class SecondaryTierManager(ABC):

210 """

211 return

212

212 - def on_schedule_end(self) -> None:

213 + def on_schedule_end(self, context: ScheduleEndContext) -> None:

214 """Called once at the end of each scheduler step.

215

216 Secondary tiers may override this for per-step cleanup or

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

+2 -1

@@ -38,6 +38,7 @@ from vllm.v1.kv_offload.tiering.base import (

38 JobMetadata,

39 JobResult,

40 RequestOffloadingContext,

41 + ScheduleEndContext,

42 SecondaryTierManager,

43 )

44 from vllm.v1.kv_offload.tiering.fs.io import load_block, store_block

@@ -201,7 +202,7 @@ class FileSystemTierManager(SecondaryTierManager):

202 self._lookup_manager.cleanup(req_context.req_id)

203

204 @override

204 - def on_schedule_end(self) -> None:

205 + def on_schedule_end(self, context: ScheduleEndContext) -> None:

206 self._lookup_manager.flush()

207

208 @override

vllm/v1/kv_offload/tiering/manager.py

+3 -2

@@ -40,6 +40,7 @@ from vllm.v1.kv_offload.base import (

40 PrepareStoreOutput,

41 ReqContext,

42 RequestOffloadingContext,

43 + ScheduleEndContext,

44 )

45 from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec

46 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager

@@ -603,7 +604,7 @@ class TieringOffloadingManager(OffloadingManager):

604 del self._req_state[req_id]

605

606 @override

606 - def on_schedule_end(self) -> None:

607 + def on_schedule_end(self, context: ScheduleEndContext) -> None:

608 """End-of-schedule hook: process finished jobs, flush deferred

609 promotions, and reset the per-step gate.

610

@@ -614,7 +615,7 @@ class TieringOffloadingManager(OffloadingManager):

615 self._processed_jobs_this_step = False

616 self._flush_pending_promotions()

617 for tier in self.secondary_tiers:

617 - tier.on_schedule_end()

618 + tier.on_schedule_end(context)

619

620 @override

621 def has_pending_work(self) -> bool:

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

+2 -1

@@ -17,6 +17,7 @@ from vllm.v1.kv_offload.tiering.base import (

17 JobMetadata,

18 JobResult,

19 RequestOffloadingContext,

20 + ScheduleEndContext,

21 SecondaryTierManager,

22 )

23 from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig

@@ -242,7 +243,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

243 def on_request_finished(self, req_context: ReqContext) -> None:

244 self._lookup_manager.cleanup(req_context.req_id)

245

245 - def on_schedule_end(self) -> None:

246 + def on_schedule_end(self, context: ScheduleEndContext) -> None:

247 self._lookup_manager.flush()

248

249 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:

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

+2 -1

@@ -26,6 +26,7 @@ from vllm.v1.kv_offload.file_mapper import FileMapper

26 from vllm.v1.kv_offload.tiering.base import (

27 JobMetadata,

28 JobResult,

29 + ScheduleEndContext,

30 SecondaryTierManager,

31 )

32 from vllm.v1.kv_offload.tiering.p2p.control import ControlTransport, ZmqTransport

@@ -426,7 +427,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):

427 time.sleep(_DRAIN_SLEEP_S)

428

429 @override

429 - def on_schedule_end(self) -> None:

430 + def on_schedule_end(self, context: ScheduleEndContext) -> None:

431 return

432

433 # ------------------------------------------------------------------

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