[Bugfix] Reclaim aborted chunk SHM and consumed connector state (#8082)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Gagan Dhakrey <gagandhakrey@gmail.com>

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import threading
+import uuid
 from collections import deque
 from types import MethodType, SimpleNamespace
 from unittest.mock import Mock, call, patch
@@ -18,6 +19,7 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
 from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt
+from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
 from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
@@ -27,6 +29,129 @@ from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.fixture
+def shm_sender(build_adapter):
+    adapter, _ = build_adapter(stage_id=0)
+    connector = SharedMemoryConnector({"stage_id": 0})
+    adapter.connector = connector
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
+        codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
+    )
+    yield adapter, connector
+    connector.close()
+
+
+def test_shm_abort_cleans_only_owned_chunks_before_id_reuse(shm_sender):
+    adapter, connector = shm_sender
+    ext_id = f"abort_{uuid.uuid4().hex}"
+    request = _req("old", RequestStatus.RUNNING, external_req_id=ext_id)
+    sibling_key = f"{ext_id}_1_2_0_0"
+    assert connector.put("0", "1", sibling_key, "sibling")[0]
+    adapter.save_async(None, request)
+    adapter._send_single_request(adapter._pending_save_reqs.popleft())
+
+    adapter.finish_requests([request.request_id], RequestStatus.FINISHED_ABORTED, {request.request_id: request})
+    # Queue a replacement before cleanup runs. The same save loop must reclaim
+    # the old generation before its key can be written by the new generation.
+    replacement = _req("new", RequestStatus.RUNNING, external_req_id=ext_id)
+    adapter.save_async(None, replacement)
+    adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    assert connector.get("0", "1", f"{ext_id}_0_0") is None
+    adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    assert connector.get("0", "1", f"{ext_id}_0_0") is not None
+    assert connector.get("0", "1", sibling_key)[0] == "sibling"
+
+
+def test_shm_abort_reclaims_inflight_put_without_blocking(shm_sender, monkeypatch):
+    adapter, connector = shm_sender
+    request = _req("inflight", RequestStatus.RUNNING, external_req_id=f"inflight_{uuid.uuid4().hex}")
+    put_started = threading.Event()
+    release_put = threading.Event()
+    real_put = connector.put
+
+    def blocking_put(**kwargs):
+        put_started.set()
+        assert release_put.wait(timeout=5)
+        return real_put(**kwargs)
+
+    monkeypatch.setattr(connector, "put", blocking_put)
+    adapter.save_async(None, request)
+    task = adapter._pending_save_reqs.popleft()
+    sender = threading.Thread(target=adapter._send_single_request, args=(task,))
+    sender.start()
+    try:
+        assert put_started.wait(timeout=5)
+        adapter.finish_requests([request.request_id], RequestStatus.FINISHED_ABORTED, {request.request_id: request})
+        assert sender.is_alive()
+    finally:
+        release_put.set()
+        sender.join(timeout=5)
+    assert not sender.is_alive()
+    while adapter._pending_save_reqs:
+        adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    assert not connector._pending_keys
+    assert connector.get("0", "1", f"{request.external_req_id}_0_0") is None
+    assert not adapter._sender_tokens
+    assert not adapter.code_prompt_token_ids
+
+
+@pytest.mark.parametrize("fail_terminal", [False, True])
+def test_shm_terminal_keeps_payload_only_after_success(shm_sender, monkeypatch, fail_terminal):
+    adapter, connector = shm_sender
+    request = _req("terminal", RequestStatus.FINISHED_STOPPED, external_req_id=f"terminal_{uuid.uuid4().hex}")
+    if fail_terminal:
+        real_put = connector.put
+
+        def failed_put(**kwargs):
+            real_put(**kwargs)
+            return False, 0, None
+
+        monkeypatch.setattr(connector, "put", failed_put)
+    adapter.save_async(None, request)
+    adapter.cleanup_sender(request.external_req_id)
+    while adapter._pending_save_reqs:
+        adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    result = connector.get("0", "1", f"{request.external_req_id}_0_0")
+    assert (result is None) == fail_terminal
+    assert not adapter._sender_tokens
+
+
+@pytest.mark.parametrize("stage_id", [0, 1])
+def test_receiver_cleanup_does_not_accumulate_tombstones(build_adapter, stage_id):
+    adapter, _ = build_adapter(stage_id=stage_id)
+    for index in range(256):
+        adapter.cleanup_receiver(f"completed-{index}")
+    assert not getattr(adapter, "_cancelled_load_reqs", ())
+    assert not adapter._registered_load_entries
+
+
+def test_recv_loop_does_not_recreate_mapping_for_cancelled_entry(build_adapter):
+    adapter, connector = build_adapter(stage_id=1)
+    request = _req("cancelled", RequestStatus.RUNNING, external_req_id="external")
+    adapter.load_async(request)
+    adapter.cleanup_receiver(request.request_id)
+    # Stop after the loop has drained the stale queue entry.
+    adapter._recv_cond.wait = lambda **kwargs: adapter.stop_event.set()
+    adapter.recv_loop()
+    connector.get.assert_not_called()
+    assert not adapter.request_ids_mapping
+    assert not adapter._pending_load_reqs
+
+
+def test_idle_save_loop_reaps_consumed_shm(shm_sender):
+    adapter, connector = shm_sender
+    key = f"idle_{uuid.uuid4().hex}"
+    assert connector.put("0", "1", key, "payload")[0]
+    receiver = SharedMemoryConnector({})
+    try:
+        assert receiver.get("0", "1", key)[0] == "payload"
+        adapter._save_cond.wait = lambda **kwargs: adapter.stop_event.set()
+        adapter.save_loop()
+        assert not connector._pending_keys
+    finally:
+        receiver.close()
+
+
 class DummyWaitingQueue(list):
     def prepend_requests(self, requests):
         self[:0] = list(requests)
@@ -583,7 +708,7 @@ def test_load_poll_uses_request_scoped_payload_sender_endpoint(build_adapter):
     connector.get.assert_called_once_with(
         "1",
         "2",
-        "req-1_1_0",
+        "external-1_1_0",
         {"source_host": "10.0.0.1", "source_port": 50051},
     )
 
@@ -1910,7 +2035,7 @@ def test_cleanup_clears_all_state(build_adapter):
     assert req_id not in adapter.get_req_chunk
     assert req_id not in adapter.requests_with_ready_chunks
     assert req_id not in adapter.request_ids_mapping
-    assert req_id in adapter._cancelled_load_reqs
+    assert req_id not in adapter._registered_load_entries
     assert req_id not in adapter._finished_load_reqs
     assert req_id not in adapter._pending_ar_prompt_updates
     assert req_id not in adapter._streaming_condition_lengths
@@ -3229,12 +3354,11 @@ def test_process_pending_chunks_purges_zombies_in_running_deque(
     # 2. Live request is still in the deque.
     assert live_req in adapter.waiting_for_chunk_running_requests
     # 3. ``cleanup_receiver`` ran for the zombie (drops origin-status mapping
-    #    and registers the id as cancelled so a late load/poll is dropped too).
+    #    and unregisters it so a late load/poll is dropped too).
     assert zombie_req.request_id not in adapter.requests_origin_status
-    assert zombie_req.request_id in adapter._cancelled_load_reqs
+    assert zombie_req.request_id not in adapter._registered_load_entries
     # 4. Live request's bookkeeping is untouched.
     assert adapter.requests_origin_status[live_req.request_id] == RequestStatus.RUNNING
-    assert live_req.request_id not in adapter._cancelled_load_reqs
 
     # 5. ``restore_queues`` (which the scheduler runs in its ``finally``
     #    clause) now only re-injects the live request -- the zombie is
@@ -3267,7 +3391,7 @@ def test_process_pending_chunks_purges_zombies_in_waiting_deque(build_adapter):
 
     assert live_req in adapter.waiting_for_chunk_waiting_requests
     assert zombie_req not in adapter.waiting_for_chunk_waiting_requests
-    assert zombie_req.request_id in adapter._cancelled_load_reqs
+    assert zombie_req.request_id not in adapter._registered_load_entries
 
     adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
     assert waiting_queue == [live_req]
@@ -3295,8 +3419,8 @@ def test_purge_preserves_live_order_with_interleaved_zombies(build_adapter):
     adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
 
     assert list(adapter.waiting_for_chunk_running_requests) == [live1, live2]
-    assert zombie1.request_id in adapter._cancelled_load_reqs
-    assert zombie2.request_id in adapter._cancelled_load_reqs
+    assert zombie1.request_id not in adapter._registered_load_entries
+    assert zombie2.request_id not in adapter._registered_load_entries
 
 
 def test_restore_queues_purges_late_aborts_after_process_pending_chunks(
@@ -3327,7 +3451,7 @@ def test_restore_queues_purges_late_aborts_after_process_pending_chunks(
     # instead of blindly extending it onto running_queue.
     adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
     assert running_queue == []
-    assert req.request_id in adapter._cancelled_load_reqs
+    assert req.request_id not in adapter._registered_load_entries
 
 
 def test_purge_is_noop_on_empty_deques(build_adapter):

tests/distributed/omni_connectors/test_shm_connector.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for SharedMemoryConnector focusing on TP / CFG / metadata fallback."""
 
 import fcntl
 import os
 import time
+import uuid
 
 import pytest
 import torch
@@ -212,7 +213,7 @@ class TestCleanup:
         connector.put("s0", "s1", "cleanup_req_42", data)
         assert "cleanup_req_42" in connector._pending_keys
 
-        connector.cleanup("req_42")
+        connector.cleanup("cleanup_req_42")
         assert "cleanup_req_42" not in connector._pending_keys
 
         result = connector.get("s0", "s1", "cleanup_req_42", metadata=None)
@@ -223,7 +224,7 @@ class TestCleanup:
         connector.put("s0", "s1", "consumed_req_99", data)
         connector.get("s0", "s1", "consumed_req_99", metadata=None)
 
-        connector.cleanup("req_99")
+        connector.cleanup("consumed_req_99")
         assert "consumed_req_99" not in connector._pending_keys
 
     def test_close_cleans_all_pending(self, connector):
@@ -233,3 +234,63 @@ class TestCleanup:
         assert len(connector._pending_keys) == 3
         connector.close()
         assert len(connector._pending_keys) == 0
+
+
+@pytest.mark.parametrize("suffix", ["_1_2", "_1", "_suffix"])
+def test_cleanup_preserves_other_request_keys(connector, suffix):
+    request_id = f"ownership_{uuid.uuid4().hex}"
+    own_key = f"{request_id}_0_0"
+    sibling_key = f"{request_id}{suffix}_0_0"
+    assert connector.put("0", "1", own_key, "own")[0]
+    assert connector.put("0", "1", sibling_key, "sibling")[0]
+
+    # A raw request id is not an exact key and must not match any chunks.
+    connector.cleanup(request_id)
+    assert connector.get("0", "1", own_key)[0] == "own"
+    connector.cleanup(own_key)
+    assert connector.get("0", "1", sibling_key)[0] == "sibling"
+
+
+def test_sender_reaps_keys_consumed_by_another_connector(connector):
+    receiver = SharedMemoryConnector({})
+    prefix = f"reap_{uuid.uuid4().hex}"
+    try:
+        for index in range(256):
+            key = f"{prefix}_{index}"
+            assert connector.put("0", "1", key, index)[0]
+            assert receiver.get("0", "1", key)[0] == index
+            assert len(connector._pending_keys) <= 1
+        connector.reap_consumed()
+        assert not connector._pending_keys
+    finally:
+        receiver.close()
+
+
+def test_reap_rotates_past_unread_keys(connector):
+    receiver = SharedMemoryConnector({})
+    prefix = f"rotate_{uuid.uuid4().hex}"
+    try:
+        for index in range(130):
+            assert connector.put("0", "1", f"{prefix}_{index}", index)[0]
+        for index in range(65, 130):
+            assert receiver.get("0", "1", f"{prefix}_{index}")[0] == index
+        for _ in range(3):
+            connector.reap_consumed()
+        assert len(connector._pending_keys) == 65
+        assert connector.get("0", "1", f"{prefix}_0")[0] == 0
+    finally:
+        receiver.close()
+
+
+def test_consumed_bad_payload_removes_lock_file(connector, monkeypatch):
+    key = f"bad_payload_{uuid.uuid4().hex}"
+    assert connector.put("0", "1", key, "payload")[0]
+
+    def fail_deserialize(_data):
+        raise ValueError("invalid payload")
+
+    monkeypatch.setattr(connector, "deserialize_obj", fail_deserialize)
+    assert connector.get("0", "1", key) is None
+    assert not os.path.exists(f"/dev/shm/shm_{key}_lockfile.lock")
+    connector.reap_consumed()
+    assert key not in connector._pending_keys

vllm_omni/distributed/omni_connectors/connectors/base.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from abc import ABC, abstractmethod
 from typing import Any
@@ -75,6 +75,9 @@ class OmniConnectorBase(ABC):
     def abandon_get(self, get_key: str) -> None:
         """Retire unresolved discovery attempts; never cancel an active DMA READ."""
 
+    def reap_consumed(self) -> None:
+        """Release producer bookkeeping for payloads consumed by another process."""
+
     @abstractmethod
     def health(self) -> dict[str, Any]:
         """Return health status and metrics."""

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -3,7 +3,9 @@
 
 import fcntl
 import os
+import threading
 import time
+from collections import OrderedDict
 from multiprocessing import shared_memory as shm_pkg
 from typing import Any
 
@@ -33,7 +35,8 @@ class SharedMemoryConnector(OmniConnectorBase):
     def __init__(self, config: dict[str, Any]):
         self.config = config
         self.stage_id = config.get("stage_id", -1)
-        self._pending_keys: set[str] = set()
+        self._pending_keys: OrderedDict[str, None] = OrderedDict()
+        self._pending_keys_lock = threading.Lock()
         self._metrics = {
             "puts": 0,
             "gets": 0,
@@ -48,6 +51,7 @@ class SharedMemoryConnector(OmniConnectorBase):
         data: Any,
     ) -> tuple[bool, int, dict[str, Any] | None]:
         try:
+            self.reap_consumed()
             payload = self.serialize_obj(data)
             size = len(payload)
 
@@ -59,7 +63,8 @@ class SharedMemoryConnector(OmniConnectorBase):
 
             # meta contains {'name': ..., 'size': ...}
             metadata = {"shm": meta, "size": size}
-            self._pending_keys.add(put_key)
+            with self._pending_keys_lock:
+                self._pending_keys[put_key] = None
 
             self._metrics["puts"] += 1
             self._metrics["bytes_transferred"] += size
@@ -71,15 +76,15 @@ class SharedMemoryConnector(OmniConnectorBase):
             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)
+                consumed = True
                 fcntl.flock(lockf, fcntl.LOCK_UN)
             obj = self.deserialize_obj(data_bytes)
             result = (obj, int(shm_handle.get("size", 0)))
-            deserialized = True
             return result
         except BlockingIOError:
             return None
@@ -87,7 +92,7 @@ class SharedMemoryConnector(OmniConnectorBase):
             logger.error(f"SharedMemoryConnector shm get failed for req : {e}")
             return None
         finally:
-            if deserialized:
+            if consumed:
                 try:
                     os.remove(lock_file)
                 except FileNotFoundError:
@@ -104,7 +109,8 @@ class SharedMemoryConnector(OmniConnectorBase):
             shm_handle = {"name": get_key, "size": shm.size}
             result = self._get_data_with_lock(lock_file, shm_handle)
             if result is not None:
-                self._pending_keys.discard(get_key)
+                with self._pending_keys_lock:
+                    self._pending_keys.pop(get_key, None)
             return result
         except FileNotFoundError:
             return None
@@ -139,7 +145,8 @@ class SharedMemoryConnector(OmniConnectorBase):
                 lock_file = f"/dev/shm/shm_{shm_handle['name']}_lockfile.lock"
                 result = self._get_data_with_lock(lock_file, shm_handle)
                 if result is not None:
-                    self._pending_keys.discard(get_key)
+                    with self._pending_keys_lock:
+                        self._pending_keys.pop(get_key, None)
             else:
                 # Missing or non-SHM metadata falls back to key-based lookup.
                 result = self._get_by_key(get_key)
@@ -151,20 +158,10 @@ class SharedMemoryConnector(OmniConnectorBase):
         return result
 
     def cleanup(self, request_id: str) -> None:
-        """Best-effort cleanup of unconsumed SHM segments for *request_id*.
-
-        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.
-        """
-        stale = [
-            k
-            for k in self._pending_keys
-            if k == request_id or k.startswith(request_id + "_") or k.endswith("_" + request_id)
-        ]
-        for key in stale:
-            self._pending_keys.discard(key)
+        """Unlink the exact key passed to ``put()``, never a request-id prefix."""
+        key = request_id
+        with self._pending_keys_lock:
+            self._pending_keys.pop(key, None)
             try:
                 seg = shm_pkg.SharedMemory(name=key)
                 seg.close()
@@ -183,20 +180,18 @@ class SharedMemoryConnector(OmniConnectorBase):
 
     def close(self) -> None:
         """Unlink all remaining tracked SHM segments."""
-        for key in list(self._pending_keys):
-            try:
-                seg = shm_pkg.SharedMemory(name=key)
-                seg.close()
-                seg.unlink()
-            except Exception:
-                pass
-            lock_file = f"/dev/shm/shm_{key}_lockfile.lock"
-            if os.path.exists(lock_file):
-                try:
-                    os.remove(lock_file)
-                except OSError:
-                    pass
-        self._pending_keys.clear()
+        with self._pending_keys_lock:
+            keys = list(self._pending_keys)
+        for key in keys:
+            self.cleanup(key)
+
+    def reap_consumed(self) -> None:
+        """Bounded round-robin sweep; receivers unlink SHM in another process."""
+        with self._pending_keys_lock:
+            for _ in range(min(64, len(self._pending_keys))):
+                key, _ = self._pending_keys.popitem(last=False)
+                if os.path.exists(f"/dev/shm/{key}"):
+                    self._pending_keys[key] = None
 
     def health(self) -> dict[str, Any]:
         return {"status": "healthy", **self._metrics}

vllm_omni/distributed/omni_connectors/transfer_adapter/base.py

@@ -25,7 +25,6 @@ class OmniTransferAdapterBase:
         self._pending_load_reqs = deque()
         # Requests that have successfully retrieved data
         self._finished_load_reqs = set()
-        self._cancelled_load_reqs: set[str] = set()
 
         # Requests that are waiting to be saved
         self._pending_save_reqs = deque()
@@ -65,10 +64,6 @@ class OmniTransferAdapterBase:
                     break
                 request = self._pending_load_reqs.popleft()
                 request_id = request.request_id
-                if request_id in self._cancelled_load_reqs:
-                    self._cancelled_load_reqs.discard(request_id)
-                    continue
-                self.request_ids_mapping[request_id] = request.external_req_id
                 try:
                     is_success = self._poll_single_request(request)
                     if is_success:
@@ -136,6 +131,8 @@ class OmniTransferAdapterBase:
                     logger.error("Send gave up for %s: %s", failed, e)
                     self.record_send_failure(failed, f"{type(e).__name__}: {e}")
 
+            if self.connector is not None:
+                self.connector.reap_consumed()
             with self._save_cond:
                 if not self._pending_save_reqs and not self.stop_event.is_set():
                     self._save_cond.wait(timeout=0.1)

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -18,6 +18,7 @@ from vllm.v1.utils import ConstantList
 from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload
 
 from ..adapter import construct_next_stage_streaming_input_prompt
+from ..connectors.shm_connector import SharedMemoryConnector
 from ..factory import OmniConnectorFactory
 from ..utils.config import ConnectorSpec, stage_receives_chunks
 from ..utils.initialization import resolve_connector_spec
@@ -44,6 +45,8 @@ class _SenderGeneration:
         # was up. The terminal task owes that cleanup once it is done with the
         # connector -- on every exit path, including a failed or raising put.
         self.cleanup_deferred = False
+        self.num_puts = 0
+        self.terminal_sent = False
 
 
 class _LoadEntry:
@@ -332,7 +335,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             if request_id in self._registered_load_entries:
                 return
             entry = _LoadEntry(request)
-            self._cancelled_load_reqs.discard(request_id)
             self._registered_load_entries[request_id] = entry
             self._pending_load_reqs.append(entry)
         with self._recv_cond:
@@ -497,8 +499,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         with self._receiver_state_lock:
             if self._registered_load_entries.get(req_id) is not entry:
                 return True
+            external_req_id = entry.external_req_id
+            self.request_ids_mapping[req_id] = external_req_id
             chunk_id = self.get_req_chunk[req_id]
-            external_req_id = self.request_ids_mapping.get(req_id, req_id)
         connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}"
 
         # Use timeout=0 for non-blocking poll
@@ -674,6 +677,12 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         return False
 
     def _send_single_request(self, task: dict):
+        if "cleanup_shm" in task:
+            assert isinstance(self.connector, SharedMemoryConnector)
+            key_prefix, num_puts = task["cleanup_shm"]
+            for chunk_id in range(num_puts):
+                self.connector.cleanup(f"{key_prefix}_{chunk_id}")
+            return
         request = task["request"]
         external_req_id = request.external_req_id
         sender_token = task.get("sender_token")
@@ -710,6 +719,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 if self._sender_tokens.get(external_req_id) is sender_token:
                     sender_token.in_flight = False
                     if sender_token.cancelled and not sender_token.terminal_pending:
+                        self._queue_shm_cleanup_locked(external_req_id, sender_token)
                         self._sender_tokens.pop(external_req_id, None)
                         self._clear_sender_state_locked(external_req_id)
 
@@ -813,6 +823,10 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             logger.debug("Skipping cancelled chunk for request %s before connector put", external_req_id)
             return
 
+        if sender_token is not None:
+            # Include the in-flight key even if cancellation wins before put
+            # returns or the connector raises after creating the segment.
+            sender_token.num_puts = self.put_req_chunk[external_req_id] + 1
         success, size, metadata = self.connector.put(
             from_stage=str(stage_id),
             to_stage=str(next_stage_id),
@@ -822,6 +836,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         with self._sender_state_lock:
             if sender_token is not None:
+                if success and is_finished:
+                    sender_token.terminal_sent = True
                 still_current = self._sender_tokens.get(external_req_id) is sender_token
                 retiring_without_terminal = sender_token.cancelled and not (
                     is_finished or sender_token.terminal_pending
@@ -941,7 +957,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self._discard_from_chunk_deque(self.waiting_for_chunk_running_requests, request_id)
         self._discard_from_chunk_deque(self._held_non_active, request_id)
 
-        self._cancelled_load_reqs.add(request_id)
         self._finished_load_reqs.discard(request_id)
         self._waiting_since.pop(request_id, None)
 
@@ -987,9 +1002,25 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 # in-flight code can mutate this request's state. Let it own
                 # cleanup without making the scheduler wait for connector I/O.
                 return
+            self._queue_shm_cleanup_locked(external_req_id, sender_token)
             self._sender_tokens.pop(external_req_id, None)
             self._clear_sender_state_locked(external_req_id)
 
+    def _queue_shm_cleanup_locked(self, external_req_id: str, sender_token: _SenderGeneration) -> None:
+        # Successful terminal chunks must remain readable by downstream. On
+        # abort, queue exact keys on the same save thread as puts: the scheduler
+        # never performs SHM I/O, and cleanup precedes reuse of this external id.
+        if (
+            isinstance(self.connector, SharedMemoryConnector)
+            and sender_token.num_puts
+            and not sender_token.terminal_sent
+        ):
+            self._pending_save_reqs.append(
+                {"cleanup_shm": (f"{external_req_id}_{self.connector.stage_id}", sender_token.num_puts)}
+            )
+            with self._save_cond:
+                self._save_cond.notify()
+
     def _clear_sender_state_locked(self, external_req_id: str) -> None:
         """Clear sender state while ``_sender_state_lock`` is held."""
         self.put_req_chunk.pop(external_req_id, None)