[Bugfix] Reclaim unconsumed SHM segments and sender state on request abort

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -2327,6 +2327,140 @@ def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapt
     assert request.request_id not in adapter.requests_origin_status
 
 
+def test_finish_requests_enqueues_connector_cleanup_for_aborts(build_adapter):
+    """Aborting a live request must sweep its connector-side resources
+    (e.g. unconsumed SHM segments — see the /dev/shm leak in
+    vllm-project/vllm-omni C1). The sweep is enqueued on the save queue —
+    not run inline — so it executes on the save thread strictly after any
+    ``put()`` already queued for the same request.
+    """
+    adapter, connector = build_adapter(stage_id=1)
+    request = _req("req-abort", RequestStatus.RUNNING, external_req_id="ext-abort")
+    adapter.put_req_chunk["ext-abort"] = 3
+    adapter.request_payload["ext-abort"] = object()
+
+    adapter.finish_requests(["req-abort"], RequestStatus.FINISHED_ABORTED, {"req-abort": request})
+
+    tasks = list(adapter._pending_save_reqs)
+    assert len(tasks) == 1
+    assert tasks[0] == {"connector_cleanup": "ext-abort"}
+
+    adapter._send_single_request(tasks[0])
+    connector.cleanup.assert_called_once_with("ext-abort")
+    assert "ext-abort" not in adapter.put_req_chunk
+    assert "ext-abort" not in adapter.request_payload
+
+
+def test_finish_requests_connector_cleanup_ordered_after_queued_saves(build_adapter):
+    """FIFO guarantee: a save enqueued before the abort must be processed
+    before the connector sweep, otherwise the late chunk would be written
+    after its segment sweep and leak.
+    """
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-fifo", RequestStatus.RUNNING, external_req_id="ext-fifo")
+    adapter.save_async(multimodal_output=None, request=request)
+
+    adapter.finish_requests(["req-fifo"], RequestStatus.FINISHED_ABORTED, {"req-fifo": request})
+
+    kinds = ["cleanup" if "connector_cleanup" in task else "save" for task in adapter._pending_save_reqs]
+    assert kinds == ["save", "cleanup"]
+
+
+def test_connector_cleanup_does_not_retire_successor_generation(build_adapter):
+    """``finish_requests`` retires the aborted sender generation inline; the
+    queued sweep runs later on the save thread. A request that reuses the
+    external id in between owns a fresh generation, and the sweep must only
+    unlink connector segments -- not cancel that successor's sends.
+    """
+    adapter, connector = build_adapter(stage_id=1)
+    aborted = _req("req-old", RequestStatus.RUNNING, external_req_id="ext-reused")
+    adapter.finish_requests([aborted.request_id], RequestStatus.FINISHED_ABORTED, {aborted.request_id: aborted})
+    sweep = adapter._pending_save_reqs.popleft()
+    assert sweep == {"connector_cleanup": "ext-reused"}
+
+    successor = _req("req-new", RequestStatus.RUNNING, external_req_id="ext-reused")
+    adapter.save_async(multimodal_output=None, request=successor)
+    successor_token = adapter._sender_tokens["ext-reused"]
+
+    adapter._send_single_request(sweep)
+
+    connector.cleanup.assert_called_once_with("ext-reused")
+    assert adapter._sender_tokens.get("ext-reused") is successor_token
+    assert not successor_token.cancelled
+
+
+def test_finish_requests_skips_connector_cleanup_for_finished_or_unknown(build_adapter):
+    """A request that already finished naturally may have an unconsumed
+    terminal chunk downstream; its segments must NOT be swept, or the next
+    stage waits for its finish marker forever. Unknown ids are skipped too.
+    """
+    adapter, _ = build_adapter(stage_id=1)
+    finished = _req("req-done", RequestStatus.FINISHED_STOPPED, external_req_id="ext-done")
+
+    adapter.finish_requests(
+        ["req-done", "req-ghost"],
+        RequestStatus.FINISHED_ABORTED,
+        {"req-done": finished},
+    )
+
+    assert len(adapter._pending_save_reqs) == 0
+
+
+def test_finish_requests_skips_connector_cleanup_for_non_abort_status(build_adapter):
+    """The sweep is gated on ``finished_status``, not on "request is still
+    live". A live request finished as FINISHED_STOPPED (e.g. a resumable
+    segment stop) keeps an unconsumed terminal chunk downstream; unlinking
+    it would strand the next stage waiting for a finish marker.
+    """
+    adapter, connector = build_adapter(stage_id=1)
+    request = _req("req-stop", RequestStatus.RUNNING, external_req_id="ext-stop")
+
+    adapter.finish_requests(["req-stop"], RequestStatus.FINISHED_STOPPED, {"req-stop": request})
+
+    assert len(adapter._pending_save_reqs) == 0
+    connector.cleanup.assert_not_called()
+
+
+@pytest.mark.parametrize("finished_status", [RequestStatus.FINISHED_ABORTED, RequestStatus.FINISHED_ERROR])
+def test_finish_requests_sweeps_sender_only_stage_zero(build_adapter, finished_status):
+    """Stage 0 is sender-only: ``process_pending_chunks`` returns early, so
+    the request never lands in a chunk deque and no terminal chunk is ever
+    sent on abort. ``finish_requests`` is the only place its segments can
+    be swept. Regression guard for the talker-side leak in #6352.
+    """
+    adapter, connector = build_adapter(stage_id=0)
+    request = _req("req-talker", RequestStatus.RUNNING, external_req_id="ext-talker")
+    adapter.put_req_chunk["ext-talker"] = 7
+    adapter.requests_num_chunks_sent["ext-talker"] = 7
+
+    adapter.finish_requests(["req-talker"], finished_status, {"req-talker": request})
+
+    tasks = list(adapter._pending_save_reqs)
+    assert tasks == [{"connector_cleanup": "ext-talker"}]
+
+    adapter._send_single_request(tasks[0])
+    connector.cleanup.assert_called_once_with("ext-talker")
+    assert "ext-talker" not in adapter.put_req_chunk
+    assert "ext-talker" not in adapter.requests_num_chunks_sent
+
+
+def test_terminal_send_does_not_sweep_connector_state(build_adapter):
+    """Natural finish: the terminal chunk just sent must stay readable for
+    the downstream consumer — ``connector.cleanup`` must not run on this
+    path (only sender-side dict state is reclaimed).
+    """
+    adapter, connector = build_adapter(stage_id=1)
+    request = _req("req-nat", RequestStatus.FINISHED_STOPPED, external_req_id="ext-nat")
+
+    adapter._send_single_request(
+        {"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
+    )
+
+    assert connector.put.called
+    connector.cleanup.assert_not_called()
+    assert "ext-nat" not in adapter.put_req_chunk
+
+
 def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_adapter):
     adapter, _ = build_adapter(stage_id=1)
     zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)

tests/distributed/omni_connectors/test_shm_connector.py

@@ -208,23 +208,24 @@ def test_sender_writes_multiple_receiver_keys(self, connector):
 
 class TestCleanup:
     def test_cleanup_removes_unconsumed_segment(self, connector):
-        data = {"leak": True}
-        connector.put("s0", "s1", "cleanup_req_42", data)
-        assert "cleanup_req_42" in connector._pending_keys
+        """KV key shape ``omni_{from}_to_{to}_kv_cache_{request_id}``."""
+        key = "omni_s0_to_s1_kv_cache_req_42"
+        connector.put("s0", "s1", key, {"leak": True})
+        assert key in connector._pending_keys
 
         connector.cleanup("req_42")
-        assert "cleanup_req_42" not in connector._pending_keys
+        assert key not in connector._pending_keys
 
-        result = connector.get("s0", "s1", "cleanup_req_42", metadata=None)
+        result = connector.get("s0", "s1", key, metadata=None)
         assert result is None
 
     def test_cleanup_noop_for_consumed_segment(self, connector):
-        data = {"consumed": True}
-        connector.put("s0", "s1", "consumed_req_99", data)
-        connector.get("s0", "s1", "consumed_req_99", metadata=None)
+        key = "omni_s0_to_s1_kv_cache_req_99"
+        connector.put("s0", "s1", key, {"consumed": True})
+        connector.get("s0", "s1", key, metadata=None)
 
         connector.cleanup("req_99")
-        assert "consumed_req_99" not in connector._pending_keys
+        assert key not in connector._pending_keys
 
     def test_close_cleans_all_pending(self, connector):
         for i in range(3):
@@ -233,3 +234,81 @@ def test_close_cleans_all_pending(self, connector):
         assert len(connector._pending_keys) == 3
         connector.close()
         assert len(connector._pending_keys) == 0
+
+    def test_cleanup_sweeps_chunk_keys_and_lock_files(self, connector):
+        """Abort sweep: chunk-style keys ``{ext_req_id}_{stage}_{chunk}`` are
+        matched by external request id; both the SHM segments and their
+        lock files must be removed, leaving other requests untouched.
+        """
+        for chunk in range(2):
+            ok, _, _ = connector.put("s0", "s1", f"req-c1_0_{chunk}", {"chunk": chunk})
+            assert ok
+        ok, _, _ = connector.put("s0", "s1", "other-req_0_0", {"keep": True})
+        assert ok
+
+        connector.cleanup("req-c1")
+
+        for chunk in range(2):
+            assert connector.get("s0", "s1", f"req-c1_0_{chunk}", metadata=None) is None
+            assert not os.path.exists(f"/dev/shm/shm_req-c1_0_{chunk}_lockfile.lock")
+        assert not any(k.startswith("req-c1_") for k in connector._pending_keys)
+
+        result = connector.get("s0", "s1", "other-req_0_0", metadata=None)
+        assert result is not None
+        assert result[0] == {"keep": True}
+
+    def test_cleanup_does_not_touch_id_prefixed_sibling(self, connector):
+        """``abc`` must not sweep ``abc_def``. A bare ``request_id + "_"``
+        prefix match would unlink a live sibling's chunks, stranding its
+        consumer on a chunk that no longer exists — strictly worse than
+        the leak this sweep fixes.
+        """
+        connector.put("s0", "s1", "abc_0_0", {"mine": True})
+        connector.put("s0", "s1", "abc_def_0_0", {"sibling": True})
+
+        connector.cleanup("abc")
+
+        assert "abc_0_0" not in connector._pending_keys
+        assert "abc_def_0_0" in connector._pending_keys
+        result = connector.get("s0", "s1", "abc_def_0_0", metadata=None)
+        assert result is not None
+        assert result[0] == {"sibling": True}
+
+    def test_cleanup_does_not_match_bare_id_suffix(self, connector):
+        """Numeric request ids are common, and a bare ``"_" + request_id``
+        suffix match would make ``cleanup("0")`` unlink chunk 0 of every
+        request in flight.
+        """
+        connector.put("s0", "s1", "other-req_0_0", {"keep": True})
+        connector.put("s0", "s1", "0_0_0", {"mine": True})
+
+        connector.cleanup("0")
+
+        assert "0_0_0" not in connector._pending_keys
+        assert "other-req_0_0" in connector._pending_keys
+        result = connector.get("s0", "s1", "other-req_0_0", metadata=None)
+        assert result is not None
+        assert result[0] == {"keep": True}
+
+    def test_cleanup_sweeps_rank_aware_kv_keys(self, connector):
+        """Rank-aware KV shape ``{req}_{from_stage}_{chunk}_{from}_{to}``
+        stays covered; stage names need not be numeric.
+        """
+        connector.put("s0", "s1", "req-tp_s0_0_0_1", {"rank": 1})
+        connector.put("s0", "s1", "req-tp_s0_0_1_0", {"rank": 0})
+
+        connector.cleanup("req-tp")
+
+        assert not any(k.startswith("req-tp") for k in connector._pending_keys)
+
+    def test_get_removes_lock_file_even_for_falsy_payload(self, connector):
+        """The lock file must be tied to segment consumption, not payload
+        truthiness — an empty-dict payload previously leaked its lock file.
+        """
+        ok, _, _ = connector.put("s0", "s1", "falsy_key_1", {})
+        assert ok
+
+        result = connector.get("s0", "s1", "falsy_key_1", metadata=None)
+        assert result is not None
+        assert result[0] == {}
+        assert not os.path.exists("/dev/shm/shm_falsy_key_1_lockfile.lock")

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -71,23 +71,24 @@ def put(
             return False, 0, None
 
     def _get_data_with_lock(self, lock_file: str, shm_handle: dict[str, Any]) -> tuple[Any, int] | None:
-        deserialized = False
+        consumed = False
         try:
             with open(lock_file, "rb+") as lockf:
                 fcntl.flock(lockf, fcntl.LOCK_EX | fcntl.LOCK_NB)
                 data_bytes = shm_read_bytes(shm_handle)
                 fcntl.flock(lockf, fcntl.LOCK_UN)
+            # The segment is already unlinked and no retry can succeed, so
+            # the lock file must go even if deserialization then fails.
+            consumed = True
             obj = self.deserialize_obj(data_bytes)
-            result = (obj, int(shm_handle.get("size", 0)))
-            deserialized = True
-            return result
+            return obj, int(shm_handle.get("size", 0))
         except BlockingIOError:
             return None
         except Exception as e:
             logger.error(f"SharedMemoryConnector shm get failed for req : {e}")
             return None
         finally:
-            if deserialized:
+            if consumed:
                 try:
                     os.remove(lock_file)
                 except FileNotFoundError:
@@ -150,19 +151,37 @@ def get(
             self._metrics["gets"] += 1
         return result
 
+    @staticmethod
+    def _key_belongs_to(key: str, request_id: str) -> bool:
+        """Whether *key* is one of *request_id*'s own connector keys.
+
+        Shapes: ``{id}``, ``{id}_{stage}_{chunk}``,
+        ``{id}_{stage}_{chunk}_{from_rank}_{to_rank}``, and
+        ``omni_{from}_to_{to}_kv_cache_{id}``. Stage names may be non-numeric.
+
+        Structural rather than a bare prefix/suffix match: ``abc`` must not
+        unlink live ``abc_def`` segments, nor ``0`` chunk 0 of every request
+        in flight. Unlinking a live segment strands its consumer.
+        """
+        if key == request_id:
+            return True
+        if key.endswith(f"_kv_cache_{request_id}"):
+            return True
+        prefix = f"{request_id}_"
+        if not key.startswith(prefix):
+            return False
+        fields = key[len(prefix) :].split("_")
+        # Two trailing fields for a chunk key, four for a rank-aware KV key.
+        return len(fields) in (2, 4) and fields[-1].isdecimal()
+
     def cleanup(self, request_id: str) -> None:
-        """Best-effort cleanup of unconsumed SHM segments for *request_id*.
+        """Best-effort unlink of *request_id*'s unconsumed segments + lock files.
 
-        Matches pending keys where *request_id* appears as the full key,
-        as a ``_``-delimited prefix, or as a ``_``-delimited suffix.
-        If ``get()`` was never called, we unlink it here so /dev/shm
-        doesn't leak.
+        An aborted request has no downstream consumer left to call ``get()``,
+        so without this sweep its segments hold /dev/shm until process exit.
         """
-        stale = [
-            k
-            for k in self._pending_keys
-            if k == request_id or k.startswith(request_id + "_") or k.endswith("_" + request_id)
-        ]
+        # Snapshot: get() discards keys from the recv thread mid-iteration.
+        stale = [k for k in list(self._pending_keys) if self._key_belongs_to(k, request_id)]
         for key in stale:
             self._pending_keys.discard(key)
             try: