[Core] Reduce local stage payload overhead with shared host rings

Signed-off-by: Sy03 <1370724210@qq.com>
This commit is contained in:
Sy03 committed 2026-10-05 21:56:29 +08:00
1 parent 023716ecc4
commit 06202d8f2a
9 files changed
+1108 -12

No files matched your search

@@ -0,0 +1,79 @@
# 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.
`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.
@@ -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"
@@ -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()
@@ -0,0 +1,346 @@
# 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
import pytest
import torch
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.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))
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_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": {1, 2}}) is None
assert sender.put("0", "1", "falsey", False)[0]
assert receiver.get("0", "1", "falsey")[0] is False
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)
def test_deadline_receive_never_waits_for_another_process_ring_lock(edge):
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, 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)[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_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()
@@ -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(
@@ -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.
@@ -201,7 +211,29 @@ class SharedMemoryConnector(OmniConnectorBase):
) -> tuple[bool, int, dict[str, Any] | None]:
try:
self.reap_consumed()
payload = self.serialize_obj(data)
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 +246,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
@@ -286,6 +320,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 +363,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 +397,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 +413,7 @@ class SharedMemoryConnector(OmniConnectorBase):
def reap_consumed(self) -> None:
"""Bounded round-robin sweep; receivers unlink SHM in another process."""
self._host_ring.reap()
with self._pending_keys_lock:
for _ in range(min(64, len(self._pending_keys))):
key, _ = self._pending_keys.popitem(last=False)
@@ -0,0 +1,431 @@
# 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 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])
@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:
os.unlink(self.registry)
except FileNotFoundError:
pass
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._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.
marker = os.open(channel.registry, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
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()
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
self._reap_readers()
# Re-discover only on a miss, retaining first-message delivery and
# support for newly started producer processes without a poll timer.
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:
os.unlink(path)
except FileNotFoundError:
pass
continue
except (BlockingIOError, ValueError):
continue
result = self._receive(channel, key, None)
if result is not None:
return result
return None
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()
try:
os.rmdir(f"/dev/shm/{self.directory}")
except (FileNotFoundError, OSError):
# Other producers in this deployment may still own rings.
pass
@@ -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
@@ -350,6 +350,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()
@@ -0,0 +1,174 @@
# 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.
The format uses tagged nodes rather than 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
from dataclasses import dataclass
import msgspec
import torch
_LENGTH = struct.Struct("<I")
_ENCODER = msgspec.msgpack.Encoder()
_DECODER = msgspec.msgpack.Decoder()
_SCALAR, _DICT, _LIST, _TENSOR = range(4)
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 | None:
"""Prepare a plain tree containing dense CPU tensors; otherwise return None.
Retaining tensors is enough because the caller writes before put() returns.
This is not an asynchronous snapshot or a CUDA source ownership protocol.
"""
tensors: list[torch.Tensor] = []
offsets: list[int] = []
descriptors: list[tuple[str, tuple[int, ...], int]] = []
data_size = 0
def visit(value: object, depth: int) -> list:
nonlocal data_size
if depth > 64:
raise _UnsupportedTreeError
if type(value) in (type(None), bool, int, float, str, bytes):
return [_SCALAR, value]
if type(value) is torch.Tensor:
if value.device.type != "cpu" or value.layout != torch.strided or value.is_quantized:
raise _UnsupportedTreeError
tensor = value
offset = _align(data_size)
index = len(tensors)
tensors.append(tensor)
offsets.append(offset)
descriptors.append((str(tensor.dtype).removeprefix("torch."), tuple(tensor.shape), offset))
data_size = offset + tensor.nbytes
if max_bytes is not None and data_size > max_bytes:
raise _UnsupportedTreeError
return [_TENSOR, index]
if type(value) is dict and all(type(key) is str for key in value):
return [_DICT, {key: visit(item, depth + 1) for key, item in value.items()}]
if isinstance(value, (list, tuple)) and type(value) in (list, tuple):
# msgpack represents both as arrays; match OmniSerializer's wire type.
return [_LIST, [visit(item, depth + 1) for item in value]]
if isinstance(value, msgspec.Struct):
try:
fields = msgspec.to_builtins(value, builtin_types=(torch.Tensor, bytes))
except TypeError:
raise _UnsupportedTreeError from None
return visit(fields, depth + 1)
raise _UnsupportedTreeError
try:
tree = visit(payload, 0)
if not tensors:
return None
header = _ENCODER.encode([descriptors, tree])
except (_UnsupportedTreeError, OverflowError):
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 tree
# is known to fit. Large and unsupported payloads 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 = _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(node: list) -> object:
tag, value = node
if tag == _SCALAR:
return value
if tag == _TENSOR:
return tensors[value]
if tag == _DICT:
return {key: restore(item) for key, item in value.items()}
if tag == _LIST:
return [restore(item) for item in value]
raise ValueError("unknown tensor frame node")
return restore(tree)