[NIXL] Ignore abort on already-finished request

This situation can occur when the API server receives a client
disconnect (and thus sends an abort) around the same time a prefill
completes and we keep the blocks (delay_free_blocks) around for a
remote decode. We should assume the blocks may be used, and so
we ignore the abort. If they are not used, they should be freed
by the connector after a timeout.

The original error was:

```
[scheduler.py:1183] Finished sending KV transfer for request cmpl-37c560d3-5680-4bd1-97f9-7ed31a56de60-0
  File "/opt/vllm-source/vllm/v1/engine/core.py", line 292, in step
     engine_core_outputs = self.scheduler.update_from_output(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/vllm-source/vllm/v1/core/sched/scheduler.py", line 893, in update_from_output
    self._update_from_kv_xfer_finished(
  File "/opt/vllm-source/vllm/v1/core/sched/scheduler.py", line 1184, in _update_from_kv_xfer_finish>
    self._free_blocks(self.requests[req_id])
                      ~~~~~~~~~~~~~^^^^^^^^

KeyError: 'cmpl-37c560d3-5680-4bd1-97f9-7ed31a56de60-0'
```

But since vllm-project#25844 we would log a warning. This fix makes it so
that situation in `_update_from_kv_xfer_finish()` should never
occur.

Signed-off-by: Mark McLoughlin <markmc@redhat.com>

Author: markmc

tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py

@@ -43,6 +43,7 @@ def test_basic_lifecycle():
     # STEP (1): Prefill.
     # (1a): schedule()
     scheduler_output = scheduler.schedule()
+    assert len(scheduler.requests) == 1
     assert len(scheduler.running) == 1
     assert len(scheduler_output.scheduled_new_reqs) == 1
 
@@ -67,6 +68,7 @@ def test_basic_lifecycle():
     assert len(scheduler.waiting) == 0
 
     # ... but blocks should not be freed.
+    assert len(scheduler.requests) == 1
     blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
         0
     ].req_to_blocks[request_id]
@@ -76,6 +78,7 @@ def test_basic_lifecycle():
     # STEP (2): Send Finished to PB.
     # (2a): schedule() - pass finished request to PB.
     scheduler_output = scheduler.schedule()
+    assert len(scheduler.requests) == 1
     assert len(scheduler.running) == 0
     assert len(scheduler_output.finished_req_ids) == 1
     assert request_id in scheduler_output.finished_req_ids
@@ -92,6 +95,7 @@ def test_basic_lifecycle():
     # STEP (3): Finished sending.
     # (3a): schedule() - pass finished request to PB.
     scheduler_output = scheduler.schedule()
+    assert len(scheduler.requests) == 1
     assert len(scheduler.running) == 0
     assert len(scheduler_output.finished_req_ids) == 0
     assert len(scheduler_output.scheduled_new_reqs) == 0
@@ -133,6 +137,7 @@ def test_short_prompt_lifecycle():
     # STEP (1): Prefill.
     # (1a): schedule()
     scheduler_output = scheduler.schedule()
+    assert len(scheduler.requests) == 1
     assert len(scheduler.running) == 1
     assert len(scheduler_output.scheduled_new_reqs) == 1
 
@@ -178,7 +183,7 @@ def test_prefix_cache_lifecycle():
         reqs=[request_normal], use_eos=True
     )
     scheduler.update_from_output(scheduler_output, model_runner_output)
-    scheduler.schedule()
+    scheduler_output = scheduler.schedule()
     scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
 
     #####################
@@ -213,3 +218,45 @@ def test_prefix_cache_lifecycle():
     )
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert_scheduler_empty(scheduler)
+
+
+def test_abort_during_kv_transfer():
+    """Test aborting request does not release blocks for remote decode."""
+
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+
+    # Prime the KVCache.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_EXTERNAL_FULL_BLOCKS = 2
+    NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
+
+    request = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        do_remote_decode=True,
+    )
+
+    scheduler.add_request(request)
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(reqs=[request])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    scheduler_output = scheduler.schedule()
+    scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
+
+    # Request removed from PB but blocks should not be freed.
+    assert len(scheduler.requests) == 1
+
+    # Abort the request, and check the blocks are still not freed
+    scheduler.finish_requests([request.request_id], RequestStatus.FINISHED_ABORTED)
+    assert len(scheduler.requests) == 1
+
+    # Simulate a finished sending notification
+    scheduler_output = scheduler.schedule()
+    model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    model_runner_output.kv_connector_output = KVConnectorOutput(
+        finished_sending=[request.request_id]
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert_scheduler_empty(scheduler)

vllm/v1/core/sched/scheduler.py

@@ -1191,7 +1191,7 @@ def finish_requests(
         # First pass: collect requests to remove from queues
         for req_id in request_ids:
             request = self.requests.get(req_id)
-            if request is None:
+            if request is None or request.is_finished():
                 # Invalid request ID.
                 continue
 
@@ -1369,14 +1369,8 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):
             self.finished_recving_kv_req_ids.add(req_id)
         for req_id in kv_connector_output.finished_sending or ():
             logger.debug("Finished sending KV transfer for request %s", req_id)
-            if req_id not in self.requests:
-                logger.warning(
-                    "Got finished sending KV transfer for request %s,"
-                    "but the request is already freed.",
-                    req_id,
-                )
-            else:
-                self._free_blocks(self.requests[req_id])
+            assert req_id in self.requests
+            self._free_blocks(self.requests[req_id])
 
     def _update_requests_with_invalid_blocks(
         self, requests: Iterable[Request], invalid_block_ids: set[int]