From 06202d8f2a1b82d40670422e66b4305cb563b512 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Mon, 5 Oct 2026 04:06:23 +0800 Subject: [PATCH] [Core] Reduce local stage payload overhead with shared host rings Signed-off-by: Sy03 <1370724210@qq.com> --- docs/design/feature/stage_host_transport.md | 79 ++++ .../test_chunk_transfer_adapter.py | 9 +- .../omni_connectors/test_shm_connector.py | 4 +- .../omni_connectors/test_shm_ring.py | 346 ++++++++++++++ .../omni_connectors/connectors/base.py | 4 +- .../connectors/shm_connector.py | 67 ++- .../omni_connectors/connectors/shm_ring.py | 431 ++++++++++++++++++ .../omni_connectors/utils/serialization.py | 6 +- .../omni_connectors/utils/tensor_frame.py | 174 +++++++ 9 files changed, 1108 insertions(+), 12 deletions(-) create mode 100644 docs/design/feature/stage_host_transport.md create mode 100644 tests/distributed/omni_connectors/test_shm_ring.py create mode 100644 vllm_omni/distributed/omni_connectors/connectors/shm_ring.py create mode 100644 vllm_omni/distributed/omni_connectors/utils/tensor_frame.py diff --git a/docs/design/feature/stage_host_transport.md b/docs/design/feature/stage_host_transport.md new file mode 100644 index 000000000..cddfa340d --- /dev/null +++ b/docs/design/feature/stage_host_transport.md @@ -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. diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index f0b8d45b5..0eeedfcfa 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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" diff --git a/tests/distributed/omni_connectors/test_shm_connector.py b/tests/distributed/omni_connectors/test_shm_connector.py index d0efd06f5..5ba462fcb 100644 --- a/tests/distributed/omni_connectors/test_shm_connector.py +++ b/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() diff --git a/tests/distributed/omni_connectors/test_shm_ring.py b/tests/distributed/omni_connectors/test_shm_ring.py new file mode 100644 index 000000000..4b24523de --- /dev/null +++ b/tests/distributed/omni_connectors/test_shm_ring.py @@ -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() diff --git a/vllm_omni/distributed/omni_connectors/connectors/base.py b/vllm_omni/distributed/omni_connectors/connectors/base.py index a67a35a80..4fdf1adbf 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/base.py +++ b/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( diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index d0f2cdb1a..769643070 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/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. @@ -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) diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_ring.py b/vllm_omni/distributed/omni_connectors/connectors/shm_ring.py new file mode 100644 index 000000000..7719f3693 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_ring.py @@ -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(" 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(" 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(" _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 diff --git a/vllm_omni/distributed/omni_connectors/utils/serialization.py b/vllm_omni/distributed/omni_connectors/utils/serialization.py index 60fc725d6..dd5c0e5c9 100644 --- a/vllm_omni/distributed/omni_connectors/utils/serialization.py +++ b/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 @@ -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() diff --git a/vllm_omni/distributed/omni_connectors/utils/tensor_frame.py b/vllm_omni/distributed/omni_connectors/utils/tensor_frame.py new file mode 100644 index 000000000..5bc7ea82f --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/utils/tensor_frame.py @@ -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(" 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)