From d486bf097a8d2b31c30a105785c1de21f256a263 Mon Sep 17 00:00:00 2001 From: NATURE Date: Thu, 24 Sep 2026 17:59:43 +0800 Subject: [PATCH] [Bugfix] Reclaim aborted chunk SHM and consumed connector state (#8082) Signed-off-by: natureofnature Co-authored-by: Gagan Dhakrey --- .../test_chunk_transfer_adapter.py | 142 ++++++++++++++++-- .../omni_connectors/test_shm_connector.py | 67 ++++++++- .../omni_connectors/connectors/base.py | 5 +- .../connectors/shm_connector.py | 65 ++++---- .../omni_connectors/transfer_adapter/base.py | 7 +- .../chunk_transfer_adapter.py | 37 ++++- 6 files changed, 267 insertions(+), 56 deletions(-) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 774257334..919b9ab2b 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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): diff --git a/tests/distributed/omni_connectors/test_shm_connector.py b/tests/distributed/omni_connectors/test_shm_connector.py index bb502253a..8d9ad0270 100644 --- a/tests/distributed/omni_connectors/test_shm_connector.py +++ b/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 diff --git a/vllm_omni/distributed/omni_connectors/connectors/base.py b/vllm_omni/distributed/omni_connectors/connectors/base.py index 92257d9de..a67a35a80 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/base.py +++ b/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.""" diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index ae45a1a3c..3ad69e661 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/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} diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py index 3420dcb94..8b64e3f24 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py +++ b/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) diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 13d3dca3e..a07e2824c 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/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)