[Bugfix] Preserve host transport fallback when discovery cleanup fails

Signed-off-by: Sy03 <1370724210@qq.com>

docs/design/feature/stage_host_transport.md

@@ -0,0 +1,84 @@
+# Local stage host transport
+
+`SharedMemoryConnector` uses a persistent host ring for ordinary local stage
+payloads by default. Models keep the existing `put`/`get` and typed payload
+contracts; no model-specific sender, receiver or deployment change is required.
+This applies to both V1 and MRV2 paths using this connector, including CPU codec
+codes and CPU hidden states. Remote connectors use their existing transports.
+
+Each producer process and destination stage lazily allocates a 4 MiB ring.
+Small CPU tensor trees use a msgpack header and raw tensor regions. Plain
+serialized payloads can use the same ring. Frames larger than a quarter of the
+ring, and frames arriving while it is full, use the existing per-key SHM path.
+The fallback preserves delivery without blocking a model step on ring capacity.
+Large or unsupported tensor trees do not allocate an unused contiguous snapshot.
+
+## Ownership and delivery
+
+The producer copies the payload before `put` returns. Readers hold the ring's
+file lock while copying into independently owned CPU storage, then mark the
+frame consumed. Multiple receiving processes may compete for a key; one reader
+claims it. Existing TP fanout remains responsible for distributing that owned
+payload to peer ranks. This transport does not provide multicast.
+
+Small tensor leaves share one owned allocation. Large multi-leaf messages use
+separate allocations, so retaining small metadata does not retain a large
+hidden state. No view of ring storage escapes the receive boundary.
+
+Cancellation retires exact keys or tracked numeric chunk suffixes. Replacing
+an unclaimed key retires its previous frame, including transitions between
+the ring and per-key SHM. Descriptors include monotonic frame positions; an old
+descriptor cannot read a newer publication after wraparound. Keys must have
+one active publisher, as assumed by the existing stage request routing.
+
+Only producers register ring allocations with Python's shared-memory resource
+tracker. Readers map files without taking ownership. Producer process identity
+includes its start time, and available Linux pidfds detect producer exit, so
+old frames cannot be delivered after a producer restart. Normal close removes
+owned allocations and discovery markers. Resource tracking covers allocation
+cleanup when the owning application exits; an abnormal exit can leave small
+discovery markers, which contain no payload data.
+Live receivers release cached mappings after a producer exits and prune stale
+markers during discovery within their deployment scope.
+
+## Configuration and observation
+
+Connector extras are shared by the endpoints of an edge. For example:
+
+```yaml
+connectors:
+  shm:
+    name: SharedMemoryConnector
+    extra:
+      host_ring_bytes: 4194304
+```
+
+Set `host_ring_bytes: 0` at both endpoints to use the per-key path, including
+controlled comparisons or deployments with older receivers. Nonzero capacities
+must be between 4 KiB and 64 MiB. All ring endpoints must support this protocol.
+The launcher supplies a common deployment scope; discovery scans only that
+scope's directory. Receive attempts do not wait for another process's file lock.
+An unchanged empty channel skips its frame lock. Channel registration and
+retirement advance a serialized discovery stamp, so empty polling does not
+rescan the directory and new producers remain discoverable without a timer.
+Ring endpoints must share a Linux PID namespace. Deployments that share SHM
+across different PID namespaces must use `host_ring_bytes: 0`.
+
+`health()` exposes `host_ring_puts`, `host_ring_gets` and `host_ring_fallbacks`
+alongside the existing counters. Ring setup failure logs once per edge and uses
+per-key SHM. CPU payload size and the cost of serialization determine the benefit;
+reduced transfer latency does not imply higher model throughput.
+
+## Scope
+
+This mechanism reduces host serialization copies and per-message segment work.
+It does not retain worker outputs on the GPU, allocate a CUDA arena, or replace
+GPU stream synchronization. It also does not impose a global byte budget on
+queued snapshots, downstream model caches or client audio delivery. The bounded
+ring with a compatibility fallback is not end-to-end streaming backpressure.
+
+Regression tests are in
+`tests/distributed/omni_connectors/test_shm_ring.py`, collected by the existing
+`core_model and cpu` CI lane. They cover ownership, tensor types, typed payloads,
+pressure fallback, cancellation, overwrite, deadlines, competing processes and
+producer restart.

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -31,10 +31,10 @@ 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):
+@pytest.fixture(params=[0, 4 << 20], ids=["per-key-shm", "host-ring"])
+def shm_sender(build_adapter, request):
     adapter, _ = build_adapter(stage_id=0)
-    connector = SharedMemoryConnector({"stage_id": 0})
+    connector = SharedMemoryConnector({"stage_id": 0, "host_ring_bytes": request.param})
     adapter.connector = connector
     adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
         codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
@@ -111,13 +111,14 @@ def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender):
 
     adapter.release_shm_resources(ext_id)
     # The release is queued behind sends; nothing is unlinked on the caller.
-    assert f"{ext_id}_0_1" in connector._pending_keys
+    assert adapter._reclaimed_shm_total == 0
     while adapter._pending_save_reqs:
         adapter._send_single_request(adapter._pending_save_reqs.popleft())
 
     assert adapter._reclaimed_shm_total == 2
     assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys)
     assert connector.get("0", "1", f"{ext_id}_0_1") is None
+    assert connector.get("0", "1", f"{ext_id}_0_2") is None
     assert connector.get("0", "1", sibling_key)[0] == "sibling"
 
 

tests/distributed/omni_connectors/test_shm_connector.py

@@ -19,7 +19,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 @pytest.fixture()
 def connector():
-    c = SharedMemoryConnector({})
+    # These cases inspect the per-key segment/lock protocol. The enabled ring
+    # protocol and its pressure fallback are covered in test_shm_ring.py.
+    c = SharedMemoryConnector({"host_ring_bytes": 0})
     yield c
     c.close()
 

tests/distributed/omni_connectors/test_shm_ring.py

@@ -0,0 +1,527 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cross-process ownership, pressure and compatibility of the default host path."""
+
+import fcntl
+import multiprocessing
+import os
+import time
+import uuid
+from dataclasses import asdict
+from multiprocessing import shared_memory
+from pathlib import Path
+
+import msgspec
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from vllm.outputs import CompletionOutput, RequestOutput
+
+from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
+from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
+from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
+from vllm_omni.distributed.omni_connectors.utils.tensor_frame import prepare_tensor_frame
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def edge():
+    scope = uuid.uuid4().hex
+    sender = SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope, "host_ring_bytes": 4096})
+    receiver = SharedMemoryConnector({"stage_id": 1, "wakeup_scope": scope, "host_ring_bytes": 4096})
+    yield sender, receiver
+    receiver.close()
+    sender.close()
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+        torch.complex64,
+    ],
+)
+def test_owned_tensor_tree_preserves_values_shapes_and_types(edge, dtype):
+    sender, receiver = edge
+    tensor = torch.arange(24).to(dtype).reshape(4, 6).T
+    data = {
+        "codes": {"audio": tensor},
+        "other": (tensor[0], [torch.empty(0, dtype=dtype), torch.tensor(1)]),
+        "__tensor__": False,
+        "meta": {"finished": torch.tensor(False)},
+    }
+    ok, size, metadata = sender.put("0", "1", "tree", data)
+    assert ok and "host_ring" in metadata
+    output, read_size = receiver.get("0", "1", "tree", metadata)
+    assert read_size == size and output["__tensor__"] is False
+    assert type(output["other"]) is list and type(output["other"][1]) is list
+    assert output["other"][1][0].shape == (0,)
+    assert output["other"][1][1].shape == ()
+    assert output["codes"]["audio"].dtype == dtype
+    assert torch.equal(output["codes"]["audio"], tensor)
+    tensor.fill_(0)
+    assert torch.count_nonzero(output["codes"]["audio"]) > 0
+    for index in range(100):
+        sender.put("0", "1", f"wrap{index}", {"codes": torch.full((56,), index)})
+        assert receiver.get("0", "1", f"wrap{index}")[0]["codes"][0] == index
+    assert torch.count_nonzero(output["codes"]["audio"]) > 0
+
+
+def test_conjugate_tensor_and_scalar_finish(edge):
+    sender, receiver = edge
+    tensor = torch.tensor([1 + 2j, 3 - 4j]).conj()
+    sender.put("0", "1", "conjugate", {"condition": tensor, "finished": torch.tensor(True)})
+    result = receiver.get("0", "1", "conjugate")[0]
+    assert torch.equal(result["condition"], tensor)
+    assert result["finished"].shape == () and result["finished"].item() is True
+
+
+def test_shared_typed_payload_keeps_omitted_defaults_and_cpu_tensors(edge):
+    sender, receiver = edge
+    payload = OmniPayloadStruct(
+        codes=CodesStruct(audio=torch.arange(28).reshape(1, 28)),
+        meta=MetaStruct(finished=torch.tensor(False), left_context_size=0),
+        kv_metadata={"binary_handle": b"\x00\xff\x10"},
+    )
+    ok, _, metadata = sender.put("0", "1", "typed", payload)
+    assert ok and "host_ring" in metadata
+    result = receiver.get("0", "1", "typed")[0]
+    assert set(result) == {"codes", "meta", "kv_metadata"}
+    assert set(result["codes"]) == {"audio"}
+    assert torch.equal(result["codes"]["audio"], payload.codes.audio)
+    assert result["meta"]["left_context_size"] == 0
+    assert result["kv_metadata"]["binary_handle"] == b"\x00\xff\x10"
+
+
+def test_source_can_be_overwritten_as_soon_as_put_returns(edge):
+    sender, receiver = edge
+    source = torch.arange(56)
+    sender.put("0", "1", "source-reuse", source)
+    source.fill_(-1)
+    assert receiver.get("0", "1", "source-reuse")[0].tolist() == list(range(56))
+
+
+@pytest.mark.parametrize("dtype", [np.bool_, np.int32, np.float16, np.float64, np.complex64])
+def test_mixed_numpy_image_and_tensor_keep_shared_wire_semantics(edge, dtype):
+    sender, receiver = edge
+    array = np.arange(24).astype(dtype).reshape(4, 6)[::-1, ::2]
+    image = Image.fromarray(np.arange(18, dtype=np.uint8).reshape(2, 3, 3))
+    payload = {"codes": torch.arange(28), "array": array, "image": image, "slice": slice(1, 7, 2)}
+    _, _, metadata = sender.put("0", "1", "mixed-native", payload)
+    assert "host_ring" in metadata
+    output = receiver.get("0", "1", "mixed-native", metadata)[0]
+    assert np.array_equal(output["array"], array)
+    assert output["array"].dtype == array.dtype
+    assert not output["array"].flags.writeable
+    assert output["image"].mode == image.mode and output["image"].tobytes() == image.tobytes()
+    assert output["slice"] == [1, 7, 2]
+    assert torch.equal(output["codes"], payload["codes"])
+    array.fill(0)
+    image.paste(0, (0, 0, 3, 2))
+    assert np.count_nonzero(output["array"]) and any(output["image"].tobytes())
+
+
+def test_large_array_rejects_unused_framing_snapshot(monkeypatch):
+    def unexpected_snapshot(value):
+        pytest.fail("a too-large array must not be snapshotted while testing ring admission")
+
+    monkeypatch.setattr(OmniSerializer.encoder, "_enc_hook", unexpected_snapshot)
+    assert prepare_tensor_frame({"array": np.zeros(1024)}, max_bytes=1024) is None
+
+
+def test_native_scalar_restore_preserves_nested_markers_and_input_containers():
+    array = np.arange(3)
+    image = Image.fromarray(np.array([[1, 2]], dtype=np.uint8))
+    tensor = torch.arange(2)
+    payload = {"ids": list(range(8192)), "mixed": [1.5, None, (array, {"image": image, "tensor": tensor})]}
+    decoded = msgspec.msgpack.decode(OmniSerializer.serialize(payload))
+    output = OmniSerializer.restore(decoded)
+    assert output is not decoded and output["ids"] is not decoded["ids"]
+    assert output["ids"] == payload["ids"]
+    assert np.array_equal(output["mixed"][2][0], array)
+    assert output["mixed"][2][1]["image"].tobytes() == image.tobytes()
+    assert torch.equal(output["mixed"][2][1]["tensor"], tensor)
+    assert isinstance(decoded["mixed"][2][0], dict)
+
+
+def test_fallback_miss_does_not_register_an_unclaimed_allocation(edge, monkeypatch, caplog):
+    sender, receiver = edge
+    key = "fallback-claim-" + uuid.uuid4().hex
+    _, _, metadata = sender.put("0", "1", key, torch.arange(4096))
+    assert "shm" in metadata
+    original = shared_memory.SharedMemory
+    calls = []
+
+    def tracked_open(*args, **kwargs):
+        calls.append((args, kwargs))
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(shared_memory, "SharedMemory", tracked_open)
+    with open(f"/dev/shm/shm_{key}_lockfile.lock", "rb+") as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        assert receiver.get("0", "1", key) is None
+        assert not calls
+    assert torch.equal(receiver.get("0", "1", key)[0], torch.arange(4096))
+    assert receiver.get("0", "1", key, metadata) is None
+    assert not any(record.levelname == "ERROR" for record in caplog.records)
+
+
+def test_tensor_frames_preserve_shared_output_reconstruction(edge):
+    sender, receiver = edge
+    completion = asdict(CompletionOutput(index=0, text="", token_ids=[1], cumulative_logprob=0.0, logprobs=None))
+    completion["multimodal_output"] = {"codes": torch.arange(28)}
+    payload = {
+        "request_id": "output",
+        "prompt": None,
+        "prompt_token_ids": [1],
+        "prompt_logprobs": None,
+        "outputs": [completion],
+        "finished": False,
+    }
+    _, _, metadata = sender.put("0", "1", "output", payload)
+    assert "host_ring" in metadata
+    result = receiver.get("0", "1", "output", metadata)[0]
+    assert isinstance(result, RequestOutput)
+    assert isinstance(result.outputs[0], CompletionOutput)
+    assert torch.equal(result.outputs[0].multimodal_output["codes"], torch.arange(28))
+
+
+def test_large_tensor_does_not_share_owner_with_small_metadata():
+    with SharedMemoryConnector({}) as connector:
+        payload = {"hidden": torch.ones(16384), "finished": torch.tensor(False)}
+        assert connector.put("0", "1", "large-owner", payload)[0]
+        output = connector.get("0", "1", "large-owner")[0]
+        assert output["finished"].untyped_storage().nbytes() == 1
+        assert output["hidden"].untyped_storage().nbytes() == 65536
+
+
+def test_full_ring_and_large_frame_use_existing_shm(edge):
+    sender, receiver = edge
+    payload = {"codes": torch.arange(56)}
+    metadata = {}
+    for index in range(12):
+        ok, _size, metadata[index] = sender.put("0", "1", f"full{index}", payload)
+        assert ok
+    assert any("host_ring" in item for item in metadata.values())
+    assert any("shm" in item for item in metadata.values())
+    # Unclaimed earlier frames pin the ring; later requests still arrive.
+    for index in reversed(range(12)):
+        assert torch.equal(receiver.get("0", "1", f"full{index}")[0]["codes"], payload["codes"])
+    ok, _size, large_metadata = sender.put("0", "1", "large", torch.arange(4096))
+    assert ok and "shm" in large_metadata
+    assert torch.equal(receiver.get("0", "1", "large", large_metadata)[0], torch.arange(4096))
+    assert "host_ring" in sender.put("0", "1", "recovered", payload)[2]
+    assert receiver.get("0", "1", "recovered") is not None
+
+
+def test_reput_across_ring_and_fallback_and_stale_descriptor(edge):
+    sender, receiver = edge
+    _, _, old = sender.put("0", "1", "reuse", torch.tensor([11]))
+    _, _, current = sender.put("0", "1", "reuse", torch.tensor([22]))
+    assert receiver.get("0", "1", "reuse", old) is None
+    assert receiver.get("0", "1", "reuse", current)[0].item() == 22
+    sender.put("0", "1", "reuse", torch.tensor([33]))
+    sender.put("0", "1", "reuse", torch.full((4096,), 44))
+    assert receiver.get("0", "1", "reuse")[0][0].item() == 44
+    sender.put("0", "1", "reuse", torch.full((4096,), 55))
+    sender.put("0", "1", "reuse", torch.tensor([66]))
+    assert receiver.get("0", "1", "reuse")[0].item() == 66
+    assert not os.path.exists("/dev/shm/reuse")
+
+
+def test_cancel_without_receiver_scan_returns_credit_and_prefix_is_exact(edge):
+    sender, receiver = edge
+    for key in ("req_0_0", "req_0_1", "req_0_10other", "request_0_0"):
+        assert sender.put("0", "1", key, torch.arange(56))[0]
+    assert sender.cleanup_prefix("req_0_") == 2
+    assert receiver.get("0", "1", "req_0_0") is None
+    assert receiver.get("0", "1", "req_0_10other") is not None
+    assert receiver.get("0", "1", "request_0_0") is not None
+    assert "host_ring" in sender.put("0", "1", "next", torch.arange(56))[2]
+
+
+def test_flat_scope_and_metadata_work_across_different_parents(monkeypatch):
+    monkeypatch.setattr(os, "getppid", lambda: 100)
+    sender = SharedMemoryConnector({"stage_id": 0, "wakeup_scope": "flat-scope"})
+    monkeypatch.setattr(os, "getppid", lambda: 200)
+    receiver = SharedMemoryConnector({"stage_id": 1, "wakeup_scope": "flat-scope"})
+    try:
+        assert sender._wake_directory == receiver._wake_directory
+        assert sender.put("0", "1", "scope", torch.tensor([7]))[0]
+        assert receiver.get("0", "1", "scope")[0].item() == 7
+    finally:
+        sender.close()
+        receiver.close()
+
+
+def test_empty_poll_reuses_discovery_and_never_waits_for_a_ring_lock(edge, monkeypatch):
+    sender, receiver = edge
+    sender.put("0", "1", "prime-poll", torch.arange(3))
+    assert receiver.get("0", "1", "prime-poll") is not None
+
+    def unexpected_work(*args, **kwargs):
+        pytest.fail("an empty, unchanged channel must not rescan discovery or acquire its frame lock")
+
+    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_ring.glob.iglob", unexpected_work)
+    monkeypatch.setattr(fcntl, "flock", unexpected_work)
+    for index in range(128):
+        assert receiver.get("0", "1", f"missing-poll-{index}") is None
+    monkeypatch.undo()
+    sender.put("0", "1", "new-poll", torch.arange(5))
+    assert receiver.get("0", "1", "new-poll")[0].tolist() == list(range(5))
+
+
+def test_new_producer_is_discovered_after_a_miss_even_when_the_clock_does_not_advance(edge, monkeypatch):
+    sender, receiver = edge
+    sender.put("0", "1", "first-producer", torch.arange(3))
+    assert receiver.get("0", "1", "first-producer") is not None
+    assert receiver.get("0", "1", "second-producer") is None
+    old_version = receiver._host_ring._discovery_versions["0", "1"]
+    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_ring.time.time_ns", lambda: 0)
+    with SharedMemoryConnector(sender.config) as second:
+        second.put("0", "1", "second-producer", torch.arange(5))
+        assert receiver.get("0", "1", "second-producer")[0].tolist() == list(range(5))
+        new_version = receiver._host_ring._discovery_versions["0", "1"]
+        assert new_version[2] == old_version[2] + 1
+
+
+def test_one_receiver_discovers_each_input_edge_at_the_same_registry_version(edge):
+    sender, receiver = edge
+    with SharedMemoryConnector({**sender.config, "stage_id": 2}) as second:
+        sender.put("0", "1", "fanin-first", torch.arange(3))
+        second.put("2", "1", "fanin-second", torch.arange(5))
+        assert receiver.get("0", "1", "fanin-first")[0].tolist() == list(range(3))
+        assert receiver.get("2", "1", "fanin-second")[0].tolist() == list(range(5))
+
+
+def test_discovery_attaches_all_producers_even_when_the_first_one_has_the_requested_key(edge, monkeypatch):
+    sender, receiver = edge
+    with SharedMemoryConnector(sender.config) as second:
+        first_metadata = sender.put("0", "1", "multi-first", torch.arange(3))[2]
+        second_metadata = second.put("0", "1", "multi-second", torch.arange(5))[2]
+        paths = [
+            f"/dev/shm/{receiver._host_ring.directory}/{metadata['host_ring']['name']}"
+            for metadata in (first_metadata, second_metadata)
+        ]
+        monkeypatch.setattr(
+            "vllm_omni.distributed.omni_connectors.connectors.shm_ring.glob.iglob", lambda _: iter(paths)
+        )
+        assert receiver.get("0", "1", "multi-first")[0].tolist() == list(range(3))
+        assert receiver.get("0", "1", "multi-second")[0].tolist() == list(range(5))
+
+
+def test_closed_channel_releases_reader_cache_while_the_producer_process_is_still_alive(edge):
+    sender, receiver = edge
+    sender.put("0", "1", "close-cache", torch.arange(3))
+    assert receiver.get("0", "1", "close-cache") is not None
+    channel = next(iter(receiver._host_ring.readers.values()))
+    fd = channel.fd
+    sender.close()
+    assert receiver.get("0", "1", "not-published-after-close") is None
+    assert not receiver._host_ring.readers
+    with pytest.raises(OSError):
+        os.fstat(fd)
+
+
+def test_disabled_ring_and_unsupported_tree_keep_legacy_serializer(edge):
+    sender, receiver = edge
+    with SharedMemoryConnector({"host_ring_bytes": 0}) as legacy:
+        _, _, metadata = legacy.put("0", "1", "legacy", {"codes": torch.arange(3)})
+        assert "shm" in metadata
+        assert receiver.get("0", "1", "legacy", metadata)[0]["codes"].tolist() == [0, 1, 2]
+    assert prepare_tensor_frame({"codes": torch.arange(3), "opaque": object()}) is None
+    assert sender.put("0", "1", "falsey", False)[0]
+    assert receiver.get("0", "1", "falsey")[0] is False
+
+
+def test_tensor_free_lists_reuse_native_serialization():
+    payload = {"prompt_token_ids": list(range(8192)), "condition": [0.25] * 8192}
+    with SharedMemoryConnector({}) as connector:
+        frame = prepare_tensor_frame(payload)
+        assert isinstance(frame, bytes)
+        assert frame == connector.serialize_obj(payload)
+        assert connector.put("0", "1", "native-lists", payload)[0]
+        assert connector.get("0", "1", "native-lists")[0] == payload
+
+
+def test_scalar_lists_and_literal_extensions_preserve_wire_semantics():
+    literal = msgspec.msgpack.Ext(42, bytes(20))
+    payload = {"condition": [0.25] * 8192, "codes": torch.arange(28), "extension": literal}
+    with SharedMemoryConnector({}) as connector:
+        _, _, metadata = connector.put("0", "1", "mixed-lists", payload)
+        assert "host_ring" in metadata
+        result = connector.get("0", "1", "mixed-lists")[0]
+        assert result["condition"] == payload["condition"]
+        assert torch.equal(result["codes"], payload["codes"])
+        assert result["extension"].code == literal.code
+        assert result["extension"].data == literal.data
+
+
+def _locked_writer(name, ready, release):
+    fd = os.open(f"/dev/shm/{name}", os.O_RDWR)
+    try:
+        fcntl.flock(fd, fcntl.LOCK_EX)
+        ready.set()
+        assert release.wait(10)
+    finally:
+        os.close(fd)
+
+
+@pytest.mark.parametrize("use_metadata", [False, True])
+def test_deadline_receive_never_waits_for_another_process_ring_lock(edge, use_metadata):
+    sender, receiver = edge
+    _, _, metadata = sender.put("0", "1", "locked", torch.arange(3))
+    ctx = multiprocessing.get_context("spawn")
+    ready, release = ctx.Event(), ctx.Event()
+    writer = ctx.Process(target=_locked_writer, args=(metadata["host_ring"]["name"], ready, release))
+    writer.start()
+    try:
+        assert ready.wait(15)
+        start = time.monotonic()
+        assert (
+            receiver.get_with_deadline("0", "1", "locked", metadata if use_metadata else None, deadline=start + 1)
+            is None
+        )
+        assert time.monotonic() - start < 1
+    finally:
+        release.set()
+        writer.join(15)
+        if writer.is_alive():
+            writer.kill()
+            writer.join()
+    assert writer.exitcode == 0
+    assert receiver.get("0", "1", "locked", metadata if use_metadata else None)[0].tolist() == [0, 1, 2]
+
+
+def _competing_reader(scope, barrier, results):
+    with SharedMemoryConnector({"stage_id": 1, "wakeup_scope": scope}) as receiver:
+        barrier.wait(timeout=20)
+        result = receiver.get("0", "1", "shared")
+        results.put(None if result is None else result[0]["codes"].tolist())
+
+
+def test_two_reader_processes_claim_exactly_once_and_close_does_not_unlink_owner():
+    scope = uuid.uuid4().hex
+    ctx = multiprocessing.get_context("spawn")
+    barrier, results = ctx.Barrier(3), ctx.Queue()
+    with SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope}) as sender:
+        _, _, metadata = sender.put("0", "1", "shared", {"codes": torch.arange(56)})
+        readers = [ctx.Process(target=_competing_reader, args=(scope, barrier, results)) for _ in range(2)]
+        for reader in readers:
+            reader.start()
+        try:
+            barrier.wait(timeout=30)
+            outputs = [results.get(timeout=20) for _ in readers]
+            assert outputs.count(None) == 1
+            assert next(output for output in outputs if output is not None) == list(range(56))
+            assert os.path.exists(f"/dev/shm/{metadata['host_ring']['name']}")
+        finally:
+            for reader in readers:
+                reader.join(15)
+                if reader.is_alive():
+                    reader.kill()
+                    reader.join()
+            results.close()
+        assert all(reader.exitcode == 0 for reader in readers)
+
+
+def test_ring_setup_failure_keeps_connector_usable(edge, monkeypatch):
+    sender, receiver = edge
+
+    def no_space(*args):
+        raise OSError(28, "No space left on device")
+
+    monkeypatch.setattr(os, "posix_fallocate", no_space)
+    ok, _, metadata = sender.put("0", "1", "no-ring-space", torch.arange(3))
+    assert ok and "shm" in metadata
+    assert receiver.get("0", "1", "no-ring-space")[0].tolist() == [0, 1, 2]
+
+
+def test_discovery_stamp_failure_keeps_fallback_and_cleanup_usable(edge, monkeypatch):
+    sender, receiver = edge
+
+    def stamp_denied(*args, **kwargs):
+        raise PermissionError("directory timestamp update denied")
+
+    monkeypatch.setattr(os, "utime", stamp_denied)
+    ok, _, metadata = sender.put("0", "1", "no-ring-stamp", torch.arange(3))
+    assert ok and "shm" in metadata
+    assert not sender._host_ring.producers
+    assert not list(Path(f"/dev/shm/{sender._host_ring.directory}").iterdir())
+    assert receiver.get("0", "1", "no-ring-stamp")[0].tolist() == [0, 1, 2]
+    sender.close()
+
+
+def test_close_is_idempotent_and_removes_owned_ring(edge):
+    sender, receiver = edge
+    _, _, metadata = sender.put("0", "1", "close", torch.arange(3))
+    receiver.get("0", "1", "close")
+    sender.close()
+    sender.close()
+    receiver.close()
+    receiver.close()
+    assert not os.path.exists(f"/dev/shm/{metadata['host_ring']['name']}")
+
+
+def _orphan_producer(scope, pipe):
+    producer = SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope})
+    metadata = producer.put("0", "1", "restart", torch.tensor([11]))[2]
+    pipe.send(metadata)
+    pipe.recv()
+    pipe.close()
+    os._exit(0)
+
+
+def test_dead_producer_payload_cannot_reach_a_reused_key():
+    scope = uuid.uuid4().hex
+    ctx = multiprocessing.get_context("spawn")
+    parent, child = ctx.Pipe()
+    process = ctx.Process(target=_orphan_producer, args=(scope, child))
+    process.start()
+    child.close()
+    old = parent.recv()
+    try:
+        with SharedMemoryConnector({"stage_id": 1, "wakeup_scope": scope}) as receiver:
+            # Cache the live channel without claiming the pending frame.
+            assert receiver.get("0", "1", "not-published") is None
+            channel = next(iter(receiver._host_ring.readers.values()))
+            fd = channel.fd
+            parent.send("exit")
+            process.join(15)
+            assert process.exitcode == 0
+            assert receiver.get("0", "1", "restart", old) is None
+            assert not receiver._host_ring.readers
+            with pytest.raises(OSError):
+                os.fstat(fd)
+            with SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope}) as replacement:
+                replacement.put("0", "1", "restart", torch.tensor([22]))
+                assert receiver.get("0", "1", "restart")[0].item() == 22
+    finally:
+        if process.is_alive():
+            process.kill()
+            process.join()
+        # The multiprocessing parent still owns this test's resource tracker.
+        # Retire only the intentionally orphaned allocation and its exact marker.
+        segment = shared_memory.SharedMemory(name=old["host_ring"]["name"])
+        segment.close()
+        segment.unlink()
+        directory = old["host_ring"]["name"].rsplit("_", 4)[0]
+        try:
+            os.unlink(f"/dev/shm/{directory}/{old['host_ring']['name']}")
+        except FileNotFoundError:
+            pass
+        try:
+            os.rmdir(f"/dev/shm/{directory}")
+        except FileNotFoundError:
+            pass
+        parent.close()

vllm_omni/distributed/omni_connectors/connectors/base.py

@@ -56,8 +56,8 @@ class OmniConnectorBase(ABC):
         pass
 
     @abstractmethod
-    def cleanup(self, request_id: str) -> None:
-        """Clean up resources for a request."""
+    def cleanup(self, request_id: str) -> bool | None:
+        """Clean up resources; backends may report whether a payload was retired."""
         pass
 
     def get_with_deadline(

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -17,7 +17,10 @@ from typing import Any
 from vllm_omni.entrypoints.stage_utils import shm_read_bytes, shm_write_bytes
 
 from ..utils.logging import get_connector_logger
+from ..utils.serialization import OmniSerializer
+from ..utils.tensor_frame import TensorFrame, prepare_tensor_frame
 from .base import OmniConnectorBase
+from .shm_ring import HostRingTransport
 
 logger = get_connector_logger(__name__)
 
@@ -55,7 +58,8 @@ class SharedMemoryConnector(OmniConnectorBase):
     def __init__(self, config: dict[str, Any]):
         self.config = config
         self.stage_id = config.get("stage_id", -1)
-        scope = config.get("extra", {}).get("wakeup_scope")
+        extra = config.get("extra", config)
+        scope = extra.get("wakeup_scope")
         self._wake_directory = (
             f"/dev/shm/omni_shm_wake_{os.getuid()}_{hashlib.sha256(str(scope).encode()).hexdigest()[:24]}"
             if scope is not None
@@ -67,7 +71,13 @@ class SharedMemoryConnector(OmniConnectorBase):
             "puts": 0,
             "gets": 0,
             "bytes_transferred": 0,
+            "host_ring_puts": 0,
+            "host_ring_gets": 0,
+            "host_ring_fallbacks": 0,
         }
+        # One bounded host fast path per edge, allocated only on the first put.
+        # A zero capacity preserves the per-key protocol for controlled comparisons.
+        self._host_ring = HostRingTransport(self._wake_directory, int(extra.get("host_ring_bytes", 4 << 20)))
         # Receiver side: FIFO read end (plus a write end of our own, so the
         # FIFO never reports EOF when no sender has it open) and a counter of
         # drained wakeups. Each receiver owns its path, including across restarts.
@@ -200,8 +210,32 @@ class SharedMemoryConnector(OmniConnectorBase):
         data: Any,
     ) -> tuple[bool, int, dict[str, Any] | None]:
         try:
-            self.reap_consumed()
-            payload = self.serialize_obj(data)
+            # Ring put() reclaims its own edge under the publication lock.
+            # Keep the per-key sweep here without taking every ring lock twice.
+            self._reap_segments()
+            use_ring = self._host_ring.capacity and str(from_stage).isdigit() and str(to_stage).isdigit()
+            frame = (
+                prepare_tensor_frame(data, max_bytes=self._host_ring.capacity // 4 - 16 - len(put_key.encode()))
+                if use_ring
+                else None
+            )
+            payload: TensorFrame | bytes = frame if frame is not None else self.serialize_obj(data)
+            ring_meta = self._host_ring.put(str(from_stage), str(to_stage), put_key, payload)
+            if ring_meta is not None:
+                # A previous put may have used the per-key fallback.
+                with self._pending_keys_lock:
+                    had_segment = put_key in self._pending_keys
+                if had_segment:
+                    self._cleanup_segment(put_key)
+                size = ring_meta["size"]
+                self._wake_receiver(to_stage)
+                self._metrics["puts"] += 1
+                self._metrics["host_ring_puts"] += 1
+                self._metrics["bytes_transferred"] += size
+                return True, size, {"host_ring": ring_meta, "size": size}
+            self._metrics["host_ring_fallbacks"] += 1
+            if isinstance(payload, TensorFrame):
+                payload = self.serialize_obj(data)
             size = len(payload)
 
             lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock"
@@ -214,6 +248,8 @@ class SharedMemoryConnector(OmniConnectorBase):
             metadata = {"shm": meta, "size": size}
             with self._pending_keys_lock:
                 self._pending_keys[put_key] = None
+            # A successful fallback replaces an unclaimed ring version, too.
+            self._host_ring.cancel(key=put_key)
             self._wake_receiver(to_stage)
 
             self._metrics["puts"] += 1
@@ -236,7 +272,7 @@ class SharedMemoryConnector(OmniConnectorBase):
             obj = self.deserialize_obj(data_bytes)
             result = (obj, int(shm_handle.get("size", 0)))
             return result
-        except BlockingIOError:
+        except (BlockingIOError, FileNotFoundError):
             return None
         except Exception as e:
             logger.error(f"SharedMemoryConnector shm get failed for req : {e}")
@@ -250,13 +286,15 @@ class SharedMemoryConnector(OmniConnectorBase):
 
     def _get_by_key(self, get_key: str) -> tuple[Any, int] | None:
         """Read a SHM segment addressed purely by *get_key*."""
-        shm = None
         try:
-            shm = shm_pkg.SharedMemory(name=get_key)
-            if shm is None or shm.size == 0:
+            # Inspect size without registering a reader before it owns the
+            # delivery lock. Competing readers must not re-register an already
+            # consumed allocation with the application's resource tracker.
+            size = os.stat(f"/dev/shm/{get_key}").st_size
+            if size == 0:
                 return None
             lock_file = f"/dev/shm/shm_{get_key}_lockfile.lock"
-            shm_handle = {"name": get_key, "size": shm.size}
+            shm_handle = {"name": get_key, "size": size}
             result = self._get_data_with_lock(lock_file, shm_handle)
             if result is not None:
                 with self._pending_keys_lock:
@@ -264,20 +302,9 @@ class SharedMemoryConnector(OmniConnectorBase):
             return result
         except FileNotFoundError:
             return None
-        except ValueError as e:
-            # A receiver can observe a newly-created POSIX SHM object before
-            # the writer has finished sizing it. Treat that as "not ready yet"
-            # so async polling can retry without a traceback.
-            if "empty file" in str(e):
-                return None
-            logger.debug("_get_by_key: unexpected error reading SHM segment %s", get_key, exc_info=True)
-            return None
         except Exception:
             logger.debug("_get_by_key: unexpected error reading SHM segment %s", get_key, exc_info=True)
             return None
-        finally:
-            if shm:
-                shm.close()
 
     def get(
         self,
@@ -286,6 +313,23 @@ class SharedMemoryConnector(OmniConnectorBase):
         get_key: str,
         metadata=None,
     ) -> tuple[Any, int] | None:
+        if isinstance(metadata, dict) and get_key in metadata:
+            metadata = metadata[get_key]
+        ring_meta = metadata.get("host_ring") if isinstance(metadata, dict) else None
+        ring_result = self._host_ring.get(str(from_stage), str(to_stage), get_key, ring_meta)
+        if ring_result is not None:
+            obj, size = ring_result
+            if isinstance(obj, bytes):
+                obj = self.deserialize_obj(obj)
+            else:
+                # Preserve the centralized serializer's output reconstruction,
+                # including RequestOutput/CompletionOutput dictionaries.
+                obj = OmniSerializer.restore(obj)
+            self._metrics["gets"] += 1
+            self._metrics["host_ring_gets"] += 1
+            return obj, size
+        if ring_meta is not None:
+            return None
         if metadata is not None:
             if isinstance(metadata, dict) and get_key in metadata:
                 metadata = metadata.get(get_key)
@@ -312,7 +356,10 @@ class SharedMemoryConnector(OmniConnectorBase):
 
         Returns True when an unconsumed segment was actually unlinked.
         """
-        key = request_id
+        ring_cancelled = self._host_ring.cancel(key=request_id)
+        return self._cleanup_segment(request_id) or bool(ring_cancelled)
+
+    def _cleanup_segment(self, key: str) -> bool:
         unlinked = False
         with self._pending_keys_lock:
             self._pending_keys.pop(key, None)
@@ -343,10 +390,14 @@ class SharedMemoryConnector(OmniConnectorBase):
         """
         with self._pending_keys_lock:
             keys = [k for k in self._pending_keys if k.startswith(key_prefix) and k[len(key_prefix) :].isdigit()]
-        return sum(self.cleanup(key) for key in keys)
+        return sum(self._cleanup_segment(key) for key in keys) + self._host_ring.cancel(prefix=key_prefix)
 
     def close(self) -> None:
         """Unlink all remaining tracked SHM segments."""
+        # A constructor failure may be followed by the base destructor.
+        if not hasattr(self, "_host_ring"):
+            return
+        self._host_ring.close()
         with self._pending_keys_lock:
             keys = list(self._pending_keys)
         for key in keys:
@@ -355,6 +406,10 @@ class SharedMemoryConnector(OmniConnectorBase):
 
     def reap_consumed(self) -> None:
         """Bounded round-robin sweep; receivers unlink SHM in another process."""
+        self._host_ring.reap()
+        self._reap_segments()
+
+    def _reap_segments(self) -> None:
         with self._pending_keys_lock:
             for _ in range(min(64, len(self._pending_keys))):
                 key, _ = self._pending_keys.popitem(last=False)

vllm_omni/distributed/omni_connectors/connectors/shm_ring.py

@@ -0,0 +1,489 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Bounded host frames shared by local stage processes.
+
+Each producer/destination edge owns one mmap, with a lock on that same file.
+Readers copy before marking a frame consumed, and all readers participate in
+the same lock. Pressure returns False to the existing per-key SHM path; this
+fast path does not implement end-to-end stream admission or GPU ownership.
+"""
+
+from __future__ import annotations
+
+import fcntl
+import glob
+import hashlib
+import mmap
+import os
+import select
+import struct
+import threading
+import time
+import uuid
+from collections import deque
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from multiprocessing import shared_memory as shm_pkg
+
+from ..utils.logging import get_connector_logger
+from ..utils.tensor_frame import TensorFrame, read_tensor_frame
+
+logger = get_connector_logger(__name__)
+
+_HEADER_SIZE = 64
+_HEADER = struct.Struct("<8sQQQQQ")  # magic, capacity, head, tail, creator pid, process start
+_MAGIC = b"OMSHMR01"
+_FRAME = struct.Struct("<IIBBHI")  # total, payload size, kind, state, key length, reserved
+_ALIGN = 16
+_READY, _DONE = 0, 1
+_BYTES, _TENSORS, _PADDING = 0, 1, 2
+_STATE_OFFSET = 9
+_DEFAULT_BYTES = 4 << 20
+
+
+def _align(size: int) -> int:
+    return (size + _ALIGN - 1) & ~(_ALIGN - 1)
+
+
+def _process_start(pid: int) -> int:
+    with open(f"/proc/{pid}/stat") as file:
+        return int(file.read().rsplit(")", 1)[1].split()[19])
+
+
+@contextmanager
+def _registry_update(directory: str) -> Iterator[int]:
+    """Publish a strictly changing discovery stamp at channel birth/retirement.
+
+    Kernel directory mtimes can coincide within a clock tick. Serialize marker
+    updates and preserve the previous stamp before the filesystem changes it.
+    Receivers can then cache discovery without a timer or a per-message scan.
+    """
+    fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY)
+    try:
+        fcntl.flock(fd, fcntl.LOCK_EX)
+        previous = os.fstat(fd).st_mtime_ns
+        yield fd
+        stamp = max(previous + 1, time.time_ns())
+        os.utime(fd, ns=(stamp, stamp))
+    finally:
+        os.close(fd)
+
+
+@dataclass
+class _Channel:
+    name: str
+    fd: int
+    mapping: mmap.mmap
+    capacity: int
+    owner: bool
+    creator_pid: int
+    creator_start: int
+    segment: shm_pkg.SharedMemory | None = None
+    registry: str | None = None
+    pidfd: int | None = None
+    cursor: int = 0
+    index: dict[str, int] = field(default_factory=dict)
+    positions: deque[tuple[int, str]] = field(default_factory=deque)
+
+    @classmethod
+    def open(cls, name: str, capacity: int | None = None) -> _Channel:
+        owner = capacity is not None
+        segment = None
+        mapping = None
+        pidfd = None
+        fd = -1
+        try:
+            if capacity is not None:
+                # Only the producer registers the allocation with Python's
+                # resource tracker. Readers map the existing file directly.
+                segment = shm_pkg.SharedMemory(name=name, create=True, size=_HEADER_SIZE + capacity)
+            fd = os.open(f"/dev/shm/{name}", os.O_RDWR | os.O_NOFOLLOW)
+            fcntl.flock(fd, fcntl.LOCK_EX | (0 if owner else fcntl.LOCK_NB))
+            try:
+                if capacity is not None:
+                    os.posix_fallocate(fd, 0, _HEADER_SIZE + capacity)
+                size = os.fstat(fd).st_size
+                if size < _HEADER_SIZE:
+                    raise ValueError("incomplete SHM ring header")
+                mapping = mmap.mmap(fd, size)
+                if capacity is not None:
+                    # Fault once per edge, outside the per-message path.
+                    mapping[:] = bytes(size)
+                    _HEADER.pack_into(mapping, 0, _MAGIC, capacity, 0, 0, os.getpid(), _process_start(os.getpid()))
+                magic, actual_capacity, head, tail, creator_pid, creator_start = _HEADER.unpack_from(mapping)
+                if magic != _MAGIC or actual_capacity != size - _HEADER_SIZE or not 0 <= tail <= head:
+                    raise ValueError("invalid SHM ring header")
+                if not owner:
+                    try:
+                        pidfd = os.pidfd_open(creator_pid)
+                    except (AttributeError, OSError):
+                        # Older Linux kernels use the process-start check on receive.
+                        pidfd = None
+                    if _process_start(creator_pid) != creator_start:
+                        raise FileNotFoundError("SHM ring producer was replaced")
+                return cls(
+                    name, fd, mapping, actual_capacity, owner, creator_pid, creator_start, segment=segment, pidfd=pidfd
+                )
+            finally:
+                fcntl.flock(fd, fcntl.LOCK_UN)
+        except BaseException:
+            if mapping is not None:
+                mapping.close()
+            if fd >= 0:
+                os.close(fd)
+            if segment is not None:
+                segment.close()
+                segment.unlink()
+            if pidfd is not None:
+                os.close(pidfd)
+            raise
+
+    @contextmanager
+    def locked(self, *, blocking: bool = True) -> Iterator[None]:
+        fcntl.flock(self.fd, fcntl.LOCK_EX | (0 if blocking else fcntl.LOCK_NB))
+        try:
+            yield
+        finally:
+            fcntl.flock(self.fd, fcntl.LOCK_UN)
+
+    def bounds(self) -> tuple[int, int]:
+        _magic, _capacity, head, tail, _pid, _start = _HEADER.unpack_from(self.mapping)
+        return head, tail
+
+    def producer_alive(self) -> bool:
+        if self.pidfd is not None:
+            return not select.select([self.pidfd], [], [], 0)[0]
+        try:
+            return _process_start(self.creator_pid) == self.creator_start
+        except FileNotFoundError:
+            return False
+
+    def frame(self, position: int) -> tuple[int, int, int, int, int, int]:
+        frame = _FRAME.unpack_from(self.mapping, _HEADER_SIZE + position % self.capacity)
+        total, length, kind, _state, key_length, _reserved = frame
+        if total < _FRAME.size or total % _ALIGN or position % self.capacity + total > self.capacity:
+            raise ValueError("invalid SHM ring frame size")
+        if kind != _PADDING and _FRAME.size + length + key_length > total:
+            raise ValueError("invalid SHM ring payload bounds")
+        return frame
+
+    def scan(self) -> None:
+        head, tail = self.bounds()
+        self.cursor = max(self.cursor, tail)
+        while self.cursor < head:
+            position = self.cursor
+            total, _length, kind, state, key_length, _reserved = self.frame(position)
+            self.cursor += total
+            if kind != _PADDING and state == _READY:
+                start = _HEADER_SIZE + position % self.capacity + _FRAME.size
+                key = self.mapping[start : start + key_length].decode()
+                self.index[key] = position
+                self.positions.append((position, key))
+        # Remove positions retired by another reader or overwritten after wrap.
+        while self.positions and self.positions[0][0] < tail:
+            position, key = self.positions.popleft()
+            if self.index.get(key) == position:
+                del self.index[key]
+
+    def consume(self, position: int) -> None:
+        self.mapping[_HEADER_SIZE + position % self.capacity + _STATE_OFFSET] = _DONE
+
+    def reclaim(self) -> None:
+        head, tail = self.bounds()
+        while tail < head:
+            total, _length, kind, state, _key_length, _reserved = self.frame(tail)
+            if kind != _PADDING and state != _DONE:
+                break
+            tail += total
+        struct.pack_into("<Q", self.mapping, 24, tail)
+
+    def cancel(self, key: str | None = None, prefix: str | None = None) -> int:
+        self.scan()
+        count = 0
+        candidates = (
+            (((key, self.index[key]),) if key in self.index else ()) if key is not None else tuple(self.index.items())
+        )
+        for candidate, position in candidates:
+            matches = candidate == key or (
+                prefix is not None and candidate.startswith(prefix) and candidate[len(prefix) :].isdigit()
+            )
+            if matches:
+                if self.frame(position)[3] == _READY:
+                    self.consume(position)
+                    count += 1
+                del self.index[candidate]
+        self.reclaim()
+        return count
+
+    def close(self) -> None:
+        self.mapping.close()
+        os.close(self.fd)
+        if self.pidfd is not None:
+            os.close(self.pidfd)
+        if self.segment is not None:
+            self.segment.close()
+            try:
+                self.segment.unlink()
+            except FileNotFoundError:
+                pass
+        if self.registry is not None:
+            try:
+                with _registry_update(os.path.dirname(self.registry)) as directory_fd:
+                    os.unlink(os.path.basename(self.registry), dir_fd=directory_fd)
+            except FileNotFoundError:
+                pass
+            except OSError as error:
+                # The allocation is already retired. A failed discovery stamp
+                # must not mask setup failure or prevent the per-key fallback.
+                logger.debug("Failed to retire host ring marker %s: %s", self.registry, error)
+
+
+class HostRingTransport:
+    """Private SharedMemoryConnector fast path; no new model-specific API."""
+
+    def __init__(self, scope: str, capacity: int = _DEFAULT_BYTES):
+        if capacity < 0 or capacity and not 4096 <= capacity <= 64 << 20:
+            raise ValueError("host_ring_bytes must be zero or between 4096 and 67108864")
+        self.capacity = _align(capacity)
+        scope_id = hashlib.sha256(scope.encode()).hexdigest()[:24]
+        self.directory = f"omni_shm_ring_v1_{os.getuid()}_{scope_id}"
+        self.prefix = f"{self.directory}_"
+        self.producers: dict[tuple[str, str], _Channel] = {}
+        self._disabled_edges: set[tuple[str, str]] = set()
+        self.readers: dict[str, _Channel] = {}
+        self._discovery_versions: dict[tuple[str, str], tuple[int, int, int]] = {}
+        self._lock = threading.RLock()
+        self._closed = False
+
+    def put(self, source: str, destination: str, key: str, payload: bytes | TensorFrame) -> dict | None:
+        """Return an immutable descriptor, or None without consuming ring credit."""
+        encoded_key = key.encode()
+        length = payload.size if isinstance(payload, TensorFrame) else len(payload)
+        total = _align(_FRAME.size + len(encoded_key) + length)
+        if not self.capacity or not source.isdigit() or not destination.isdigit():
+            return None
+        # Keep one exceptional payload from monopolizing a channel.
+        if total > self.capacity // 4 or len(encoded_key) > 65535:
+            return None
+        with self._lock:
+            if self._closed:
+                raise RuntimeError("host ring is closed")
+            if (source, destination) in self._disabled_edges:
+                return None
+            channel = self.producers.get((source, destination))
+            if channel is None:
+                name = f"{self.prefix}{source}_{destination}_{os.getpid()}_{uuid.uuid4().hex}"
+                try:
+                    os.makedirs(f"/dev/shm/{self.directory}", mode=0o700, exist_ok=True)
+                    channel = _Channel.open(name, self.capacity)
+                    channel.registry = f"/dev/shm/{self.directory}/{name}"
+                    # Discovery is confined to this deployment directory. The
+                    # marker appears only after the allocation is ready.
+                    with _registry_update(os.path.dirname(channel.registry)) as directory_fd:
+                        marker = os.open(name, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600, dir_fd=directory_fd)
+                        os.close(marker)
+                except OSError as error:
+                    if channel is not None:
+                        channel.close()
+                    self._disabled_edges.add((source, destination))
+                    logger.warning(
+                        "Host ring unavailable on edge %s -> %s; using per-key SHM: %s", source, destination, error
+                    )
+                    return None
+                self.producers[source, destination] = channel
+            with channel.locked():
+                channel.reclaim()
+                head, tail = channel.bounds()
+                position = head % channel.capacity
+                padding = channel.capacity - position if position + total > channel.capacity else 0
+                if padding + total > channel.capacity - (head - tail):
+                    return None
+                if padding:
+                    _FRAME.pack_into(channel.mapping, _HEADER_SIZE + position, padding, 0, _PADDING, _DONE, 0, 0)
+                    head += padding
+                    position = 0
+                start = _HEADER_SIZE + position + _FRAME.size
+                channel.mapping[start : start + len(encoded_key)] = encoded_key
+                start += len(encoded_key)
+                if isinstance(payload, TensorFrame):
+                    view = memoryview(channel.mapping)
+                    try:
+                        payload.write(view, start)
+                    finally:
+                        view.release()
+                else:
+                    channel.mapping[start : start + length] = payload
+                _FRAME.pack_into(
+                    channel.mapping,
+                    _HEADER_SIZE + position,
+                    total,
+                    length,
+                    _TENSORS if isinstance(payload, TensorFrame) else _BYTES,
+                    _READY,
+                    len(encoded_key),
+                    0,
+                )
+                # Retire an earlier version while holding the same lock as all readers.
+                channel.cancel(key=key)
+                struct.pack_into("<Q", channel.mapping, 16, head + total)
+                channel.index[key] = head
+                channel.positions.append((head, key))
+                channel.cursor = head + total
+            return {"name": channel.name, "position": head, "size": length}
+
+    def _reader(self, name: str) -> _Channel:
+        if not name.startswith(f"omni_shm_ring_v1_{os.getuid()}_") or not name.replace("_", "").isalnum():
+            raise ValueError("invalid local SHM ring name")
+        channel = self.readers.get(name)
+        if channel is None:
+            channel = _Channel.open(name)
+            self.readers[name] = channel
+        return channel
+
+    def _receive(self, channel: _Channel, key: str, descriptor: dict | None) -> tuple[object, int] | None:
+        if not channel.producer_alive():
+            self.readers.pop(channel.name, None)
+            channel.close()
+            self._discovery_versions.clear()
+            return None
+        # This optimistic check only rejects a miss; successful reads still
+        # take the shared lock. An unchanged publication cursor cannot have
+        # introduced a new key since this receiver's last locked scan.
+        if key not in channel.index and struct.unpack_from("<Q", channel.mapping, 16)[0] == channel.cursor:
+            return None
+        try:
+            fcntl.flock(channel.fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except BlockingIOError:
+            return None
+        try:
+            channel.scan()
+            position = channel.index.get(key)
+            if position is None or descriptor is not None and position != descriptor["position"]:
+                return None
+            _total, length, kind, state, key_length, _reserved = channel.frame(position)
+            del channel.index[key]
+            if state != _READY:
+                return None
+            start = _HEADER_SIZE + position % channel.capacity + _FRAME.size + key_length
+            view = memoryview(channel.mapping)[start : start + length]
+            try:
+                # No reference to ring storage escapes this locked boundary.
+                result = read_tensor_frame(view) if kind == _TENSORS else bytes(view)
+            finally:
+                view.release()
+                channel.consume(position)
+            return result, length
+        finally:
+            fcntl.flock(channel.fd, fcntl.LOCK_UN)
+
+    def get(self, source: str, destination: str, key: str, descriptor: dict | None) -> tuple[object, int] | None:
+        with self._lock:
+            if self._closed:
+                return None
+            if descriptor is not None:
+                try:
+                    result = self._receive(self._reader(descriptor["name"]), key, descriptor)
+                    if result is None:
+                        self._reap_readers()
+                    return result
+                except (FileNotFoundError, BlockingIOError):
+                    return None
+            if not self.capacity or not source.isdigit() or not destination.isdigit():
+                return None
+            edge_prefix = f"{self.prefix}{source}_{destination}_"
+            for name, channel in tuple(self.readers.items()):
+                if name.startswith(edge_prefix):
+                    result = self._receive(channel, key, None)
+                    if result is not None:
+                        return result
+            directory = f"/dev/shm/{self.directory}"
+            try:
+                status = os.stat(directory)
+            except FileNotFoundError:
+                self._discovery_versions.clear()
+                self._reap_readers()
+                return None
+            version = (status.st_dev, status.st_ino, status.st_mtime_ns)
+            edge = (source, destination)
+            if version == self._discovery_versions.get(edge):
+                return None
+            self._reap_readers()
+            # A new producer advances the stamp before publishing any frames,
+            # retaining first-message discovery without scanning every miss.
+            discovered_result = None
+            discovery_complete = True
+            for path in glob.iglob(f"/dev/shm/{self.directory}/{edge_prefix}*"):
+                name = os.path.basename(path)
+                if name in self.readers:
+                    continue
+                try:
+                    channel = self._reader(name)
+                except FileNotFoundError:
+                    # A published marker whose allocation or producer is gone
+                    # cannot become valid again: allocation names are unique.
+                    try:
+                        with _registry_update(directory) as directory_fd:
+                            os.unlink(name, dir_fd=directory_fd)
+                    except FileNotFoundError:
+                        pass
+                    except OSError as error:
+                        logger.debug("Failed to prune host ring marker %s: %s", path, error)
+                    continue
+                except (BlockingIOError, ValueError):
+                    # Attachment may race a writer holding the frame lock.
+                    # Retry discovery after it releases, even at this stamp.
+                    discovery_complete = False
+                    continue
+                if discovered_result is None:
+                    discovered_result = self._receive(channel, key, None)
+            # Attach every producer before caching this edge's stamp. Returning
+            # as soon as one frame is found would hide its unvisited peers.
+            if discovery_complete:
+                self._discovery_versions[edge] = version
+            else:
+                self._discovery_versions.pop(edge, None)
+            return discovered_result
+
+    def cancel(self, key: str | None = None, prefix: str | None = None) -> int:
+        with self._lock:
+            count = 0
+            for channel in self.producers.values():
+                with channel.locked():
+                    count += channel.cancel(key, prefix)
+            return count
+
+    def reap(self) -> None:
+        with self._lock:
+            for channel in self.producers.values():
+                with channel.locked():
+                    channel.reclaim()
+            self._reap_readers()
+
+    def _reap_readers(self) -> None:
+        for name, channel in tuple(self.readers.items()):
+            if not channel.producer_alive() or not os.path.exists(f"/dev/shm/{name}"):
+                channel.close()
+                del self.readers[name]
+
+    def close(self) -> None:
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+            for channel in self.producers.values():
+                with channel.locked():
+                    channel.scan()
+                    for position in channel.index.values():
+                        channel.consume(position)
+                channel.close()
+            for channel in self.readers.values():
+                channel.close()
+            self.producers.clear()
+            self.readers.clear()
+            self._discovery_versions.clear()
+            try:
+                os.rmdir(f"/dev/shm/{self.directory}")
+            except (FileNotFoundError, OSError):
+                # Other producers in this deployment may still own rings.
+                pass

vllm_omni/distributed/omni_connectors/utils/serialization.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 collections.abc import Mapping
 from dataclasses import asdict, fields, is_dataclass
@@ -205,7 +205,7 @@ class OmniMsgpackDecoder:
                 return self._decode_pil_image(obj)
 
             # Process values recursively first
-            processed = {k: self._post_process(v) for k, v in obj.items()}
+            processed = {k: self._post_process(v) if isinstance(v, (dict, list, tuple)) else v for k, v in obj.items()}
 
             # Check if this looks like an OmniRequestOutput (check before RequestOutput
             # since OmniRequestOutput may also have some RequestOutput-like fields)
@@ -223,10 +223,13 @@ class OmniMsgpackDecoder:
             return processed
 
         if isinstance(obj, list):
-            return [self._post_process(item) for item in obj]
+            # Native msgpack already owns scalar values. Only containers can
+            # contain wire markers; avoid a recursive call for every token or
+            # float in long model payloads, while preserving new-list semantics.
+            return [self._post_process(item) if isinstance(item, (dict, list, tuple)) else item for item in obj]
 
         if isinstance(obj, tuple):
-            return tuple(self._post_process(item) for item in obj)
+            return tuple(self._post_process(item) if isinstance(item, (dict, list, tuple)) else item for item in obj)
 
         return obj
 
@@ -350,6 +353,10 @@ class OmniSerde:
         """Deserialize bytes to an object."""
         return self.decoder.decode(data)
 
+    def restore(self, obj: object) -> object:
+        """Restore wire objects from an already decoded transport tree."""
+        return self.decoder._post_process(obj)
+
 
 # Global instance for simple interface
 OmniSerializer = OmniSerde()

vllm_omni/distributed/omni_connectors/utils/tensor_frame.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Owned CPU tensor trees with a msgpack header and raw tensor regions.
+
+Tensor references use message-scoped msgpack extensions, preserving literal
+extensions and native scalar containers without reserved dictionary keys. It never
+exposes a view of transport storage to a caller. Unsupported objects retain
+the ordinary OmniSerializer path.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import math
+import struct
+import uuid
+from dataclasses import dataclass
+
+import msgspec
+import numpy as np
+import torch
+from PIL import Image
+
+from .serialization import OmniSerializer
+
+_LENGTH = struct.Struct("<I")
+_ENCODER = msgspec.msgpack.Encoder()
+_REFERENCE = struct.Struct("<16sI")
+_REFERENCE_CODE = 42
+_DECODER = msgspec.msgpack.Decoder()
+
+
+def _align(size: int) -> int:
+    return (size + 15) & ~15
+
+
+class _UnsupportedTreeError(Exception):
+    pass
+
+
+@dataclass
+class TensorFrame:
+    header: bytes
+    tensors: list[torch.Tensor]
+    offsets: list[int]
+    size: int
+
+    def write(self, buffer: memoryview, start: int) -> None:
+        """Copy directly into caller-owned writable storage, while it is locked."""
+        _LENGTH.pack_into(buffer, start, len(self.header))
+        buffer[start + 4 : start + 4 + len(self.header)] = self.header
+        data_start = start + _align(4 + len(self.header))
+        # c_char is an anchor only; all bounds are checked by the ring writer.
+        anchor = ctypes.c_char.from_buffer(buffer)
+        base = ctypes.addressof(anchor) + data_start
+        for tensor, offset in zip(self.tensors, self.offsets, strict=True):
+            if tensor.nbytes:
+                ctypes.memmove(base + offset, tensor.data_ptr(), tensor.nbytes)
+
+
+def prepare_tensor_frame(payload: object, *, max_bytes: int | None = None) -> TensorFrame | bytes | None:
+    """Encode native containers once, or prepare their dense CPU tensor regions.
+
+    Scalar/token lists stay in msgpack's C encoder. Tensor-free messages reuse
+    those encoded bytes, rather than constructing and discarding a tagged tree.
+    Sources remain caller-owned until the synchronous put() returns.
+    """
+    tensors: list[torch.Tensor] = []
+    offsets: list[int] = []
+    descriptors: list[tuple[str, tuple[int, ...], int]] = []
+    data_size = 0
+    nonce: bytes | None = None
+
+    def encode_tensor(value: object) -> object:
+        nonlocal data_size, nonce
+        if type(value) is not torch.Tensor:
+            # Reuse the common wire representations for mixed model payloads.
+            # Large arrays/images keep the ordinary path without first making
+            # an extra byte snapshot just to discover that the frame cannot fit.
+            if max_bytes is not None:
+                if isinstance(value, np.ndarray) and value.nbytes > max_bytes:
+                    raise _UnsupportedTreeError
+                if isinstance(value, Image.Image) and value.width * value.height * len(value.getbands()) > max_bytes:
+                    raise _UnsupportedTreeError
+            return OmniSerializer.encoder._enc_hook(value)
+        if value.device.type != "cpu" or value.layout != torch.strided or value.is_quantized:
+            raise _UnsupportedTreeError
+        offset = _align(data_size)
+        index = len(tensors)
+        tensors.append(value)
+        offsets.append(offset)
+        descriptors.append((str(value.dtype).removeprefix("torch."), tuple(value.shape), offset))
+        data_size = offset + value.nbytes
+        if max_bytes is not None and data_size > max_bytes:
+            raise _UnsupportedTreeError
+        if nonce is None:
+            nonce = uuid.uuid4().bytes
+        return msgspec.msgpack.Ext(_REFERENCE_CODE, _REFERENCE.pack(nonce, index))
+
+    try:
+        tree = msgspec.msgpack.Encoder(enc_hook=encode_tensor).encode(payload)
+        if not tensors:
+            return tree
+        header = _ENCODER.encode([descriptors, tree, nonce])
+    except (_UnsupportedTreeError, TypeError, OverflowError, RecursionError):
+        return None
+    size = _align(4 + len(header)) + data_size
+    if max_bytes is not None and size > max_bytes:
+        return None
+    # Only materialize non-contiguous/flagged storage after the entire message
+    # is known to fit. Unsupported and large frames avoid an unused copy.
+    for index, tensor in enumerate(tensors):
+        if tensor.is_conj():
+            tensor = tensor.resolve_conj()
+        if tensor.is_neg():
+            tensor = tensor.resolve_neg()
+        if not tensor.is_contiguous():
+            tensor = tensor.contiguous()
+        tensors[index] = tensor
+    return TensorFrame(header, tensors, offsets, size)
+
+
+def read_tensor_frame(buffer: memoryview) -> object:
+    """Decode into independent CPU allocations before returning ring credit."""
+    if len(buffer) < 4:
+        raise ValueError("truncated tensor frame")
+    header_size = _LENGTH.unpack_from(buffer)[0]
+    data_start = _align(4 + header_size)
+    if data_start > len(buffer):
+        raise ValueError("tensor frame header exceeds payload")
+    header_view = buffer[4 : 4 + header_size]
+    try:
+        descriptors, tree, nonce = _DECODER.decode(header_view)
+    finally:
+        header_view.release()
+    # All leaves share one independently owned allocation. This amortizes CPU
+    # allocation/copy dispatch across small codec groups and hidden-state leaves.
+    data_size = len(buffer) - data_start
+    # A small metadata tensor must not keep a large, already consumed hidden
+    # state alive. Pack small trees; give large multi-leaf trees separate owners.
+    storage = torch.empty(data_size, dtype=torch.uint8) if data_size <= 65536 or len(descriptors) == 1 else None
+    anchor = ctypes.c_char.from_buffer(buffer)
+    if data_size and storage is not None:
+        ctypes.memmove(storage.data_ptr(), ctypes.addressof(anchor) + data_start, data_size)
+    tensors = []
+    for name, shape, offset in descriptors:
+        dtype = getattr(torch, name, None)
+        if not isinstance(dtype, torch.dtype):
+            raise ValueError("unknown tensor frame dtype")
+        if any(type(dim) is not int or dim < 0 for dim in shape) or type(offset) is not int or offset < 0:
+            raise ValueError("invalid tensor frame region")
+        count = math.prod(shape)
+        size = count * dtype.itemsize
+        start = data_start + offset
+        if start + size > len(buffer):
+            raise ValueError("tensor frame region exceeds payload")
+        if size and storage is not None:
+            tensor = storage[offset : offset + size].view(dtype).reshape(shape)
+        else:
+            tensor = torch.empty(shape, dtype=dtype)
+            if size:
+                ctypes.memmove(tensor.data_ptr(), ctypes.addressof(anchor) + start, size)
+        tensors.append(tensor)
+
+    def restore_reference(code: int, data: memoryview) -> object:
+        if code == _REFERENCE_CODE and len(data) == _REFERENCE.size:
+            reference_nonce, index = _REFERENCE.unpack(data)
+            if reference_nonce == nonce:
+                if index >= len(tensors):
+                    raise ValueError("invalid tensor frame reference")
+                return tensors[index]
+        # Literal user extensions keep their existing msgpack wire semantics.
+        return msgspec.msgpack.Ext(code, bytes(data))
+
+    return msgspec.msgpack.decode(tree, ext_hook=restore_reference)