From 4bfa0f2b1458be320fa39c6fa54be5f83cef2444 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Thu, 28 May 2026 19:00:18 +0300 Subject: [PATCH] [KV Offload] Rename `SecondaryTierManager.get_finished()` to `get_finished_jobs()` (#43870) Signed-off-by: Ronen Schaffer --- tests/v1/kv_offload/test_fs_tier.py | 4 ++-- tests/v1/kv_offload/test_tiering_offloading.py | 2 +- vllm/v1/kv_offload/tiering/base.py | 8 ++++---- vllm/v1/kv_offload/tiering/example/manager.py | 4 ++-- vllm/v1/kv_offload/tiering/fs/manager.py | 4 ++-- vllm/v1/kv_offload/tiering/manager.py | 8 ++++---- 6 files changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/v1/kv_offload/test_fs_tier.py b/tests/v1/kv_offload/test_fs_tier.py index e9e999449d0..70ecc7943bd 100644 --- a/tests/v1/kv_offload/test_fs_tier.py +++ b/tests/v1/kv_offload/test_fs_tier.py @@ -72,14 +72,14 @@ def make_job( def drain(tier: FileSystemTierManager, max_rounds: int = 40) -> list: """ - Call get_finished() repeatedly until no new results arrive for 5 + Call get_finished_jobs() repeatedly until no new results arrive for 5 consecutive rounds or max_rounds is reached. """ results = [] idle = 0 for _ in range(max_rounds): time.sleep(0.01) - new = list(tier.get_finished()) + new = list(tier.get_finished_jobs()) results.extend(new) if new: idle = 0 diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/test_tiering_offloading.py index 61f7cdb7606..7359a39384f 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/test_tiering_offloading.py @@ -184,7 +184,7 @@ class TestTieringOffloadingManager: # End of step 1: _maybe_process_finished_jobs() was already called by # prepare_store() above (setting the per-step flag), so take_events() - # does NOT poll get_finished() again — cascade completions remain + # does NOT poll get_finished_jobs() again — cascade completions remain # unprocessed until the next step. list(self.manager.take_events()) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 8014ac9b0ce..7d4d6f031da 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -50,7 +50,7 @@ class SecondaryTierManager(ABC): IMPORTANT: All methods run in the Scheduler process and must be lightweight and non-blocking. submit_load() and submit_store() submit - async jobs; get_finished() polls for completion. + async jobs; get_finished_jobs() polls for completion. """ def __init__( @@ -106,7 +106,7 @@ class SecondaryTierManager(ABC): 3. Allocating space in this tier 4. Submitting the async transfer (read from primary via block_ids) - Report completion via ``get_finished()``. + Report completion via ``get_finished_jobs()``. Args: job_metadata: Job metadata including job_id, keys, and block_ids @@ -131,7 +131,7 @@ class SecondaryTierManager(ABC): The implementation must copy data from this tier into the primary-tier slots identified by ``block_ids``. - Report completion via ``get_finished()``. + Report completion via ``get_finished_jobs()``. Args: job_metadata: Job metadata including job_id, keys, and block_ids @@ -140,7 +140,7 @@ class SecondaryTierManager(ABC): pass @abstractmethod - def get_finished(self) -> Iterable[JobResult]: + def get_finished_jobs(self) -> Iterable[JobResult]: """ Return all jobs (loads and stores) that completed since the last call. diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index 65d519e46bb..933a8aa8c83 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -61,7 +61,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): # key -> True (only care about presence) self.blocks: dict[OffloadKey, bool] = {} - # Completed jobs waiting to be retrieved by get_finished() + # Completed jobs waiting to be retrieved by get_finished_jobs() self.completed_jobs: list[JobResult] = [] def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: @@ -120,7 +120,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) - def get_finished(self) -> Iterable[JobResult]: + def get_finished_jobs(self) -> Iterable[JobResult]: """ Poll for finished jobs. diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 25318b760d9..e5921d6ffd8 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -47,7 +47,7 @@ class FileSystemTierManager(SecondaryTierManager): queue, so neither starves. submit_store / submit_load are non-blocking: they enqueue tasks and return. - get_finished() polls job completion and returns completed JobResults. + get_finished_jobs() polls job completion and returns completed JobResults. """ @@ -131,7 +131,7 @@ class FileSystemTierManager(SecondaryTierManager): ) self._pool.enqueue_load(job_metadata.job_id, len(job_metadata.keys), tasks) - def get_finished(self) -> Iterable[JobResult]: + def get_finished_jobs(self) -> Iterable[JobResult]: """ Collect completed jobs from the finished-jobs queue. """ diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 4a2d79254ea..36fc27e48de 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -113,7 +113,7 @@ class TieringOffloadingManager(OffloadingManager): Key internal state: - Minimal state tracking; relies on secondary tiers to report completion - via get_finished() + via get_finished_jobs() - Secondary tiers return JobResult objects containing all necessary information - job_id_counter: monotonically increasing counter for job IDs @@ -182,14 +182,14 @@ class TieringOffloadingManager(OffloadingManager): Unconditionally poll all secondary tiers for completed jobs. This method: - 1. Calls get_finished() on each secondary tier + 1. Calls get_finished_jobs() on each secondary tier 2. For completed stores (primary→secondary): calls primary.complete_read() to decrement ref_cnt 3. For completed loads (secondary→primary): calls primary.complete_write() to make blocks available """ for i, tier in enumerate(self.secondary_tiers): - for completed_job in tier.get_finished(): + for completed_job in tier.get_finished_jobs(): job_id = completed_job.job_id job_metadata = self._transfer_jobs.pop(job_id, None) assert job_metadata is not None, ( @@ -464,7 +464,7 @@ class TieringOffloadingManager(OffloadingManager): tier.submit_store(job_metadata) # Note: The async transfers are now in flight. Their completion is - # tracked via get_finished() / _maybe_process_finished_jobs(). + # tracked via get_finished_jobs() / _maybe_process_finished_jobs(). def take_events(self) -> Iterable[OffloadingEvent]: """