From a3ed5ab10ca127f2dc8a87c6a23962fd5bec8caf Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 28 May 2026 23:45:18 +0300 Subject: [PATCH] [KV Offload] Add per-request offloading policy via `on_new_request` lifecycle hook (#43205) Signed-off-by: Ronen Schaffer Co-authored-by: Or Ozeri Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../offloading_connector/test_scheduler.py | 28 +++--- .../unit/offloading_connector/utils.py | 2 + .../v1/kv_offload/test_tiering_offloading.py | 74 +++++++++++++++ .../kv_connector/v1/offloading/scheduler.py | 62 +++++++----- .../kv_connector/v1/offloading_connector.py | 4 + vllm/v1/kv_offload/base.py | 37 ++++++++ vllm/v1/kv_offload/cpu/manager.py | 4 + vllm/v1/kv_offload/tiering/base.py | 24 ++++- vllm/v1/kv_offload/tiering/example/manager.py | 5 +- vllm/v1/kv_offload/tiering/fs/manager.py | 4 + vllm/v1/kv_offload/tiering/manager.py | 95 ++++++++++++++++++- 11 files changed, 294 insertions(+), 45 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 0cfa1fc168..ff30a6f42f 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -24,7 +24,9 @@ from vllm.v1.kv_cache_interface import ( from vllm.v1.kv_offload.base import ( OffloadingEvent, OffloadingManager, + OffloadPolicy, ReqContext, + RequestOffloadingContext, get_offload_block_hash, ) from vllm.v1.request import RequestStatus @@ -748,12 +750,8 @@ class TestSlidingWindowLookup: @pytest.mark.parametrize("async_scheduling", [True, False]) -def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bool): - """With do_remote_decode=True, after loading prefix blocks from CPU, - all blocks must be re-stored — not just the newly computed ones. - - This supports P/D disaggregation where the prefill instance offloads the - complete KV cache so a remote decode node can consume it.""" +def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool): + """With REQUEST_LEVEL policy, all blocks are stored — including prefix hits.""" gpu_block_size = 4 block_size_factor = 3 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 # Reset GPU prefix cache so the next request must load from CPU. runner.scheduler.reset_prefix_cache() - # New request with do_remote_decode=True and 2 offloaded blocks. - # The first offloaded block matches what we stored in CPU. - runner.new_request( - token_ids=[0] * offloaded_block_size * 2, - kv_transfer_params={"do_remote_decode": True}, + # Manager returns REQUEST_LEVEL for the next request. + runner.manager.on_new_request.return_value = RequestOffloadingContext( + policy=OffloadPolicy.REQUEST_LEVEL ) + + # New request with 2 offloaded blocks; first matches what's in CPU. + runner.new_request(token_ids=[0] * offloaded_block_size * 2) runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) @@ -957,9 +956,10 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool): ) def setup(r, max_offload_tokens): - r.new_request(token_ids=[0] * offloaded_block_size * 3) - req = r.scheduler.requests[str(r.req_id)] - req.kv_transfer_params = {"max_offload_tokens": max_offload_tokens} + r.new_request( + token_ids=[0] * offloaded_block_size * 3, + kv_transfer_params={"max_offload_tokens": max_offload_tokens}, + ) r.manager.prepare_store.side_effect = ( lambda keys, req_context: generate_store_output(keys) ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index ad4d1ae249..4248a89bf6 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -43,6 +43,7 @@ from vllm.v1.kv_offload.base import ( OffloadingSpec, OffloadKey, PrepareStoreOutput, + RequestOffloadingContext, make_offload_key, ) from vllm.v1.kv_offload.worker.worker import ( @@ -119,6 +120,7 @@ class MockOffloadingSpec(OffloadingSpec): self.manager.lookup.return_value = 0 self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys) self.manager.lookup.return_value = False + self.manager.on_new_request.return_value = RequestOffloadingContext() self.handler = MockOffloadingHandler() def get_manager(self) -> OffloadingManager: diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index 7359a39384..a64caddb2b 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -19,7 +19,9 @@ import torch from vllm.v1.kv_offload.base import ( OffloadKey, + OffloadPolicy, ReqContext, + RequestOffloadingContext, make_offload_key, ) from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager @@ -412,6 +414,78 @@ class TestTieringOffloadingManager: job_metadata = self.secondary_tier1.submit_store.call_args.args[0] assert job_metadata.req_context is ctx + def test_on_new_request_lifecycle(self, manager_setup): + """Policy defaults to BLOCK_LEVEL, escalates when a tier requests it, + and is cleaned up on on_request_finished.""" + # Default: all tiers return BLOCK_LEVEL + ctx = ReqContext(req_id="req_policy_lifecycle") + result = self.manager.on_new_request(ctx) + assert result.policy == OffloadPolicy.BLOCK_LEVEL + self.manager.on_request_finished(ctx) + + # Escalate: tier1 requests REQUEST_LEVEL + self.secondary_tier1.on_new_request = ( + lambda req_context: RequestOffloadingContext( + policy=OffloadPolicy.REQUEST_LEVEL + ) + ) + + ctx = ReqContext(req_id="req_policy_lifecycle_2") + result = self.manager.on_new_request(ctx) + assert result.policy == OffloadPolicy.REQUEST_LEVEL + assert ctx.req_id in self.manager._request_level_tiers + + # Cleanup + self.manager.on_request_finished(ctx) + assert ctx.req_id not in self.manager._request_level_tiers + + def test_prepare_store_cascades_existing_blocks_to_request_level_tiers( + self, manager_setup + ): + """prepare_store cascades hit blocks to request-level tiers only.""" + # Store some blocks to primary first + existing_blocks = to_keys(range(3)) + result = self.manager.prepare_store(existing_blocks, _CTX) + assert result is not None + self.manager.complete_store(existing_blocks, _CTX, success=True) + # Drain cascade completions + list(self.manager.take_events()) + list(self.manager.take_events()) + + # Make tier1 request-level, tier2 stays block-level + self.secondary_tier1.on_new_request = ( + lambda req_context: RequestOffloadingContext( + policy=OffloadPolicy.REQUEST_LEVEL + ) + ) + + ctx = ReqContext(req_id="req_cascade") + self.manager.on_new_request(ctx) + + # Spy on submit_store + self.secondary_tier1.submit_store = MagicMock( + wraps=self.secondary_tier1.submit_store + ) + self.secondary_tier2.submit_store = MagicMock( + wraps=self.secondary_tier2.submit_store + ) + + # Call prepare_store with existing + new blocks + new_blocks = to_keys(range(3, 5)) + all_blocks = existing_blocks + new_blocks + result = self.manager.prepare_store(all_blocks, ctx) + assert result is not None + assert set(result.keys_to_store) == set(new_blocks) + + # Only tier1 (request-level) should get existing blocks cascaded now. + # New blocks are cascaded to ALL tiers later via complete_store(). + self.secondary_tier1.submit_store.assert_called_once() + job_metadata = self.secondary_tier1.submit_store.call_args.args[0] + assert set(job_metadata.keys) == set(existing_blocks) + + # tier2 (block-level) does not get existing blocks here. + self.secondary_tier2.submit_store.assert_not_called() + class TestTieringOffloadingWithoutSecondaryTiers: """Test TieringOffloadingManager with no secondary tiers (backward compat).""" 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 c7d10a3e39..e336ed94ff 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -29,7 +29,9 @@ from vllm.v1.kv_offload.base import ( OffloadingManager, OffloadingSpec, OffloadKey, + OffloadPolicy, ReqContext, + RequestOffloadingContext, get_offload_block_hash, make_offload_key, ) @@ -172,8 +174,9 @@ class RequestGroupState: class RequestOffloadState: config: SchedulerOffloadConfig req: Request + req_context: ReqContext + offloading_context: RequestOffloadingContext group_states: tuple[RequestGroupState, ...] = field(init=False) - req_context: ReqContext = field(init=False) # upper bound on tokens to offload for this request; None means no cap max_offload_tokens: int | None = None # number of hits in the GPU cache @@ -186,10 +189,6 @@ class RequestOffloadState: self.group_states = tuple( RequestGroupState() for _ in self.config.kv_group_configs ) - self.req_context = ReqContext( - req_id=self.req.request_id, - kv_transfer_params=self.req.kv_transfer_params, - ) params = self.req.kv_transfer_params # NOTE: This field is experimental and subject to change in the future. @@ -248,6 +247,13 @@ class RequestOffloadState: ) +def _create_req_context(req: Request) -> ReqContext: + return ReqContext( + req_id=req.request_id, + kv_transfer_params=req.kv_transfer_params, + ) + + class OffloadingConnectorScheduler: """Implementation of Scheduler side methods""" @@ -512,6 +518,18 @@ class OffloadingConnectorScheduler: return num_hit_tokens + def on_new_request(self, request: Request) -> None: + """Called when a new request is added to the scheduler.""" + req_context = _create_req_context(request) + offloading_context = self.manager.on_new_request(req_context) + req_status = RequestOffloadState( + config=self.config, + req=request, + req_context=req_context, + offloading_context=offloading_context, + ) + self._req_status[request.request_id] = req_status + def get_num_new_matched_tokens( self, request: Request, num_computed_tokens: int ) -> tuple[int | None, bool]: @@ -534,24 +552,15 @@ class OffloadingConnectorScheduler: - `True` if tokens will be loaded asynchronously (between scheduler steps). """ - is_new_request = False - if req_status := self._req_status.get(request.request_id): - # make sure block IDs are cleared - for group_state in req_status.group_states: - group_state.block_ids.clear() - else: - is_new_request = True - req_status = RequestOffloadState(config=self.config, req=request) - self._req_status[request.request_id] = req_status + req_status = self._req_status[request.request_id] + for group_state in req_status.group_states: + group_state.block_ids.clear() req_status.update_offload_keys() req_status.num_locally_computed_tokens = num_computed_tokens num_hit_tokens = self._lookup(req_status) - if is_new_request: - req_status.update_num_hit_blocks( - num_computed_tokens + (num_hit_tokens or 0) - ) + req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0)) self._touch(req_status) @@ -568,9 +577,6 @@ class OffloadingConnectorScheduler: num_locally_computed_tokens = req_status.num_locally_computed_tokens num_cached_tokens = num_locally_computed_tokens + num_external_tokens - params = req_status.req_context.kv_transfer_params - do_remote_decode = params is not None and params.get("do_remote_decode") - keys_to_load: list[OffloadKey] = [] dst_block_ids: list[int] = [] # per group @@ -624,10 +630,10 @@ class OffloadingConnectorScheduler: group_sizes.append(num_pending_gpu_blocks) block_indices.append(num_locally_computed_gpu_blocks) - if not do_remote_decode: - # For P/D prefill requests (do_remote_decode=True), we do - # NOT skip saving the hit prefix, as we need to stream the - # entire KV cache so a remote decode node can consume it. + # Skip prefix-hit blocks for block-level policy; for + # request-level, next_stored_block_idx stays at 0 so all + # blocks (including hits) are offloaded. + if req_status.offloading_context.policy == OffloadPolicy.BLOCK_LEVEL: group_state.next_stored_block_idx = num_blocks # Fence dst blocks against finished-request pending stores. @@ -950,6 +956,12 @@ class OffloadingConnectorScheduler: # TODO(orozery): possibly kickoff offload for last block # which may have been deferred due to async scheduling req_status = self._req_status.get(request.request_id) + + req_context = ( + req_status.req_context if req_status else _create_req_context(request) + ) + self.manager.on_request_finished(req_context) + if req_status is None: return False, None if not req_status.transfer_jobs: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 6c75bda0c4..20888c71f8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -124,6 +124,10 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): return self.connector_worker.build_connector_worker_meta() return None + def on_new_request(self, request: "Request") -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.on_new_request(request) + def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int ) -> tuple[int | None, bool]: diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index de65be1c05..d932fd2d1b 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -7,6 +7,7 @@ Core abstractions for KV cache offloading in vLLM v1. from abc import ABC, abstractmethod from collections.abc import Collection, Iterable, Iterator, Sequence from dataclasses import dataclass +from enum import Enum from typing import TYPE_CHECKING, Any, NewType import numpy as np @@ -49,6 +50,20 @@ class ReqContext: kv_transfer_params: dict[str, Any] | None = None +class OffloadPolicy(Enum): + # Offload only newly-computed blocks as they arrive; prefix-hit + # blocks (already offloaded by a prior request) are skipped. + BLOCK_LEVEL = "block_level" + # Offload all blocks for the request, including prefix hits. + # Used by tiers that need the complete KV context for a request. + REQUEST_LEVEL = "request_level" + + +@dataclass +class RequestOffloadingContext: + policy: OffloadPolicy = OffloadPolicy.BLOCK_LEVEL + + class LoadStoreSpec(ABC): """ Abstract metadata that encapsulates information allowing a worker @@ -210,6 +225,28 @@ class OffloadingManager(ABC): """ return + @abstractmethod + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + """ + Called when a new request is first seen by the scheduler. + + Returns a RequestOffloadingContext indicating how this request's + blocks should be offloaded. + + Args: + req_context: per-request context. + """ + pass + + def on_request_finished(self, req_context: ReqContext) -> None: + """ + Called when a request has finished. + + Args: + req_context: per-request context. + """ + return + def take_events(self) -> Iterable[OffloadingEvent]: """ Take the offloading events from the manager. diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 39e64933b5..4700e250c5 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -11,6 +11,7 @@ from vllm.v1.kv_offload.base import ( OffloadKey, PrepareStoreOutput, ReqContext, + RequestOffloadingContext, ) from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy @@ -94,6 +95,9 @@ class CPUOffloadingManager(OffloadingManager): # --- OffloadingManager interface --- + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + return RequestOffloadingContext() + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: if self.counts is not None: if key in self.counts: diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 7d4d6f031d..b824eec1c0 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING import numpy as np -from vllm.v1.kv_offload.base import OffloadKey, ReqContext +from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext if TYPE_CHECKING: from vllm.v1.kv_offload.base import OffloadingSpec @@ -163,6 +163,28 @@ class SecondaryTierManager(ABC): """ return + @abstractmethod + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + """ + Called when a new request is first seen by the scheduler. + + Returns a RequestOffloadingContext expressing this tier's preference + for how blocks should be offloaded for this request. + + Args: + req_context: Per-request context. + """ + pass + + def on_request_finished(self, req_context: ReqContext) -> None: + """ + Called when a request has finished. + + Args: + req_context: per-request context. + """ + return + def shutdown(self) -> None: """Release resources held by this tier (threads, connections, etc.).""" return diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index 933a8aa8c8..29ec627ca6 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -13,7 +13,7 @@ import logging from collections.abc import Iterable from typing import TYPE_CHECKING -from vllm.v1.kv_offload.base import OffloadKey, ReqContext +from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext from vllm.v1.kv_offload.tiering.base import ( JobMetadata, JobResult, @@ -132,6 +132,9 @@ class ExampleSecondaryTierManager(SecondaryTierManager): self.completed_jobs = [] return result + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + return RequestOffloadingContext() + def get_num_blocks(self) -> int: """Get the number of blocks currently stored in this tier.""" return len(self.blocks) diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index e5921d6ffd..1d37983d55 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -27,6 +27,7 @@ from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.base import ( JobMetadata, JobResult, + RequestOffloadingContext, SecondaryTierManager, ) from vllm.v1.kv_offload.tiering.fs.io import load_block, store_block @@ -100,6 +101,9 @@ class FileSystemTierManager(SecondaryTierManager): thread_name_prefix="vllm_kv_py_fs", ) + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + return RequestOffloadingContext() + def lookup( self, key: OffloadKey, req_context: ReqContext | None = None ) -> bool | None: diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 36fc27e48d..ae8121b4aa 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -20,7 +20,8 @@ Key Design Principles: protecting blocks from eviction until complete_read() is called """ -from collections.abc import Collection, Iterable +from collections import defaultdict +from collections.abc import Collection, Iterable, Sequence from dataclasses import dataclass, field import numpy as np @@ -31,8 +32,10 @@ from vllm.v1.kv_offload.base import ( OffloadingEvent, OffloadingManager, OffloadKey, + OffloadPolicy, PrepareStoreOutput, ReqContext, + RequestOffloadingContext, ) from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager @@ -158,6 +161,13 @@ class TieringOffloadingManager(OffloadingManager): # Reset at the end of each step in take_events(). self._processed_jobs_this_step: bool = False + # Per-request set of secondary tiers that requested REQUEST_LEVEL + # policy. Populated in on_new_request(), + # cleaned up in on_request_finished(). + self._request_level_tiers: defaultdict[str, set[SecondaryTierManager]] = ( + defaultdict(set) + ) + def _next_job_id(self) -> JobId: """Generate a unique job ID for async transfer tracking.""" job_id = self._job_id_counter @@ -387,6 +397,9 @@ class TieringOffloadingManager(OffloadingManager): that any completed async transfers have their ref_cnt decremented before the primary tier makes eviction decisions. + For request-level tiers, blocks already present in the primary tier + are immediately cascaded via submit_store(). + Args: keys: Blocks to prepare for storing. req_context: Per-request context. @@ -400,14 +413,63 @@ class TieringOffloadingManager(OffloadingManager): # successfully transferred to secondary tiers. self._maybe_process_finished_jobs() - # Step 2: Store to primary tier + # Step 2: Store to primary tier (new blocks only). + # Cascading of these newly-stored blocks to ALL secondary tiers + # happens later in complete_store(), after the GPU→Primary transfer + # completes. primary_result = self.primary_tier.prepare_store(keys, req_context) - # Note: Secondary tier cascading will happen in complete_store() - # after the GPU→Primary transfer completes and blocks are ready. + if primary_result is None: + return None + + # Step 3: For request-level tiers, cascade blocks already in primary + request_level_tiers = self._request_level_tiers.get(req_context.req_id) + if request_level_tiers is not None: + keys_to_store_set = set(primary_result.keys_to_store) + keys_already_in_primary = tuple( + k for k in keys if k not in keys_to_store_set + ) + if keys_already_in_primary: + self._cascade_existing_blocks_to_request_level_tiers( + keys_already_in_primary, req_context, request_level_tiers + ) return primary_result + def _cascade_existing_blocks_to_request_level_tiers( + self, + keys: Sequence[OffloadKey], + req_context: ReqContext, + request_level_tiers: set[SecondaryTierManager], + ) -> None: + """ + For tiers that requested request-level policy, submit_store() for + blocks that are already present in the primary tier. + """ + # Filter out keys that are not ready in primary (e.g. in-flight) + ready_keys = tuple( + k for k in keys if self.primary_tier.lookup(k, req_context) is True + ) + if not ready_keys: + return + + for tier in request_level_tiers: + primary_blocks_spec = self.primary_tier.prepare_read( + ready_keys, req_context + ) + + job_id = self._next_job_id() + assert isinstance(primary_blocks_spec, CPULoadStoreSpec) + job_metadata = JobMetadata( + job_id=job_id, + keys=ready_keys, + block_ids=primary_blocks_spec.block_ids, + is_promotion=False, + req_context=req_context, + ) + self._transfer_jobs[job_id] = job_metadata + tier.submit_store(job_metadata) + def complete_store( self, keys: Collection[OffloadKey], @@ -466,6 +528,31 @@ class TieringOffloadingManager(OffloadingManager): # Note: The async transfers are now in flight. Their completion is # tracked via get_finished_jobs() / _maybe_process_finished_jobs(). + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + """ + Query each secondary tier for its offload policy preference. + + Returns REQUEST_LEVEL if ANY secondary tier wants request-level. + Only stores REQUEST_LEVEL tier decisions for use in prepare_store. + """ + for tier in self.secondary_tiers: + tier_ctx = tier.on_new_request(req_context) + if tier_ctx.policy == OffloadPolicy.REQUEST_LEVEL: + self._request_level_tiers[req_context.req_id].add(tier) + + policy = ( + OffloadPolicy.REQUEST_LEVEL + if req_context.req_id in self._request_level_tiers + else OffloadPolicy.BLOCK_LEVEL + ) + return RequestOffloadingContext(policy=policy) + + def on_request_finished(self, req_context: ReqContext) -> None: + self.primary_tier.on_request_finished(req_context) + for tier in self.secondary_tiers: + tier.on_request_finished(req_context) + self._request_level_tiers.pop(req_context.req_id, None) + def take_events(self) -> Iterable[OffloadingEvent]: """ End-of-step hook: flush deferred work, yield events, reset per-step state.