@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a3ed5ab10c
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 12:15 UTC

[KV Offload] Add per-request offloading policy via `on_new_request` lifecycle hook (#43205)

Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com> Co-authored-by: Or Ozeri <or@ozery.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Ronen Schaffer committed May 28, 2026 at 23:45 UTC a3ed5ab10ca127f2dc8a87c6a23962fd5bec8caf
11 files changed +294 -45
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+14 -14
@@ -24,7 +24,9 @@ from vllm.v1.kv_cache_interface import (
24 from vllm.v1.kv_offload.base import (
25 OffloadingEvent,
26 OffloadingManager,
27 + OffloadPolicy,
28 ReqContext,
29 + RequestOffloadingContext,
30 get_offload_block_hash,
31 )
32 from vllm.v1.request import RequestStatus
@@ -748,12 +750,8 @@ class TestSlidingWindowLookup:
750
751
752 @pytest.mark.parametrize("async_scheduling", [True, False])
751 -def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bool):
752 - """With do_remote_decode=True, after loading prefix blocks from CPU,
753 - all blocks must be re-stored — not just the newly computed ones.
754 -
755 - This supports P/D disaggregation where the prefill instance offloads the
756 - complete KV cache so a remote decode node can consume it."""
753 +def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool):
754 + """With REQUEST_LEVEL policy, all blocks are stored — including prefix hits."""
755 gpu_block_size = 4
756 block_size_factor = 3
757 offloaded_block_size = gpu_block_size * block_size_factor
@@ -780,12 +778,13 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo
778 # Reset GPU prefix cache so the next request must load from CPU.
779 runner.scheduler.reset_prefix_cache()
780
783 - # New request with do_remote_decode=True and 2 offloaded blocks.
784 - # The first offloaded block matches what we stored in CPU.
785 - runner.new_request(
786 - token_ids=[0] * offloaded_block_size * 2,
787 - kv_transfer_params={"do_remote_decode": True},
781 + # Manager returns REQUEST_LEVEL for the next request.
782 + runner.manager.on_new_request.return_value = RequestOffloadingContext(
783 + policy=OffloadPolicy.REQUEST_LEVEL
784 )
785 +
786 + # New request with 2 offloaded blocks; first matches what's in CPU.
787 + runner.new_request(token_ids=[0] * offloaded_block_size * 2)
788 runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
789 runner.manager.prepare_store.side_effect = lambda keys, req_context: (
790 generate_store_output(keys)
@@ -957,9 +956,10 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
956 )
957
958 def setup(r, max_offload_tokens):
960 - r.new_request(token_ids=[0] * offloaded_block_size * 3)
961 - req = r.scheduler.requests[str(r.req_id)]
962 - req.kv_transfer_params = {"max_offload_tokens": max_offload_tokens}
959 + r.new_request(
960 + token_ids=[0] * offloaded_block_size * 3,
961 + kv_transfer_params={"max_offload_tokens": max_offload_tokens},
962 + )
963 r.manager.prepare_store.side_effect = (
964 lambda keys, req_context: generate_store_output(keys)
965 )
tests/v1/kv_connector/unit/offloading_connector/utils.py
+2
@@ -43,6 +43,7 @@ from vllm.v1.kv_offload.base import (
43 OffloadingSpec,
44 OffloadKey,
45 PrepareStoreOutput,
46 + RequestOffloadingContext,
47 make_offload_key,
48 )
49 from vllm.v1.kv_offload.worker.worker import (
@@ -119,6 +120,7 @@ class MockOffloadingSpec(OffloadingSpec):
120 self.manager.lookup.return_value = 0
121 self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)
122 self.manager.lookup.return_value = False
123 + self.manager.on_new_request.return_value = RequestOffloadingContext()
124 self.handler = MockOffloadingHandler()
125
126 def get_manager(self) -> OffloadingManager:
tests/v1/kv_offload/test_tiering_offloading.py
+74
@@ -19,7 +19,9 @@ import torch
19
20 from vllm.v1.kv_offload.base import (
21 OffloadKey,
22 + OffloadPolicy,
23 ReqContext,
24 + RequestOffloadingContext,
25 make_offload_key,
26 )
27 from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
@@ -412,6 +414,78 @@ class TestTieringOffloadingManager:
414 job_metadata = self.secondary_tier1.submit_store.call_args.args[0]
415 assert job_metadata.req_context is ctx
416
417 + def test_on_new_request_lifecycle(self, manager_setup):
418 + """Policy defaults to BLOCK_LEVEL, escalates when a tier requests it,
419 + and is cleaned up on on_request_finished."""
420 + # Default: all tiers return BLOCK_LEVEL
421 + ctx = ReqContext(req_id="req_policy_lifecycle")
422 + result = self.manager.on_new_request(ctx)
423 + assert result.policy == OffloadPolicy.BLOCK_LEVEL
424 + self.manager.on_request_finished(ctx)
425 +
426 + # Escalate: tier1 requests REQUEST_LEVEL
427 + self.secondary_tier1.on_new_request = (
428 + lambda req_context: RequestOffloadingContext(
429 + policy=OffloadPolicy.REQUEST_LEVEL
430 + )
431 + )
432 +
433 + ctx = ReqContext(req_id="req_policy_lifecycle_2")
434 + result = self.manager.on_new_request(ctx)
435 + assert result.policy == OffloadPolicy.REQUEST_LEVEL
436 + assert ctx.req_id in self.manager._request_level_tiers
437 +
438 + # Cleanup
439 + self.manager.on_request_finished(ctx)
440 + assert ctx.req_id not in self.manager._request_level_tiers
441 +
442 + def test_prepare_store_cascades_existing_blocks_to_request_level_tiers(
443 + self, manager_setup
444 + ):
445 + """prepare_store cascades hit blocks to request-level tiers only."""
446 + # Store some blocks to primary first
447 + existing_blocks = to_keys(range(3))
448 + result = self.manager.prepare_store(existing_blocks, _CTX)
449 + assert result is not None
450 + self.manager.complete_store(existing_blocks, _CTX, success=True)
451 + # Drain cascade completions
452 + list(self.manager.take_events())
453 + list(self.manager.take_events())
454 +
455 + # Make tier1 request-level, tier2 stays block-level
456 + self.secondary_tier1.on_new_request = (
457 + lambda req_context: RequestOffloadingContext(
458 + policy=OffloadPolicy.REQUEST_LEVEL
459 + )
460 + )
461 +
462 + ctx = ReqContext(req_id="req_cascade")
463 + self.manager.on_new_request(ctx)
464 +
465 + # Spy on submit_store
466 + self.secondary_tier1.submit_store = MagicMock(
467 + wraps=self.secondary_tier1.submit_store
468 + )
469 + self.secondary_tier2.submit_store = MagicMock(
470 + wraps=self.secondary_tier2.submit_store
471 + )
472 +
473 + # Call prepare_store with existing + new blocks
474 + new_blocks = to_keys(range(3, 5))
475 + all_blocks = existing_blocks + new_blocks
476 + result = self.manager.prepare_store(all_blocks, ctx)
477 + assert result is not None
478 + assert set(result.keys_to_store) == set(new_blocks)
479 +
480 + # Only tier1 (request-level) should get existing blocks cascaded now.
481 + # New blocks are cascaded to ALL tiers later via complete_store().
482 + self.secondary_tier1.submit_store.assert_called_once()
483 + job_metadata = self.secondary_tier1.submit_store.call_args.args[0]
484 + assert set(job_metadata.keys) == set(existing_blocks)
485 +
486 + # tier2 (block-level) does not get existing blocks here.
487 + self.secondary_tier2.submit_store.assert_not_called()
488 +
489
490 class TestTieringOffloadingWithoutSecondaryTiers:
491 """Test TieringOffloadingManager with no secondary tiers (backward compat)."""
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+37 -25
@@ -29,7 +29,9 @@ from vllm.v1.kv_offload.base import (
29 OffloadingManager,
30 OffloadingSpec,
31 OffloadKey,
32 + OffloadPolicy,
33 ReqContext,
34 + RequestOffloadingContext,
35 get_offload_block_hash,
36 make_offload_key,
37 )
@@ -172,8 +174,9 @@ class RequestGroupState:
174 class RequestOffloadState:
175 config: SchedulerOffloadConfig
176 req: Request
177 + req_context: ReqContext
178 + offloading_context: RequestOffloadingContext
179 group_states: tuple[RequestGroupState, ...] = field(init=False)
176 - req_context: ReqContext = field(init=False)
180 # upper bound on tokens to offload for this request; None means no cap
181 max_offload_tokens: int | None = None
182 # number of hits in the GPU cache
@@ -186,10 +189,6 @@ class RequestOffloadState:
189 self.group_states = tuple(
190 RequestGroupState() for _ in self.config.kv_group_configs
191 )
189 - self.req_context = ReqContext(
190 - req_id=self.req.request_id,
191 - kv_transfer_params=self.req.kv_transfer_params,
192 - )
192 params = self.req.kv_transfer_params
193
194 # NOTE: This field is experimental and subject to change in the future.
@@ -248,6 +247,13 @@ class RequestOffloadState:
247 )
248
249
250 +def _create_req_context(req: Request) -> ReqContext:
251 + return ReqContext(
252 + req_id=req.request_id,
253 + kv_transfer_params=req.kv_transfer_params,
254 + )
255 +
256 +
257 class OffloadingConnectorScheduler:
258 """Implementation of Scheduler side methods"""
259
@@ -512,6 +518,18 @@ class OffloadingConnectorScheduler:
518
519 return num_hit_tokens
520
521 + def on_new_request(self, request: Request) -> None:
522 + """Called when a new request is added to the scheduler."""
523 + req_context = _create_req_context(request)
524 + offloading_context = self.manager.on_new_request(req_context)
525 + req_status = RequestOffloadState(
526 + config=self.config,
527 + req=request,
528 + req_context=req_context,
529 + offloading_context=offloading_context,
530 + )
531 + self._req_status[request.request_id] = req_status
532 +
533 def get_num_new_matched_tokens(
534 self, request: Request, num_computed_tokens: int
535 ) -> tuple[int | None, bool]:
@@ -534,24 +552,15 @@ class OffloadingConnectorScheduler:
552 - `True` if tokens will be loaded asynchronously
553 (between scheduler steps).
554 """
537 - is_new_request = False
538 - if req_status := self._req_status.get(request.request_id):
539 - # make sure block IDs are cleared
540 - for group_state in req_status.group_states:
541 - group_state.block_ids.clear()
542 - else:
543 - is_new_request = True
544 - req_status = RequestOffloadState(config=self.config, req=request)
545 - self._req_status[request.request_id] = req_status
555 + req_status = self._req_status[request.request_id]
556 + for group_state in req_status.group_states:
557 + group_state.block_ids.clear()
558
559 req_status.update_offload_keys()
560 req_status.num_locally_computed_tokens = num_computed_tokens
561
562 num_hit_tokens = self._lookup(req_status)
551 - if is_new_request:
552 - req_status.update_num_hit_blocks(
553 - num_computed_tokens + (num_hit_tokens or 0)
554 - )
563 + req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0))
564
565 self._touch(req_status)
566
@@ -568,9 +577,6 @@ class OffloadingConnectorScheduler:
577 num_locally_computed_tokens = req_status.num_locally_computed_tokens
578 num_cached_tokens = num_locally_computed_tokens + num_external_tokens
579
571 - params = req_status.req_context.kv_transfer_params
572 - do_remote_decode = params is not None and params.get("do_remote_decode")
573 -
580 keys_to_load: list[OffloadKey] = []
581 dst_block_ids: list[int] = []
582 # per group
@@ -624,10 +630,10 @@ class OffloadingConnectorScheduler:
630 group_sizes.append(num_pending_gpu_blocks)
631 block_indices.append(num_locally_computed_gpu_blocks)
632
627 - if not do_remote_decode:
628 - # For P/D prefill requests (do_remote_decode=True), we do
629 - # NOT skip saving the hit prefix, as we need to stream the
630 - # entire KV cache so a remote decode node can consume it.
633 + # Skip prefix-hit blocks for block-level policy; for
634 + # request-level, next_stored_block_idx stays at 0 so all
635 + # blocks (including hits) are offloaded.
636 + if req_status.offloading_context.policy == OffloadPolicy.BLOCK_LEVEL:
637 group_state.next_stored_block_idx = num_blocks
638
639 # Fence dst blocks against finished-request pending stores.
@@ -950,6 +956,12 @@ class OffloadingConnectorScheduler:
956 # TODO(orozery): possibly kickoff offload for last block
957 # which may have been deferred due to async scheduling
958 req_status = self._req_status.get(request.request_id)
959 +
960 + req_context = (
961 + req_status.req_context if req_status else _create_req_context(request)
962 + )
963 + self.manager.on_request_finished(req_context)
964 +
965 if req_status is None:
966 return False, None
967 if not req_status.transfer_jobs:
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+4
@@ -124,6 +124,10 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
124 return self.connector_worker.build_connector_worker_meta()
125 return None
126
127 + def on_new_request(self, request: "Request") -> None:
128 + assert self.connector_scheduler is not None
129 + self.connector_scheduler.on_new_request(request)
130 +
131 def get_num_new_matched_tokens(
132 self, request: "Request", num_computed_tokens: int
133 ) -> tuple[int | None, bool]:
vllm/v1/kv_offload/base.py
+37
@@ -7,6 +7,7 @@ Core abstractions for KV cache offloading in vLLM v1.
7 from abc import ABC, abstractmethod
8 from collections.abc import Collection, Iterable, Iterator, Sequence
9 from dataclasses import dataclass
10 +from enum import Enum
11 from typing import TYPE_CHECKING, Any, NewType
12
13 import numpy as np
@@ -49,6 +50,20 @@ class ReqContext:
50 kv_transfer_params: dict[str, Any] | None = None
51
52
53 +class OffloadPolicy(Enum):
54 + # Offload only newly-computed blocks as they arrive; prefix-hit
55 + # blocks (already offloaded by a prior request) are skipped.
56 + BLOCK_LEVEL = "block_level"
57 + # Offload all blocks for the request, including prefix hits.
58 + # Used by tiers that need the complete KV context for a request.
59 + REQUEST_LEVEL = "request_level"
60 +
61 +
62 +@dataclass
63 +class RequestOffloadingContext:
64 + policy: OffloadPolicy = OffloadPolicy.BLOCK_LEVEL
65 +
66 +
67 class LoadStoreSpec(ABC):
68 """
69 Abstract metadata that encapsulates information allowing a worker
@@ -210,6 +225,28 @@ class OffloadingManager(ABC):
225 """
226 return
227
228 + @abstractmethod
229 + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
230 + """
231 + Called when a new request is first seen by the scheduler.
232 +
233 + Returns a RequestOffloadingContext indicating how this request's
234 + blocks should be offloaded.
235 +
236 + Args:
237 + req_context: per-request context.
238 + """
239 + pass
240 +
241 + def on_request_finished(self, req_context: ReqContext) -> None:
242 + """
243 + Called when a request has finished.
244 +
245 + Args:
246 + req_context: per-request context.
247 + """
248 + return
249 +
250 def take_events(self) -> Iterable[OffloadingEvent]:
251 """
252 Take the offloading events from the manager.
vllm/v1/kv_offload/cpu/manager.py
+4
@@ -11,6 +11,7 @@ from vllm.v1.kv_offload.base import (
11 OffloadKey,
12 PrepareStoreOutput,
13 ReqContext,
14 + RequestOffloadingContext,
15 )
16 from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
17 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
@@ -94,6 +95,9 @@ class CPUOffloadingManager(OffloadingManager):
95
96 # --- OffloadingManager interface ---
97
98 + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
99 + return RequestOffloadingContext()
100 +
101 def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
102 if self.counts is not None:
103 if key in self.counts:
vllm/v1/kv_offload/tiering/base.py
+23 -1
@@ -11,7 +11,7 @@ from typing import TYPE_CHECKING
11
12 import numpy as np
13
14 -from vllm.v1.kv_offload.base import OffloadKey, ReqContext
14 +from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext
15
16 if TYPE_CHECKING:
17 from vllm.v1.kv_offload.base import OffloadingSpec
@@ -163,6 +163,28 @@ class SecondaryTierManager(ABC):
163 """
164 return
165
166 + @abstractmethod
167 + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
168 + """
169 + Called when a new request is first seen by the scheduler.
170 +
171 + Returns a RequestOffloadingContext expressing this tier's preference
172 + for how blocks should be offloaded for this request.
173 +
174 + Args:
175 + req_context: Per-request context.
176 + """
177 + pass
178 +
179 + def on_request_finished(self, req_context: ReqContext) -> None:
180 + """
181 + Called when a request has finished.
182 +
183 + Args:
184 + req_context: per-request context.
185 + """
186 + return
187 +
188 def shutdown(self) -> None:
189 """Release resources held by this tier (threads, connections, etc.)."""
190 return
vllm/v1/kv_offload/tiering/example/manager.py
+4 -1
@@ -13,7 +13,7 @@ import logging
13 from collections.abc import Iterable
14 from typing import TYPE_CHECKING
15
16 -from vllm.v1.kv_offload.base import OffloadKey, ReqContext
16 +from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext
17 from vllm.v1.kv_offload.tiering.base import (
18 JobMetadata,
19 JobResult,
@@ -132,6 +132,9 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
132 self.completed_jobs = []
133 return result
134
135 + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
136 + return RequestOffloadingContext()
137 +
138 def get_num_blocks(self) -> int:
139 """Get the number of blocks currently stored in this tier."""
140 return len(self.blocks)
vllm/v1/kv_offload/tiering/fs/manager.py
+4
@@ -27,6 +27,7 @@ from vllm.v1.kv_offload.file_mapper import FileMapper
27 from vllm.v1.kv_offload.tiering.base import (
28 JobMetadata,
29 JobResult,
30 + RequestOffloadingContext,
31 SecondaryTierManager,
32 )
33 from vllm.v1.kv_offload.tiering.fs.io import load_block, store_block
@@ -100,6 +101,9 @@ class FileSystemTierManager(SecondaryTierManager):
101 thread_name_prefix="vllm_kv_py_fs",
102 )
103
104 + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
105 + return RequestOffloadingContext()
106 +
107 def lookup(
108 self, key: OffloadKey, req_context: ReqContext | None = None
109 ) -> bool | None:
vllm/v1/kv_offload/tiering/manager.py
+91 -4
@@ -20,7 +20,8 @@ Key Design Principles:
20 protecting blocks from eviction until complete_read() is called
21 """
22
23 -from collections.abc import Collection, Iterable
23 +from collections import defaultdict
24 +from collections.abc import Collection, Iterable, Sequence
25 from dataclasses import dataclass, field
26
27 import numpy as np
@@ -31,8 +32,10 @@ from vllm.v1.kv_offload.base import (
32 OffloadingEvent,
33 OffloadingManager,
34 OffloadKey,
35 + OffloadPolicy,
36 PrepareStoreOutput,
37 ReqContext,
38 + RequestOffloadingContext,
39 )
40 from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
41 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
@@ -158,6 +161,13 @@ class TieringOffloadingManager(OffloadingManager):
161 # Reset at the end of each step in take_events().
162 self._processed_jobs_this_step: bool = False
163
164 + # Per-request set of secondary tiers that requested REQUEST_LEVEL
165 + # policy. Populated in on_new_request(),
166 + # cleaned up in on_request_finished().
167 + self._request_level_tiers: defaultdict[str, set[SecondaryTierManager]] = (
168 + defaultdict(set)
169 + )
170 +
171 def _next_job_id(self) -> JobId:
172 """Generate a unique job ID for async transfer tracking."""
173 job_id = self._job_id_counter
@@ -387,6 +397,9 @@ class TieringOffloadingManager(OffloadingManager):
397 that any completed async transfers have their ref_cnt decremented
398 before the primary tier makes eviction decisions.
399
400 + For request-level tiers, blocks already present in the primary tier
401 + are immediately cascaded via submit_store().
402 +
403 Args:
404 keys: Blocks to prepare for storing.
405 req_context: Per-request context.
@@ -400,14 +413,63 @@ class TieringOffloadingManager(OffloadingManager):
413 # successfully transferred to secondary tiers.
414 self._maybe_process_finished_jobs()
415
403 - # Step 2: Store to primary tier
416 + # Step 2: Store to primary tier (new blocks only).
417 + # Cascading of these newly-stored blocks to ALL secondary tiers
418 + # happens later in complete_store(), after the GPU→Primary transfer
419 + # completes.
420 primary_result = self.primary_tier.prepare_store(keys, req_context)
421
406 - # Note: Secondary tier cascading will happen in complete_store()
407 - # after the GPU→Primary transfer completes and blocks are ready.
422 + if primary_result is None:
423 + return None
424 +
425 + # Step 3: For request-level tiers, cascade blocks already in primary
426 + request_level_tiers = self._request_level_tiers.get(req_context.req_id)
427 + if request_level_tiers is not None:
428 + keys_to_store_set = set(primary_result.keys_to_store)
429 + keys_already_in_primary = tuple(
430 + k for k in keys if k not in keys_to_store_set
431 + )
432 + if keys_already_in_primary:
433 + self._cascade_existing_blocks_to_request_level_tiers(
434 + keys_already_in_primary, req_context, request_level_tiers
435 + )
436
437 return primary_result
438
439 + def _cascade_existing_blocks_to_request_level_tiers(
440 + self,
441 + keys: Sequence[OffloadKey],
442 + req_context: ReqContext,
443 + request_level_tiers: set[SecondaryTierManager],
444 + ) -> None:
445 + """
446 + For tiers that requested request-level policy, submit_store() for
447 + blocks that are already present in the primary tier.
448 + """
449 + # Filter out keys that are not ready in primary (e.g. in-flight)
450 + ready_keys = tuple(
451 + k for k in keys if self.primary_tier.lookup(k, req_context) is True
452 + )
453 + if not ready_keys:
454 + return
455 +
456 + for tier in request_level_tiers:
457 + primary_blocks_spec = self.primary_tier.prepare_read(
458 + ready_keys, req_context
459 + )
460 +
461 + job_id = self._next_job_id()
462 + assert isinstance(primary_blocks_spec, CPULoadStoreSpec)
463 + job_metadata = JobMetadata(
464 + job_id=job_id,
465 + keys=ready_keys,
466 + block_ids=primary_blocks_spec.block_ids,
467 + is_promotion=False,
468 + req_context=req_context,
469 + )
470 + self._transfer_jobs[job_id] = job_metadata
471 + tier.submit_store(job_metadata)
472 +
473 def complete_store(
474 self,
475 keys: Collection[OffloadKey],
@@ -466,6 +528,31 @@ class TieringOffloadingManager(OffloadingManager):
528 # Note: The async transfers are now in flight. Their completion is
529 # tracked via get_finished_jobs() / _maybe_process_finished_jobs().
530
531 + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
532 + """
533 + Query each secondary tier for its offload policy preference.
534 +
535 + Returns REQUEST_LEVEL if ANY secondary tier wants request-level.
536 + Only stores REQUEST_LEVEL tier decisions for use in prepare_store.
537 + """
538 + for tier in self.secondary_tiers:
539 + tier_ctx = tier.on_new_request(req_context)
540 + if tier_ctx.policy == OffloadPolicy.REQUEST_LEVEL:
541 + self._request_level_tiers[req_context.req_id].add(tier)
542 +
543 + policy = (
544 + OffloadPolicy.REQUEST_LEVEL
545 + if req_context.req_id in self._request_level_tiers
546 + else OffloadPolicy.BLOCK_LEVEL
547 + )
548 + return RequestOffloadingContext(policy=policy)
549 +
550 + def on_request_finished(self, req_context: ReqContext) -> None:
551 + self.primary_tier.on_request_finished(req_context)
552 + for tier in self.secondary_tiers:
553 + tier.on_request_finished(req_context)
554 + self._request_level_tiers.pop(req_context.req_id, None)
555 +
556 def take_events(self) -> Iterable[OffloadingEvent]:
557 """
558 End-of-step hook: flush deferred work, yield events, reset per-step state.