[Refactor][KV Offloading] Consolidate duplicated cleanup and clamp logic

- Extract _calc_num_offloadable_tokens: eliminates 8-line clamp
  duplication between _build_store_jobs and request_finished
- Extract _drop_store_job: eliminates pop+cleanup+discard teardown
  duplication between build_connector_meta and test harness
- Upgrade _remove_pending_job to safe .get()+.discard() for all
  contexts
- Remove stale TODO(orozery) comment (now implemented)
- Remove dead code: unreachable 'if job_id not in store_jobs' guard
- Merge test_partial_block_not_offloaded into
  test_last_block_offload_at_request_finish via prompt_offset
  parametrization

Signed-off-by: Alex <alex.tech.lab@outlook.com>

Author: Alex-ai-future

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

@@ -2762,20 +2762,25 @@ def capture_fence():
 
 
 @pytest.mark.parametrize("async_scheduling", [True, False])
-def test_last_block_offloaded_on_request_finish(request_runner, async_scheduling: bool):
+@pytest.mark.parametrize("prompt_offset", [-1, -2])
+def test_last_block_offload_at_request_finish(
+    request_runner, async_scheduling: bool, prompt_offset: int
+):
     """The last block that fills at request finish time is offloaded.
 
     Regression test for orozery's TODO: when the finishing token (EOS)
     completes a block, that block's KV must still be offloaded. Before
     schedule time the block wasn't full, so _build_store_jobs skipped it;
     only request_finished has the complete token count.
 
-    Scenario (block_size=4, factor=1):
-    - prompt = 3 tokens (not a full block)
-    - decode EOS -> 4 tokens = 1 full block
+    prompt_offset=-1: prompt has block_size-1 tokens, decode EOS fills
+    the block → store job created → offloaded.
+    prompt_offset=-2: prompt has block_size-2 tokens, decode EOS leaves
+    the block partial → no store job → not offloaded.
     """
     block_size = 4
     block_size_factor = 1
+    should_offload = prompt_offset == -1
 
     runner = request_runner(
         block_size=block_size,
@@ -2787,9 +2792,8 @@ def test_last_block_offloaded_on_request_finish(request_runner, async_scheduling
         generate_store_output(keys)
     )
 
-    runner.new_request(token_ids=[0] * (block_size - 1))
+    runner.new_request(token_ids=[0] * (block_size + prompt_offset))
 
-    # Track _finished_req_store_jobs state during the run
     cs = runner.connector_scheduler
     finished_req_store_jobs_seen = []
 
@@ -2804,49 +2808,20 @@ def patched_update(*args, **kwargs):
 
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored=(0,),
-        drain_finished_req_stores=True,
-    )
-
-    # Verify _finished_req_store_jobs was non-empty at some point
-    non_empty_states = [s for s in finished_req_store_jobs_seen if s]
-    assert non_empty_states, (
-        "_finished_req_store_jobs was never non-empty. "
-        "Expected request_finished to create a store job for the last block. "
-        f"All states seen: {finished_req_store_jobs_seen}"
-    )
-
-    assert runner.manager.complete_store.call_count == 1
-
-
-def test_partial_block_not_offloaded(request_runner):
-    """Tokens that do not form a complete block are not offloaded.
-
-    Scenario:
-    - prompt has block_size - 2 tokens
-    - decode 1 token
-    - total is block_size - 1 tokens, not enough for a full block
-    """
-    block_size = 4
-    block_size_factor = 1
-
-    runner = request_runner(
-        block_size=block_size,
-        num_gpu_blocks=10,
-        async_scheduling=False,
-        block_size_factor=block_size_factor,
-    )
-    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
-        generate_store_output(keys)
-    )
-
-    runner.new_request(token_ids=[0] * (block_size - 2))
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored=(),
+        expected_stored=(0,) if should_offload else (),
+        drain_finished_req_stores=should_offload,
     )
 
-    assert runner.manager.complete_store.call_count == 0
+    if should_offload:
+        non_empty_states = [s for s in finished_req_store_jobs_seen if s]
+        assert non_empty_states, (
+            "_finished_req_store_jobs was never non-empty. "
+            "Expected request_finished to create a store job for the last block. "
+            f"All states seen: {finished_req_store_jobs_seen}"
+        )
+        assert runner.manager.complete_store.call_count == 1
+    else:
+        assert runner.manager.complete_store.call_count == 0
 
 
 @pytest.mark.parametrize("async_scheduling", [True, False])

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -477,25 +477,15 @@ def _discard_finished_req_store_jobs(self):
         """
         cs = self.connector_scheduler
         for job_id in list(cs._finished_req_store_jobs.keys()):
-            job_status = cs._jobs.pop(job_id, None)
+            job_status = cs._drop_store_job(job_id)
             if job_status is not None:
                 req_status = cs._req_status.get(job_status.req_id)
-                if req_status is not None:
-                    req_status.transfer_jobs.discard(job_id)
-                    if not req_status.transfer_jobs and req_status.req.is_finished():
-                        del cs._req_status[job_status.req_id]
-                for bid in job_status.non_sliding_window_block_ids or ():
-                    pending = cs._block_id_to_pending_jobs.get(bid)
-                    if pending:
-                        pending.discard(job_id)
-                        if not pending:
-                            del cs._block_id_to_pending_jobs[bid]
-                for bid in job_status.sliding_window_block_ids or ():
-                    pending = cs._block_id_to_pending_jobs.get(bid)
-                    if pending:
-                        pending.discard(job_id)
-                        if not pending:
-                            del cs._block_id_to_pending_jobs[bid]
+                if (
+                    req_status is not None
+                    and not req_status.transfer_jobs
+                    and req_status.req.is_finished()
+                ):
+                    del cs._req_status[job_status.req_id]
         cs._finished_req_store_jobs.clear()
 
     def _run(

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

@@ -430,10 +430,33 @@ def _generate_job_id(self) -> int:
 
     def _remove_pending_job(self, job_id: int, block_ids: list[int] | None) -> None:
         for bid in block_ids or ():
-            pending = self._block_id_to_pending_jobs[bid]
-            pending.remove(job_id)
-            if not pending:
-                del self._block_id_to_pending_jobs[bid]
+            pending = self._block_id_to_pending_jobs.get(bid)
+            if pending is not None:
+                pending.discard(job_id)
+                if not pending:
+                    del self._block_id_to_pending_jobs[bid]
+
+    def _calc_num_offloadable_tokens(
+        self, req_status: RequestOffloadState, num_computed_tokens: int
+    ) -> int:
+        num = min(num_computed_tokens, req_status.req.num_tokens)
+        max_offload_tokens = req_status.max_offload_tokens
+        if max_offload_tokens is not None:
+            num = min(num, max_offload_tokens)
+        if self.config.offload_prompt_only:
+            num = min(num, req_status.req.num_prompt_tokens)
+        return num
+
+    def _drop_store_job(self, job_id: int) -> TransferJobStatus | None:
+        job = self._jobs.pop(job_id, None)
+        if job is None:
+            return None
+        self._remove_pending_job(job_id, job.sliding_window_block_ids)
+        self._remove_pending_job(job_id, job.non_sliding_window_block_ids)
+        req_status = self._req_status.get(job.req_id)
+        if req_status is not None:
+            req_status.transfer_jobs.discard(job_id)
+        return job
 
     def _maximal_prefix_lookup(
         self, keys: Iterable[OffloadKey], req_context: ReqContext
@@ -1079,20 +1102,9 @@ def _build_store_jobs(
 
             num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
             num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens
-            # with async scheduling, some tokens may be missing
-            num_offloadable_tokens = min(num_tokens_after_batch, req.num_tokens)
-            max_offload_tokens = req_status.max_offload_tokens
-            if max_offload_tokens is not None:
-                num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens)
-
-            # Skip decode-phase blocks: clamp to the prompt length so only
-            # prefill (prompt) blocks become eligible for store. next_stored_idx
-            # never advances past this boundary, so decode blocks are never
-            # queued in this or any later step.
-            if self.config.offload_prompt_only:
-                num_offloadable_tokens = min(
-                    num_offloadable_tokens, req.num_prompt_tokens
-                )
+            num_offloadable_tokens = self._calc_num_offloadable_tokens(
+                req_status, num_tokens_after_batch
+            )
 
             self._build_store_jobs_for_req(
                 req_status, num_offloadable_tokens, store_jobs
@@ -1151,8 +1163,6 @@ def build_connector_meta(
         if self._current_batch_allocated_block_ids:
             jobs_to_drop = set()
             for job_id in finished_req_job_ids:
-                if job_id not in store_jobs:
-                    continue  # Already removed or not present
                 job = self._jobs.get(job_id)
                 if job is not None and job.is_store:
                     job_blocks = set(job.non_sliding_window_block_ids or [])
@@ -1162,24 +1172,7 @@ def build_connector_meta(
             for job_id in jobs_to_drop:
                 del store_jobs[job_id]
                 self._current_batch_jobs_to_flush.discard(job_id)
-                job = self._jobs.pop(job_id, None)
-                if job is not None:
-                    for bid in job.sliding_window_block_ids or []:
-                        pjs = self._block_id_to_pending_jobs.get(bid)
-                        if pjs is not None:
-                            pjs.discard(job_id)
-                            if not pjs:
-                                del self._block_id_to_pending_jobs[bid]
-                    for bid in job.non_sliding_window_block_ids or []:
-                        pjs = self._block_id_to_pending_jobs.get(bid)
-                        if pjs is not None:
-                            pjs.discard(job_id)
-                            if not pjs:
-                                del self._block_id_to_pending_jobs[bid]
-                    # Also remove from req_status.transfer_jobs
-                    req_status = self._req_status.get(job.req_id)
-                    if req_status is not None:
-                        req_status.transfer_jobs.discard(job_id)
+                self._drop_store_job(job_id)
 
         meta = OffloadingConnectorMetadata(
             load_jobs=self._current_batch_load_jobs,
@@ -1308,8 +1301,6 @@ def request_finished(
             Optional KVTransferParams to be included in the request outputs
             returned by the engine.
         """
-        # 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)
 
         if req_status is None:
@@ -1329,14 +1320,9 @@ def request_finished(
         # append_output_token_ids → update_block_hashes, but offload_keys
         # is only refreshed in _update_req_states / build_connector_meta.
         req_status.update_offload_keys()
-        num_offloadable_tokens = request.num_tokens
-        if self.config.offload_prompt_only:
-            num_offloadable_tokens = min(
-                num_offloadable_tokens, request.num_prompt_tokens
-            )
-        max_offload_tokens = req_status.max_offload_tokens
-        if max_offload_tokens is not None:
-            num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens)
+        num_offloadable_tokens = self._calc_num_offloadable_tokens(
+            req_status, request.num_tokens
+        )
         self._build_store_jobs_for_req(
             req_status, num_offloadable_tokens, self._finished_req_store_jobs
         )