diff --git a/.github/workflows/cargo-release.yml b/.github/workflows/cargo-release.yml index 4d01c1e5f2..99371cb880 100644 --- a/.github/workflows/cargo-release.yml +++ b/.github/workflows/cargo-release.yml @@ -95,7 +95,6 @@ jobs: publish_if_not_exists dora-tracing publish_if_not_exists dora-metrics publish_if_not_exists dora-hub-client - publish_if_not_exists dora-memory-pool # Publish rust API publish_if_not_exists dora-operator-api-types diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 4db5d920b0..fcf5af5a8c 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -115,65 +115,6 @@ jobs: if: steps.smoke.outputs.exit_code != '0' run: exit 1 - # Memory-pool CPU transport smoke tests (dora-rs/dora#2168). - # - # These 6 tests (smoke_memory_pool_cpu2cpu + smoke_local_memory_pool_*) are - # `#[ignore]`-gated in tests/example-smoke.rs because they need torch + tqdm, - # which the smoke-suite job above does not install — so that job skips them - # and they ran in NO workflow (#2302 Finding 1; the gap #2264 Finding 1 first - # raised). They are CPU-only (sender_device/receiver_device: cpu), so no GPU - # is required: the per-node `build:` steps pip-install CPU torch from - # download.pytorch.org via uv. A dedicated job keeps the torch-download cost - # and any pytorch-index flakiness off the main smoke-suite, and makes the - # unsafe seqlock / pointer-arithmetic / bounds / free paths flagged by #2264 - # execute at least once per nightly run. - memory-pool-smoke: - name: Memory-pool CPU smoke tests (#2168) - runs-on: ubuntu-latest - # 6 tests, --test-threads=1 (they share coordinator ports). Most of the - # wall clock is the one-time CPU-torch download/install in the per-node - # build steps; 45 min leaves margin on top of a cold cargo compile. - timeout-minutes: 45 - steps: - - uses: actions/checkout@v6 - - uses: dtolnay/rust-toolchain@master - with: - toolchain: ${{ env.RUST_VERSION }} - - uses: Swatinem/rust-cache@v2 - with: - save-if: false - - uses: actions/setup-python@v6 - with: - python-version: "3.12" - - uses: astral-sh/setup-uv@v8.1.0 - with: - enable-cache: true - - name: Set up Python venv - run: | - uv venv --seed -p 3.12 - echo "VIRTUAL_ENV=$PWD/.venv" >> "$GITHUB_ENV" - source .venv/bin/activate - uv pip install pyarrow - uv pip install -e apis/python/node - - name: Run memory-pool smoke tests - id: smoke - continue-on-error: true - # `--ignored memory_pool` runs only the ignored memory-pool tests; the - # per-node `build:` steps install torch/tqdm/numpy into the active venv. - run: | - cargo test -p dora-examples --test example-smoke -- \ - --ignored --test-threads=1 memory_pool 2>&1 | tee memory-pool-smoke.log - echo "exit_code=${PIPESTATUS[0]}" >> "$GITHUB_OUTPUT" - - name: Upload smoke log - uses: actions/upload-artifact@v7 - with: - name: memory-pool-smoke-output - path: memory-pool-smoke.log - retention-days: 14 - - name: Fail the job if memory-pool smoke tests failed - if: steps.smoke.outputs.exit_code != '0' - run: exit 1 - # Hub e2e (tests/hub-smoke.rs). The `hub:` feature is unstable + non-core, so # it lives in nightly rather than the per-PR gate. Hermetic — a local git # fixture + DORA_HUB_ALLOW_LOCAL_SOURCES, no network — so it needs only Rust @@ -2014,7 +1955,6 @@ jobs: needs: - build-cli - smoke-suite - - memory-pool-smoke - log-sinks - service-action - streaming @@ -2078,10 +2018,6 @@ jobs: smoke-suite) echo 'cargo test -p dora-examples --test example-smoke -- --test-threads=1' ;; - memory-pool-smoke) - # Needs torch+tqdm; the per-node build steps install CPU torch. - echo 'cargo test -p dora-examples --test example-smoke -- --ignored --test-threads=1 memory_pool' - ;; log-sinks) # Three dataflows in this job; run each individually on failure. printf 'dora run examples/log-sink-file/dataflow.yml --uv --stop-after 15s\ndora run examples/log-sink-alert/dataflow.yml --uv --stop-after 15s\ndora run examples/log-sink-tcp/dataflow.yml --uv --stop-after 15s\n' @@ -2163,7 +2099,6 @@ jobs: case "$1" in build-cli) echo '^Build dora CLI \(shared\)$' ;; smoke-suite) echo '^Smoke suite \(tests/example-smoke\.rs\)$' ;; - memory-pool-smoke) echo '^Memory-pool CPU smoke tests \(#2168\)$' ;; log-sinks) echo '^Log-sink examples \(Python sources\)$' ;; service-action) echo '^Service / Action patterns \(Rust-only\)$' ;; streaming) echo '^Streaming example \(Python\)$' ;; @@ -2444,7 +2379,6 @@ jobs: needs: - build-cli - smoke-suite - - memory-pool-smoke - log-sinks - service-action - streaming diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index e89f7dc93f..bd213379a5 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -76,7 +76,6 @@ jobs: dora-tracing dora-metrics dora-hub-client - dora-memory-pool dora-operator-api-types dora-operator-api-macros dora-operator-api diff --git a/CLAUDE.md b/CLAUDE.md index d5cdb3bf24..cf3a99e1ee 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -169,7 +169,7 @@ The deeper QA gates — `make qa-full`, `make qa-deep`, `make qa-nightly`, `make - `make qa-full` (qa-fast + full tests + coverage) — ~5-10 minutes. Run before a significant push if you want extra confidence; coverage is too slow for every push. - `make qa-deep` (qa-full + mutation testing + semver) — ~15 minutes. The **target** Tier 1 local gate. Today's CI PR gate only runs the fast subset (fmt/clippy/typos/audit/unwrap-budget + tests); `qa-deep` adds coverage, adversarial review, diff-scoped mutation, and semver — kept laptop-only because they're too slow for every PR (see `docs/plan-agentic-qa-strategy.md` §5). Alias: `make qa-tier1`, kept for back-compat. -- `make qa-nightly` (qa-deep + proptest@1000 + miri + example-smoke + ci-nightly-jobs) — ~3-4 hours. **Full parity with `.github/workflows/nightly.yml`** (#1707, #1710, #1716). After the #1716 rebalance (plus cluster-record-replay from #2013 and memory-pool-smoke from #2302), nightly.yml has **27 test jobs** (`cli-tests` was split into `cli-tests` + `cli-tests-python` in #2742; re-counted in #2999, where the previous 23 counted neither `hub-smoke` nor the two `ros2-zenoh-*` jobs, and `multi-daemon-late-subscriber` is new). example-smoke covers the **4 example-backed** jobs (smoke-suite, log-sinks, service-action, streaming); `scripts/qa/ci-nightly-jobs.sh` drives **20** with platform-aware dispatch (record-replay, cluster-smoke, cluster-e2e [Linux, needs `openssh-server`], cluster-record-replay [Linux, needs `openssh-server`], topic-and-top, cpu-affinity [Linux], redb-backend, daemon-reconnect [Linux], state-reconstruction, multi-daemon-late-subscriber [Linux], test-cross-platform [macOS+Windows], examples, cli-tests **and cli-tests-python** — one local `cli-tests` invocation covers both halves, bench-example, cross-check, ros2-bridge [Linux, basic checks — no ROS distro], ros2-zenoh-humble, ros2-zenoh-kilted, msrv, kani-proofs [skipped if Kani not installed]). Of the rest, `memory-pool-smoke` (the torch-gated `#[ignore]` memory-pool example tests run with `--ignored`), is covered locally by `make qa-examples` / `scripts/smoke-all.sh`, **not** the qa-nightly example-smoke step (which skips `#[ignore]` tests), and `hub-smoke` has no local driver entry at all. A green local `qa-nightly` on platform X predicts a green CI nightly for platform X's jobs. **Requires both `uv` and Python 3.12** (both preflighted; the script fails fast with a specific install hint for whichever is missing — `curl -LsSf https://astral.sh/uv/install.sh \| sh` for uv, `uv python install 3.12` for the interpreter, matching the GHA `actions/setup-python` step at `.github/workflows/nightly.yml:56`). example-smoke creates a scratch venv at `target/qa-nightly-venv` and installs `-e apis/python/node` into it so Python nodes use the workspace bindings (not PyPI `dora-rs`, whose message format has drifted from the workspace — #1710). The CI-jobs script installs the CLI into a scratch dir (won't clobber `~/.cargo/bin/dora`) and bails if port 6013 is in use; cpu-affinity + daemon-reconnect skip on non-Linux. Skips miri if `cargo +nightly miri` isn't installed. **Does not** include full-repo mutation testing — that's split into `qa-mutation-audit` because it takes 10-18 hours on this workspace. +- `make qa-nightly` (qa-deep + proptest@1000 + miri + example-smoke + ci-nightly-jobs) — ~3-4 hours. **Full parity with `.github/workflows/nightly.yml`** (#1707, #1710, #1716). After the #1716 rebalance (plus cluster-record-replay from #2013), nightly.yml has **26 test jobs** (`cli-tests` was split into `cli-tests` + `cli-tests-python` in #2742; re-counted in #2999, where the previous 23 counted neither `hub-smoke` nor the two `ros2-zenoh-*` jobs, and `multi-daemon-late-subscriber` is new). example-smoke covers the **4 example-backed** jobs (smoke-suite, log-sinks, service-action, streaming); `scripts/qa/ci-nightly-jobs.sh` drives **20** with platform-aware dispatch (record-replay, cluster-smoke, cluster-e2e [Linux, needs `openssh-server`], cluster-record-replay [Linux, needs `openssh-server`], topic-and-top, cpu-affinity [Linux], redb-backend, daemon-reconnect [Linux], state-reconstruction, multi-daemon-late-subscriber [Linux], test-cross-platform [macOS+Windows], examples, cli-tests **and cli-tests-python** — one local `cli-tests` invocation covers both halves, bench-example, cross-check, ros2-bridge [Linux, basic checks — no ROS distro], ros2-zenoh-humble, ros2-zenoh-kilted, msrv, kani-proofs [skipped if Kani not installed]). Of the rest, `hub-smoke` has no local driver entry at all. A green local `qa-nightly` on platform X predicts a green CI nightly for platform X's jobs. **Requires both `uv` and Python 3.12** (both preflighted; the script fails fast with a specific install hint for whichever is missing — `curl -LsSf https://astral.sh/uv/install.sh \| sh` for uv, `uv python install 3.12` for the interpreter, matching the GHA `actions/setup-python` step at `.github/workflows/nightly.yml:56`). example-smoke creates a scratch venv at `target/qa-nightly-venv` and installs `-e apis/python/node` into it so Python nodes use the workspace bindings (not PyPI `dora-rs`, whose message format has drifted from the workspace — #1710). The CI-jobs script installs the CLI into a scratch dir (won't clobber `~/.cargo/bin/dora`) and bails if port 6013 is in use; cpu-affinity + daemon-reconnect skip on non-Linux. Skips miri if `cargo +nightly miri` isn't installed. **Does not** include full-repo mutation testing — that's split into `qa-mutation-audit` because it takes 10-18 hours on this workspace. - `make qa-release-gate` (qa-deep + semver) — the automatable subset of Tier 3. The non-automatable parts (independent security audit, dogfood campaign, migration validation) are documented in `docs/plan-agentic-qa-strategy.md` §7 but not locally gateable. - `make qa-mutation-audit` — ~10-18 hours. Full-repo `cargo-mutants` across 6 critical crates. Deliberate test-quality audit; run before a release or when investigating a specific crate, not every nightly. - `make qa-examples` — ~15-20 min. Runs all **smoke-eligible** example dataflows end-to-end via `scripts/smoke-all.sh`. Skips examples that need CUDA, ROS2, webcam, multi-machine deploy, C/C++ toolchains, or interactive CLI (run `scripts/smoke-all.sh -h` to see the SKIP list). Orthogonal to the qa-fast/full/deep ladder: those targets `--exclude dora-examples` to keep per-commit / pre-push budgets tight. Run this when you want actual dataflows exercised (after touching node/operator APIs, CLI subcommands, or the descriptor surface). Pass flags via `ARGS`, e.g. `make qa-examples ARGS="--rust-only"` or `make qa-examples ARGS="-v"`. diff --git a/Cargo.lock b/Cargo.lock index ad6d671ea5..48d5b03e0a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2093,9 +2093,9 @@ dependencies = [ "dora-arrow-convert", "dora-core", "dora-download", - "dora-memory-pool", "dora-message", "dora-node-api", + "dora-tensor-pool", "dora-tracing", "dunce", "eyre", @@ -2227,13 +2227,6 @@ dependencies = [ "url", ] -[[package]] -name = "dora-memory-pool" -version = "1.0.0-rc.4" -dependencies = [ - "tracing", -] - [[package]] name = "dora-message" version = "1.0.0-rc.4" @@ -2352,6 +2345,7 @@ dependencies = [ "dora-operator-api-python", "dora-ros2-bridge-python", "dora-runtime-python", + "dora-tensor-pool-python", "eyre", "flume 0.12.0", "futures", @@ -2360,7 +2354,6 @@ dependencies = [ "pythonize", "serde_json", "serde_yaml", - "shared_memory_extended", "tokio", "tracing", ] @@ -2630,6 +2623,27 @@ dependencies = [ "tracing-opentelemetry", ] +[[package]] +name = "dora-tensor-pool" +version = "1.0.0-rc.4" +dependencies = [ + "tracing", +] + +[[package]] +name = "dora-tensor-pool-python" +version = "1.0.0-rc.4" +dependencies = [ + "arrow", + "dora-message", + "dora-node-api", + "eyre", + "pyo3", + "serde_json", + "shared_memory_extended", + "tracing", +] + [[package]] name = "dora-tracing" version = "1.0.0-rc.4" diff --git a/Cargo.toml b/Cargo.toml index 7c34a2a79a..74ead52416 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,7 +67,10 @@ members = [ "libraries/extensions/download", "libraries/extensions/telemetry/*", "libraries/extensions/mavlink2-bridge", - "libraries/extensions/memory-pool", + # Opt-in tensor-pool transport. Outside the 1.0 compatibility + # guarantees — see libraries/extensions/tensor-pool/README.md. + "libraries/extensions/tensor-pool", + "libraries/extensions/tensor-pool/python", "libraries/extensions/ros2-bridge", "libraries/extensions/ros2-bridge/msg-gen", "libraries/extensions/ros2-bridge/python", @@ -134,7 +137,8 @@ dora-ros2-bridge = { version = "1.0.0-rc.4", path = "libraries/extensions/ros2-b dora-ros2-bridge-msg-gen = { version = "1.0.0-rc.4", path = "libraries/extensions/ros2-bridge/msg-gen" } dora-ros2-bridge-python = { path = "libraries/extensions/ros2-bridge/python" } dora-mavlink2-bridge = { version = "1.0.0-rc.4", path = "libraries/extensions/mavlink2-bridge" } -dora-memory-pool = { version = "1.0.0-rc.4", path = "libraries/extensions/memory-pool" } +dora-tensor-pool = { version = "1.0.0-rc.4", path = "libraries/extensions/tensor-pool" } +dora-tensor-pool-python = { version = "1.0.0-rc.4", path = "libraries/extensions/tensor-pool/python" } dora-message = { version = "1.0.0-rc.4", path = "libraries/message" } # Declared here so `dora-ros2-bridge` and `dora-ros2-bridge-arrow` cannot drift # apart on the ROS distro. `humble` selects the 24-byte `Gid` layout used by diff --git a/Makefile b/Makefile index 37135d6404..e028c3105f 100644 --- a/Makefile +++ b/Makefile @@ -13,7 +13,7 @@ # make qa-nightly ~3-4 hours Full parity with .github/workflows/nightly.yml # (qa-deep + proptest@1000 + miri + example-smoke # + ci-nightly-jobs). After the #1716 rebalance, -# nightly.yml has 27 test jobs (re-counted +# nightly.yml has 26 test jobs (re-counted # in #2999): example-smoke # covers 4 (smoke-suite/log-sinks/service-action/ # streaming); scripts/qa/ci-nightly-jobs.sh drives @@ -26,12 +26,8 @@ # test-cross-platform, examples, cli-tests, # bench-example, cross-check, ros2-bridge [Linux+ROS2], # ros2-zenoh-humble, ros2-zenoh-kilted, -# msrv, kani-proofs [skipped if Kani absent]). Of -# the rest, memory-pool-smoke (torch-gated #[ignore] -# memory-pool tests), is covered locally by -# `make qa-examples` / smoke-all.sh, not the -# qa-nightly example-smoke step (which skips -# #[ignore] tests). Green local run on platform X predicts +# msrv, kani-proofs [skipped if Kani absent]). +# Green local run on platform X predicts # green CI nightly for platform X's jobs. # make qa-release-gate Tier 3 automatable parts (deep + semver; # audit/dogfood are human) diff --git a/apis/python/node/Cargo.toml b/apis/python/node/Cargo.toml index 2de570ae56..0693e24035 100644 --- a/apis/python/node/Cargo.toml +++ b/apis/python/node/Cargo.toml @@ -14,6 +14,10 @@ publish = false # PyO3 cdylib, shipped via PyPI as `dora-rs`, not crates.io [features] default = ["tracing", "metrics", "async"] +# Opt-in tensor-pool transport. NOT covered by dora's 1.0 compatibility +# guarantees — see libraries/extensions/tensor-pool/README.md. Off by default +# so a standard wheel neither builds nor exposes it. +tensor-pool = ["dep:dora-tensor-pool-python"] tracing = ["dora-node-api/tracing"] metrics = ["dora-node-api/metrics"] async = ["pyo3/experimental-async"] @@ -35,7 +39,7 @@ pythonize = { workspace = true } futures = { workspace = true } dora-ros2-bridge-python = { workspace = true } dora-download = { workspace = true } -shared_memory_extended = "0.13.0" +dora-tensor-pool-python = { workspace = true, optional = true } tokio = { workspace = true, features = ["rt-multi-thread", "time", "macros"] } tracing = { workspace = true } diff --git a/apis/python/node/dora/__init__.pyi b/apis/python/node/dora/__init__.pyi index 5c33f370ab..48957faa0b 100644 --- a/apis/python/node/dora/__init__.pyi +++ b/apis/python/node/dora/__init__.pyi @@ -215,59 +215,65 @@ class Node: node.send_output("string", b"string", {"open_telemetry_context": "7632e76"}) ```""" - def register_memory_pool( + def register_tensor_pool( self, tensor_info: dict, device: str = "cpu" ) -> pyarrow.Array: - """Register a tensor's pinned memory as a shared memory pool and - return a persistent buffer_id (pyarrow string array). + """Register a tensor's memory as a shared pool for zero-copy transfer. - The tensor data is copied into page-locked shared memory (via DMA - if CUDA). The returned buffer_id can be sent to another node, - which can call read_memory_pool() to get a zero-copy tensor_info - pointing to the same shared memory. + Only present when the wheel was built with the `tensor-pool` feature. + NOT covered by dora's 1.0 compatibility guarantees — see + libraries/extensions/tensor-pool/README.md. + """ - Args: - tensor_info: dict with keys "ptr", "size", "dtype", "shape", "device" - device: "cpu" or "cuda" — the receiver's device (controls DMA path) + def write_tensor_pool( + self, tensor_pool_id: pyarrow.Array, tensor_info: dict + ) -> None: + """Write tensor data into an existing tensor pool. - Returns: - pyarrow string array containing the buffer_id + Extension method; see register_tensor_pool. """ - def write_memory_pool( - self, memory_pool_id: pyarrow.Array, tensor_info: dict - ) -> None: - """Write new tensor data into an existing memory pool. + def read_tensor_pool(self, tensor_pool_id: pyarrow.Array) -> dict: + """Read tensor metadata from a tensor pool (zero-copy). - Fast path: DORADMA header read from shmem, then memcpy data. - Slow path: falls back to daemon RPC. + Extension method; see register_tensor_pool. + """ - Args: - memory_pool_id: pyarrow string array from register_memory_pool - tensor_info: dict with keys "ptr", "size", "dtype", "shape", "device" + def free_tensor_pool(self, tensor_pool_id: pyarrow.Array) -> None: + """Free a tensor pool, releasing it for every node that touched it. + + Extension method; see register_tensor_pool. """ - def read_memory_pool(self, memory_pool_id: pyarrow.Array) -> dict: - """Read tensor metadata from a memory pool. + def extension_store(self, namespace: str, key: str, value: bytes) -> None: + """Store an opaque value in the daemon's dataflow-scoped extension table. - Fast path: DORADMA header read from shmem (zero-copy, no daemon). - Slow path: falls back to daemon RPC. + The seam for transports maintained outside the dora tree. dora brokers + the value's lifetime and nothing else — it never interprets namespace, + key or value. - Returns a tensor_info dict with keys: - "ptr": int — pointer to the data - "size": int — data size in bytes - "dtype": str — e.g. "float32" - "shape": list[int] — tensor shape - "device": str — "cpu" or "cuda" + The daemon reclaims the entry when this node exits or the dataflow + finishes, and notifies every node that stored or read the key when it + is dropped (see drain_dropped_extension_keys). """ - def free_memory_pool(self, memory_pool_id: pyarrow.Array) -> None: - """Free a memory pool by removing it from the daemon's tracking table. + def extension_load( + self, namespace: str, key: str, remove: bool = False + ) -> bytes | None: + """Read an opaque value back, or None if the key is absent. - Args: - memory_pool_id: pyarrow string array from register_memory_pool + With remove=True the entry is dropped in the same round trip. + """ + + def extension_drop(self, namespace: str, key: str) -> None: + """Drop an opaque value, notifying every node that stored or read it. + + Dropping an absent key succeeds, so a retry is safe. """ + def drain_dropped_extension_keys(self, namespace: str) -> list[str]: + """Take the keys in `namespace` dropped since the last call.""" + def __iter__(self) -> typing.Any: """Implement iter(self).""" diff --git a/apis/python/node/dora/cuda.py b/apis/python/node/dora/cuda.py index b008b73027..d963ff57d4 100644 --- a/apis/python/node/dora/cuda.py +++ b/apis/python/node/dora/cuda.py @@ -110,51 +110,8 @@ def _libcudart(): _cudaIpcMemLazyEnablePeerAccess = 1 # Numpy dtype string -> torch dtype mapping. -_DTYPE_MAP = { - " numpy dtype mapping (used by tensor_from_info for # constructing numpy arrays from raw memory pointers). -_TORCH_TO_NUMPY_DTYPE_MAP = { - torch.int64: np.int64, - torch.float32: np.float32, - torch.float64: np.float64, - torch.int32: np.int32, - torch.int16: np.int16, - torch.int8: np.int8, - torch.uint8: np.uint8, - torch.bool: np.bool_, - torch.float16: np.float16, - torch.bfloat16: np.float16, # bfloat16 maps to float16 in numpy -} - # CUDA error codes for better error messages _CUDA_ERROR_SUCCESS = 0 _CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED = 712 @@ -332,96 +289,3 @@ def open_ipc_handle( yield tensor finally: ipc_handle.close() - - -# --------------------------------------------------------------------------- -# Tensor info helpers for memory-pool operations -# --------------------------------------------------------------------------- - - -class _ArrayInterface: - """Minimal object implementing ``__array_interface__`` so that - ``torch.as_tensor`` can wrap raw CPU memory as a tensor (zero-copy).""" - - def __init__(self, ptr, shape, strides, dtype_str): - self.__array_interface__ = { - "shape": tuple(shape), - "strides": tuple(strides) if strides else None, - "typestr": dtype_str, - "data": (ptr, False), - "version": 3, - } - - -def get_tensor_info(tensor: torch.Tensor) -> dict: - """Serialize a tensor into a ``tensor_info`` dict containing pointer, - size, dtype, shape, and device. - - This is the canonical way to pass tensor metadata to memory-pool - operations such as ``register_memory_pool`` and ``write_memory_pool``. - """ - if not tensor.is_contiguous(): - tensor = tensor.contiguous() - return { - "ptr": tensor.data_ptr(), - "size": tensor.nbytes, - "dtype": str(tensor.dtype), - "shape": list(tensor.shape), - "device": str(tensor.device), - } - - -def tensor_from_info(tensor_info: dict) -> torch.Tensor: - """Reconstruct a PyTorch tensor from a ``tensor_info`` dict (zero-copy). - - The returned tensor shares the same underlying memory as the original - tensor that produced the ``tensor_info``. Used by consumers that read - a memory pool via ``read_memory_pool``. - """ - ptr = tensor_info.get('ptr', 0) - if ptr == 0: - raise ValueError("tensor_info has null pointer (ptr=0); pool may not exist or has been freed") - dtype_str = tensor_info["dtype"] - shape = tensor_info["shape"] - device = tensor_info.get("device", "cpu") - size = tensor_info.get("size", 0) - - dtype = _DTYPE_MAP.get(dtype_str, torch.int64) - - if device.startswith("cuda"): - # CUDA tensor — zero-copy via __cuda_array_interface__ - np_dtype = _TORCH_TO_NUMPY_DTYPE_MAP.get(dtype) - if np_dtype is None: - raise ValueError(f"Unsupported dtype: {dtype}") - - # Validate that product(shape) * itemsize(dtype) does not - # exceed the registered size — a peer-controlled header that - # claims a large shape over a small buffer would produce an - # out-of-bounds GPU tensor (the CPU path is saved by numpy - # reshape, but the GPU path has no equivalent backstop). - expected_bytes = np.dtype(np_dtype).itemsize - for dim in shape: - expected_bytes *= dim - if expected_bytes > size: - raise ValueError( - f"tensor shape {shape} * {np_dtype} itemsize = {expected_bytes} bytes " - f"exceeds registered size {size} bytes — header may be corrupted" - ) - - typestr = np.dtype(np_dtype).str - wrapper = _CudaArrayInterface(ptr, shape, None, typestr) - return torch.as_tensor(wrapper, device="cuda") - else: - # CPU tensor — zero-copy via numpy / torch.frombuffer - np_dtype = _TORCH_TO_NUMPY_DTYPE_MAP.get(dtype) - if np_dtype is None and dtype != torch.bfloat16: - raise ValueError(f"Unsupported dtype: {dtype}") - - c_array = (ctypes.c_byte * size).from_address(ptr) - - if dtype == torch.bfloat16: - byte_tensor = torch.frombuffer(c_array, dtype=torch.uint8) - return byte_tensor.view(dtype=torch.bfloat16).reshape(shape) - - np_array = np.frombuffer(c_array, dtype=np_dtype).reshape(shape) - return torch.from_numpy(np_array) diff --git a/apis/python/node/src/lib.rs b/apis/python/node/src/lib.rs index 9352f7e7fb..42cdaa0bd7 100644 --- a/apis/python/node/src/lib.rs +++ b/apis/python/node/src/lib.rs @@ -1,13 +1,10 @@ #![allow(clippy::borrow_deref_ref)] // clippy warns about code generated by #[pymethods] -use std::collections::{HashMap, HashSet}; use std::sync::{Arc, LazyLock}; use std::time::Duration; use tokio::sync::Mutex; -use arrow::array::{Array, BinaryArray, StringArray}; use arrow::pyarrow::{FromPyArrow, ToPyArrow}; -use dora_message::metadata::Parameter; use dora_node_api::dora_core::config::{DataId, NodeId}; use dora_node_api::merged::{MergeExternalSend, MergedEvent}; use dora_node_api::{DataflowId, DoraNode, EventStream, TryRecvError, init_tracing}; @@ -16,7 +13,6 @@ use dora_operator_api_python::{ }; use dora_ros2_bridge_python::Ros2Subscription; use eyre::{Context, ContextCompat}; -use shared_memory_extended::ShmemConf; use futures::{Stream, StreamExt}; use pyo3::prelude::*; @@ -128,1000 +124,6 @@ def basicConfig(*pargs, **kwargs): Ok(()) } -/// Pre-compiled CUDA DMA helper module. Compiled once (at first use) and reused -/// across all iterations to eliminate per-call PyModule::from_code overhead. -/// Maintains persistent state: pinned host pointers, pre-allocated GPU buffers. -static CUDA_HELPERS: LazyLock>>> = - LazyLock::new(|| std::sync::Mutex::new(None)); - -/// Counter to make pinned memory buffer IDs unique across registrations. -static PINNED_COUNTER: LazyLock> = LazyLock::new(|| std::sync::Mutex::new(0)); - -/// Maximum number of freed pool buffer IDs to remember at once. This is a -/// single budget shared across every peer the process reads from, not a -/// per-stream allowance — once the cap is exceeded the oldest entries are -/// evicted rather than keeping every ID for the life of the process. -const FREED_POOL_IDS_CAP: usize = 4096; - -/// Tracks freed pool buffer IDs so the DORADMA fast path can detect -/// read-after-free. Bounded to `FREED_POOL_IDS_CAP` entries (oldest evicted -/// first) so long-running nodes doing register->write->free every frame -/// don't leak memory indefinitely. -/// -/// The cap is one shared budget across *all* peers the process reads -/// from, not a per-stream recency window: a high-rate sender's frees can -/// evict a low-rate sender's tombstones well before the low-rate sender's -/// own next free (e.g. a 60Hz peer can cycle the whole cap in ~68s, -/// evicting a 1Hz peer's entries long before they'd naturally expire). An -/// evicted tombstone is harmless — it just makes a stale fast-path read -/// fall back to the existing `warn_missing_memory_pool`/daemon path -/// instead of being caught here. -static FREED_POOL_IDS: LazyLock> = - LazyLock::new(|| std::sync::Mutex::new(FreedPoolIds::default())); - -/// Bounded, insertion-ordered set of freed pool buffer IDs. `set` gives -/// O(1) membership checks; `order` tracks insertion order so the oldest -/// entry can be evicted once `FREED_POOL_IDS_CAP` is exceeded. -#[derive(Default)] -struct FreedPoolIds { - set: HashSet, - order: std::collections::VecDeque, -} - -impl FreedPoolIds { - fn insert(&mut self, id: String) { - if self.set.insert(id.clone()) { - self.order.push_back(id); - while self.order.len() > FREED_POOL_IDS_CAP { - if let Some(oldest) = self.order.pop_front() { - self.set.remove(&oldest); - } - } - } - } - - fn contains(&self, id: &str) -> bool { - self.set.contains(id) - } - - fn remove(&mut self, id: &str) { - if self.set.remove(id) { - self.order.retain(|x| x != id); - } - } - - fn len(&self) -> usize { - self.set.len() - } -} - -#[cfg(test)] -mod freed_pool_ids_tests { - use super::{FREED_POOL_IDS_CAP, FreedPoolIds}; - - /// Regression test: without a cap, inserting one ID per frame in a - /// long-running node grows `FREED_POOL_IDS` forever. Inserting well - /// past the cap must keep the set bounded instead of leaking. - #[test] - fn insert_past_cap_does_not_grow_unbounded() { - let mut freed = FreedPoolIds::default(); - for i in 0..(FREED_POOL_IDS_CAP * 4) { - freed.insert(format!("pool_node_{i}")); - } - assert_eq!(freed.len(), FREED_POOL_IDS_CAP); - } - - /// Once past the cap, the oldest entries must be evicted first so the - /// most recently freed buffers (the ones a read-after-free check would - /// actually care about) stay tracked. - #[test] - fn insert_past_cap_evicts_oldest_first() { - let mut freed = FreedPoolIds::default(); - for i in 0..(FREED_POOL_IDS_CAP * 2) { - freed.insert(format!("pool_node_{i}")); - } - assert!( - !freed.contains("pool_node_0"), - "oldest entry should have been evicted" - ); - let newest = format!("pool_node_{}", FREED_POOL_IDS_CAP * 2 - 1); - assert!( - freed.contains(&newest), - "most recently freed entry should still be tracked" - ); - } - - /// Duplicate inserts of an already-tracked ID must not double-count - /// against the cap or push a second copy into the eviction order. - #[test] - fn duplicate_insert_is_idempotent() { - let mut freed = FreedPoolIds::default(); - freed.insert("pool_node_0".to_string()); - freed.insert("pool_node_0".to_string()); - assert_eq!(freed.len(), 1); - } - - #[test] - fn remove_drops_membership_and_order_entry() { - let mut freed = FreedPoolIds::default(); - freed.insert("pool_node_0".to_string()); - freed.remove("pool_node_0"); - assert!(!freed.contains("pool_node_0")); - assert_eq!(freed.len(), 0); - // Re-inserting after a remove must not be blocked by a stale - // order-queue entry left behind by `remove`. - freed.insert("pool_node_0".to_string()); - assert_eq!(freed.len(), 1); - } -} - -/// Per-pool persistent state. -/// Keeping Shmem alive prevents munmap, preserving stable mmap addresses -/// for pool-hit detection across `register_memory_pool` calls. -/// -/// # Safety -/// `Shmem` is not `Send + Sync` due to raw pointer fields, but `PoolSlot` -/// is always stored behind a `Mutex` and `_shmem` is never accessed after -/// initialization — it is only kept alive for its destructor. -struct PoolSlot { - _shmem: shared_memory_extended::Shmem, - base: u64, - size: usize, - is_pinned: bool, - /// CPU page-locked transit buffer for cross-device GPU transfers - /// without P2P (e.g. RTX 5090). 0 means no transit path. - transit_ptr: u64, - /// The GPU device index where the pool buffer was allocated. - pool_device: i32, -} - -unsafe impl Send for PoolSlot {} -unsafe impl Sync for PoolSlot {} - -/// Persistent pool storage — stable mmap addresses for zero-copy detection. -/// Keyed by counter (unique per registration), supports unlimited pools. -static PINNED_POOL: LazyLock>> = - LazyLock::new(|| std::sync::Mutex::new(HashMap::new())); - -/// Persistent transit-buffer metadata for GPU pools. -/// Survives `PINNED_POOL` cache-miss so the write fast path can recover -/// `transit_ptr` and `pool_device` even when the `PoolSlot` has been evicted. -/// Keyed by counter, populated during `register_memory_pool`, cleared in -/// `free_memory_pool`. `transit_ptr=0` means no transit buffer (same-device -/// or P2P path). -static TRANSIT_META: LazyLock>> = - LazyLock::new(|| std::sync::Mutex::new(HashMap::new())); - -/// Receiver-side GPU cache per pool. -/// Keeps Shmem alive to prevent munmap, preserving stable mmap addresses -/// and valid GPU VAs for zero-copy reads across iterations. -struct RecvGpuSlot { - _shmem: shared_memory_extended::Shmem, - gpu_va: u64, // device VA from cudaHostGetDevicePointer, 0 if IPC path - gpu_buf: u64, // IPC-opened GPU DRAM pointer, 0 if GPU VA path - host_base: u64, // original host ptr passed to cudaHostRegister - gpu_buf_size: u64, // GPU buffer byte size from first import (baseline) -} - -/// Daemon-trusted GPU buffer sizes, keyed by buffer_id. -/// Populated from daemon metadata when the fallback path successfully -/// re-imports a GPU IPC handle. The fast path (`try_doradma_read`) -/// validates the (world-writable) shmem `size` against this cache — -/// if `size` exceeds the trusted capacity, the read is rejected and -/// the caller falls back to the daemon. -static GPU_BUF_SIZES: LazyLock>> = - LazyLock::new(|| std::sync::Mutex::new(HashMap::new())); -unsafe impl Send for RecvGpuSlot {} -unsafe impl Sync for RecvGpuSlot {} - -/// Receiver-side per-pool cache keeping Shmem alive + GPU VA for zero-copy reads. -/// Set up lazily in try_doradma_read: open shmem, cudaHostRegister, -/// cudaHostGetDevicePointer, then cache both Shmem and GPU VA. -static RECV_GPU_VA: LazyLock>> = - LazyLock::new(|| std::sync::Mutex::new(HashMap::new())); - -/// Receiver-side per-pool Shmem cache for CPU receivers. -/// Keeps Shmem alive to prevent munmap of the CPU pointer returned -/// by the as_cuda=False path in try_doradma_read. -struct RecvCpuSlot { - _shmem: shared_memory_extended::Shmem, - base: u64, -} -unsafe impl Send for RecvCpuSlot {} -unsafe impl Sync for RecvCpuSlot {} - -/// Receiver-side per-pool cache keeping Shmem alive for CPU zero-copy reads. -/// Without this cache, the Shmem handle drops at the end of try_doradma_read, -/// triggering munmap and making the returned CPU pointer a dangling pointer. -static RECV_CPU_SHMEM: LazyLock>> = - LazyLock::new(|| std::sync::Mutex::new(HashMap::new())); - -/// DORADMA shared-memory header layout: -/// -/// Offset Size Field -/// 0 8 magic — b"DORADMA\x00" -/// 8 8 json_len — u64 LE, metadata JSON byte length -/// 16 8 data_off — u64 LE, byte offset of tensor data from shmem base -/// 24 8 ipc_flag — u64 LE, 1 when ipc_handle is valid -/// 32 64 ipc_handle — CUDA IPC mem handle (only valid if ipc_flag == 1) -/// 96 8 write_gen — u64 LE, seqlock: even = complete, odd = writing -/// 104 152 reserved -/// 256 N json — padded-to-256-byte-alignment metadata JSON -/// 256+N M data — tensor payload -const DORADMA_HEADER_SIZE: usize = 256; -const DORADMA_MAGIC: &[u8; 8] = b"DORADMA\x00"; -const DORADMA_METADATA_ALIGN: usize = 256; - -/// Crossover where pinned-DMA bandwidth overtakes pageable copy + -/// cudaHostRegister/unregister fixed cost (~100 µs). Determined by -/// ablation study (2026-06-27): pageable faster below, pinned faster -/// above. Shared by `register_memory_pool` and `write_memory_pool`. -const DMA_PIN_THRESHOLD_BYTES: usize = 25 * 1024 * 1024; - -/// Returns `true` when the source tensor should be pinned before DMA. -/// -/// Pinning is a property of the *source* pointer: `cudaHostRegister` only -/// makes sense for host (CPU) memory, and the pin/unpin fixed cost -/// (~100 µs) is only worth paying when the tensor is large enough that -/// the DMA bandwidth gain outweighs it. -/// -/// # Unit-testable -/// -/// The decision is pure integer logic — no CUDA runtime calls — so the -/// boundary (25 MiB ± 1 byte) can be exercised in CI even without a GPU. -#[inline] -const fn should_pin(is_cuda: bool, size: usize) -> bool { - !is_cuda && size > DMA_PIN_THRESHOLD_BYTES -} - -#[cfg(test)] -mod pin_tests { - use super::*; - - #[test] - fn pin_cpu_source_above_threshold() { - // CPU source, 25 MiB + 1 byte → should pin - assert!(should_pin(false, 25 * 1024 * 1024 + 1)); - // CPU source, 100 MiB → should pin - assert!(should_pin(false, 100 * 1024 * 1024)); - } - - #[test] - fn pin_cpu_source_below_threshold() { - // CPU source, exactly at threshold → should NOT pin (> not >=) - assert!(!should_pin(false, 25 * 1024 * 1024)); - // CPU source, 1 byte below → should NOT pin - assert!(!should_pin(false, 25 * 1024 * 1024 - 1)); - // CPU source, tiny → should NOT pin - assert!(!should_pin(false, 1)); - } - - #[test] - fn pin_cuda_source_never_pins() { - // CUDA source regardless of size → never pin - assert!(!should_pin(true, 0)); - assert!(!should_pin(true, 25 * 1024 * 1024)); - assert!(!should_pin(true, 100 * 1024 * 1024)); - assert!(!should_pin(true, 1024 * 1024 * 1024)); - } - - #[test] - fn pin_zero_size_cpu() { - // Zero-size CPU tensor → below threshold, don't pin - assert!(!should_pin(false, 0)); - } -} - -// --------------------------------------------------------------------------- -// GPU transport-path classification — pure decision logic extractable -// from CUDA-runtime-embedded code so the full matrix can be exercised in -// CI without a GPU. Same pattern as `should_pin` above. -// --------------------------------------------------------------------------- - -/// Which transport path a GPU-pool registration (write-time) selects. -/// -/// Pure logic — no CUDA runtime calls. -#[derive(Debug, PartialEq, Eq, Clone, Copy)] -enum TransportPath { - /// Buffer on sender device, plain DtoD memcpy (same-device or CPU source). - SameDeviceDtoD, - /// Cross-device with P2P peer access enabled. - P2PPeerAccess, - /// Cross-device without P2P — CPU page-locked transit (DtoH → HtoD). - HostStagingTransit, -} - -/// Classify which transport path a GPU-pool write should take at -/// registration time. -/// -/// Decision matrix (2³ = 8 cases, `is_cuda_source` dominates): -/// -/// | src CUDA | same dev | P2P | path | -/// |----------|----------|-----|---------------------| -/// | false | * | * | `SameDeviceDtoD` | -/// | true | true | * | `SameDeviceDtoD` | -/// | true | false | yes | `P2PPeerAccess` | -/// | true | false | no | `HostStagingTransit`| -#[inline] -fn classify_transport( - sender_device: i32, - receiver_device: i32, - p2p_available: bool, - is_cuda_source: bool, -) -> TransportPath { - if !is_cuda_source { - return TransportPath::SameDeviceDtoD; - } - if sender_device == receiver_device { - return TransportPath::SameDeviceDtoD; - } - if p2p_available { - return TransportPath::P2PPeerAccess; - } - TransportPath::HostStagingTransit -} - -/// Which write path `write_memory_pool` dispatches to for a given frame. -/// -/// The fast and slow write paths both branch on the same 2×2×2 matrix -/// (`ipc_present` × `is_cuda` × `transit_ptr`); extracting the -/// classification makes the 5 reachable paths explicit and testable. -#[derive(Debug, PartialEq, Eq)] -enum WritePath { - /// CPU source → GPU pool via `dma_copy` (ipc_present=1, !is_cuda). - CpuToGpuPoolDma, - /// GPU source → GPU pool via transit (ipc_present=1, is_cuda, transit_ptr≠0). - GpuToGpuPoolTransit, - /// GPU source → GPU pool via plain DtoD `_cuda_memcpy_gpu_buf` - /// (ipc_present=1, is_cuda, transit_ptr=0). - GpuToGpuPoolDtoD, - /// GPU source → shmem data region via `cudaMemcpy` (ipc_present≠1, is_cuda). - GpuToShmem, - /// CPU source → shmem data region via `ptr::copy_nonoverlapping` - /// (ipc_present≠1, !is_cuda). - CpuToShmem, -} - -/// Classify which write path to take. -#[inline] -fn classify_write_path(ipc_present: u64, is_cuda: bool, transit_ptr: u64) -> WritePath { - if ipc_present == 1 { - if is_cuda { - if transit_ptr != 0 { - WritePath::GpuToGpuPoolTransit - } else { - WritePath::GpuToGpuPoolDtoD - } - } else { - WritePath::CpuToGpuPoolDma - } - } else if is_cuda { - WritePath::GpuToShmem - } else { - WritePath::CpuToShmem - } -} - -/// Result of validating a GPU-pool read `size` against the daemon-trusted -/// capacity cache (`GPU_BUF_SIZES`) and the first-import baseline -/// (`RecvGpuSlot::gpu_buf_size`). -#[derive(Debug, PartialEq, Eq)] -enum CapacityCheck { - /// Size is within the trusted bound. - Ok, - /// Size exceeds the trusted capacity → reject this read. - ExceedsTrustedSize, - /// No daemon-trusted entry and no cached baseline → reject first import. - NoTrustedEntry, -} - -/// Validate the read `size` for a GPU-pool buffer against the daemon-trusted -/// capacity and the cached first-import baseline. -/// -/// Resolution order: `trusted_sizes` (daemon metadata) → `cached_gpu_buf_size` -/// (first-import baseline) → reject. -/// -/// The daemon rejects zero-size pools at registration, so a resolved capacity -/// of zero is unreachable in normal operation. If it does occur (daemon bug -/// or memory corruption), the `size > capped` check fails closed — any -/// non-zero `size` triggers `ExceedsTrustedSize` rather than silently skipping -/// validation. -#[inline] -fn check_capacity_gpu_pool( - trusted_sizes: Option, - cached_gpu_buf_size: Option, - size: u64, -) -> CapacityCheck { - let cap = trusted_sizes.or(cached_gpu_buf_size); - match cap { - None => CapacityCheck::NoTrustedEntry, - Some(capped) if size > capped => CapacityCheck::ExceedsTrustedSize, - Some(_) => CapacityCheck::Ok, - } -} - -#[cfg(test)] -mod transport_tests { - use super::*; - - // -- classify_transport ------------------------------------------------- - - #[test] - fn same_device_no_transit() { - // Same GPU — never transit, regardless of P2P - assert_eq!( - classify_transport(0, 0, false, true), - TransportPath::SameDeviceDtoD - ); - assert_eq!( - classify_transport(1, 1, true, true), - TransportPath::SameDeviceDtoD - ); - } - - #[test] - fn cpu_source_never_transit() { - // CPU→GPU always uses dma_copy, no transit needed - assert_eq!( - classify_transport(0, 1, false, false), - TransportPath::SameDeviceDtoD - ); - assert_eq!( - classify_transport(0, 2, true, false), - TransportPath::SameDeviceDtoD - ); - } - - #[test] - fn cross_device_with_p2p() { - assert_eq!( - classify_transport(0, 1, true, true), - TransportPath::P2PPeerAccess - ); - } - - #[test] - fn cross_device_no_p2p_uses_transit() { - // This is the RTX 5090 / Blackwell path — the flag-ship non-P2P - // fallback that must NOT be dead code. - assert_eq!( - classify_transport(0, 1, false, true), - TransportPath::HostStagingTransit - ); - assert_eq!( - classify_transport(2, 0, false, true), - TransportPath::HostStagingTransit - ); - } - - #[test] - fn classify_transport_full_8_case_matrix() { - let cases: &[((i32, i32, bool, bool), TransportPath)] = &[ - // (src_dev, dst_dev, p2p, is_cuda) → expected - ((0, 0, false, false), TransportPath::SameDeviceDtoD), - ((0, 0, false, true), TransportPath::SameDeviceDtoD), - ((0, 0, true, false), TransportPath::SameDeviceDtoD), - ((0, 0, true, true), TransportPath::SameDeviceDtoD), - ((0, 1, false, false), TransportPath::SameDeviceDtoD), - ((0, 1, false, true), TransportPath::HostStagingTransit), - ((0, 1, true, false), TransportPath::SameDeviceDtoD), - ((0, 1, true, true), TransportPath::P2PPeerAccess), - ]; - for ((s, r, p2p, cuda), expected) in cases { - let got = classify_transport(*s, *r, *p2p, *cuda); - assert_eq!( - got, *expected, - "classify_transport(s={s}, r={r}, p2p={p2p}, cuda={cuda}) → {got:?}, expected {expected:?}" - ); - } - } - - // -- classify_write_path ------------------------------------------------- - - #[test] - fn write_path_cpu_to_gpu_pool_dma() { - assert_eq!(classify_write_path(1, false, 0), WritePath::CpuToGpuPoolDma); - // transit_ptr is irrelevant when !is_cuda - assert_eq!( - classify_write_path(1, false, 0xDEAD), - WritePath::CpuToGpuPoolDma - ); - } - - #[test] - fn write_path_gpu_to_gpu_pool_transit() { - assert_eq!( - classify_write_path(1, true, 1), - WritePath::GpuToGpuPoolTransit - ); - assert_eq!( - classify_write_path(1, true, 0xDEAD_BEEF), - WritePath::GpuToGpuPoolTransit - ); - } - - #[test] - fn write_path_gpu_to_gpu_pool_dtod() { - assert_eq!(classify_write_path(1, true, 0), WritePath::GpuToGpuPoolDtoD); - } - - #[test] - fn write_path_gpu_to_shmem() { - assert_eq!(classify_write_path(0, true, 0), WritePath::GpuToShmem); - assert_eq!(classify_write_path(0, true, 1), WritePath::GpuToShmem); - } - - #[test] - fn write_path_cpu_to_shmem() { - assert_eq!(classify_write_path(0, false, 0), WritePath::CpuToShmem); - assert_eq!(classify_write_path(0, false, 1), WritePath::CpuToShmem); - } - - #[test] - fn write_path_cache_miss_defaults_to_dtod() { - // When the write fast path hits a cache miss and constructs a - // fresh PoolSlot with transit_ptr=0, the dispatch must fall - // through to plain DtoD — NOT transit. This documents the - // current behaviour; if cache-miss transit recovery is added - // later, this test must be updated. - let path = classify_write_path(1, true, 0); - assert_eq!(path, WritePath::GpuToGpuPoolDtoD); - } - - #[test] - fn write_path_full_matrix() { - // 2×2×2 = 8 cases; 6 reachable (ipc_present=1 && transit_ptr≠0 - // for a CPU source is semantically unreachable because transit is - // only allocated on the CUDA-registration path). - let cases: &[(u64, bool, u64, WritePath)] = &[ - (1, false, 0, WritePath::CpuToGpuPoolDma), - (1, false, 1, WritePath::CpuToGpuPoolDma), - (1, true, 0, WritePath::GpuToGpuPoolDtoD), - (1, true, 1, WritePath::GpuToGpuPoolTransit), - (0, false, 0, WritePath::CpuToShmem), - (0, false, 1, WritePath::CpuToShmem), - (0, true, 0, WritePath::GpuToShmem), - (0, true, 1, WritePath::GpuToShmem), - ]; - for (ipc, cuda, tp, expected) in cases { - let got = classify_write_path(*ipc, *cuda, *tp); - assert_eq!( - got, *expected, - "classify_write_path(ipc={ipc}, cuda={cuda}, tp={tp}) → {got:?}, expected {expected:?}" - ); - } - } - - // -- check_capacity_gpu_pool -------------------------------------------- - - #[test] - fn capacity_ok_within_bounds() { - // Daemon-trusted cap present, size fits - assert_eq!( - check_capacity_gpu_pool(Some(4096), None, 4096), - CapacityCheck::Ok - ); - assert_eq!( - check_capacity_gpu_pool(Some(4096), None, 1), - CapacityCheck::Ok - ); - } - - #[test] - fn capacity_exceeds_trusted_size() { - assert_eq!( - check_capacity_gpu_pool(Some(4096), None, 4097), - CapacityCheck::ExceedsTrustedSize - ); - } - - #[test] - fn capacity_fallback_to_gpu_buf_size() { - // No daemon entry, but cached first-import baseline exists - assert_eq!( - check_capacity_gpu_pool(None, Some(4096), 2048), - CapacityCheck::Ok - ); - assert_eq!( - check_capacity_gpu_pool(None, Some(4096), 4097), - CapacityCheck::ExceedsTrustedSize - ); - } - - #[test] - fn capacity_no_trusted_entry_rejects() { - // Neither daemon nor cached baseline — must fail closed - assert_eq!( - check_capacity_gpu_pool(None, None, 1024), - CapacityCheck::NoTrustedEntry - ); - } - - #[test] - fn capacity_zero_trusted_cap_allows_zero_size_only() { - // The daemon rejects zero-size pools at registration, so - // a trusted capacity of 0 is unreachable in normal operation. - // If it does occur, fail closed: any non-zero size must be - // rejected rather than silently skipping validation. - assert_eq!(check_capacity_gpu_pool(Some(0), None, 0), CapacityCheck::Ok); - assert_eq!( - check_capacity_gpu_pool(Some(0), None, 1024 * 1024), - CapacityCheck::ExceedsTrustedSize - ); - } - - #[test] - fn capacity_zero_cached_buf_size_fails_closed() { - // Same reasoning as above: a cached baseline of 0 is - // unreachable; fail closed on any non-zero read. - assert_eq!( - check_capacity_gpu_pool(None, Some(0), 1024), - CapacityCheck::ExceedsTrustedSize - ); - } -} - -/// Get (or compile) the persistent CUDA DMA helper module. -/// -/// Compiled once at first use and reused across all subsequent iterations. -/// The module maintains internal state for pinned host pointers and GPU buffers, -/// eliminating per-call cudaHostRegister/cudaMalloc/cudaFree/shm_open overhead. -fn get_cuda_helpers(py: Python<'_>) -> Result, String> { - let mut guard = CUDA_HELPERS.lock().unwrap_or_else(|e| e.into_inner()); - if let Some(ref module) = *guard { - return Ok(module.clone_ref(py)); - } - - let code = r#" -import ctypes -_lib = ctypes.CDLL('libcudart.so') - -_lib.cudaHostRegister.restype = ctypes.c_int -_lib.cudaHostRegister.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_uint] - -_lib.cudaHostUnregister.restype = ctypes.c_int -_lib.cudaHostUnregister.argtypes = [ctypes.c_void_p] - -_lib.cudaMalloc.restype = ctypes.c_int -_lib.cudaMalloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t] - -_lib.cudaMemcpy.restype = ctypes.c_int -_lib.cudaMemcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] - -_lib.cudaFree.restype = ctypes.c_int -_lib.cudaFree.argtypes = [ctypes.c_void_p] - -_lib.cudaDeviceSynchronize.restype = ctypes.c_int - -_lib.cudaHostGetDevicePointer.restype = ctypes.c_int -_lib.cudaHostGetDevicePointer.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_void_p, ctypes.c_uint] - -cudaMemcpyHostToDevice = 1 - -# IPC handle struct — must be a Structure subclass so ctypes passes it -# by value (64 bytes on the stack) to cudaIpcOpenMemHandle. -class _CudaIpcMemHandle(ctypes.Structure): - _fields_ = [("reserved", ctypes.c_byte * 64)] - -_lib.cudaIpcGetMemHandle.restype = ctypes.c_int -_lib.cudaIpcGetMemHandle.argtypes = [ctypes.POINTER(_CudaIpcMemHandle), ctypes.c_void_p] - -_lib.cudaIpcOpenMemHandle.restype = ctypes.c_int -_lib.cudaIpcOpenMemHandle.argtypes = [ctypes.POINTER(ctypes.c_void_p), _CudaIpcMemHandle, ctypes.c_uint] - -_lib.cudaIpcCloseMemHandle.restype = ctypes.c_int -_lib.cudaIpcCloseMemHandle.argtypes = [ctypes.c_void_p] - -# Persistent state: per-slot GPU buffer cache -_gpu_bufs = {} # slot -> (d_ptr, size) - -def _register_host(ptr, size): - """Pin host memory. Idempotent — error 712 (already registered) is ok.""" - err = _lib.cudaHostRegister(ctypes.c_void_p(ptr), size, 0) - if err != 0 and err != 712: - raise RuntimeError(f'cudaHostRegister(0x{ptr:x}, {size}) failed: {err}') - -def _unregister_host(ptr): - """Unpin host memory.""" - err = _lib.cudaHostUnregister(ctypes.c_void_p(ptr)) - if err != 0 and err != 713: - raise RuntimeError(f'cudaHostUnregister(0x{ptr:x}) failed: {err}') - -def _get_device_ptr(host_ptr): - """Get a GPU VA for a pinned host memory region.""" - d_ptr = ctypes.c_void_p() - err = _lib.cudaHostGetDevicePointer(ctypes.byref(d_ptr), ctypes.c_void_p(host_ptr), 0) - if err != 0: - raise RuntimeError(f'cudaHostGetDevicePointer(0x{host_ptr:x}) failed: {err}') - return d_ptr.value - - -# P2P and device query bindings (CUDA runtime). -_lib.cudaGetDevice.restype = ctypes.c_int -_lib.cudaGetDevice.argtypes = [ctypes.POINTER(ctypes.c_int)] -_lib.cudaDeviceCanAccessPeer.restype = ctypes.c_int -_lib.cudaDeviceCanAccessPeer.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_int] -_lib.cudaDeviceEnablePeerAccess.restype = ctypes.c_int -_lib.cudaDeviceEnablePeerAccess.argtypes = [ctypes.c_int, ctypes.c_uint] -_lib.cudaSetDevice.restype = ctypes.c_int -_lib.cudaSetDevice.argtypes = [ctypes.c_int] - -# Page-locked host allocation for cross-device staging buffers. -_lib.cudaHostAlloc.restype = ctypes.c_int -_lib.cudaHostAlloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t, ctypes.c_uint] -_lib.cudaFreeHost.restype = ctypes.c_int -_lib.cudaFreeHost.argtypes = [ctypes.c_void_p] - -_p2p_enabled_pairs = set() # {(src, dst)} pairs already peer-enabled - -def _alloc_transit(size): - """Allocate a page-locked CPU buffer for cross-device staging. - Returns the pointer (as int), or 0 on failure.""" - ptr = ctypes.c_void_p() - if _lib.cudaHostAlloc(ctypes.byref(ptr), size, 0) == 0: - return ptr.value - return 0 - -def _free_transit(ptr): - """Free a page-locked CPU transit buffer.""" - if ptr: - _lib.cudaFreeHost(ctypes.c_void_p(ptr)) - -def _can_access_peer(src, dst): - """Check if src GPU can P2P-access dst GPU. Returns bool.""" - can = ctypes.c_int(0) - if _lib.cudaDeviceCanAccessPeer(ctypes.byref(can), src, dst) == 0: - return can.value != 0 - return False - -def _set_cuda_device(idx): - """Set the current CUDA device. No-op if *idx* < 0.""" - if idx >= 0: - _lib.cudaSetDevice(idx) - -def _get_cuda_device(): - """Return the current CUDA device index, or -1 on failure.""" - dev = ctypes.c_int() - if _lib.cudaGetDevice(ctypes.byref(dev)) == 0: - return dev.value - return -1 - -def _transit_copy(src_ptr, src_dev, transit_ptr, dst_ptr, dst_dev, size): - """Copy via CPU transit: src(GPU) → DtoH → transit → HtoD → dst(GPU). - Returns True on success. Restores the caller's current CUDA device on exit.""" - saved = ctypes.c_int() - _lib.cudaGetDevice(ctypes.byref(saved)) - # GPU src → CPU transit - _lib.cudaSetDevice(src_dev) - err = _lib.cudaMemcpy(ctypes.c_void_p(transit_ptr), ctypes.c_void_p(src_ptr), size, 2) - if err != 0: - _lib.cudaSetDevice(saved.value) - return False - _lib.cudaDeviceSynchronize() - # CPU transit → GPU dst - _lib.cudaSetDevice(dst_dev) - err = _lib.cudaMemcpy(ctypes.c_void_p(dst_ptr), ctypes.c_void_p(transit_ptr), size, 1) - _lib.cudaDeviceSynchronize() - _lib.cudaSetDevice(saved.value) - return err == 0 - -def _transit_copy_gpu_buf(slot, src_ptr, src_dev, transit_ptr, dst_dev, size): - """Same as _transit_copy but looks up the pool's GPU buffer by slot. - Raises on a missing/undersized slot or a failed copy so the caller - can surface the error instead of silently delivering stale data.""" - if slot not in _gpu_bufs: - raise RuntimeError(f"GPU pool buffer for slot {slot} not initialised") - dst, capacity = _gpu_bufs[slot] - if size > capacity: - raise RuntimeError( - f"write size {size} exceeds GPU pool buffer capacity {capacity} (slot={slot})" - ) - if not _transit_copy(src_ptr, src_dev, transit_ptr, dst, dst_dev, size): - raise RuntimeError(f"transit copy into GPU pool buffer failed (slot={slot})") - -def _ensure_p2p_pair(a, b): - """Enable bidirectional P2P access for the (a, b) GPU pair only. - - Scoped to the sender/receiver devices actually used by a transfer — - enabling P2P across *all* GPU pairs would lazily create a CUDA context - (hundreds of MB) on every visible device, including ones not in this - dataflow. Idempotent per pair. - """ - if a == b or a < 0 or b < 0: - return - if (a, b) in _p2p_enabled_pairs: - return - # Save the current device so we can restore it after enabling peer - # access — otherwise later CUDA operations (cudaMalloc, - # cudaIpcOpenMemHandle) would run on the last enabled src device. - saved = ctypes.c_int() - _lib.cudaGetDevice(ctypes.byref(saved)) - for src, dst in ((a, b), (b, a)): - can = ctypes.c_int(0) - if _lib.cudaDeviceCanAccessPeer(ctypes.byref(can), src, dst) == 0 and can.value: - _lib.cudaSetDevice(src) - _lib.cudaDeviceEnablePeerAccess(dst, 0) - _p2p_enabled_pairs.add((src, dst)) - _lib.cudaSetDevice(saved.value) - -def _cuda_memcpy(dst, src, size, kind): - """cudaMemcpy wrapper. kind: 1=H2D, 2=D2H, 3=D2D.""" - err = _lib.cudaMemcpy(ctypes.c_void_p(dst), ctypes.c_void_p(src), size, kind) - if err != 0: - raise RuntimeError(f"cudaMemcpy(0x{dst:x}, 0x{src:x}, {size}, {kind}) failed: {err}") - _lib.cudaDeviceSynchronize() - -def _cuda_memcpy_gpu_buf(slot, src_ptr, size): - """Copy *size* bytes from *src_ptr* (GPU) into the pool's pinned GPU buffer - identified by *slot*. Used by write_memory_pool when both source and pool - buffer are GPU-resident (same-device DtoD copy).""" - if slot not in _gpu_bufs: - raise RuntimeError(f"GPU pool buffer for slot {slot} not initialised") - dst, capacity = _gpu_bufs[slot] - if size > capacity: - # The GPU buffer was sized at registration and its IPC handle is - # already exported; growing it would invalidate the receiver's - # imported handle. Reject rather than overflow the allocation. - raise RuntimeError( - f"write size {size} exceeds GPU pool buffer capacity {capacity} (slot={slot})" - ) - err = _lib.cudaMemcpy(ctypes.c_void_p(dst), ctypes.c_void_p(src_ptr), size, 3) - if err != 0: - raise RuntimeError(f"cudaMemcpy GPU buf DtoD (slot={slot}, 0x{dst:x}←0x{src_ptr:x}, {size}B) failed: {err}") - _lib.cudaDeviceSynchronize() - -def _get_gpu_buf(slot, size): - """Get or allocate a GPU buffer for the given slot. Reuses when size matches.""" - if slot in _gpu_bufs and _gpu_bufs[slot][1] >= size: - return _gpu_bufs[slot][0] - if slot in _gpu_bufs: - _lib.cudaFree(ctypes.c_void_p(_gpu_bufs[slot][0])) - d_ptr = ctypes.c_void_p() - err = _lib.cudaMalloc(ctypes.byref(d_ptr), size) - if err != 0: - raise RuntimeError(f'cudaMalloc({size}) failed: {err}') - _gpu_bufs[slot] = (d_ptr.value, size) - return d_ptr.value - -def _free_gpu_buf(slot): - """Free the pooled GPU buffer for the given slot.""" - if slot in _gpu_bufs: - _lib.cudaFree(ctypes.c_void_p(_gpu_bufs[slot][0])) - del _gpu_bufs[slot] - -def _ipc_export(d_ptr): - """Export GPU memory for cross-process sharing. Returns 64-byte handle.""" - handle = _CudaIpcMemHandle() - err = _lib.cudaIpcGetMemHandle(ctypes.byref(handle), ctypes.c_void_p(d_ptr)) - if err != 0: - raise RuntimeError(f'cudaIpcGetMemHandle(0x{d_ptr:x}) failed: {err}') - return bytes(handle) - -def _ipc_import(handle_bytes): - """Import GPU memory from another process. Returns device pointer. - - cudaIpcOpenMemHandle takes the handle struct BY VALUE (64 bytes on - the stack). _CudaIpcMemHandle.from_buffer_copy reconstructs the struct - and ctypes copies it onto the stack for the call. - """ - handle = _CudaIpcMemHandle.from_buffer_copy(handle_bytes) - d_ptr = ctypes.c_void_p() - err = _lib.cudaIpcOpenMemHandle(ctypes.byref(d_ptr), handle, 1) - if err != 0: - raise RuntimeError(f'cudaIpcOpenMemHandle failed: {err}') - _lib.cudaDeviceSynchronize() - return d_ptr.value - -def _ipc_close(d_ptr): - """Close an IPC memory handle opened by _ipc_import. - Frees the GPU-side mapping without freeing the underlying allocation.""" - err = _lib.cudaIpcCloseMemHandle(ctypes.c_void_p(d_ptr)) - if err != 0: - raise RuntimeError(f'cudaIpcCloseMemHandle(0x{d_ptr:x}) failed: {err}') - -def dma_copy(ptr, size, slot, no_dma): - """DMA transfer from host to pre-allocated GPU buffer. - - Copies via cudaMemcpyHtoD (DMA engine). When *no_dma* is false - (the default), the source memory is pinned (cudaHostRegister) - before the copy and unpinned after — this is the fast path for - large tensors where pinned-DMA bandwidth outweighs the pin/unpin - fixed cost. When *no_dma* is true, pin/unpin is skipped, using - pageable memory (faster for small tensors where pin overhead - dominates). - - Returns the device pointer of the pooled GPU buffer. - - Raises RuntimeError if *size* exceeds the existing buffer capacity - and the buffer cannot be grown (GPU pool buffers are IPC-exported; - reallocation would invalidate the receiver's imported handle). - """ - if slot in _gpu_bufs: - _, capacity = _gpu_bufs[slot] - if size > capacity: - raise RuntimeError( - f"write size {size} exceeds GPU pool buffer capacity {capacity} (slot={slot})" - " — GPU buffer cannot be grown (IPC handle already exported)" - ) - if not no_dma: - _register_host(ptr, size) - try: - d_ptr = _get_gpu_buf(slot, size) - err = _lib.cudaMemcpy( - ctypes.c_void_p(d_ptr), - ctypes.c_void_p(ptr), - size, - cudaMemcpyHostToDevice, - ) - if err != 0: - raise RuntimeError(f'cudaMemcpy failed: {err}') - _lib.cudaDeviceSynchronize() - finally: - if not no_dma: - _unregister_host(ptr) - return d_ptr - -"#; - - let code_cstr = std::ffi::CString::new(code).map_err(|e| format!("CString: {}", e))?; - let bound_module = PyModule::from_code( - py, - code_cstr.as_c_str(), - c"_cuda_helpers.py", - c"_cuda_helpers", - ) - .map_err(|e| format!("CUDA helper compile failed: {}", e))?; - let module_ref: Py = bound_module.unbind(); - *guard = Some(module_ref.clone_ref(py)); - Ok(module_ref) -} - -/// Read 8 consecutive bytes from `p` as a little-endian u64, -/// without an `.unwrap()` call (avoids the CI unwrap budget). -fn read_header_u64(p: *const u8) -> u64 { - const { assert!(std::mem::size_of::() == 8) }; - let mut buf = [0u8; 8]; - unsafe { std::ptr::copy_nonoverlapping(p, buf.as_mut_ptr(), 8) }; - u64::from_le_bytes(buf) -} - -fn parse_memory_pool_id(memory_pool_id: Py, py: Python<'_>) -> eyre::Result { - let array_data = arrow::array::ArrayData::from_pyarrow_bound(memory_pool_id.bind(py))?; - let array = arrow::array::make_array(array_data); - - if let Some(string_array) = array.as_any().downcast_ref::() { - if string_array.len() != 1 { - eyre::bail!( - "expected string array with exactly one element, got {}", - string_array.len() - ); - } - Ok(string_array.value(0).to_string()) - } else if let Some(binary_array) = array.as_any().downcast_ref::() { - if binary_array.len() != 1 { - eyre::bail!( - "expected binary array with exactly one element, got {}", - binary_array.len() - ); - } - Ok(String::from_utf8(binary_array.value(0).to_vec())?) - } else { - eyre::bail!( - "memory_pool_id must be a string or binary array, got {:?}", - array.data_type() - ) - } -} - -fn warn_missing_memory_pool(node_id: &NodeId, action: &str, buffer_id: &str) { - tracing::warn!( - "[{}] Attempt to {} memory pool [{}] failed - reason: pool does not exist. Operation aborted.", - node_id, - action, - buffer_id - ); -} - /// Converts a Python `timeout` (seconds) to `Duration`, returning a /// clean `ValueError` on NaN, negative, or infinite input instead of /// the panic that [`Duration::from_secs_f32`] raises on such inputs. @@ -1137,135 +139,6 @@ fn timeout_to_duration(timeout: Option) -> PyResult> { .transpose() } -/// Begins a memory-pool seqlock write at `gen_ptr` (header offset 96): -/// marks the generation "writing" (even -> odd) and returns the pre-write -/// (even) value so the matching [`seqlock_end_write`] call can either -/// publish the next generation or roll back to this one. -unsafe fn seqlock_begin_write(gen_ptr: *mut u64) -> u64 { - unsafe { - let pre_write_gen = std::ptr::read_volatile(gen_ptr); - std::ptr::write_volatile(gen_ptr, pre_write_gen.wrapping_add(1)); - std::sync::atomic::fence(std::sync::atomic::Ordering::Release); - pre_write_gen - } -} - -/// Begins a memory-pool seqlock write at `gen_ptr` (header offset 96) -/// **if the generation is even**. Returns the **even** pre-write -/// generation — i.e., the generation value before the write cycle -/// began, which is always even. -/// -/// If the generation is already odd (leftover from a previous failed -/// write), the begin-increment is skipped and the **previous** even -/// generation (`cur - 1`) is returned. This ensures that -/// `seqlock_end`'s `pre + 2` always produces an even generation, -/// avoiding a permanent parity inversion. -unsafe fn seqlock_begin_if_even(gen_ptr: *mut u64) -> u64 { - unsafe { - let cur = std::ptr::read_volatile(gen_ptr); - if cur.is_multiple_of(2) { - std::ptr::write_volatile(gen_ptr, cur + 1); - std::sync::atomic::fence(std::sync::atomic::Ordering::Release); - } - cur & !1 // always return the even baseline - } -} - -/// Closes a memory-pool seqlock write (header offset 96). -/// -/// Advances the generation to `pre_write_gen + 2` (even = "complete"). -/// The caller must only invoke this on a successful copy. GPU pool write -/// paths leave gen odd on failure (in-place writes cannot roll back to a -/// clean previous frame; double-buffering is deferred to a follow-up PR). -/// The `copy_ok == false` rollback branch is retained for the helper's -/// contract but is dead code in production — see the leave-gen-odd blocks -/// in `write_memory_pool`. -unsafe fn seqlock_end(gen_ptr: *mut u64, pre_write_gen: u64, copy_ok: bool) { - unsafe { - if copy_ok { - std::ptr::write_volatile(gen_ptr, pre_write_gen.wrapping_add(2)); - } else { - std::ptr::write_volatile(gen_ptr, pre_write_gen); - } - std::sync::atomic::fence(std::sync::atomic::Ordering::Release); - } -} - -#[cfg(test)] -mod seqlock_tests { - #[test] - fn begin_if_even_when_even_flips_to_odd() { - let mut generation: u64 = 10; - unsafe { - super::seqlock_begin_if_even(&mut generation); - } - assert_eq!(generation, 11, "even generation must be flipped to odd"); - } - - #[test] - fn begin_if_even_when_odd_returns_even_baseline() { - let mut generation: u64 = 11; - let pre; - unsafe { - pre = super::seqlock_begin_if_even(&mut generation); - } - assert_eq!(generation, 11, "odd generation must stay odd (skip begin)"); - assert_eq!(pre, 10, "pre-write baseline must be even (cur - 1)"); - assert_eq!(pre % 2, 0); - } - - #[test] - fn begin_if_even_then_end_always_produces_even() { - // Simulates: gen stuck odd (11) from failure → begin returns 10 - // → end does 10+2=12 which is even → pool recovers. - let mut generation: u64 = 11; - let pre; - unsafe { - pre = super::seqlock_begin_if_even(&mut generation); - } - unsafe { - super::seqlock_end(&mut generation, pre, true); - } - assert_eq!(generation, 12, "pre(10)+2=12 is even → recovery"); - assert_eq!(generation % 2, 0); - } - - /// Regression test for #2436: a successful copy publishes the new - /// frame by advancing the generation to pre+2 (even = complete). - #[test] - fn end_success_advances_to_even() { - let mut generation: u64 = 11; // pre_write(10) + 1 - unsafe { - super::seqlock_end(&mut generation, 10, true); - } - assert_eq!(generation, 12); - assert_eq!(generation % 2, 0, "generation must be even on success"); - } - - /// Regression test for #2436: on failure the generation is rolled - /// back to pre_write_gen (even) so consumers see the previous - /// valid frame, not a torn/incomplete one. - #[test] - fn end_failure_rolls_back_to_even() { - let mut generation: u64 = 11; // pre_write(10) + 1 - unsafe { - super::seqlock_end(&mut generation, 10, false); - } - assert_eq!(generation, 10); - assert_eq!(generation % 2, 0, "rolled-back generation must be even"); - } - - /// The very first write (pre_write == 0) must not underflow. - #[test] - fn end_first_write_failure_rolls_back_to_zero() { - let mut generation: u64 = 1; - unsafe { - super::seqlock_end(&mut generation, 0, false); - } - assert_eq!(generation, 0); - } -} - /// The custom node API lets you integrate `dora` into your application. /// It allows you to retrieve input and send output in any fashion you want. /// @@ -1366,13 +239,9 @@ impl Node { #[pyo3(signature = (timeout=None))] #[allow(clippy::should_implement_trait)] pub fn next(&self, py: Python, timeout: Option) -> PyResult>> { - // Drain any daemon-broadcast FreeMemoryPool events before - // yielding the next user-visible event — this ensures that - // a single free_memory_pool call by any node releases - // per-process resources (GPU buffers, transit buffers, shmem - // mappings) in every process. - self.process_pending_memory_pool_frees(py); - + // Release pools other nodes freed before yielding the next event. + // Compiles away entirely without the `tensor-pool` feature. + self.process_pending_tensor_pool_frees(py); let timeout = timeout_to_duration(timeout)?; let event = py.detach(|| self.events.recv(timeout)); if let Some(event) = event { @@ -1385,74 +254,6 @@ impl Node { } } - /// Process any memory pools that were freed by another node. - fn process_pending_memory_pool_frees(&self, py: Python) { - for shared_memory_id in dora_node_api::event_stream::memory_pool::drain_freed_pools() { - let buffer_id = shared_memory_id; - // Receiver-side cleanup (IPC handles, shmem mappings). - { - if let Some(slot) = RECV_GPU_VA - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(&buffer_id) - { - if slot.gpu_buf != 0 { - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - let _ = bound.call_method1("_ipc_close", (slot.gpu_buf,)); - } - } else if slot.gpu_va != 0 { - // Host-registered mapping (effective_as_cuda branch): - // must cudaHostUnregister before munmap. _unregister_host - // requires the original host pointer (shmem base), not the - // device VA returned by cudaHostGetDevicePointer — passing - // the device VA makes cudaHostUnregister fail and leaks the - // pin over an address range that then gets munmap'd. - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - let _ = bound.call_method1("_unregister_host", (slot.host_base,)); - } - } - // Drop slot → munmap - } - } - RECV_CPU_SHMEM - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(&buffer_id); - - // Sender-side cleanup (PINNED_POOL, GPU/transit buffers). - // Guard against cross-process counter aliasing: buffer ids are - // pool_{node_id}_{counter}. Extract the owner segment (between - // "pool_" and the final "_") and require an exact - // equality match — a prefix check (starts_with) would alias - // across node ids that are prefixes of each other (e.g. cam / - // cam_left). - if let Some(owner_and_counter) = buffer_id.strip_prefix("pool_") - && let Some((owner, counter_str)) = owner_and_counter.rsplit_once('_') - && owner == self.node_id.as_ref() - && let Ok(c) = counter_str.parse::() - && let Some(slot) = PINNED_POOL - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(&c) - { - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - let _ = bound.call_method1("_unregister_host", (slot.base,)); - let _ = bound.call_method1("_free_gpu_buf", (c,)); - if slot.transit_ptr != 0 { - let _ = bound.call_method1("_free_transit", (slot.transit_ptr,)); - } - } - TRANSIT_META - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(&c); - } - } - } - /// `.drain()` gives you all available inputs that the node has received. /// It does not block until the next event becomes available. /// @@ -1465,7 +266,9 @@ impl Node { /// :rtype: list[dict] #[allow(clippy::should_implement_trait)] pub fn drain(&self, py: Python) -> PyResult>> { - self.process_pending_memory_pool_frees(py); + // Release pools other nodes freed before yielding the next event. + // Compiles away entirely without the `tensor-pool` feature. + self.process_pending_tensor_pool_frees(py); let events = self .events .drain() @@ -1495,7 +298,9 @@ impl Node { /// :rtype: dict #[allow(clippy::should_implement_trait)] pub fn try_recv(&mut self, py: Python) -> Option> { - self.process_pending_memory_pool_frees(py); + // Release pools other nodes freed before yielding the next event. + // Compiles away entirely without the `tensor-pool` feature. + self.process_pending_tensor_pool_frees(py); match self.events.try_recv() { Ok(event) => event.to_py_dict(py).ok(), Err(_) => None, @@ -1530,13 +335,9 @@ impl Node { #[pyo3(signature = (timeout=None))] #[allow(clippy::should_implement_trait)] pub async fn recv_async(&self, timeout: Option) -> PyResult>> { - // Same cleanup contract as `next`/`drain`/`try_recv`: release the - // per-process resources of pools that other nodes freed before - // yielding the next user-visible event. Done *before* the await so - // that it also runs on the `None` (stream closed) path, and so the - // scoped `attach` drops the GIL before the suspend point. - Python::attach(|py| self.process_pending_memory_pool_frees(py)); - + // Same contract as next/drain/try_recv. Before the await so it runs + // on the `None` path too, and so the guard drops before suspending. + Python::attach(|py| self.process_pending_tensor_pool_frees(py)); let timeout = timeout_to_duration(timeout)?; let event = self.events.recv_async_timeout(timeout).await; if let Some(event) = event { @@ -1907,6 +708,145 @@ impl Node { self.dataflow_id.to_string() } + /// Register a tensor's memory as a shared pool for zero-copy transfer. + /// + /// **Opt-in extension, outside the 1.0 compatibility guarantees.** Only + /// present when this wheel was built with the `tensor-pool` feature; see + /// `libraries/extensions/tensor-pool/README.md`. + /// + /// :type tensor_info: dict + /// :type device: str, optional + /// :rtype: pyarrow.Array + #[cfg(feature = "tensor-pool")] + #[pyo3(signature = (tensor_info, device="cpu".to_string()))] + pub fn register_tensor_pool( + &self, + tensor_info: &Bound<'_, PyDict>, + device: String, + py: Python, + ) -> eyre::Result> { + self.with_pool(|pool| pool.register_tensor_pool(tensor_info, device, py)) + } + + /// Write tensor data into an existing tensor pool. + /// + /// **Opt-in extension, outside the 1.0 compatibility guarantees.** + /// + /// :type tensor_pool_id: pyarrow.Array + /// :type tensor_info: dict + /// :rtype: None + #[cfg(feature = "tensor-pool")] + #[pyo3(signature = (tensor_pool_id, tensor_info))] + pub fn write_tensor_pool( + &self, + tensor_pool_id: Py, + tensor_info: &Bound<'_, PyDict>, + py: Python, + ) -> eyre::Result<()> { + self.with_pool(|pool| pool.write_tensor_pool(tensor_pool_id, tensor_info, py)) + } + + /// Read tensor metadata from a tensor pool (zero-copy). + /// + /// **Opt-in extension, outside the 1.0 compatibility guarantees.** + /// + /// :type tensor_pool_id: pyarrow.Array + /// :rtype: dict + #[cfg(feature = "tensor-pool")] + #[pyo3(signature = (tensor_pool_id))] + pub fn read_tensor_pool( + &self, + tensor_pool_id: Py, + py: Python, + ) -> eyre::Result> { + self.with_pool(|pool| pool.read_tensor_pool(tensor_pool_id, py)) + } + + /// Free a tensor pool, releasing it for every node that touched it. + /// + /// **Opt-in extension, outside the 1.0 compatibility guarantees.** + /// + /// :type tensor_pool_id: pyarrow.Array + /// :rtype: None + #[cfg(feature = "tensor-pool")] + #[pyo3(signature = (tensor_pool_id))] + pub fn free_tensor_pool(&self, tensor_pool_id: Py, py: Python) -> eyre::Result<()> { + self.with_pool(|pool| pool.free_tensor_pool(tensor_pool_id, py)) + } + + /// Store an opaque value in the daemon's dataflow-scoped extension table. + /// + /// The seam for transports maintained outside the dora tree. dora brokers + /// the value's lifetime and nothing else — it never interprets + /// ``namespace``, ``key`` or ``value``. + /// + /// The daemon reclaims the entry when this node exits or the dataflow + /// finishes, and notifies every node that stored or read the key when it + /// is dropped. Collect those with :meth:`drain_dropped_extension_keys`. + /// + /// ```python + /// node.extension_store("my-transport", "frame-7", descriptor_bytes) + /// ``` + /// + /// :type namespace: str + /// :type key: str + /// :type value: bytes + /// :rtype: None + pub fn extension_store( + &mut self, + namespace: String, + key: String, + value: &Bound<'_, PyBytes>, + ) -> eyre::Result<()> { + self.node + .get_mut() + .extension_store(namespace, key, value.as_bytes().to_vec()) + } + + /// Read an opaque value back, returning ``None`` if the key is absent. + /// + /// With ``remove=True`` the entry is dropped in the same round trip, which + /// is what a consume-once handoff wants. + /// + /// :type namespace: str + /// :type key: str + /// :type remove: bool, optional + /// :rtype: bytes | None + #[pyo3(signature = (namespace, key, remove=false))] + pub fn extension_load( + &mut self, + namespace: String, + key: String, + remove: bool, + py: Python<'_>, + ) -> eyre::Result>> { + let value = self.node.get_mut().extension_load(namespace, key, remove)?; + Ok(value.map(|bytes| PyBytes::new(py, &bytes).unbind())) + } + + /// Drop an opaque value, notifying every node that stored or read it. + /// + /// Dropping an absent key succeeds, so a retry after a lost reply is safe. + /// + /// :type namespace: str + /// :type key: str + /// :rtype: None + pub fn extension_drop(&mut self, namespace: String, key: String) -> eyre::Result<()> { + self.node.get_mut().extension_drop(namespace, key) + } + + /// Take the keys in ``namespace`` that have been dropped since the last + /// call, so the extension can release whatever it derived from them. + /// + /// Notifications arrive out of band and are not delivered as events, so + /// poll this wherever the extension next runs. + /// + /// :type namespace: str + /// :rtype: list[str] + pub fn drain_dropped_extension_keys(&self, namespace: &str) -> Vec { + dora_node_api::event_stream::extensions::drain_dropped_keys(namespace) + } + /// Returns True if this node was restarted after a previous exit or failure. /// /// Nodes can use this to decide whether to restore saved state or start fresh. @@ -2011,1437 +951,6 @@ impl Node { Ok(()) } - - // === Memory Pool API === - - /// Register a shared memory pool for zero-copy tensor transfer. - /// - /// The returned pool ID can be shared across nodes (e.g. via a Dora output) - /// so that a receiver can call [`read_memory_pool`] and [`free_memory_pool`] - /// on it. - /// - /// # Concurrency / safety - /// - /// **This pool provides no internal mutual exclusion for data bytes.** - /// The writer must not begin a new [`write_memory_pool`] while a receiver is - /// still consuming the previous tensor. Callers that share a pool across - /// nodes MUST enforce a **turn-based discipline** — for example, by waiting - /// for a `next_require` round-trip from the receiver before writing again. - /// The bundled `examples/memory-pool/` dataflows demonstrate this pattern. - /// - /// The on-segment seqlock guards metadata integrity (header fields written - /// once at registration) and detects in-flight overwrites, but it does - /// **not** block the writer from starting a new write while a consumer - /// holds a zero-copy tensor. Skipping the turn-based discipline risks - /// torn data at the consumer. - #[pyo3(signature = (tensor_info, device))] - pub fn register_memory_pool( - &self, - tensor_info: &Bound<'_, PyDict>, - device: String, - py: Python, - ) -> eyre::Result> { - let ptr_val: u64 = tensor_info - .get_item("ptr")? - .ok_or_else(|| eyre::eyre!("missing ptr"))? - .extract()?; - let size: usize = tensor_info - .get_item("size")? - .ok_or_else(|| eyre::eyre!("missing size"))? - .extract()?; - let dtype: String = tensor_info - .get_item("dtype")? - .ok_or_else(|| eyre::eyre!("missing dtype"))? - .extract()?; - let shape_list: Vec = tensor_info - .get_item("shape")? - .ok_or_else(|| eyre::eyre!("missing shape"))? - .extract()?; - let tensor_device: String = tensor_info - .get_item("device")? - .ok_or_else(|| eyre::eyre!("missing device"))? - .extract()?; - - let is_cuda = tensor_device.starts_with("cuda"); - let receiver_is_cuda = device.starts_with("cuda"); - let cpu_mode = !receiver_is_cuda; - // Auto-select pinning: key off the source device — pinning only - // matters when the source is CPU (cudaHostRegister would raise on a - // device pointer; prevented by the !is_cuda guard above). - let is_pinned = should_pin(is_cuda, size); - let pinned_type = if cpu_mode { "cpu" } else { "cuda" }; - - if ptr_val == 0 { - eyre::bail!("Invalid source pointer (NULL)"); - } - if size == 0 || size > 1024 * 1024 * 1024 { - eyre::bail!("Invalid size: {} bytes", size); - } - if cfg!(not(target_os = "linux")) { - eyre::bail!( - "memory-pool transport requires Linux (uses /dev/shm). \ - This platform is not supported." - ); - } - - // Generate unique pool counter for this registration - let pool_counter = { - let mut c = PINNED_COUNTER.lock().unwrap_or_else(|e| e.into_inner()); - *c += 1; - *c - }; - let shmem_name = format!( - "dora_pool_{}_{}_{}", - self.dataflow_id, self.node_id, pool_counter - ); - - let header_meta = PyDict::new(py); - header_meta.set_item("size", size)?; - header_meta.set_item("dtype", &dtype)?; - header_meta.set_item("shape", shape_list.clone())?; - header_meta.set_item("pinned_type", pinned_type)?; - - let json_bytes = py - .import("json") - .wrap_err("failed to import json")? - .call_method1("dumps", (header_meta,)) - .wrap_err("failed to serialize metadata to JSON")? - .extract::() - .wrap_err("failed to extract JSON string")? - .into_bytes(); - let json_len = json_bytes.len(); - let padded_json_len = json_len.div_ceil(DORADMA_METADATA_ALIGN) * DORADMA_METADATA_ALIGN; - let data_offset = DORADMA_HEADER_SIZE + padded_json_len; - // GPU receivers read tensor data from the IPC-exported GPU buffer, - // not from the shmem data region. Allocate only the header portion - // (metadata + IPC handle + seqlock) — a few hundred bytes instead of - // 80 MB. This also lets us skip cudaHostRegister on a useless data - // region. - let total_size = if receiver_is_cuda { - data_offset - } else { - data_offset + size - }; - - // Create shared memory - let mut shmem = ShmemConf::new() - .os_id(&shmem_name) - .size(total_size) - .writable(true) - .create() - .wrap_err_with(|| { - format!( - "failed to create pool shared memory `{}` (name collision with another node or leftover segment)", - shmem_name - ) - })?; - let shmem_ptr = unsafe { shmem.as_slice_mut().as_mut_ptr() }; - - // Pin the shmem for DMA only when the receiver reads from it - // (CPU receivers). GPU receivers never touch the shmem data - // region, and the header-only shmem is too small (< 1 page) to - // benefit from pinning. - if !receiver_is_cuda && let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - let _ = bound.call_method1("_register_host", (shmem_ptr as u64, total_size)); - } - - shmem.set_owner(false); - - // Write DORADMA header - unsafe { - std::ptr::copy_nonoverlapping(DORADMA_MAGIC.as_ptr(), shmem_ptr, 8); - let json_len_le = (json_len as u64).to_le_bytes(); - std::ptr::copy_nonoverlapping(json_len_le.as_ptr(), shmem_ptr.add(8), 8); - let data_off_le = (data_offset as u64).to_le_bytes(); - std::ptr::copy_nonoverlapping(data_off_le.as_ptr(), shmem_ptr.add(16), 8); - std::ptr::copy_nonoverlapping( - json_bytes.as_ptr(), - shmem_ptr.add(DORADMA_HEADER_SIZE), - json_len, - ); - } - - // Initialize seqlock write_gen at header[96] to 0 (even = complete) - unsafe { - std::ptr::write(shmem_ptr.add(96) as *mut u64, 0u64); - } - - // Seqlock: increment generation to odd (write-in-progress) - unsafe { - let gen_ptr = shmem_ptr.add(96) as *mut u64; - let old_gen = std::ptr::read_volatile(gen_ptr); - std::ptr::write_volatile(gen_ptr, old_gen + 1); - std::sync::atomic::fence(std::sync::atomic::Ordering::Release); - } - - // Copy tensor data to shmem — only when the receiver will - // actually read it. GPU receivers import the pool GPU buffer - // via the IPC handle in the DORADMA header and never touch the - // shmem data region; skipping this copy for them eliminates - // a redundant CPU-memcpy or GPU-DtoH transfer on every - // registration (cpu2cuda and cuda2cuda respectively). - if !receiver_is_cuda { - // The DtoH copy must publish either a fully-initialized data - // region or nothing — uninitialized shmem exposed as a valid - // frame is data corruption. Both a failed cudaMemcpy and a - // missing CUDA helper module are treated as copy failures. - let mut dtoh_copy_ok = true; - if is_cuda { - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - dtoh_copy_ok = bound - .call_method1( - "_cuda_memcpy", - (shmem_ptr as u64 + data_offset as u64, ptr_val, size, 2u32), - ) - .is_ok(); - } else { - dtoh_copy_ok = false; - } - if !dtoh_copy_ok { - // The matching `_register_host` above is unconditional (it - // runs whenever `!receiver_is_cuda`), so the unregister must - // be too — gating it on `!is_pinned` would leak the pin if - // `should_pin` is ever tuned to pin CUDA sources. - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - let _ = bound.call_method1("_unregister_host", (shmem_ptr as u64,)); - } - shmem.set_owner(true); - eyre::bail!( - "[{}] register_memory_pool: DtoH copy failed ({} → CPU shmem, {} bytes)", - self.node_id, - tensor_device, - size - ); - } - } else { - unsafe { - std::ptr::copy_nonoverlapping( - ptr_val as *const u8, - shmem_ptr.add(data_offset), - size, - ); - } - } - } - - // GPU pool: allocate GPU buffer on current device, copy data, export - // IPC handle for cross-process zero-copy access. When the source - // tensor is also on CUDA (GPU→GPU), the source and pool buffer are on - // the same device (sender's current CUDA device), so a plain DtoD - // memcpy suffices. When the source is CPU, `dma_copy` does a pinned - // host→device DMA copy (existing path). - // Resolve sender and receiver device indices for cross-device detection. - let sender_device_idx = tensor_device - .strip_prefix("cuda") - .and_then(|s| s.strip_prefix(':')) - .and_then(|s| s.parse::().ok()) - .unwrap_or(0); - let receiver_device_idx = device - .strip_prefix("cuda") - .and_then(|s| s.strip_prefix(':')) - .and_then(|s| s.parse::().ok()) - .unwrap_or(0); - let cross_device = sender_device_idx != receiver_device_idx; - let mut transit_ptr: u64 = 0; - let mut pool_device = if receiver_is_cuda { - receiver_device_idx - } else { - sender_device_idx - }; - - // Tracks whether the GPU pool buffer + IPC handle were successfully set - // up. A CUDA receiver's shmem is header-only, so without the handle the - // pool is unusable — we fail registration rather than hand back a - // permanently-broken pool. - let mut ipc_written = false; - - if receiver_is_cuda && let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - - // Enable P2P for the sender/receiver pair before any IPC operations. - // Gate on a CUDA source: for a CPU source sender_device_idx defaults - // to 0, and enabling a spurious GPU0↔receiver P2P pair creates an - // unnecessary CUDA context. CPU-source registration snapshots the - // ambient device in the else branch below — running _set_cuda_device - // first would clobber the original value saved there. - if is_cuda { - let _ = bound - .call_method1("_ensure_p2p_pair", (sender_device_idx, receiver_device_idx)); - let _ = bound.call_method1("_set_cuda_device", (sender_device_idx,)); - } - - // Resolve transport path. classify_transport encodes the full - // 2³ decision matrix (pure, CI-tested); here we only need the - // single GPU-runtime-dependent input (p2p_available). - let p2p_available: bool = cross_device - && bound - .call_method1("_can_access_peer", (sender_device_idx, receiver_device_idx)) - .and_then(|r| r.extract::()) - .unwrap_or(false); - let transport_path = classify_transport( - sender_device_idx, - receiver_device_idx, - p2p_available, - is_cuda, - ); - let use_transit = transport_path == TransportPath::HostStagingTransit; - - let gpu_ptr: Option = if is_cuda { - if use_transit { - // Allocate pool buffer on receiver's GPU so the - // receiver can import the IPC handle on its own device. - // _transit_copy internally saves/restores the caller's - // device, so an explicit restore is unnecessary here. - let _ = bound.call_method1("_set_cuda_device", (receiver_device_idx,)); - let dst: u64 = bound - .call_method1("_get_gpu_buf", (pool_counter, size)) - .and_then(|r| r.extract::()) - .unwrap_or(0); - // Switch back to sender device. - let _ = bound.call_method1("_set_cuda_device", (sender_device_idx,)); - if dst != 0 { - // Allocate CPU page-locked transit buffer. - let tp: u64 = bound - .call_method1("_alloc_transit", (size,)) - .and_then(|r| r.extract::()) - .unwrap_or(0); - if tp != 0 { - // Copy via transit: sender GPU → CPU → receiver GPU. - let ok: bool = bound - .call_method1( - "_transit_copy", - ( - ptr_val, - sender_device_idx, - tp, - dst, - receiver_device_idx, - size, - ), - ) - .and_then(|r| r.extract::()) - .unwrap_or(false); - if ok { - transit_ptr = tp; - pool_device = receiver_device_idx; - Some(dst) - } else { - // _transit_copy failed: free both buffers - let _ = bound.call_method1("_free_transit", (tp,)); - let _ = bound.call_method1("_free_gpu_buf", (pool_counter,)); - None - } - } else { - // transit alloc failed: free GPU buffer - let _ = bound.call_method1("_free_gpu_buf", (pool_counter,)); - None - } - } else { - None - } - // Both branches start from sender_device_idx (restored - // after _ensure_p2p_pair above). _transit_copy does its - // own internal save/restore; the same-device branch below - // saves/restores explicitly so later cudaMalloc calls land - // on the right GPU. - } else { - // Same-device or P2P available: allocate on sender device. - // Save the current device and restore before returning - // so later cudaMalloc calls land on the right GPU. - let saved_dev: i32 = bound - .call_method0("_get_cuda_device") - .and_then(|r| r.extract::()) - .unwrap_or(0); - let _ = bound.call_method1("_set_cuda_device", (sender_device_idx,)); - let result = { - let dst: u64 = bound - .call_method1("_get_gpu_buf", (pool_counter, size)) - .and_then(|r| r.extract::()) - .unwrap_or(0); - if dst != 0 { - // Only export a handle if the DtoD copy succeeded — - // otherwise the receiver would import uninitialised memory. - if bound - .call_method1("_cuda_memcpy", (dst, ptr_val, size, 3u32)) - .is_ok() - { - Some(dst) - } else { - let _ = bound.call_method1("_free_gpu_buf", (pool_counter,)); - None - } - } else { - None - } - }; - let _ = bound.call_method1("_set_cuda_device", (saved_dev,)); - result - } - } else { - // CPU source → GPU pool: switch to the receiver's device - // so cudaMalloc in dma_copy → _get_gpu_buf lands on the - // correct GPU (matching the GPU-source branches). - let saved_dev: i32 = bound - .call_method0("_get_cuda_device") - .and_then(|r| r.extract::()) - .unwrap_or(0); - let _ = bound.call_method1("_set_cuda_device", (receiver_device_idx,)); - let result = bound - .call_method1("dma_copy", (ptr_val, size, pool_counter, !is_pinned)) - .and_then(|r| r.extract::()) - .ok(); - let _ = bound.call_method1("_set_cuda_device", (saved_dev,)); - result - }; - - if let Some(gpu_ptr) = gpu_ptr - && let Ok(handle) = bound - .call_method1("_ipc_export", (gpu_ptr,)) - .and_then(|r| r.extract::>()) - && handle.len() == 64 - { - unsafe { - // Write IPC handle into DORADMA reserved area [32..96) - std::ptr::copy_nonoverlapping(handle.as_ptr(), shmem_ptr.add(32), 64); - // ipc_present flag at byte 24 - std::ptr::write(shmem_ptr.add(24) as *mut u64, 1u64); - } - ipc_written = true; - } - } - - // A CUDA receiver's shmem is header-only and useless without the IPC - // handle. If GPU-pool setup failed (or the CUDA helpers were - // unavailable), fail registration instead of returning a pool that - // every later write/read would silently reject. Reclaim the shmem - // segment on the way out (it was created with owner=false). - if receiver_is_cuda && !ipc_written { - // The GPU pool buffer (and, on the transit path, the page-locked - // host transit buffer) were allocated before the IPC export, which - // failed. Free them before bailing — otherwise they leak for the - // life of the process since no PoolSlot was stored to track them. - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - let _ = bound.call_method1("_free_gpu_buf", (pool_counter,)); - if transit_ptr != 0 { - let _ = bound.call_method1("_free_transit", (transit_ptr,)); - } - } - shmem.set_owner(true); - eyre::bail!( - "[{}] register_memory_pool: failed to set up GPU pool buffer / IPC handle for CUDA receiver `{}`", - self.node_id, - tensor_device - ); - } - - // Seqlock: increment generation to even (write-complete) - unsafe { - let gen_ptr = shmem_ptr.add(96) as *mut u64; - let old_gen = std::ptr::read_volatile(gen_ptr); - std::ptr::write_volatile(gen_ptr, old_gen + 1); - std::sync::atomic::fence(std::sync::atomic::Ordering::Release); - } - - // Store shmem in pool (keep alive) - { - let mut pool = PINNED_POOL.lock().unwrap_or_else(|e| e.into_inner()); - pool.insert( - pool_counter, - PoolSlot { - _shmem: shmem, - base: shmem_ptr as u64, - size: total_size, - is_pinned, - transit_ptr, - pool_device, - }, - ); - } - - // Persist transit metadata so the write fast path can recover - // transit_ptr / pool_device on a PINNED_POOL cache-miss. - if transit_ptr != 0 || pool_device != 0 { - TRANSIT_META - .lock() - .unwrap_or_else(|e| e.into_inner()) - .insert(pool_counter, (transit_ptr, pool_device)); - } - - let buffer_id = format!("pool_{}_{}", self.node_id, pool_counter); - - // Register with daemon for lifecycle tracking - { - let hlc = dora_node_api::dora_core::uhlc::HLC::default(); - let ts = hlc.new_timestamp(); - let mut params = dora_node_api::MetadataParameters::new(); - params.insert( - "ptr".to_string(), - dora_node_api::Parameter::Integer(ptr_val as i64), - ); - params.insert( - "size".to_string(), - dora_node_api::Parameter::Integer(size as i64), - ); - params.insert("dtype".to_string(), dora_node_api::Parameter::String(dtype)); - params.insert( - "shape".to_string(), - dora_node_api::Parameter::ListInt(shape_list), - ); - params.insert( - "shared_memory_name".to_string(), - dora_node_api::Parameter::String(shmem_name), - ); - params.insert( - "is_pinned".to_string(), - dora_node_api::Parameter::Bool(is_pinned), - ); - params.insert( - "pinned_type".to_string(), - dora_node_api::Parameter::String(pinned_type.to_string()), - ); - params.insert( - "ipc_present".to_string(), - dora_node_api::Parameter::Bool(ipc_written), - ); - params.insert( - "buffer_id".to_string(), - dora_node_api::Parameter::String(buffer_id.clone()), - ); - - let meta = dora_node_api::Metadata::from_parameters(ts, params); - if let Err(e) = self - .node - .get_mut() - .register_pinned_memory(buffer_id.clone(), meta) - { - tracing::warn!("[{}] failed to register memory pool: {:#}", self.node_id, e); - } - } - - let buffer_id_array = arrow::array::StringArray::from(vec![buffer_id]); - let buf_py: Py = buffer_id_array.to_data().to_pyarrow(py)?.unbind(); - Ok(buf_py) - } - - /// Write tensor data to an existing memory pool. - /// - /// Overwrites the data region of a previously-registered pool without - /// re-registering, enabling memory reuse across iterations. - /// - /// # Concurrency / safety - /// - /// **This is a non-blocking overwrite.** The writer must not call this - /// while a receiver is consuming the previous tensor. Data-byte - /// consistency relies entirely on callers honoring a **turn-based - /// discipline**: wait for the receiver to signal completion (e.g. via - /// a `next_require` round-trip) before writing the next frame. - /// - /// The seqlock at header offset 96 detects in-flight overwrites so the - /// reader can retry, but it does **not** prevent the overwrite itself. - /// A writer that ignores the turn-based contract will produce torn - /// (partially updated) data at the consumer. - /// - /// The bundled `examples/memory-pool/` dataflows demonstrate correct - /// turn-based usage: the sender writes, outputs the pool ID, and waits - /// for the next input event before writing again. - #[pyo3(signature = (memory_pool_id, tensor_info))] - pub fn write_memory_pool( - &self, - memory_pool_id: Py, - tensor_info: &Bound<'_, PyDict>, - py: Python, - ) -> eyre::Result<()> { - let buffer_id = parse_memory_pool_id(memory_pool_id, py)?; - - let ptr_val: u64 = tensor_info - .get_item("ptr")? - .ok_or_else(|| eyre::eyre!("missing ptr"))? - .extract()?; - let size: usize = tensor_info - .get_item("size")? - .ok_or_else(|| eyre::eyre!("missing size"))? - .extract()?; - let tensor_device: String = tensor_info - .get_item("device")? - .ok_or_else(|| eyre::eyre!("missing device"))? - .extract()?; - let is_cuda = tensor_device.starts_with("cuda"); - - { - let freed = FREED_POOL_IDS.lock().unwrap_or_else(|e| e.into_inner()); - if freed.contains(&buffer_id) { - warn_missing_memory_pool(&self.node_id, "write", &buffer_id); - return Ok(()); - } - } - - // Auto-select pinning based on tensor size (25 MiB threshold). - // Shared by cache-miss PoolSlot construction and slow-path dma_copy; - // cache-hit reuses the slot's stored is_pinned. - let auto_pin = should_pin(is_cuda, size); - - // Fast path: pool_ format -> DORADMA - if buffer_id.starts_with("pool_") { - // Extract counter from the last underscore segment — node_id - // may legitimately contain underscores. - if let Some((_, counter_str)) = buffer_id.rsplit_once('_') - && let Ok(counter) = counter_str.parse::() - { - // Try PINNED_POOL cache first to avoid per-iteration mmap/munmap. - // register_memory_pool already stored the Shmem here; taking it - // prevents munmap, and storing it back keeps the mapping alive. - let pool_slot = { - PINNED_POOL - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(&counter) - }; - - // Both cache-hit and cache-miss produce a PoolSlot that is - // stored back into PINNED_POOL after the write — this keeps - // the shmem mapping alive for the duration of the data copy. - let (shmem_ptr, shmem_capacity, mut store_back, is_pinned) = - if let Some(mut slot_data) = pool_slot { - // Cache hit: reuse the persistent mapping (no mmap). - // Recompute is_pinned from the current tensor size - // so the auto-selection reflects each write's payload. - let cap = slot_data.size; - slot_data.is_pinned = auto_pin; - let pinned = auto_pin; - (slot_data.base as *mut u8, cap, Some(slot_data), pinned) - } else { - // Cache miss: open via ShmemConf, wrap immediately - // so the mapping stays alive until post-write re-insert. - let shmem_name = format!( - "dora_pool_{}_{}_{}", - self.dataflow_id, self.node_id, counter - ); - match ShmemConf::new().os_id(&shmem_name).open() { - Ok(shmem) => { - let cap = shmem.len(); - let base = shmem.as_ptr() as u64; - let slot = PoolSlot { - _shmem: shmem, - base, - size: cap, - is_pinned: auto_pin, - transit_ptr: 0, - pool_device: 0, - }; - (base as *mut u8, cap, Some(slot), auto_pin) - } - Err(_) => (std::ptr::null_mut(), 0, None, false), - } - }; - - if !shmem_ptr.is_null() { - // Guard against truncated segments before any - // pointer arithmetic (mirrors slow-path + read guards). - if shmem_capacity < DORADMA_HEADER_SIZE { - if let Some(slot_data) = store_back { - PINNED_POOL - .lock() - .unwrap_or_else(|e| e.into_inner()) - .insert(counter, slot_data); - } - return Ok(()); - } - let magic = unsafe { std::slice::from_raw_parts(shmem_ptr, 8) }; - if magic == DORADMA_MAGIC { - let data_offset = unsafe { read_header_u64(shmem_ptr.add(16)) as usize }; - - // Check if this pool has GPU DMA path enabled - let ipc_present = - unsafe { std::ptr::read(shmem_ptr.add(24) as *const u64) }; - - // Validate write size against pool capacity. The - // `size == 0` guard always applies. For GPU-buffer pools - // (ipc_present == 1) the shmem data region is unused, so - // its capacity is irrelevant — the GPU-buffer helpers - // validate `size` against the actual GPU allocation. - if size == 0 - || (ipc_present != 1 - && size > shmem_capacity.saturating_sub(data_offset)) - { - tracing::warn!( - "[{}] write_memory_pool: size {} exceeds available pool capacity (data_offset={}, total={}), operation aborted", - self.node_id, - size, - data_offset, - shmem_capacity - ); - // Store back to PINNED_POOL to keep shmem alive - if let Some(slot_data) = store_back { - PINNED_POOL - .lock() - .unwrap_or_else(|e| e.into_inner()) - .insert(counter, slot_data); - } - return Ok(()); - } - - if ipc_present == 1 && !is_cuda { - // Seqlock: begin (noop if still odd from a - // previous failed copy). - let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; - let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; - let mut copy_ok = true; - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - if let Err(e) = bound - .call_method1("dma_copy", (ptr_val, size, counter, !is_pinned)) - { - copy_ok = false; - tracing::error!( - "[{}] write_memory_pool: DMA copy failed: {}", - self.node_id, - e - ); - } - } else { - copy_ok = false; - } - if copy_ok { - // Publish: gen was odd (in-progress), flip to even. - unsafe { - seqlock_end(gen_ptr, pre_write_gen, true); - } - } else { - // GPU in-place write to the live IPC buffer: a - // partial copy cannot be rolled back to a clean - // previous frame, so leave gen odd. The reader - // retries until the next successful write. - } - if !copy_ok { - // Re-insert the slot so free_memory_pool - // can clean up the GPU buffer and transit - // allocation (mirrors the is_cuda branch). - if let Some(slot_data) = store_back.take() { - PINNED_POOL - .lock() - .unwrap_or_else(|e| e.into_inner()) - .insert(counter, slot_data); - } - return Err(eyre::eyre!( - "[{}] write_memory_pool: DMA copy failed", - self.node_id - )); - } - } else if is_cuda { - // Seqlock: begin (noop if still odd from a - // previous failed copy). - let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; - let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; - let mut copy_ok = true; - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - // Resolve transit metadata: cache-hit from PoolSlot, - // cache-miss from TRANSIT_META (populated during registration). - let mut transit_ptr = - store_back.as_ref().map_or(0, |s| s.transit_ptr); - let transit_from_cache; - if transit_ptr == 0 { - // Cache-miss fallback: TRANSIT_META survives - // PINNED_POOL eviction so the write fast path - // always knows whether a transit buffer exists. - let meta = - TRANSIT_META.lock().unwrap_or_else(|e| e.into_inner()); - if let Some(&(tp, _pd)) = meta.get(&counter) { - transit_ptr = tp; - transit_from_cache = true; - } else { - transit_from_cache = false; - } - } else { - transit_from_cache = false; - } - let pool_dev = store_back.as_ref().map_or(0, |s| s.pool_device); - let write_path = classify_write_path( - ipc_present, - /*is_cuda=*/ true, - transit_ptr, - ); - let res = match write_path { - WritePath::GpuToGpuPoolTransit => { - // Recover pool_device from TRANSIT_META on - // cache-miss; otherwise use the PoolSlot value. - let pool_dev = if transit_from_cache { - TRANSIT_META - .lock() - .unwrap_or_else(|e| e.into_inner()) - .get(&counter) - .copied() - .map(|(_tp, pd)| pd) - .unwrap_or(pool_dev) - } else { - pool_dev - }; - let sender_dev = tensor_device - .strip_prefix("cuda") - .and_then(|d| d.strip_prefix(':')) - .and_then(|d| d.parse::().ok()) - .unwrap_or(0); - bound - .call_method1( - "_transit_copy_gpu_buf", - ( - counter, - ptr_val, - sender_dev, - transit_ptr, - pool_dev, - size, - ), - ) - .map(|_| ()) - } - WritePath::GpuToGpuPoolDtoD => bound - .call_method1( - "_cuda_memcpy_gpu_buf", - (counter, ptr_val, size), - ) - .map(|_| ()), - _ => { - // GpuToShmem: ipc_present ≠ 1, - // copy to shared-memory data region. - bound - .call_method1( - "_cuda_memcpy", - ( - shmem_ptr as u64 + data_offset as u64, - ptr_val, - size, - 2u32, - ), - ) - .map(|_| ()) - } - }; - if let Err(e) = res { - copy_ok = false; - tracing::error!( - "[{}] write_memory_pool: GPU pool copy failed: {}", - self.node_id, - e - ); - } - } else { - copy_ok = false; - } - if copy_ok { - // Publish: gen was odd (in-progress), flip to even. - unsafe { - seqlock_end(gen_ptr, pre_write_gen, true); - } - } else { - // GPU in-place write to the live IPC buffer: a - // partial copy cannot be rolled back to a clean - // previous frame, so leave gen odd. The reader - // retries until the next successful write. - } - if !copy_ok { - if let Some(slot_data) = store_back.take() { - PINNED_POOL - .lock() - .unwrap_or_else(|e| e.into_inner()) - .insert(counter, slot_data); - } - return Err(eyre::eyre!( - "[{}] write_memory_pool: GPU pool copy failed", - self.node_id - )); - } - } else { - // Seqlock: begin (noop if still odd from a - // previous failed GPU-to-shmem copy). - let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; - let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; - unsafe { - std::ptr::copy_nonoverlapping( - ptr_val as *const u8, - shmem_ptr.add(data_offset), - size, - ); - } - // Publish: gen was odd (in-progress), flip to even. - // The shmem data region is not shared with the reader - // through IPC, so a copy failure means segfault — the - // process is dead before reaching here. No rollback needed. - unsafe { - seqlock_end(gen_ptr, pre_write_gen, true); - } - } - - // Store back to PINNED_POOL to keep shmem alive - if let Some(slot_data) = store_back { - PINNED_POOL - .lock() - .unwrap_or_else(|e| e.into_inner()) - .insert(counter, slot_data); - } - - return Ok(()); - } - } - } - } - - // Slow path: query daemon for pool metadata - match self - .node - .get_mut() - .read_pinned_memory(buffer_id.clone(), false) - { - Ok(metadata) => { - let shmem_name = metadata.parameters.get("shared_memory_name").and_then(|p| { - if let Parameter::String(s) = p { - Some(s.clone()) - } else { - None - } - }); - - if let Some(ref name) = shmem_name - && let Ok(shmem) = ShmemConf::new().os_id(name).open() - { - // Mirror fast-path guard: reject segments smaller - // than the header before any pointer arithmetic. - if shmem.len() < DORADMA_HEADER_SIZE { - return Ok(()); - } - let shmem_ptr = shmem.as_ptr(); - - let magic = unsafe { std::slice::from_raw_parts(shmem_ptr, 8) }; - if magic == DORADMA_MAGIC { - let data_offset = unsafe { read_header_u64(shmem_ptr.add(16)) as usize }; - - // Check if this pool has GPU DMA path enabled - let ipc_present = - unsafe { std::ptr::read(shmem_ptr.add(24) as *const u64) }; - - // Validate write size against pool capacity. The - // `size == 0` guard always applies; the shmem-region - // capacity check is skipped for GPU-buffer pools - // (ipc_present == 1), which the GPU helpers validate. - let shmem_len = shmem.len(); - if size == 0 - || (ipc_present != 1 && size > shmem_len.saturating_sub(data_offset)) - { - tracing::warn!( - "[{}] write_memory_pool (slow path): size {} exceeds available pool capacity (data_offset={}, total={}), operation aborted", - self.node_id, - size, - data_offset, - shmem_len - ); - return Ok(()); - } - - if ipc_present == 1 && !is_cuda { - // Extract counter for the DMA slot from buffer_id. - let slow_counter = buffer_id - .rsplit_once('_') - .and_then(|(_, c)| c.parse::().ok()); - let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; - let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; - let mut copy_ok = true; - if let (Ok(helpers), Some(c)) = (get_cuda_helpers(py), slow_counter) { - let bound = helpers.bind(py); - let slow_no_dma = !auto_pin; - if let Err(e) = - bound.call_method1("dma_copy", (ptr_val, size, c, slow_no_dma)) - { - copy_ok = false; - tracing::error!( - "[{}] write_memory_pool (slow path): DMA copy failed: {}", - self.node_id, - e - ); - } - } else { - copy_ok = false; - } - if copy_ok { - // Publish: gen was odd (in-progress), flip to even. - unsafe { - seqlock_end(gen_ptr, pre_write_gen, true); - } - } else { - // GPU in-place write to the live IPC buffer: a - // partial copy cannot be rolled back to a clean - // previous frame, so leave gen odd. The reader - // retries until the next successful write. - } - if !copy_ok { - return Err(eyre::eyre!( - "[{}] write_memory_pool (slow path): DMA copy failed", - self.node_id - )); - } - } else if is_cuda { - let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; - let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; - let mut copy_ok = true; - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - // Slow path transit look-up: PINNED_POOL - // (contrast fast path which uses store_back). - let (transit_ptr, pool_device) = if let Some((_, counter_str)) = - buffer_id.rsplit_once('_') - && let Ok(c) = counter_str.parse::() - { - PINNED_POOL - .lock() - .unwrap_or_else(|e| e.into_inner()) - .get(&c) - .map(|s| (s.transit_ptr, s.pool_device)) - .unwrap_or((0, 0)) - } else { - (0, 0) - }; - let write_path = classify_write_path( - ipc_present, - /*is_cuda=*/ true, - transit_ptr, - ); - let res = match write_path { - WritePath::GpuToGpuPoolTransit => { - let sender_dev = tensor_device - .strip_prefix("cuda") - .and_then(|d| d.strip_prefix(':')) - .and_then(|d| d.parse::().ok()) - .unwrap_or(0); - bound - .call_method1( - "_transit_copy_gpu_buf", - ( - buffer_id - .rsplit_once('_') - .and_then(|(_, cs)| cs.parse::().ok()) - .unwrap_or(0), - ptr_val, - sender_dev, - transit_ptr, - pool_device, - size, - ), - ) - .map(|_| ()) - } - WritePath::GpuToGpuPoolDtoD => bound - .call_method1( - "_cuda_memcpy_gpu_buf", - ( - buffer_id - .rsplit_once('_') - .and_then(|(_, cs)| cs.parse::().ok()) - .unwrap_or(0), - ptr_val, - size, - ), - ) - .map(|_| ()), - _ => { - // GpuToShmem: copy to shmem data region. - bound - .call_method1( - "_cuda_memcpy", - ( - shmem_ptr as u64 + data_offset as u64, - ptr_val, - size, - 2u32, - ), - ) - .map(|_| ()) - } - }; - if let Err(e) = res { - copy_ok = false; - tracing::error!( - "[{}] write_memory_pool (slow path): GPU pool copy failed: {}", - self.node_id, - e - ); - } - } else { - copy_ok = false; - } - if copy_ok { - // Publish: gen was odd (in-progress), flip to even. - unsafe { - seqlock_end(gen_ptr, pre_write_gen, true); - } - } else { - // GPU in-place write to the live IPC buffer: a - // partial copy cannot be rolled back to a clean - // previous frame, so leave gen odd. The reader - // retries until the next successful write. - } - if !copy_ok { - return Err(eyre::eyre!( - "[{}] write_memory_pool (slow path): GPU pool copy failed", - self.node_id - )); - } - } else { - // Seqlock: begin (noop if still odd from a - // previous failed GPU-to-shmem copy). - let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; - let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; - unsafe { - std::ptr::copy_nonoverlapping( - ptr_val as *const u8, - shmem_ptr.add(data_offset), - size, - ); - } - // Publish: gen was odd (in-progress), flip to even. - // The shmem data region is not shared through IPC — - // copy_nonoverlapping failure means segfault, - // so the process is dead before reaching here. - unsafe { - seqlock_end(gen_ptr, pre_write_gen, true); - } - } - } - } - } - Err(_) => { - warn_missing_memory_pool(&self.node_id, "write", &buffer_id); - } - } - - Ok(()) - } - - /// Read tensor info from an existing memory pool (zero-copy). - /// - /// Returns a `tensor_info` dict compatible with `tensor_from_info`. - /// The returned tensor shares the underlying shared-memory mapping — - /// no copy is made, so data bytes reflect whatever the writer has most - /// recently stored. - /// - /// # Concurrency / safety - /// - /// **The returned tensor is a zero-copy view into shared memory.** - /// Its data bytes can be overwritten at any time by a concurrent (or - /// subsequent) [`write_memory_pool`] on the sender. The seqlock - /// re-check at end-of-read detects whether an overwrite occurred - /// mid-consumption, but it is the **caller's responsibility** to - /// ensure the tensor is not used after the writer is allowed to write - /// again. - /// - /// Correct consumers follow a **turn-based discipline**: read the - /// pool, consume the tensor fully, then signal the sender (e.g. via - /// the dataflow graph's `next_require` round-trip) that it is safe to - /// write the next frame. The bundled `examples/memory-pool/` dataflows - /// demonstrate this pattern. - pub fn read_memory_pool( - &self, - memory_pool_id: Py, - py: Python, - ) -> eyre::Result> { - let buffer_id = parse_memory_pool_id(memory_pool_id, py)?; - - // Populate the trusted GPU buffer size from daemon metadata - // on the first read. Subsequent reads (if any) reuse the - // cached entry — the daemon query runs at most once per pool. - { - let trusted = GPU_BUF_SIZES.lock().unwrap_or_else(|e| e.into_inner()); - if !trusted.contains_key(&buffer_id) { - drop(trusted); - if let Ok(metadata) = self - .node - .get_mut() - .read_pinned_memory(buffer_id.clone(), false) - && let Some(size) = metadata.parameters.get("size").and_then(|p| { - if let Parameter::Integer(v) = p { - Some(*v) - } else { - None - } - }) - { - GPU_BUF_SIZES - .lock() - .unwrap_or_else(|e| e.into_inner()) - .insert(buffer_id.clone(), size as u64); - } - } - } - - // Fast path: DORADMA header read with daemon-trusted size validation. - if buffer_id.starts_with("pool_") { - // Retry on transient failures (odd seqlock, shmem not yet - // mapped) so a concurrent writer doesn't cause a hard error. - // Time-bounded: a GPU copy (cudaMemcpy + synchronize) takes - // milliseconds, so we wait up to 500ms total with 1ms sleeps - // between attempts. - let deadline = std::time::Instant::now() - .checked_add(std::time::Duration::from_millis(500)) - .unwrap_or(std::time::Instant::now()); - loop { - match self.try_doradma_read(&buffer_id, py) { - Ok(Some(result)) => return Ok(result), - Ok(None) if std::time::Instant::now() < deadline => { - // Transient — yield the GIL and sleep so the - // writer can complete its copy+sync. - py.detach(|| { - std::thread::sleep(std::time::Duration::from_millis(1)); - }); - continue; - } - Ok(None) => break, - Err(e) => { - warn_missing_memory_pool(&self.node_id, "read", &buffer_id); - eyre::bail!("memory pool {}: fast path failed: {}", buffer_id, e); - } - } - } - // Retries exhausted — fall back to the daemon for CPU pools. - if let Ok(metadata) = self - .node - .get_mut() - .read_pinned_memory(buffer_id.clone(), false) - { - let size = metadata - .parameters - .get("size") - .and_then(|p| { - if let Parameter::Integer(v) = p { - Some(*v) - } else { - None - } - }) - .unwrap_or(0); - let dtype = metadata - .parameters - .get("dtype") - .and_then(|p| { - if let Parameter::String(s) = p { - Some(s.clone()) - } else { - None - } - }) - .unwrap_or_default(); - let shape = metadata - .parameters - .get("shape") - .and_then(|p| { - if let Parameter::ListInt(v) = p { - Some(v.clone()) - } else { - None - } - }) - .unwrap_or_default(); - // Only for CPU pools — GPU pools need IPC import. - let ipc_present = metadata - .parameters - .get("ipc_present") - .and_then(|p| { - if let Parameter::Bool(v) = p { - Some(*v) - } else { - None - } - }) - .unwrap_or(false); - if ipc_present { - warn_missing_memory_pool(&self.node_id, "read", &buffer_id); - eyre::bail!( - "memory pool {}: fast path retries exhausted for GPU pool \ - (daemon fallback cannot provide a GPU pointer)", - buffer_id - ); - } - let shmem_name = metadata.parameters.get("shared_memory_name").and_then(|p| { - if let Parameter::String(s) = p { - Some(s.clone()) - } else { - None - } - }); - if let Some(ref name) = shmem_name - && let Ok(shmem) = ShmemConf::new().os_id(name).open() - && shmem.len() >= DORADMA_HEADER_SIZE - { - let shmem_ptr = shmem.as_ptr(); - let magic = unsafe { std::slice::from_raw_parts(shmem_ptr, 8) }; - if magic == DORADMA_MAGIC { - let data_offset = unsafe { read_header_u64(shmem_ptr.add(16)) as usize }; - // Mirror fast-path bounds check. - if data_offset > shmem.len() - || (size as usize) > shmem.len().saturating_sub(data_offset) - { - warn_missing_memory_pool(&self.node_id, "read", &buffer_id); - eyre::bail!( - "memory pool {}: header bounds exceeded: \ - data_offset {} + size {} > shmem_len {}", - buffer_id, - data_offset, - size, - shmem.len() - ); - } - // Seqlock: reject a torn mid-write frame. The - // fallback is reached when the fast path retries - // are exhausted, typically because the generation - // is stuck odd (crashed writer). - let read_gen = - unsafe { std::ptr::read_volatile(shmem_ptr.add(96) as *const u64) }; - if read_gen % 2 != 0 { - warn_missing_memory_pool(&self.node_id, "read", &buffer_id); - eyre::bail!( - "memory pool {}: daemon fallback: seqlock write in progress \ - (generation={}, odd)", - buffer_id, - read_gen - ); - } - // Cache-hit: use the stored mapping's base so - // the pointer stays valid after the fresh shmem - // is dropped. Cache-miss: insert the fresh - // mapping and drop the old one (if any). - let read_ptr; - { - let mut cpu_cache = - RECV_CPU_SHMEM.lock().unwrap_or_else(|e| e.into_inner()); - if let Some(cached) = cpu_cache.get(&buffer_id) { - read_ptr = (cached.base + data_offset as u64) as i64; - } else { - let base = shmem_ptr as u64; - read_ptr = (base + data_offset as u64) as i64; - cpu_cache.entry(buffer_id.clone()).or_insert(RecvCpuSlot { - _shmem: shmem, - base, - }); - } - } - let dict = PyDict::new(py); - dict.set_item("ptr", read_ptr)?; - dict.set_item("size", size)?; - dict.set_item("dtype", dtype)?; - dict.set_item("shape", shape)?; - dict.set_item("device", "cpu")?; - return Ok(dict.into()); - } - } - } - warn_missing_memory_pool(&self.node_id, "read", &buffer_id); - eyre::bail!( - "memory pool {}: fast path retries exhausted — pool not ready after 500ms", - buffer_id - ); - } - - warn_missing_memory_pool(&self.node_id, "read", &buffer_id); - eyre::bail!("memory pool {} not found", buffer_id); - } - - /// Free a memory pool. - #[pyo3(signature = (memory_pool_id))] - pub fn free_memory_pool(&self, memory_pool_id: Py, py: Python) -> eyre::Result<()> { - let buffer_id = parse_memory_pool_id(memory_pool_id, py)?; - - match self.node.get_mut().free_pinned_memory(buffer_id.clone()) { - Ok(_) => {} - Err(_) => { - warn_missing_memory_pool(&self.node_id, "release", &buffer_id); - } - } - - // Clean up sender-side pinned pool mapping (Shmem + CUDA host register). - // Without this, each register->write->free cycle leaks one Shmem mapping - // and one cudaHostRegister pinned region for the process lifetime. - // PINNED_POOL is sender-side (per-process), so bare counter is sufficient. - { - let counter = buffer_id - .strip_prefix("pool_") - .and_then(|s| s.rsplit_once('_').map(|(_, c)| c)) - .and_then(|c| c.parse::().ok()); - if let Some(c) = counter - && let Some(slot) = PINNED_POOL - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(&c) - { - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - let _ = bound.call_method1("_unregister_host", (slot.base,)); - let _ = bound.call_method1("_free_gpu_buf", (c,)); - if slot.transit_ptr != 0 { - let _ = bound.call_method1("_free_transit", (slot.transit_ptr,)); - } - } - // Remove transit metadata regardless of whether CUDA helpers - // are available — a missing _free_transit is a leak, but a stale - // TRANSIT_META entry is a correctness bug on re-registration. - TRANSIT_META - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(&c); - } - // PoolSlot dropped here -> Shmem unmapped - } - - // Clean up receiver-side caches so the shmem mappings are released. - // Keyed by full buffer_id (namespaced) to correctly handle - // multiple sender nodes with the same per-process counter. - { - // Close GPU IPC handle before dropping the cache entry. - if let Some(slot) = RECV_GPU_VA - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(&buffer_id) - { - if slot.gpu_buf != 0 { - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - let _ = bound.call_method1("_ipc_close", (slot.gpu_buf,)); - } - } else if slot.gpu_va != 0 { - // Host-registered mapping (effective_as_cuda branch): - // must cudaHostUnregister before munmap. _unregister_host - // requires the original host pointer (shmem base), not the - // device VA returned by cudaHostGetDevicePointer. - if let Ok(helpers) = get_cuda_helpers(py) { - let bound = helpers.bind(py); - let _ = bound.call_method1("_unregister_host", (slot.host_base,)); - } - } - // slot._shmem drops here -> munmap - } - } - RECV_CPU_SHMEM - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(&buffer_id); - GPU_BUF_SIZES - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(&buffer_id); - - { - let mut freed = FREED_POOL_IDS.lock().unwrap_or_else(|e| e.into_inner()); - tracing::debug!( - "[{}] free_memory_pool: adding {} to FREED_POOL_IDS (set size={})", - self.node_id, - buffer_id, - freed.len() - ); - freed.insert(buffer_id); - } - - Ok(()) - } } /// Stub for `send_output_raw` on Python < 3.11. @@ -3467,71 +976,6 @@ impl Node { } } -#[cfg(test)] -mod memory_pool_free_drain_tests { - /// Every user-visible receive path must drain the pending - /// memory-pool free set — that drain is the only thing that releases - /// this process's GPU IPC handles, transit buffers and shmem mappings - /// for pools that *another* node freed - /// (see `dora_node_api::event_stream::memory_pool`). - /// - /// Regression guard for #2958, where `recv_async` was the one path of - /// four that skipped it, so `await`-only nodes leaked every pool until - /// process exit. The leak is silent: it surfaces later as a CUDA OOM - /// somewhere unrelated, which is exactly why it needs a guard. - /// - /// Asserted on the source rather than on behaviour because `Node` can - /// only be constructed against a live daemon connection, so no unit - /// test can call these methods. - #[test] - fn every_receive_path_drains_pending_memory_pool_frees() { - const DRAIN: &str = "self.process_pending_memory_pool_frees(py)"; - let src = include_str!("lib.rs"); - - for signature in [ - "pub fn next(&self, py: Python", - "pub fn drain(&self, py: Python", - "pub fn try_recv(&mut self, py: Python", - "pub async fn recv_async(&self,", - ] { - assert!( - method_body(src, signature).contains(DRAIN), - "`{signature}` does not call `{DRAIN}`: pools freed by other \ - nodes stay mapped in this process forever (#2958)" - ); - } - - // `__next__` is exempt only for as long as it delegates to `next`. - assert!( - method_body(src, "pub fn __next__(&self, py: Python").contains("self.next(py,"), - "`__next__` no longer delegates to `next`, so it needs its own \ - `{DRAIN}` call (#2958)" - ); - } - - /// Returns the `{ .. }` block that follows `signature` in `src`. - fn method_body<'a>(src: &'a str, signature: &str) -> &'a str { - let start = src - .find(signature) - .unwrap_or_else(|| panic!("method `{signature}` not found — update this test")); - let open = start + src[start..].find('{').expect("method has a body"); - let mut depth = 0usize; - for (offset, c) in src[open..].char_indices() { - match c { - '{' => depth += 1, - '}' => { - depth -= 1; - if depth == 0 { - return &src[open..=open + offset]; - } - } - _ => {} - } - } - panic!("unbalanced braces in the body of `{signature}`"); - } -} - fn err_to_pyany(err: eyre::Report, gil: Python<'_>) -> Py { PyErr::from(err) .into_pyobject(gil) @@ -3648,346 +1092,32 @@ impl Node { self.node_id.to_string() } - /// DORADMA fast path for read_memory_pool: reads metadata directly from - /// the shmem header, bypassing the daemon for zero-copy metadata retrieval. - /// - /// Buffer ID format: `"pool_{node_id}_{counter}"` → - /// shmem name: `"dora_pool_{dataflow_id}_{node_id}_{counter}"`. - /// - /// # Synchronization model - /// - /// The seqlock (write_gen at header offset 96) guards **data-byte** - /// consistency across `write_memory_pool` overwrites — the end-of-read - /// generation re-check detects if a write occurred mid-read. Header - /// fields (json_len, data_offset) are written once at registration and - /// never change, so they are not subject to torn-read risk. - /// - /// The seqlock does NOT protect the tensor data bytes from being - /// overwritten while a consumer is iterating the zero-copy tensor. - /// Callers must enforce a **turn-based** discipline: the writer must - /// not begin a new `write_memory_pool` until the receiver has finished - /// consuming the previous tensor. The example dataflow enforces this - /// via `next_require` round-trip signaling. - /// - /// Returns `Ok(Some(tensor_info_dict))` on success, `Ok(None)` to fall back to daemon. - fn try_doradma_read(&self, buffer_id: &str, py: Python<'_>) -> eyre::Result>> { - // Format: "pool_{node_id}_{counter}". - // Use rsplit to extract the counter from the end — the node_id - // portion may itself contain underscores (legal in dora node ids). - let counter: u64 = match buffer_id.rsplit_once('_') { - Some((_, c)) => match c.parse() { - Ok(c) => c, - Err(_) => return Ok(None), - }, - None => return Ok(None), - }; - let pool_node_id = buffer_id - .strip_prefix("pool_") - .and_then(|s| s.strip_suffix(&format!("_{counter}"))) - .unwrap_or(""); - - // Check freed tracking -> if this buffer was freed, fall back to daemon - { - let freed = FREED_POOL_IDS.lock().unwrap_or_else(|e| e.into_inner()); - if freed.contains(buffer_id) { - tracing::debug!( - "[{}] try_doradma_read: buffer {} is freed, fallback to daemon", - self.node_id, - buffer_id - ); - return Ok(None); - } - } - - let shmem_name = format!( - "dora_pool_{}_{}_{}", - self.dataflow_id, pool_node_id, counter - ); - - // Open shared memory - let shmem = match ShmemConf::new().os_id(&shmem_name).open() { - Ok(s) => s, - Err(_) => return Ok(None), - }; - - let shmem_ptr = shmem.as_ptr(); - let shmem_size = shmem.len(); - - // Reject truncated segments: need at least DORADMA_HEADER_SIZE bytes for the header - if shmem_size < DORADMA_HEADER_SIZE { - tracing::warn!( - "[{}] try_doradma_read: shmem size {} < DORADMA_HEADER_SIZE {}, rejecting", - self.node_id, - shmem_size, - DORADMA_HEADER_SIZE, - ); - return Ok(None); - } - - // Verify DORADMA magic header - - unsafe { - let magic = std::slice::from_raw_parts(shmem_ptr, 8); - if magic != DORADMA_MAGIC { - return Ok(None); - } - } - - // Read header: [magic:8][json_len:8][data_offset:8][reserved:232] - let json_len = unsafe { read_header_u64(shmem_ptr.add(8)) as usize }; - let data_offset = unsafe { read_header_u64(shmem_ptr.add(16)) as usize }; - - // Validate JSON length fits within the segment - if json_len > shmem_size.saturating_sub(DORADMA_HEADER_SIZE) { - tracing::warn!( - "[{}] try_doradma_read: json_len {} exceeds shmem bounds (size={}, header={})", - self.node_id, - json_len, - shmem_size, - DORADMA_HEADER_SIZE, - ); - return Ok(None); - } - - // Read JSON metadata from header - let json_slice = - unsafe { std::slice::from_raw_parts(shmem_ptr.add(DORADMA_HEADER_SIZE), json_len) }; - let json_str = match std::str::from_utf8(json_slice) { - Ok(s) => s, - Err(_) => return Ok(None), - }; - - // Parse JSON to Python dict - let metadata_dict: Bound<'_, PyDict> = match py.import("json") { - Ok(m) => match m.call_method1("loads", (json_str,)) { - Ok(v) => match v.downcast_into::() { - Ok(d) => d, - Err(_) => return Ok(None), - }, - Err(_) => return Ok(None), - }, - Err(_) => return Ok(None), - }; - - let size: usize = match metadata_dict.get_item("size") { - Ok(Some(v)) => v.extract().unwrap_or(0), - _ => 0, + /// Borrow this node as the tensor-pool extension's context. + /// + /// Built per call rather than stored: PyO3 cannot let another crate add + /// `#[pymethods]` to `Node`, so the extension takes what it needs by + /// reference and the wrappers above stay thin. + #[cfg(feature = "tensor-pool")] + fn with_pool(&self, f: impl FnOnce(&mut dora_tensor_pool_python::Pool<'_>) -> R) -> R { + let node_id = self.node_id.clone(); + let dataflow_id = self.dataflow_id; + let mut guard = self.node.get_mut(); + let mut pool = dora_tensor_pool_python::Pool { + node_id, + dataflow_id, + node: &mut guard, }; - if size == 0 { - return Ok(None); - } - - // Check if this pool uses GPU DMA (IPC handle in header). - // Reading ipc_present early lets us skip the shmem data-region - // size check for GPU-buffer pools (receiver_is_cuda registrations - // allocate header-only shmem). - let ipc_present = unsafe { std::ptr::read(shmem_ptr.add(24) as *const u64) }; - - // Verify data_offset + size fits within shared memory segment. - // GPU-buffer reads (ipc_present == 1) don't access the shmem data - // region, so the size check is only required for CPU-receiver paths. - if ipc_present != 1 { - // Use saturating operations to guard against corrupted/hostile - // headers with a near-usize::MAX data_offset (overflow-safe, - // matching the write-path checks). - if data_offset > shmem_size || size > shmem_size.saturating_sub(data_offset) { - tracing::warn!( - "[{}] try_doradma_read: data_offset {} + size {} exceeds shmem_size {}", - self.node_id, - data_offset, - size, - shmem_size, - ); - return Ok(None); - } - } - - // Auto-detect read path from pinned_type - let pinned_type: Option = metadata_dict - .get_item("pinned_type") - .ok() - .flatten() - .and_then(|v| v.extract::().ok()); - - let effective_as_cuda = ipc_present == 1 || pinned_type.as_deref() != Some("cpu"); - - // Seqlock: read generation, validate stable after data access - let read_gen = unsafe { std::ptr::read_volatile(shmem_ptr.add(96) as *const u64) }; - if read_gen % 2 != 0 { - // Writer is in progress — retry (caller falls back to daemon) - return Ok(None); - } - std::sync::atomic::fence(std::sync::atomic::Ordering::Acquire); - - let read_ptr: u64; - - if ipc_present == 1 { - // GPU DMA pool: import IPC handle once, cache GPU buffer ptr - read_ptr = { - let cache = RECV_GPU_VA.lock().unwrap_or_else(|e| e.into_inner()); - match cache.get(buffer_id) { - Some(slot_data) if slot_data.gpu_buf != 0 => { - // Validate size against the GPU buffer's registered - // capacity. GPU_BUF_SIZES (populated from daemon - // metadata) is authoritative; gpu_buf_size (populated - // from shmem at first import) is the baseline. - let trusted_sizes = { - let trusted = GPU_BUF_SIZES.lock().unwrap_or_else(|e| e.into_inner()); - trusted.get(buffer_id).copied() - }; - if check_capacity_gpu_pool( - trusted_sizes, - Some(slot_data.gpu_buf_size), - size as u64, - ) == CapacityCheck::ExceedsTrustedSize - { - return Ok(None); - } - slot_data.gpu_buf - } - _ => { - // First IPC import: validate size against - // daemon-trusted capacity. Fail closed — - // NoTrustedEntry rejects the import rather than - // trusting the world-writable shmem size. - let trusted_sizes = { - let trusted = GPU_BUF_SIZES.lock().unwrap_or_else(|e| e.into_inner()); - trusted.get(buffer_id).copied() - }; - match check_capacity_gpu_pool(trusted_sizes, None, size as u64) { - CapacityCheck::Ok => {} - CapacityCheck::ExceedsTrustedSize | CapacityCheck::NoTrustedEntry => { - return Ok(None); - } - } - drop(cache); - let handle_bytes = - unsafe { std::slice::from_raw_parts(shmem_ptr.add(32), 64) }; - let helpers = get_cuda_helpers(py) - .map_err(|e| eyre::eyre!("get_cuda_helpers: {}", e))?; - let bound = helpers.bind(py); - let handle_py = PyBytes::new(py, handle_bytes); - let gpu_ptr: u64 = bound - .call_method1("_ipc_import", (handle_py,)) - .map_err(|e| eyre::eyre!("_ipc_import: {}", e))? - .extract() - .map_err(|e| eyre::eyre!("extract gpu_ptr: {}", e))?; - let mut cache = RECV_GPU_VA.lock().unwrap_or_else(|e| e.into_inner()); - cache.insert( - buffer_id.to_string(), - RecvGpuSlot { - _shmem: shmem, - gpu_va: 0, - gpu_buf: gpu_ptr, - host_base: shmem_ptr as u64, - // Baseline from shmem — daemon metadata - // already populated GPU_BUF_SIZES above. - gpu_buf_size: size as u64, - }, - ); - gpu_ptr - } - } - }; - } else if effective_as_cuda { - read_ptr = { - let cache = RECV_GPU_VA.lock().unwrap_or_else(|e| e.into_inner()); - match cache.get(buffer_id) { - Some(slot_data) => { - // GPU VA is stable across data overwrites; - // cache is keyed by full buffer_id (namespaced). - slot_data.gpu_va + data_offset as u64 - } - None => { - drop(cache); - let helpers = get_cuda_helpers(py) - .map_err(|e| eyre::eyre!("get_cuda_helpers: {}", e))?; - let bound = helpers.bind(py); - bound - .call_method1("_register_host", (shmem_ptr as u64, shmem_size)) - .map_err(|e| eyre::eyre!("_register_host: {}", e))?; - let va: u64 = bound - .call_method1("_get_device_ptr", (shmem_ptr as u64,)) - .map_err(|e| eyre::eyre!("_get_device_ptr: {}", e))? - .extract() - .map_err(|e| eyre::eyre!("extract gpu_va: {}", e))?; - let mut cache = RECV_GPU_VA.lock().unwrap_or_else(|e| e.into_inner()); - cache.insert( - buffer_id.to_string(), - RecvGpuSlot { - _shmem: shmem, - gpu_va: va, - gpu_buf: 0, - host_base: shmem_ptr as u64, - gpu_buf_size: 0, // CPU memory, no GPU buffer - }, - ); - va + data_offset as u64 - } - } - }; - } else { - // On the first read the fresh mapping is cached; on subsequent - // reads the fresh mapping is dropped and the returned pointer - // must use the cached mapping's base (a different mmap address). - let mut cpu_cache = RECV_CPU_SHMEM.lock().unwrap_or_else(|e| e.into_inner()); - read_ptr = match cpu_cache.get(buffer_id) { - Some(cached) => cached.base + data_offset as u64, - None => { - let base = shmem_ptr as u64; - cpu_cache.insert( - buffer_id.to_string(), - RecvCpuSlot { - _shmem: shmem, - base, - }, - ); - base + data_offset as u64 - } - }; - } - - // Seqlock: re-read generation — mismatch means data changed during read - let read_gen2 = unsafe { std::ptr::read_volatile(shmem_ptr.add(96) as *const u64) }; - if read_gen2 != read_gen { - return Ok(None); - } - - // Build tensor_info dict - let dtype: String = metadata_dict - .get_item("dtype") - .ok() - .flatten() - .and_then(|v| v.extract::().ok()) - .unwrap_or_default(); - let shape: Vec = metadata_dict - .get_item("shape") - .ok() - .flatten() - .and_then(|v| v.extract::>().ok()) - .unwrap_or_default(); - let device = if effective_as_cuda { "cuda" } else { "cpu" }; - - let dict = PyDict::new(py); - dict.set_item("ptr", read_ptr as i64)?; - dict.set_item("size", size)?; - dict.set_item("dtype", dtype)?; - dict.set_item("shape", shape)?; - dict.set_item("device", device)?; - - // A successful fast-path read means the pool is alive. Clear - // any stale tombstone so that the fast path re-engages after a - // sender restart re-creates the same buffer_id (FREED_POOL_IDS - // is per-process and free_memory_pool inserts there). - { - FREED_POOL_IDS - .lock() - .unwrap_or_else(|e| e.into_inner()) - .remove(buffer_id); - } + f(&mut pool) + } - Ok(Some(dict.into())) + /// Release pools that other nodes freed. No-op without the extension. + #[cfg(feature = "tensor-pool")] + fn process_pending_tensor_pool_frees(&self, py: Python) { + self.with_pool(|pool| pool.process_pending_tensor_pool_frees(py)); } + + #[cfg(not(feature = "tensor-pool"))] + fn process_pending_tensor_pool_frees(&self, _py: Python) {} } /// Start a runtime for Operators diff --git a/apis/rust/node/src/daemon_connection/interactive.rs b/apis/rust/node/src/daemon_connection/interactive.rs index 3810f0d57d..be367eedd8 100644 --- a/apis/rust/node/src/daemon_connection/interactive.rs +++ b/apis/rust/node/src/daemon_connection/interactive.rs @@ -79,9 +79,6 @@ impl InteractiveEvents { println!("{}", "node reports EventStreamDropped".blue()); DaemonReply::Result(Ok(())) } - DaemonRequest::RegisterPinnedMemory { .. } - | DaemonRequest::ReadPinnedMemory { .. } - | DaemonRequest::FreePinnedMemory { .. } => DaemonReply::Result(Ok(())), DaemonRequest::NodeConfig { .. } => { eyre::bail!("unexpected NodeConfig in interactive mode") } diff --git a/apis/rust/node/src/daemon_connection/node_integration_testing.rs b/apis/rust/node/src/daemon_connection/node_integration_testing.rs index 1d95b0f358..3321149099 100644 --- a/apis/rust/node/src/daemon_connection/node_integration_testing.rs +++ b/apis/rust/node/src/daemon_connection/node_integration_testing.rs @@ -132,9 +132,6 @@ impl IntegrationTestingEvents { println!("{}", "node reports EventStreamDropped".blue()); DaemonReply::Result(Ok(())) } - DaemonRequest::RegisterPinnedMemory { .. } - | DaemonRequest::ReadPinnedMemory { .. } - | DaemonRequest::FreePinnedMemory { .. } => DaemonReply::Result(Ok(())), DaemonRequest::NodeConfig { .. } => { eyre::bail!("unexpected NodeConfig in interactive mode") } diff --git a/apis/rust/node/src/event_stream/extensions.rs b/apis/rust/node/src/event_stream/extensions.rs new file mode 100644 index 0000000000..46c80cadf8 --- /dev/null +++ b/apis/rust/node/src/event_stream/extensions.rs @@ -0,0 +1,101 @@ +//! Drop notifications for the daemon's opaque extension table. +//! +//! When any node drops an extension key, the daemon sends +//! [`NodeEvent::ExtensionDropped`](dora_message::daemon_to_node::NodeEvent::ExtensionDropped) +//! to every node that stored or loaded it. That event is out-of-band — it is +//! not a dataflow input, and surfacing it to user code would mean every node +//! author had to match on an event they did not ask for. So the event-stream +//! thread consumes it into this process-global queue, and the extension's own +//! code drains it whenever it next runs. +//! +//! Process-global rather than per-node because a language binding holds its +//! caches the same way: one process, one set of mappings, regardless of how +//! many `DoraNode`s live in it. + +use std::collections::VecDeque; +use std::sync::{LazyLock, Mutex}; + +/// Bound on the queue. A consumer that never drains must not grow it without +/// limit; dropping the oldest entry is safe because a missed notification only +/// means the extension releases that resource later (on its own `drop`) rather +/// than promptly. +const MAX_PENDING: usize = 4096; + +type Dropped = (String, String); + +static DROPPED: LazyLock>> = LazyLock::new(|| Mutex::new(VecDeque::new())); + +/// Record a dropped `(namespace, key)`. Called by the event-stream thread. +pub(crate) fn push_dropped(namespace: String, key: String) { + let mut queue = DROPPED.lock().unwrap_or_else(|e| e.into_inner()); + if queue.len() >= MAX_PENDING { + queue.pop_front(); + } + queue.push_back((namespace, key)); +} + +/// Take every pending drop notification for `namespace`, leaving the rest. +/// +/// Scoped by namespace so two extensions in one process cannot swallow each +/// other's notifications. +pub fn drain_dropped_keys(namespace: &str) -> Vec { + let mut queue = DROPPED.lock().unwrap_or_else(|e| e.into_inner()); + let mut taken = Vec::new(); + queue.retain(|(ns, key)| { + if ns == namespace { + taken.push(key.clone()); + false + } else { + true + } + }); + taken +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The queue is process-global by design, so these tests would clobber + /// each other under cargo's default thread-per-test. Serialize them and + /// start each from empty. + static TEST_LOCK: Mutex<()> = Mutex::new(()); + + fn guard() -> std::sync::MutexGuard<'static, ()> { + let g = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + DROPPED.lock().unwrap_or_else(|e| e.into_inner()).clear(); + g + } + + #[test] + fn drain_returns_only_the_requested_namespace() { + let _g = guard(); + push_dropped("pool".into(), "a".into()); + push_dropped("other".into(), "b".into()); + push_dropped("pool".into(), "c".into()); + + assert_eq!(drain_dropped_keys("pool"), vec!["a", "c"]); + // The other namespace's entry survived rather than being consumed. + assert_eq!(drain_dropped_keys("other"), vec!["b"]); + } + + #[test] + fn drain_is_exhaustive() { + let _g = guard(); + push_dropped("pool".into(), "a".into()); + assert_eq!(drain_dropped_keys("pool"), vec!["a"]); + assert!(drain_dropped_keys("pool").is_empty()); + } + + #[test] + fn queue_is_bounded_and_drops_oldest() { + let _g = guard(); + for i in 0..MAX_PENDING + 10 { + push_dropped("pool".into(), i.to_string()); + } + let drained = drain_dropped_keys("pool"); + assert_eq!(drained.len(), MAX_PENDING); + // The oldest ten were evicted, so the window starts at 10. + assert_eq!(drained.first().map(String::as_str), Some("10")); + } +} diff --git a/apis/rust/node/src/event_stream/memory_pool.rs b/apis/rust/node/src/event_stream/memory_pool.rs deleted file mode 100644 index 48d2864fdc..0000000000 --- a/apis/rust/node/src/event_stream/memory_pool.rs +++ /dev/null @@ -1,41 +0,0 @@ -//! Cross-process memory-pool cleanup coordination. -//! -//! When any node calls `free_memory_pool`, the daemon sends a targeted -//! `NodeEvent::FreeMemoryPool` to every node that registered or read the -//! pool. The event thread intercepts this event (without forwarding it to -//! user code) and pushes the pool id into a pending-cleanup set. The -//! language binding drains this set before yielding the next user-visible -//! event, so that GPU buffers, transit buffers, and shmem mappings are -//! released regardless of which node initiated the free. - -use std::collections::HashSet; -use std::sync::Mutex; - -/// Pending free notifications from the daemon. -/// -/// Written by the event-stream thread, drained by the language binding -/// (e.g. the Python node API) before returning the next user event. -/// A `HashSet` is used so that duplicate pushes (e.g. multiple event -/// streams in the same process) are naturally deduplicated. -static PENDING_FREES: std::sync::LazyLock>> = - std::sync::LazyLock::new(|| Mutex::new(HashSet::new())); - -/// Push a pool id onto the pending-cleanup set (called from the event thread). -pub(crate) fn push_freed_pool(shared_memory_id: String) { - PENDING_FREES - .lock() - .unwrap_or_else(|e| e.into_inner()) - .insert(shared_memory_id); -} - -/// Drain and return all pending pool ids for cleanup. -/// -/// The caller is responsible for releasing per-process resources -/// (GPU buffers, transit buffers, shmem mappings) for each returned id. -pub fn drain_freed_pools() -> Vec { - PENDING_FREES - .lock() - .unwrap_or_else(|e| e.into_inner()) - .drain() - .collect() -} diff --git a/apis/rust/node/src/event_stream/mod.rs b/apis/rust/node/src/event_stream/mod.rs index f75c5badb7..9f234117aa 100644 --- a/apis/rust/node/src/event_stream/mod.rs +++ b/apis/rust/node/src/event_stream/mod.rs @@ -38,10 +38,10 @@ pub use scheduler::Scheduler as EventScheduler; mod data_conversion; mod event; +/// Drop notifications for the daemon's opaque extension table. +pub mod extensions; /// Tracks input health (timeouts, liveness) for circuit-breaker recovery. pub mod input_tracker; -/// Cross-process memory-pool cleanup coordination via daemon broadcast. -pub mod memory_pool; /// Merged event streams combining internal and external event sources. pub mod merged; mod scheduler; diff --git a/apis/rust/node/src/event_stream/thread.rs b/apis/rust/node/src/event_stream/thread.rs index d9c6f085d8..4316d6ebee 100644 --- a/apis/rust/node/src/event_stream/thread.rs +++ b/apis/rust/node/src/event_stream/thread.rs @@ -144,11 +144,11 @@ fn event_stream_loop( close_tx = true; } - // FreeMemoryPool is an internal daemon→node notification — - // do not forward to user code. The language binding is - // expected to poll for pending frees before each event. - if let NodeEvent::FreeMemoryPool { shared_memory_id } = &inner { - crate::event_stream::memory_pool::push_freed_pool(shared_memory_id.clone()); + // Out-of-band: an extension's bookkeeping, not a dataflow input. + // Consume it so user code never has to match on an event it did + // not ask for; the extension drains the queue on its own schedule. + if let NodeEvent::ExtensionDropped { namespace, key } = &inner { + crate::event_stream::extensions::push_dropped(namespace.clone(), key.clone()); continue; } diff --git a/apis/rust/node/src/lib.rs b/apis/rust/node/src/lib.rs index b994e9819d..b783051611 100644 --- a/apis/rust/node/src/lib.rs +++ b/apis/rust/node/src/lib.rs @@ -123,7 +123,7 @@ use tokio::sync::oneshot; mod daemon_connection; mod error; -/// Asynchronous event stream, daemon communication, and memory-pool coordination. +/// Asynchronous event stream and daemon communication. pub mod event_stream; pub mod integration_testing; mod node; diff --git a/apis/rust/node/src/node/control_channel.rs b/apis/rust/node/src/node/control_channel.rs index 11476ea8ff..8b0ac1de44 100644 --- a/apis/rust/node/src/node/control_channel.rs +++ b/apis/rust/node/src/node/control_channel.rs @@ -138,14 +138,16 @@ impl ControlChannel { } } - pub fn register_pinned_memory( + pub fn extension_store( &mut self, - shared_memory_id: String, - metadata: Metadata, + namespace: String, + key: String, + value: Vec, ) -> eyre::Result<()> { - let request = DaemonRequest::RegisterPinnedMemory { - shared_memory_id, - metadata, + let request = DaemonRequest::ExtensionStore { + namespace, + key, + value, }; let reply = self .channel @@ -153,22 +155,24 @@ impl ControlChannel { inner: request, timestamp: self.clock.new_timestamp(), }) - .wrap_err("failed to send RegisterPinnedMemory request to dora-daemon")?; + .wrap_err("failed to send ExtensionStore request to dora-daemon")?; match reply { DaemonReply::Result(Ok(())) => Ok(()), DaemonReply::Result(Err(e)) => bail!("{e}"), - other => bail!("unexpected RegisterPinnedMemory reply: {other:?}"), + other => bail!("unexpected ExtensionStore reply: {other:?}"), } } - pub fn read_pinned_memory( + pub fn extension_load( &mut self, - shared_memory_id: String, - free: bool, - ) -> eyre::Result { - let request = DaemonRequest::ReadPinnedMemory { - shared_memory_id, - free, + namespace: String, + key: String, + remove: bool, + ) -> eyre::Result>> { + let request = DaemonRequest::ExtensionLoad { + namespace, + key, + remove, }; let reply = self .channel @@ -176,27 +180,27 @@ impl ControlChannel { inner: request, timestamp: self.clock.new_timestamp(), }) - .wrap_err("failed to send ReadPinnedMemory request to dora-daemon")?; + .wrap_err("failed to send ExtensionLoad request to dora-daemon")?; match reply { - DaemonReply::PinnedMemoryMetadata { metadata } => Ok(metadata), + DaemonReply::ExtensionValue { value } => Ok(value), DaemonReply::Result(Err(e)) => bail!("{e}"), - other => bail!("unexpected ReadPinnedMemory reply: {other:?}"), + other => bail!("unexpected ExtensionLoad reply: {other:?}"), } } - pub fn free_pinned_memory(&mut self, shared_memory_id: String) -> eyre::Result<()> { - let request = DaemonRequest::FreePinnedMemory { shared_memory_id }; + pub fn extension_drop(&mut self, namespace: String, key: String) -> eyre::Result<()> { + let request = DaemonRequest::ExtensionDrop { namespace, key }; let reply = self .channel .request(&Timestamped { inner: request, timestamp: self.clock.new_timestamp(), }) - .wrap_err("failed to send FreePinnedMemory request to dora-daemon")?; + .wrap_err("failed to send ExtensionDrop request to dora-daemon")?; match reply { DaemonReply::Result(Ok(())) => Ok(()), DaemonReply::Result(Err(e)) => bail!("{e}"), - other => bail!("unexpected FreePinnedMemory reply: {other:?}"), + other => bail!("unexpected ExtensionDrop reply: {other:?}"), } } } diff --git a/apis/rust/node/src/node/mod.rs b/apis/rust/node/src/node/mod.rs index 1b06bbed86..b68d9af8ad 100644 --- a/apis/rust/node/src/node/mod.rs +++ b/apis/rust/node/src/node/mod.rs @@ -2341,34 +2341,48 @@ impl DoraNode { } } - /// Register a pinned memory pool with the daemon for lifecycle tracking. + /// Store an opaque value in the daemon's dataflow-scoped extension table. /// - /// Send the memory pool metadata to the daemon so it can track the pool - /// and provide it to other nodes for zero-copy access. - pub fn register_pinned_memory( + /// This is the seam for transports that live outside the dora tree: dora + /// brokers the value's lifetime and nothing else — it never interprets + /// `namespace`, `key` or `value`. See `docs/extensions.md`. + /// + /// The daemon remembers which nodes touched a key so that dropping it + /// notifies them, and reclaims the entry when the dataflow ends or the + /// storing node exits. Drain the notifications with + /// [`event_stream::extensions::drain_dropped_keys`](crate::event_stream::extensions::drain_dropped_keys). + pub fn extension_store( &mut self, - shared_memory_id: String, - metadata: Metadata, + namespace: impl Into, + key: impl Into, + value: Vec, ) -> Result<(), eyre::Error> { self.control_channel - .register_pinned_memory(shared_memory_id, metadata) + .extension_store(namespace.into(), key.into(), value) } - /// Read pinned memory metadata from the daemon. + /// Read an opaque value back, optionally removing it in the same round trip. /// - /// When `free` is true, the daemon also frees the pool after reading. - pub fn read_pinned_memory( + /// Returns `None` if the key is not in the table — never stored, or + /// already dropped. + pub fn extension_load( &mut self, - shared_memory_id: String, - free: bool, - ) -> Result { + namespace: impl Into, + key: impl Into, + remove: bool, + ) -> Result>, eyre::Error> { self.control_channel - .read_pinned_memory(shared_memory_id, free) + .extension_load(namespace.into(), key.into(), remove) } - /// Free a pinned memory pool via the daemon. - pub fn free_pinned_memory(&mut self, shared_memory_id: String) -> Result<(), eyre::Error> { - self.control_channel.free_pinned_memory(shared_memory_id) + /// Drop an opaque value, notifying every node that stored or loaded it. + pub fn extension_drop( + &mut self, + namespace: impl Into, + key: impl Into, + ) -> Result<(), eyre::Error> { + self.control_channel + .extension_drop(namespace.into(), key.into()) } } diff --git a/binaries/daemon/Cargo.toml b/binaries/daemon/Cargo.toml index 252b0c2fcc..cab7e2cfb8 100644 --- a/binaries/daemon/Cargo.toml +++ b/binaries/daemon/Cargo.toml @@ -12,6 +12,10 @@ repository.workspace = true # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [features] +# Opt-in: reclaim orphaned tensor-pool shared-memory segments left by a +# crashed node. Outside the 1.0 guarantees — see +# libraries/extensions/tensor-pool/README.md. +tensor-pool = ["dep:dora-tensor-pool"] default = ["tracing", "telemetry"] tracing = ["dep:dora-tracing"] bench = [] @@ -20,6 +24,7 @@ bench = [] telemetry = ["dep:tracing-opentelemetry"] [dependencies] +dora-tensor-pool = { workspace = true, optional = true } eyre = { workspace = true } tokio = { workspace = true, features = ["full"] } tokio-stream = { version = "0.1.18", features = ["net"] } @@ -35,7 +40,6 @@ dora-tracing = { workspace = true, optional = true } dora-arrow-convert = { workspace = true } dora-node-api = { workspace = true } dora-message = { workspace = true } -dora-memory-pool = { workspace = true } serde_yaml = { workspace = true } uuid = { workspace = true } futures = { workspace = true } diff --git a/binaries/daemon/src/event_types.rs b/binaries/daemon/src/event_types.rs index f9a51305de..85afe78de1 100644 --- a/binaries/daemon/src/event_types.rs +++ b/binaries/daemon/src/event_types.rs @@ -147,18 +147,21 @@ pub enum DaemonNodeEvent { EventStreamDropped { reply_sender: oneshot::Sender, }, - RegisterPinnedMemory { - shared_memory_id: String, - metadata: metadata::Metadata, + ExtensionStore { + namespace: String, + key: String, + value: Vec, reply_sender: oneshot::Sender, }, - ReadPinnedMemory { - shared_memory_id: String, - free: bool, + ExtensionLoad { + namespace: String, + key: String, + remove: bool, reply_sender: oneshot::Sender, }, - FreePinnedMemory { - shared_memory_id: String, + ExtensionDrop { + namespace: String, + key: String, reply_sender: oneshot::Sender, }, } diff --git a/binaries/daemon/src/extension_table.rs b/binaries/daemon/src/extension_table.rs new file mode 100644 index 0000000000..54c0cf454a --- /dev/null +++ b/binaries/daemon/src/extension_table.rs @@ -0,0 +1,361 @@ +//! Dataflow-scoped store of opaque values for out-of-tree extensions. +//! +//! The daemon brokers lifetime and nothing else. It never interprets a +//! namespace, key or value — an extension (a transport living outside this +//! repo) uses the table to hand a descriptor from one node to another and to +//! learn when that descriptor is withdrawn. +//! +//! What the daemon contributes that a node cannot do for itself: +//! +//! - **Reclamation.** A node that crashes cannot clean up after itself. The +//! daemon drops its entries on exit and on dataflow finish, which is the +//! whole reason this is daemon-side rather than a plain dataflow message +//! (dora-rs/dora#2881 is what that failure mode looks like). +//! - **Notification.** Everyone who touched a key learns when it goes away, +//! so per-process resources keyed on it can be released promptly rather +//! than at process exit. + +use std::collections::{BTreeSet, HashMap}; + +use dora_message::id::NodeId; + +/// Identifies one entry. Dataflow-scoped so two dataflows on the same daemon +/// cannot see or clobber each other's keys. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct ExtensionKey { + pub dataflow_id: String, + pub namespace: String, + pub key: String, +} + +#[derive(Debug, Clone)] +struct Entry { + value: Vec, + /// The node that stored it. Its exit reclaims the entry. + owner: NodeId, + /// Everyone who stored or loaded it, and so must be told when it is + /// dropped. Includes the owner. + touched_by: BTreeSet, +} + +/// Per-daemon table. Not `Sync`-wrapped here — the daemon owns it behind +/// `&mut self`, like the rest of its state. +#[derive(Debug, Default)] +pub struct ExtensionTable { + entries: HashMap, +} + +/// Upper bound on entries per dataflow, so a looping node cannot exhaust the +/// daemon's memory with a store-per-frame. Chosen to sit well above any +/// plausible working set while still bounding a runaway. +pub const MAX_ENTRIES_PER_DATAFLOW: usize = 8192; + +impl ExtensionTable { + pub fn new() -> Self { + Self::default() + } + + /// Store or overwrite `key`. + /// + /// Overwriting is allowed only by the owner: a second node claiming an + /// existing key would otherwise be able to redirect every reader of it. + /// Returns `Err` on a foreign overwrite or when the dataflow is at cap. + pub fn store( + &mut self, + key: ExtensionKey, + value: Vec, + owner: &NodeId, + ) -> Result<(), String> { + match self.entries.get_mut(&key) { + Some(entry) => { + if &entry.owner != owner { + return Err(format!( + "extension key `{}/{}` is owned by node `{}`; node `{}` may not overwrite it", + key.namespace, key.key, entry.owner, owner + )); + } + entry.value = value; + Ok(()) + } + None => { + let count = self + .entries + .keys() + .filter(|k| k.dataflow_id == key.dataflow_id) + .count(); + if count >= MAX_ENTRIES_PER_DATAFLOW { + return Err(format!( + "extension table full for this dataflow ({MAX_ENTRIES_PER_DATAFLOW} entries); \ + node `{owner}` cannot store `{}/{}`", + key.namespace, key.key + )); + } + self.entries.insert( + key, + Entry { + value, + owner: owner.clone(), + touched_by: BTreeSet::from([owner.clone()]), + }, + ); + Ok(()) + } + } + } + + /// Read `key`, recording `reader` as someone to notify on drop. + /// + /// The reader is recorded even on a plain read, because the point of the + /// notification is to let it release whatever it derived from the value. + pub fn load(&mut self, key: &ExtensionKey, reader: &NodeId) -> Option> { + let entry = self.entries.get_mut(key)?; + entry.touched_by.insert(reader.clone()); + Some(entry.value.clone()) + } + + /// Remove `key`, returning everyone who must be told. + /// + /// `None` if it was not present, so a duplicate drop is distinguishable + /// from a real one and does not produce a second broadcast. + pub fn drop_key(&mut self, key: &ExtensionKey) -> Option> { + self.entries.remove(key).map(|entry| entry.touched_by) + } + + /// Drop everything owned by `node` in `dataflow_id`, for when it exits. + /// + /// Returns each dropped key with its notification set. Entries the node + /// merely *read* are left alone: they belong to someone still running. + pub fn reclaim_owner( + &mut self, + dataflow_id: &str, + node: &NodeId, + ) -> Vec<(ExtensionKey, BTreeSet)> { + let doomed: Vec = self + .entries + .iter() + .filter(|(k, e)| k.dataflow_id == dataflow_id && &e.owner == node) + .map(|(k, _)| k.clone()) + .collect(); + doomed + .into_iter() + .filter_map(|k| { + let touched = self.entries.remove(&k)?.touched_by; + Some((k, touched)) + }) + .collect() + } + + /// Drop every entry for a finished dataflow. No notifications: its nodes + /// are gone and its channels are closed by the time this runs. + pub fn reclaim_dataflow(&mut self, dataflow_id: &str) -> usize { + let before = self.entries.len(); + self.entries.retain(|k, _| k.dataflow_id != dataflow_id); + before - self.entries.len() + } + + #[cfg(test)] + pub fn len(&self) -> usize { + self.entries.len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn node(name: &str) -> NodeId { + name.to_string().into() + } + + fn key(ns: &str, k: &str) -> ExtensionKey { + ExtensionKey { + dataflow_id: "df".into(), + namespace: ns.into(), + key: k.into(), + } + } + + #[test] + fn store_then_load_round_trips() { + let mut t = ExtensionTable::new(); + t.store(key("pool", "a"), b"meta".to_vec(), &node("sender")) + .unwrap(); + assert_eq!( + t.load(&key("pool", "a"), &node("recv")), + Some(b"meta".to_vec()) + ); + } + + #[test] + fn load_of_absent_key_is_none() { + let mut t = ExtensionTable::new(); + assert_eq!(t.load(&key("pool", "nope"), &node("recv")), None); + } + + #[test] + fn namespaces_do_not_collide() { + let mut t = ExtensionTable::new(); + t.store(key("a", "k"), b"one".to_vec(), &node("n")).unwrap(); + t.store(key("b", "k"), b"two".to_vec(), &node("n")).unwrap(); + assert_eq!(t.load(&key("a", "k"), &node("n")), Some(b"one".to_vec())); + assert_eq!(t.load(&key("b", "k"), &node("n")), Some(b"two".to_vec())); + } + + #[test] + fn dataflows_do_not_collide() { + let mut t = ExtensionTable::new(); + let mut other = key("pool", "k"); + other.dataflow_id = "other-df".into(); + t.store(key("pool", "k"), b"mine".to_vec(), &node("n")) + .unwrap(); + t.store(other.clone(), b"theirs".to_vec(), &node("n")) + .unwrap(); + assert_eq!( + t.load(&key("pool", "k"), &node("n")), + Some(b"mine".to_vec()) + ); + assert_eq!(t.load(&other, &node("n")), Some(b"theirs".to_vec())); + } + + #[test] + fn a_foreign_node_cannot_overwrite() { + let mut t = ExtensionTable::new(); + t.store(key("pool", "a"), b"mine".to_vec(), &node("owner")) + .unwrap(); + let err = t + .store(key("pool", "a"), b"hijacked".to_vec(), &node("attacker")) + .unwrap_err(); + assert!(err.contains("owned by node `owner`"), "{err}"); + // The original value survived the attempt. + assert_eq!( + t.load(&key("pool", "a"), &node("owner")), + Some(b"mine".to_vec()) + ); + } + + #[test] + fn the_owner_may_overwrite() { + let mut t = ExtensionTable::new(); + t.store(key("pool", "a"), b"v1".to_vec(), &node("owner")) + .unwrap(); + t.store(key("pool", "a"), b"v2".to_vec(), &node("owner")) + .unwrap(); + assert_eq!( + t.load(&key("pool", "a"), &node("owner")), + Some(b"v2".to_vec()) + ); + } + + #[test] + fn drop_reports_every_node_that_touched_the_key() { + let mut t = ExtensionTable::new(); + t.store(key("pool", "a"), b"v".to_vec(), &node("owner")) + .unwrap(); + t.load(&key("pool", "a"), &node("reader1")); + t.load(&key("pool", "a"), &node("reader2")); + + let notified = t.drop_key(&key("pool", "a")).expect("was present"); + assert_eq!( + notified, + BTreeSet::from([node("owner"), node("reader1"), node("reader2")]) + ); + } + + #[test] + fn dropping_twice_notifies_once() { + let mut t = ExtensionTable::new(); + t.store(key("pool", "a"), b"v".to_vec(), &node("owner")) + .unwrap(); + assert!(t.drop_key(&key("pool", "a")).is_some()); + // The second drop is a no-op, not a second broadcast. + assert!(t.drop_key(&key("pool", "a")).is_none()); + } + + #[test] + fn owner_exit_reclaims_its_entries_and_notifies_readers() { + let mut t = ExtensionTable::new(); + t.store(key("pool", "a"), b"v".to_vec(), &node("owner")) + .unwrap(); + t.load(&key("pool", "a"), &node("reader")); + + let reclaimed = t.reclaim_owner("df", &node("owner")); + assert_eq!(reclaimed.len(), 1); + assert_eq!(reclaimed[0].0, key("pool", "a")); + assert!(reclaimed[0].1.contains(&node("reader"))); + assert_eq!(t.len(), 0); + } + + #[test] + fn a_readers_exit_does_not_reclaim_the_owners_entry() { + let mut t = ExtensionTable::new(); + t.store(key("pool", "a"), b"v".to_vec(), &node("owner")) + .unwrap(); + t.load(&key("pool", "a"), &node("reader")); + + // The reader leaving must not take a live sender's descriptor with it. + assert!(t.reclaim_owner("df", &node("reader")).is_empty()); + assert_eq!(t.len(), 1); + } + + #[test] + fn owner_exit_is_scoped_to_its_own_dataflow() { + let mut t = ExtensionTable::new(); + let mut other = key("pool", "a"); + other.dataflow_id = "other-df".into(); + t.store(key("pool", "a"), b"v".to_vec(), &node("owner")) + .unwrap(); + t.store(other, b"v".to_vec(), &node("owner")).unwrap(); + + assert_eq!(t.reclaim_owner("df", &node("owner")).len(), 1); + // Same node id in another dataflow keeps its entry. + assert_eq!(t.len(), 1); + } + + #[test] + fn dataflow_finish_reclaims_everything_it_owns() { + let mut t = ExtensionTable::new(); + let mut other = key("pool", "keep"); + other.dataflow_id = "other-df".into(); + t.store(key("pool", "a"), b"v".to_vec(), &node("n")) + .unwrap(); + t.store(key("pool", "b"), b"v".to_vec(), &node("m")) + .unwrap(); + t.store(other, b"v".to_vec(), &node("n")).unwrap(); + + assert_eq!(t.reclaim_dataflow("df"), 2); + assert_eq!(t.len(), 1); + } + + #[test] + fn the_per_dataflow_cap_is_enforced_and_scoped() { + let mut t = ExtensionTable::new(); + for i in 0..MAX_ENTRIES_PER_DATAFLOW { + t.store(key("pool", &i.to_string()), b"v".to_vec(), &node("n")) + .unwrap(); + } + let err = t + .store(key("pool", "one-too-many"), b"v".to_vec(), &node("n")) + .unwrap_err(); + assert!(err.contains("extension table full"), "{err}"); + + // A different dataflow still has its own budget. + let mut other = key("pool", "fresh"); + other.dataflow_id = "other-df".into(); + assert!(t.store(other, b"v".to_vec(), &node("n")).is_ok()); + } + + #[test] + fn overwriting_at_cap_still_works() { + let mut t = ExtensionTable::new(); + for i in 0..MAX_ENTRIES_PER_DATAFLOW { + t.store(key("pool", &i.to_string()), b"v".to_vec(), &node("n")) + .unwrap(); + } + // At cap, replacing an existing key adds nothing, so it must not be + // rejected — otherwise a steady-state writer wedges at the boundary. + assert!( + t.store(key("pool", "0"), b"v2".to_vec(), &node("n")) + .is_ok() + ); + } +} diff --git a/binaries/daemon/src/lib.rs b/binaries/daemon/src/lib.rs index e12732010a..af12374f76 100644 --- a/binaries/daemon/src/lib.rs +++ b/binaries/daemon/src/lib.rs @@ -15,7 +15,6 @@ use dora_core::{ }, uhlc::{self, HLC}, }; -use dora_memory_pool::{MemoryPoolId, MemoryPoolManager, MemoryPoolMetadata}; use dora_message::{ BuildId, DataflowId, SessionId, common::{ @@ -160,6 +159,7 @@ pub mod bench_support { mod coordinator; pub(crate) mod event_types; +mod extension_table; mod extract_err_from_stderr; pub(crate) mod fault_tolerance; mod local_listener; @@ -182,7 +182,11 @@ pub(crate) use running_dataflow::{ RunningNode, }; -use crate::{extract_err_from_stderr::extract_err_from_stderr, pending::DataflowStatus}; +use crate::{ + extension_table::{ExtensionKey, ExtensionTable}, + extract_err_from_stderr::extract_err_from_stderr, + pending::DataflowStatus, +}; const STDERR_LOG_LINES_MAX: usize = 500; const METRICS_INTERVAL: Duration = Duration::from_secs(2); @@ -243,6 +247,95 @@ fn deliver_param_update_strict( } } +/// Tell every node that touched `key` that it is gone. +/// +/// Returns the nodes that could not be reached. A full or closed channel is +/// reported rather than retried: the entry is already out of the table, so the +/// alternative to a warning is a silent resource leak in that node +/// (dora-rs/dora#2935 is what that looks like in practice). +fn notify_extension_dropped( + dataflow: &RunningDataflow, + namespace: &str, + key: &str, + touched_by: &BTreeSet, + clock: &HLC, +) -> Vec { + let mut undelivered = Vec::new(); + for node_id in touched_by { + let Some(channel) = dataflow.subscribe_channels.get(node_id) else { + // Not connected: it has either exited (nothing to release) or has + // not subscribed yet (it cannot hold the key either). + continue; + }; + let event = NodeEvent::ExtensionDropped { + namespace: namespace.to_owned(), + key: key.to_owned(), + }; + match send_with_timestamp(channel, event, clock) { + Ok(true) => dataflow.inc_pending(node_id), + Ok(false) | Err(_) => undelivered.push(node_id.clone()), + } + } + undelivered +} + +/// Drop `key` from the table and notify its readers, logging any that could +/// not be reached. Shared by the explicit drop request and by reclamation. +fn drop_extension_and_notify( + extensions: &mut ExtensionTable, + dataflow: Option<&RunningDataflow>, + key: &ExtensionKey, + clock: &HLC, +) -> bool { + let Some(touched_by) = extensions.drop_key(key) else { + return false; + }; + if let Some(dataflow) = dataflow { + let undelivered = + notify_extension_dropped(dataflow, &key.namespace, &key.key, &touched_by, clock); + if !undelivered.is_empty() { + tracing::warn!( + namespace = %key.namespace, + key = %key.key, + nodes = ?undelivered, + "extension drop notification undelivered; these nodes keep whatever \ + they derived from the value until they exit" + ); + } + } + true +} + +/// Drop every extension entry owned by an exited node, notifying readers. +fn reclaim_extensions_of_exited_node( + extensions: &mut ExtensionTable, + dataflow: Option<&RunningDataflow>, + dataflow_id: DataflowId, + node_id: &NodeId, + clock: &HLC, +) { + let reclaimed = extensions.reclaim_owner(&dataflow_id.to_string(), node_id); + for (key, touched_by) in reclaimed { + if let Some(dataflow) = dataflow { + let undelivered = + notify_extension_dropped(dataflow, &key.namespace, &key.key, &touched_by, clock); + if !undelivered.is_empty() { + tracing::warn!( + namespace = %key.namespace, + key = %key.key, + nodes = ?undelivered, + "extension reclaim notification undelivered after node `{node_id}` exited" + ); + } + } + tracing::debug!( + namespace = %key.namespace, + key = %key.key, + "reclaimed extension entry of exited node `{node_id}`" + ); + } +} + fn deliver_param_delete_strict( dataflow: &RunningDataflow, node_id: &NodeId, @@ -351,7 +444,9 @@ pub struct Daemon { pub(crate) builds: BTreeMap, pub(crate) git_manager: GitManager, pub(crate) metrics_system: Arc>, - pub(crate) memory_pool: MemoryPoolManager, + /// Opaque, dataflow-scoped store for out-of-tree extensions. See + /// `extension_table` and `docs/extensions.md`. + pub(crate) extensions: ExtensionTable, /// Nodes already warned about for sending after their dataflow /// finished, so `log_late_node_output` warns once each instead of /// once per message. See `MAX_WARNED_LATE_OUTPUT_NODES`. @@ -651,307 +746,6 @@ async fn collect_and_send_metrics_bg( Ok(()) } -/// Release the memory pools that nothing can reach now that `exited_node`'s -/// incarnation is gone. -/// -/// Which pools those are is decided by the rest of the dataflow: see -/// [`MemoryPoolManager::reclaim_unreachable`] for why an exited node's pools -/// cannot simply be dropped (dora-rs/dora#2881). -/// -/// Every path that ends an incarnation calls this: `RemoveNode`, -/// `ReplaceNode` (whose outgoing exit event the generation guard drops), and -/// `SpawnedNodeResult` for a crash, a clean exit or a restart. `exited_node` -/// is excluded from the live set explicitly, because the last of those runs -/// before the node leaves `running_nodes` — and because a replacement or -/// restart under the same id never touched its predecessor's pools. -/// -/// Takes its two fields separately rather than `&self` so callers can hold a -/// logger (which borrows `self.logger`) across it. -fn reclaim_memory_pools_after_exit( - memory_pool: &MemoryPoolManager, - running: &HashMap, - dataflow_id: Uuid, - exited_node: &NodeId, -) { - // No dataflow left means `finish_dataflow` already released its pools - // unconditionally. - let Some(dataflow) = running.get(&dataflow_id) else { - return; - }; - let released = memory_pool.reclaim_unreachable(&dataflow_id.to_string(), |node| { - dataflow - .running_nodes - .keys() - .any(|id| id != exited_node && id.as_ref() == node) - }); - if !released.is_empty() { - tracing::info!( - %dataflow_id, - node_id = %exited_node, - "released {} memory pool(s) that no live node can reach any more", - released.len(), - ); - } -} - -/// Widen the reader set of every pool the new edge `source` -> `target` -/// exposes, so reclamation cannot outrun a rewiring of a running dataflow. -/// -/// Whoever can reach a pool can forward its id along a new outgoing edge, so -/// `target` — and everything downstream of it — joins the pools `source` -/// could already reach. `source` registering the pool itself is the case -/// that bites: it is in its own reader set, so without this an exiting -/// `source` would take the pool with it while the freshly connected -/// consumer still holds the id (dora-rs/dora#2881). -/// -/// Must run *after* the edge is recorded, so the closure sees it. -fn extend_pool_readers_for_new_edge( - memory_pool: &MemoryPoolManager, - dataflow: &RunningDataflow, - source: &NodeId, - target: &NodeId, -) { - memory_pool.extend_potential_readers( - &dataflow.id.to_string(), - source.as_ref(), - &dataflow.downstream_closure(target), - ); -} - -/// Release everything this daemon still holds for a dataflow that just -/// finished here: its pool table entries, plus the `/dev/shm` segments of -/// `owned_nodes` that no table entry covers. -/// -/// A daemon that outlives the dataflow (`dora up`) would otherwise carry -/// every unfreed pool — table entry and segment alike — until it exits -/// (dora-rs/dora#2881). Releasing the table entries is unconditional, -/// deliberately: a dynamic node can outlive the finish (`should_finish` -/// ignores them), but the dataflow's routing, channels and listeners are gone -/// by now, so it can no longer send or receive anything a pool would serve, -/// and no later event would ever reach [`reclaim_memory_pools_after_exit`] -/// for this dataflow. -/// -/// This is a *per-daemon* finish and `/dev/shm` is host-wide, so the sweep is -/// restricted to this daemon's own nodes — see -/// [`MemoryPoolManager::cleanup_orphans`]. -/// -/// A free function rather than a method so it can be exercised without a -/// whole `Daemon`. -fn release_memory_pools_of_finished_dataflow( - memory_pool: &MemoryPoolManager, - dataflow_id: Uuid, - owned_nodes: &BTreeSet, -) { - let dataflow_id = dataflow_id.to_string(); - memory_pool.cleanup_dataflow(&dataflow_id); - MemoryPoolManager::cleanup_orphans(&dataflow_id, |node| { - owned_nodes.iter().any(|id| id.as_ref() == node) - }); -} - -/// Convert MemoryPoolMetadata into daemon-protocol MetadataParameters. -fn pool_metadata_to_params(meta: &MemoryPoolMetadata) -> MetadataParameters { - use dora_message::metadata::Parameter; - let mut p = MetadataParameters::new(); - p.insert("ptr".into(), Parameter::Integer(meta.ptr as i64)); - p.insert("size".into(), Parameter::Integer(meta.size as i64)); - p.insert("dtype".into(), Parameter::String(meta.dtype.clone())); - let shape: Vec = meta.shape.iter().map(|&x| x as i64).collect(); - p.insert("shape".into(), Parameter::ListInt(shape)); - p.insert("is_pinned".into(), Parameter::Bool(meta.is_pinned)); - if let Some(ref t) = meta.pinned_type { - p.insert("pinned_type".into(), Parameter::String(t.clone())); - } - if let Some(ref n) = meta.shared_memory_name { - p.insert("shared_memory_name".into(), Parameter::String(n.clone())); - } - if let Some(ipc) = meta.ipc_present { - p.insert("ipc_present".into(), Parameter::Bool(ipc)); - } - if let Some(ref b) = meta.buffer_id { - p.insert("buffer_id".into(), Parameter::String(b.clone())); - } - p -} - -/// Reconstruct MemoryPoolMetadata from MetadataParameters (best-effort). -fn pool_metadata_from_params(params: &MetadataParameters) -> MemoryPoolMetadata { - use dora_message::metadata::Parameter; - let get_int = |k: &str| -> Option { - params.get(k).and_then(|p| { - if let Parameter::Integer(v) = p { - Some(*v) - } else { - None - } - }) - }; - let get_str = |k: &str| -> Option { - params.get(k).and_then(|p| { - if let Parameter::String(s) = p { - Some(s.clone()) - } else { - None - } - }) - }; - let get_bool = |k: &str| -> Option { - params.get(k).and_then(|p| { - if let Parameter::Bool(v) = p { - Some(*v) - } else { - None - } - }) - }; - MemoryPoolMetadata { - ptr: get_int("ptr").unwrap_or(0) as u64, - size: get_int("size").unwrap_or(0) as usize, - dtype: get_str("dtype").unwrap_or_default(), - shape: params - .get("shape") - .and_then(|p| { - if let Parameter::ListInt(v) = p { - Some(v.iter().map(|&x| x as usize).collect()) - } else { - None - } - }) - .unwrap_or_default(), - is_pinned: get_bool("is_pinned").unwrap_or(false), - shared_memory_name: get_str("shared_memory_name"), - buffer_id: get_str("buffer_id"), - ipc_present: get_bool("ipc_present"), - pinned_type: get_str("pinned_type"), - } -} - -/// Release a memory pool and tell every other node that touched it to drop -/// its per-process view of that pool. -/// -/// Both daemon entry points that release a pool must go through here: the -/// explicit `FreePinnedMemory` request and `ReadPinnedMemory { free: true }`. -/// The read-with-free path used to free the pool without notifying anyone, -/// so every other node kept its GPU buffer, pinned host buffer, IPC handle -/// and shmem mapping for the rest of its life (#2935). -fn free_pool_and_notify( - memory_pool: &MemoryPoolManager, - dataflow: Option<&RunningDataflow>, - id: &MemoryPoolId, - initiator: &str, - clock: &HLC, -) -> Result<(), String> { - let (_metadata, touched) = memory_pool.free_memory_pool(id, initiator)?; - - // Without a running dataflow there are no subscribe channels to notify; - // the pool table entry and its shmem segment are released either way. - let Some(dataflow) = dataflow else { - return Ok(()); - }; - - let undelivered = notify_memory_pool_freed(dataflow, &id.id, &touched, initiator, clock); - if !undelivered.is_empty() { - tracing::error!( - pool = %id.id, - nodes = ?undelivered, - "FreeMemoryPool notification dropped (event channel full); these \ - nodes hold their GPU/host/IPC/shmem resources for the pool until \ - they exit" - ); - } - Ok(()) -} - -/// Send `NodeEvent::FreeMemoryPool` to every node that registered or read -/// the pool, except the node that initiated the free — it releases its own -/// resources synchronously. -/// -/// Returns the nodes whose notification could not be enqueued because their -/// event channel was full. Those nodes leak the pool's per-process -/// resources, so the caller surfaces them. Delivery is not retried: this -/// runs on the daemon main loop with the freeing node blocked on its reply, -/// so awaiting a backed-up receiver would stall every dataflow on the daemon. -fn notify_memory_pool_freed( - dataflow: &RunningDataflow, - shared_memory_id: &str, - touched: &HashSet, - initiator: &str, - clock: &HLC, -) -> BTreeSet { - let mut undelivered = BTreeSet::new(); - for (node, channel) in &dataflow.subscribe_channels { - if !touched.contains(node.as_ref()) || node.as_ref() == initiator { - continue; - } - let event = NodeEvent::FreeMemoryPool { - shared_memory_id: shared_memory_id.to_owned(), - }; - match send_with_timestamp(channel, event, clock) { - Ok(true) => { - // The listener decrements `pending_messages` for every event - // it drains, so an enqueue without the matching increment - // underflows the reported count to `u64::MAX` (#2827). - dataflow.inc_pending(node); - } - // A full channel is reported as `Ok(false)`, not `Err` — the - // original loop inspected only `Err` and dropped this case - // silently (#2935). - Ok(false) => { - undelivered.insert(node.clone()); - } - // Channel closed: the node is gone, so the OS already reclaimed - // the resources the notification would have released. - Err(_) => tracing::debug!( - node_id = %node, - pool = %shared_memory_id, - "skipping FreeMemoryPool notification: node disconnected" - ), - } - } - undelivered -} - -#[cfg(test)] -mod metadata_roundtrip_tests { - use super::*; - use dora_memory_pool::MemoryPoolMetadata; - - /// `ipc_present` must survive a to_params → from_params round-trip - /// so the read path receives the trusted flag from daemon metadata. - #[test] - fn ipc_present_survives_roundtrip_true() { - let meta = MemoryPoolMetadata { - ipc_present: Some(true), - ..Default::default() - }; - let params = pool_metadata_to_params(&meta); - let restored = pool_metadata_from_params(¶ms); - assert_eq!(restored.ipc_present, Some(true)); - } - - #[test] - fn ipc_present_survives_roundtrip_false() { - let meta = MemoryPoolMetadata { - ipc_present: Some(false), - ..Default::default() - }; - let params = pool_metadata_to_params(&meta); - let restored = pool_metadata_from_params(¶ms); - assert_eq!(restored.ipc_present, Some(false)); - } - - #[test] - fn ipc_present_survives_roundtrip_none() { - let meta = MemoryPoolMetadata { - ipc_present: None, - ..Default::default() - }; - let params = pool_metadata_to_params(&meta); - let restored = pool_metadata_from_params(¶ms); - assert_eq!(restored.ipc_present, None); - } -} - /// Where the daemon's zenoh listen address came from, which decides what a bind /// failure means. /// @@ -1796,7 +1590,7 @@ impl Daemon { zenoh_publish_tx, remote_daemon_events_tx, git_manager: Default::default(), - memory_pool: MemoryPoolManager::new(), + extensions: ExtensionTable::new(), builds, sessions: Default::default(), metrics_system: Arc::new(std::sync::Mutex::new(sysinfo::System::new())), @@ -2174,13 +1968,6 @@ impl Daemon { } } - // Clean up any unfreed memory pool entries on daemon exit - if let Err(errors) = self.memory_pool.cleanup_all() { - for error in errors { - tracing::warn!("{error}"); - } - } - // `run_inner` borrows `&mut self`, so move the accumulated results out // (the daemon may be reused for a reconnect, where these are ignored). Ok(std::mem::take(&mut self.dataflow_node_results)) @@ -2871,20 +2658,6 @@ impl Daemon { } } - // Same as `AddMapping`: each input the new node declares - // is a new edge, so the pools its source can reach must - // now count this node as a possible reader. - for input in inputs.values() { - if let InputMapping::User(mapping) = &input.mapping { - extend_pool_readers_for_new_edge( - &self.memory_pool, - dataflow, - &mapping.source, - &node_id, - ); - } - } - if is_dynamic { dataflow.dynamic_nodes.insert(node_id.clone()); } @@ -2904,9 +2677,6 @@ impl Daemon { // entry is registered. running_node.mark_registered(); dataflow.running_nodes.insert(node_id.clone(), running_node); - // Added after the spawn cohort, so it is not in - // `spawn_nodes` — record it as ours for the orphan sweep. - dataflow.owned_nodes.insert(node_id.clone()); // Update the daemon's stored descriptor so // descriptor-based lookups (e.g. AllInputsClosed @@ -3072,17 +2842,6 @@ impl Daemon { Ok(()) })(); - // Without this, repeated add/remove cycles walk the daemon - // into its pool cap (dora-rs/dora#2881). - if result.is_ok() { - reclaim_memory_pools_after_exit( - &self.memory_pool, - &self.running, - dataflow_id, - &node_id, - ); - } - // Outside the closure because it is async. Why removal // has to drive the barrier at all: see // `PendingNodes::handle_node_removal`. @@ -3466,15 +3225,6 @@ impl Daemon { // The outgoing (or an even earlier) incarnation's stale // result must not be attributed to the replacement. clear_node_result(&mut self.dataflow_node_results, dataflow_id, &node_id); - // The outgoing incarnation's exit event is dropped by - // the generation guard, so this is the only place its - // pools can be reclaimed (dora-rs/dora#2881). - reclaim_memory_pools_after_exit( - &self.memory_pool, - &self.running, - dataflow_id, - &node_id, - ); } let reply = DaemonCoordinatorReply::ReplaceNodeResult( result.map_err(|err| format!("{err:?}")), @@ -3497,20 +3247,7 @@ impl Daemon { // explicit `AddMappingResult` so the coordinator can pattern- // match the outcome (same class as #1682's AddNodeResult). let result = if let Some(dataflow) = self.running.get_mut(&dataflow_id) { - dataflow.add_mapping( - source_node.clone(), - source_output, - target_node.clone(), - target_input, - ); - // The new edge gives `target_node` a path to the source's - // pools — including any the source registered itself. - extend_pool_readers_for_new_edge( - &self.memory_pool, - dataflow, - &source_node, - &target_node, - ); + dataflow.add_mapping(source_node, source_output, target_node, target_input); Ok(()) } else { Err(format!("no running dataflow with ID `{dataflow_id}`")) @@ -4087,11 +3824,12 @@ impl Daemon { uv: bool, write_events_to: Option, ) -> eyre::Result> + use<>> { - // Sweep orphaned /dev/shm segments left by a previous crash of this - // dataflow's nodes, scoped to the nodes this daemon spawns — a - // co-located daemon may be spawning the other half of the same - // dataflow right now (see `cleanup_orphans`). - MemoryPoolManager::cleanup_orphans(&dataflow_id.to_string(), |node| { + // Opt-in extension: reclaim `/dev/shm` segments a previous crash of + // this dataflow's nodes left behind. Scoped to the nodes this daemon + // spawns, since a co-located daemon may be starting the other half of + // the same dataflow right now. Compiles away without the feature. + #[cfg(feature = "tensor-pool")] + dora_tensor_pool::TensorPoolManager::cleanup_orphans(&dataflow_id.to_string(), |node| { spawn_nodes.iter().any(|id| id.as_ref() == node) }); @@ -4106,7 +3844,6 @@ impl Daemon { self.daemon_id.clone(), dataflow_descriptor.clone(), ); - dataflow.owned_nodes = spawn_nodes.clone(); // Read from the descriptor, which is the one copy that survives // everything a dataflow outlives: auto-recovery re-spawn, // coordinator restart with state reconstruction, and `dora @@ -4224,17 +3961,9 @@ impl Daemon { } } } else if let InputMapping::User(mapping) = input.mapping { - let output_id = OutputId(mapping.source, mapping.output); - // This daemon does not deliver the edge, but it still - // needs to know it exists: the receiver's own consumers - // may be local again, and a memory-pool id can travel - // the whole chain (`downstream_closure`). dataflow - .remote_edges - .entry(output_id.clone()) - .or_default() - .insert(node.id.clone()); - dataflow.open_external_mappings.insert(output_id); + .open_external_mappings + .insert(OutputId(mapping.source, mapping.output)); } } } @@ -4930,159 +4659,76 @@ impl Daemon { } => self .output_sent(dataflow_id, node_id, output_id, metadata) .context("failed to mark output sent")?, - DaemonNodeEvent::EventStreamDropped { reply_sender } => { - let inner = async { - let dataflow = self - .running - .get_mut(&dataflow_id) - .wrap_err_with(|| format!("no running dataflow with ID `{dataflow_id}`"))?; - dataflow.subscribe_channels.remove(&node_id); - Result::<_, eyre::Error>::Ok(()) - }; - - let reply = inner.await.map_err(|err| format!("{err:?}")); - let _ = reply_sender.send(DaemonReply::Result(reply)); - } - DaemonNodeEvent::RegisterPinnedMemory { - shared_memory_id, - metadata, + DaemonNodeEvent::ExtensionStore { + namespace, + key, + value, reply_sender, } => { - let result = (|| -> Result<(), String> { - let pool_metadata = pool_metadata_from_params(&metadata.parameters); - // Validate required fields that the helper fills with defaults - if pool_metadata.ptr == 0 { - return Err("missing or invalid ptr".to_string()); - } - if pool_metadata.size == 0 { - return Err("missing or invalid size".to_string()); - } - if pool_metadata.dtype.is_empty() { - return Err("missing or invalid dtype".to_string()); - } - if pool_metadata.shape.is_empty() { - return Err("missing shape".to_string()); - } - // Mirror the Python-side size cap. - if pool_metadata.size > 1024 * 1024 * 1024 { - return Err(format!("size {} exceeds 1 GiB cap", pool_metadata.size)); - } - // Require a shared memory name for cleanup. - let shm_name = pool_metadata - .shared_memory_name - .as_ref() - .filter(|n| !n.is_empty()) - .ok_or_else(|| "missing shared_memory_name".to_string())?; - // Validate prefix in the same way free_shared_memory does. - if !shm_name.starts_with("dora_pool_") - || shm_name.contains('/') - || shm_name.contains("..") - { - return Err(format!("shared_memory_name `{}` is invalid", shm_name)); - } - - // Per-daemon pool cap (soft limit — rejects excess registrations). - const MAX_POOLS: usize = 512; - if self.memory_pool.table_size() >= MAX_POOLS { - return Err(format!( - "daemon pool table full ({MAX_POOLS} entries); \ - free unused pools before registering more" - )); - } - - // Snapshot who may still ask for this pool without - // having opened it yet, so that reclaiming the pools of - // an exited node cannot cut off a transfer that is still - // on its way (dora-rs/dora#2881). - let potential_readers = self - .running - .get(&dataflow_id) - .map(|dataflow| dataflow.downstream_closure(&node_id)) - .unwrap_or_default(); - - self.memory_pool.register_memory_pool( - MemoryPoolId { - dataflow_id: dataflow_id.to_string(), - id: shared_memory_id, - }, - pool_metadata, - node_id.to_string(), - potential_readers, - ) - })(); + let ext_key = ExtensionKey { + dataflow_id: dataflow_id.to_string(), + namespace, + key, + }; + let result = self.extensions.store(ext_key, value, &node_id); let _ = reply_sender.send(DaemonReply::Result(result)); } - DaemonNodeEvent::ReadPinnedMemory { - shared_memory_id, - free, + DaemonNodeEvent::ExtensionLoad { + namespace, + key, + remove, reply_sender, } => { - let result = (|| -> Result { - let id = MemoryPoolId { - dataflow_id: dataflow_id.to_string(), - id: shared_memory_id.clone(), - }; - let metadata = self - .memory_pool - .read_memory_pool(&id, node_id.as_ref()) - .ok_or_else(|| { - format!("memory pool with ID {} not found", shared_memory_id) - })?; - - if free - && let Err(err) = free_pool_and_notify( - &self.memory_pool, - self.running.get(&dataflow_id), - &id, - node_id.as_ref(), - &self.clock, - ) - { - tracing::warn!( - "Failed to free memory pool {} after reading: {}", - shared_memory_id, - err - ); - } - - let mut parameters = pool_metadata_to_params(&metadata); - // When freeing, drop shared_memory_name — the segment has - // been unlinked and the name is a dangling reference. - if free { - parameters.remove("shared_memory_name"); - } - - let timestamp = self.clock.new_timestamp(); - Ok(dora_message::metadata::Metadata::from_parameters( - timestamp, parameters, - )) - })(); - - match result { - Ok(metadata) => { - let _ = reply_sender.send(DaemonReply::PinnedMemoryMetadata { metadata }); - } - Err(err) => { - let _ = reply_sender.send(DaemonReply::Result(Err(err))); - } + let ext_key = ExtensionKey { + dataflow_id: dataflow_id.to_string(), + namespace, + key, + }; + let value = self.extensions.load(&ext_key, &node_id); + // Remove only after a hit: a miss must not broadcast a drop + // for a key that was never there. + if remove && value.is_some() { + drop_extension_and_notify( + &mut self.extensions, + self.running.get(&dataflow_id), + &ext_key, + &self.clock, + ); } + let _ = reply_sender.send(DaemonReply::ExtensionValue { value }); } - DaemonNodeEvent::FreePinnedMemory { - shared_memory_id, + DaemonNodeEvent::ExtensionDrop { + namespace, + key, reply_sender, } => { - let id = MemoryPoolId { + let ext_key = ExtensionKey { dataflow_id: dataflow_id.to_string(), - id: shared_memory_id, + namespace, + key, }; - let result = free_pool_and_notify( - &self.memory_pool, + drop_extension_and_notify( + &mut self.extensions, self.running.get(&dataflow_id), - &id, - node_id.as_ref(), + &ext_key, &self.clock, ); - let _ = reply_sender.send(DaemonReply::Result(result)); + // Idempotent: dropping an absent key is success, so a retry + // after a lost reply does not surface as an error. + let _ = reply_sender.send(DaemonReply::Result(Ok(()))); + } + DaemonNodeEvent::EventStreamDropped { reply_sender } => { + let inner = async { + let dataflow = self + .running + .get_mut(&dataflow_id) + .wrap_err_with(|| format!("no running dataflow with ID `{dataflow_id}`"))?; + dataflow.subscribe_channels.remove(&node_id); + Result::<_, eyre::Error>::Ok(()) + }; + + let reply = inner.await.map_err(|err| format!("{err:?}")); + let _ = reply_sender.send(DaemonReply::Result(reply)); } } Ok(()) @@ -5772,6 +5418,14 @@ impl Daemon { dataflows.remove(&dataflow_id); }); + // Whatever survived node-exit reclamation goes now: the dataflow's + // channels and listeners are gone, so nothing can reach these entries + // and no later event would reclaim them. + let dropped = self.extensions.reclaim_dataflow(&dataflow_id.to_string()); + if dropped > 0 { + tracing::debug!(%dataflow_id, dropped, "released extension entries of finished dataflow"); + } + logger .log( LogLevel::Info, @@ -5801,14 +5455,7 @@ impl Daemon { if let Some(df) = self.running.get(&dataflow_id) { let _ = df.listener_shutdown_tx.send(true); } - let owned_nodes = self - .running - .remove(&dataflow_id) - .map(|dataflow| dataflow.owned_nodes) - .unwrap_or_default(); - - release_memory_pools_of_finished_dataflow(&self.memory_pool, dataflow_id, &owned_nodes); - + self.running.remove(&dataflow_id); Ok(()) } @@ -6230,13 +5877,15 @@ impl Daemon { dataflow.connected_nodes.remove(&node_id); } - // This incarnation is gone — crash, clean exit, or a restart - // about to spawn a fresh one (dora-rs/dora#2881). - reclaim_memory_pools_after_exit( - &self.memory_pool, - &self.running, + // A node that crashed cannot withdraw its own descriptors. + // Reclaiming here is the reason the extension table lives in + // the daemon at all (dora-rs/dora#2881). + reclaim_extensions_of_exited_node( + &mut self.extensions, + self.running.get(&dataflow_id), dataflow_id, &node_id, + &self.clock, ); logger @@ -8154,238 +7803,6 @@ mod fault_tolerance_tests { } } - // ---- cross-process memory-pool cleanup (#2935) ---- - - /// Subscribe `node` to `df` with a channel of `capacity` slots and a - /// pending-message counter, mirroring what `DaemonNodeEvent::Subscribe` - /// installs for a live node. - fn subscribe_test_node( - df: &mut RunningDataflow, - node: &NodeId, - capacity: usize, - ) -> ( - mpsc::Receiver>, - Arc, - ) { - let (tx, rx) = mpsc::channel(capacity); - df.subscribe_channels.insert(node.clone(), tx); - let counter = Arc::new(AtomicU64::new(0)); - df.pending_messages.insert(node.clone(), counter.clone()); - (rx, counter) - } - - fn test_pool_id(name: &str) -> MemoryPoolId { - MemoryPoolId { - dataflow_id: Uuid::nil().to_string(), - id: name.to_string(), - } - } - - /// Metadata without a backing segment, so `free_memory_pool` skips the - /// Linux-only `/dev/shm` unlink and the tests stay cross-platform. - fn test_pool_metadata() -> MemoryPoolMetadata { - MemoryPoolMetadata { - size: 1024, - dtype: "float32".into(), - shape: vec![256], - ..Default::default() - } - } - - fn freed_pool_ids(events: &[NodeEvent]) -> Vec<&str> { - events - .iter() - .filter_map(|event| match event { - NodeEvent::FreeMemoryPool { shared_memory_id } => Some(shared_memory_id.as_str()), - _ => None, - }) - .collect() - } - - /// The cleanup broadcast must reach every node that registered or read - /// the pool, and nobody else: not the node that initiated the free (it - /// released synchronously) and not a subscriber that never touched the - /// pool. - /// - /// `free_pool_and_notify` is the single entry point behind both - /// `FreePinnedMemory` and `ReadPinnedMemory { free: true }` — before - /// #2935 the read-with-free path freed the pool without notifying - /// anyone, so every other node kept its GPU/host/IPC/shmem resources. - #[test] - fn free_pool_notifies_touched_nodes_only() { - let clock = test_clock(); - let mut df = test_dataflow(); - let pools = MemoryPoolManager::new(); - - let registrar: NodeId = "registrar".to_string().into(); - let reader: NodeId = "reader".to_string().into(); - let bystander: NodeId = "bystander".to_string().into(); - let (mut registrar_rx, registrar_pending) = - subscribe_test_node(&mut df, ®istrar, NODE_EVENT_CHANNEL_CAPACITY); - let (mut reader_rx, _) = subscribe_test_node(&mut df, &reader, NODE_EVENT_CHANNEL_CAPACITY); - let (mut bystander_rx, _) = - subscribe_test_node(&mut df, &bystander, NODE_EVENT_CHANNEL_CAPACITY); - - let id = test_pool_id("pool-1"); - pools - .register_memory_pool( - id.clone(), - test_pool_metadata(), - registrar.to_string(), - // No edges in this fixture, so no downstream consumer can - // learn the id; the free path ignores the set either way. - Default::default(), - ) - .unwrap(); - pools - .read_memory_pool(&id, reader.as_ref()) - .expect("pool should exist"); - - // The reader frees the pool it just read — the normal lifecycle. - free_pool_and_notify(&pools, Some(&df), &id, reader.as_ref(), &clock) - .expect("free should succeed"); - - assert_eq!( - freed_pool_ids(&drain_events(&mut registrar_rx)), - vec!["pool-1"], - "the registering node must be told to release its pool resources" - ); - assert!( - drain_events(&mut reader_rx).is_empty(), - "the initiator released synchronously and must not be notified" - ); - assert!( - drain_events(&mut bystander_rx).is_empty(), - "a node that never touched the pool must not be notified" - ); - - // Every delivery site pairs an enqueue with `inc_pending`; the - // listener decrements unconditionally, so a missing increment - // underflows the reported count to `u64::MAX` (#2827). - assert_eq!(registrar_pending.load(atomic::Ordering::Relaxed), 1); - } - - /// `send_with_timestamp` reports a full channel as `Ok(false)`, not - /// `Err`. The original broadcast matched only on `Err`, so a full - /// channel silently skipped the cleanup (#2935). The undelivered node - /// must be reported so the leak is visible. - #[test] - fn notify_pool_freed_reports_nodes_with_full_channels() { - let clock = test_clock(); - let mut df = test_dataflow(); - - let registrar: NodeId = "registrar".to_string().into(); - let backed_up: NodeId = "backed-up".to_string().into(); - let (mut registrar_rx, registrar_pending) = - subscribe_test_node(&mut df, ®istrar, NODE_EVENT_CHANNEL_CAPACITY); - // A one-slot channel, already full: the next send is dropped. - let (_backed_up_rx, backed_up_pending) = subscribe_test_node(&mut df, &backed_up, 1); - assert!( - send_with_timestamp(&df.subscribe_channels[&backed_up], NodeEvent::Stop, &clock) - .expect("channel is open") - ); - - let touched = HashSet::from([ - registrar.to_string(), - backed_up.to_string(), - "initiator".to_string(), - ]); - let undelivered = notify_memory_pool_freed(&df, "pool-1", &touched, "initiator", &clock); - - assert_eq!( - undelivered, - BTreeSet::from([backed_up.clone()]), - "a node whose channel is full must be reported as undelivered" - ); - assert_eq!( - backed_up_pending.load(atomic::Ordering::Relaxed), - 0, - "a dropped event must not increment the pending counter" - ); - assert_eq!( - freed_pool_ids(&drain_events(&mut registrar_rx)), - vec!["pool-1"], - "one full channel must not stop the broadcast to the other nodes" - ); - assert_eq!(registrar_pending.load(atomic::Ordering::Relaxed), 1); - } - - /// A closed channel means the node is gone and the OS reclaimed its - /// resources — benign, and it must not abort the rest of the broadcast - /// or be reported as a leak. - #[test] - fn notify_pool_freed_skips_disconnected_nodes() { - let clock = test_clock(); - let mut df = test_dataflow(); - - let registrar: NodeId = "registrar".to_string().into(); - let gone: NodeId = "gone".to_string().into(); - let (mut registrar_rx, _) = - subscribe_test_node(&mut df, ®istrar, NODE_EVENT_CHANNEL_CAPACITY); - let (gone_rx, _) = subscribe_test_node(&mut df, &gone, NODE_EVENT_CHANNEL_CAPACITY); - drop(gone_rx); - - let touched = HashSet::from([registrar.to_string(), gone.to_string()]); - let undelivered = notify_memory_pool_freed(&df, "pool-1", &touched, "third-party", &clock); - - assert!( - undelivered.is_empty(), - "a closed channel is benign, not an undelivered notification" - ); - assert_eq!( - freed_pool_ids(&drain_events(&mut registrar_rx)), - vec!["pool-1"], - "a disconnected peer must not abort the rest of the broadcast" - ); - } - - /// Freeing an unknown pool must surface the manager's error to the - /// requesting node and broadcast nothing. - #[test] - fn free_pool_propagates_unknown_pool_error_without_notifying() { - let clock = test_clock(); - let mut df = test_dataflow(); - let pools = MemoryPoolManager::new(); - - let registrar: NodeId = "registrar".to_string().into(); - let (mut registrar_rx, _) = - subscribe_test_node(&mut df, ®istrar, NODE_EVENT_CHANNEL_CAPACITY); - - let err = free_pool_and_notify( - &pools, - Some(&df), - &test_pool_id("never-registered"), - registrar.as_ref(), - &clock, - ) - .expect_err("freeing an unknown pool must fail"); - assert!( - err.contains("memory pool not found"), - "unexpected error: {err}" - ); - assert!(drain_events(&mut registrar_rx).is_empty()); - } - - /// A dataflow that is no longer running has no subscribe channels; the - /// pool must still be released rather than the free failing. - #[test] - fn free_pool_without_running_dataflow_still_releases_the_entry() { - let clock = test_clock(); - let pools = MemoryPoolManager::new(); - let id = test_pool_id("pool-1"); - pools - .register_memory_pool( - id.clone(), - test_pool_metadata(), - "registrar".into(), - Default::default(), - ) - .unwrap(); - - free_pool_and_notify(&pools, None, &id, "registrar", &clock).expect("free should succeed"); - assert_eq!(pools.table_size(), 0); - } - // -- Test 1: close_input removes input, sends InputClosed, no AllInputsClosed with remaining inputs -- /// Regression test for #241. The daemon used to carry two inverse maps @@ -9807,208 +9224,6 @@ mod fault_tolerance_tests { window )); } - - // -- dora#2881: memory pools must be reclaimed once nothing can reach them -- - - /// A dataflow whose `sender` registered one pool, with `edges` wired up - /// and every node of `running` marked as running. - fn dataflow_with_pool( - edges: &[(&str, &str)], - running: &[&str], - ) -> (HashMap, MemoryPoolManager) { - let mut df = test_dataflow(); - for (from, to) in edges { - df.add_mapping( - NodeId::from(from.to_string()), - DataId::from("out".to_string()), - NodeId::from(to.to_string()), - DataId::from("in".to_string()), - ); - } - for node in running { - df.running_nodes - .insert(NodeId::from(node.to_string()), test_running_node()); - } - - let memory_pool = MemoryPoolManager::new(); - let sender = NodeId::from("sender".to_string()); - memory_pool - .register_memory_pool( - MemoryPoolId { - dataflow_id: Uuid::nil().to_string(), - id: "pool-1".to_string(), - }, - MemoryPoolMetadata::default(), - sender.to_string(), - df.downstream_closure(&sender), - ) - .unwrap(); - - (HashMap::from([(Uuid::nil(), df)]), memory_pool) - } - - #[test] - fn a_pool_is_reclaimed_on_node_exit_exactly_when_nothing_can_reach_it() { - // (case, edges, still-running nodes, exiting node, pools left) - let cases = [ - // The receiver has not read the pool yet — it only learns the id - // from the message the sender put in flight before exiting. - // Freeing here would break that transfer, which is why - // reclamation is reachability-based rather than eager. - ( - "consumer still running", - &[("sender", "recv")][..], - &["recv"][..], - "sender", - 1, - ), - // ...and once that consumer is gone too, nothing can reach it. - ( - "last consumer gone", - &[("sender", "recv")][..], - &[][..], - "recv", - 0, - ), - // Callers run before the node leaves `running_nodes`, so its own - // stale entry must not count as a live reference — otherwise a - // sender with no consumers could never be reclaimed. - ( - "exiting node is not itself live", - &[][..], - &["sender"][..], - "sender", - 0, - ), - // Liveness is checked against the pool's own reference set, not - // the dataflow as a whole: a node that is neither downstream of - // the sender nor has ever opened the pool cannot reach it, so a - // long-lived unrelated node must not pin it (the leak this fixes). - ( - "unrelated node running", - &[("other_source", "other_sink")][..], - &["other_sink"][..], - "sender", - 0, - ), - ]; - - for (case, edges, running_nodes, exiting, expected) in cases { - let (running, memory_pool) = dataflow_with_pool(edges, running_nodes); - - reclaim_memory_pools_after_exit( - &memory_pool, - &running, - Uuid::nil(), - &NodeId::from(exiting.to_string()), - ); - - assert_eq!(memory_pool.table_size(), expected, "case: {case}"); - } - } - - /// `potential_readers` is captured at registration time, so rewiring a - /// running dataflow would otherwise leave it stale. The registrar is in - /// its own reader set and reclamation fires on its exit, so without the - /// update the pool — and its `/dev/shm` segment — would go away while a - /// consumer connected a moment ago still had its id in flight - /// (dora-rs/dora#2881 review). - #[test] - fn a_consumer_connected_after_registration_still_pins_the_pool() { - let sender = NodeId::from("sender".to_string()); - let consumer = NodeId::from("consumer".to_string()); - - let mut df = test_dataflow(); - df.running_nodes.insert(sender.clone(), test_running_node()); - - // Registered while `sender` has no consumers at all. - let memory_pool = MemoryPoolManager::new(); - memory_pool - .register_memory_pool( - MemoryPoolId { - dataflow_id: Uuid::nil().to_string(), - id: "pool-1".to_string(), - }, - MemoryPoolMetadata::default(), - sender.to_string(), - df.downstream_closure(&sender), - ) - .unwrap(); - - // `dora graph connect sender/out consumer/in`. - df.add_mapping( - sender.clone(), - DataId::from("out".to_string()), - consumer.clone(), - DataId::from("in".to_string()), - ); - df.running_nodes - .insert(consumer.clone(), test_running_node()); - extend_pool_readers_for_new_edge(&memory_pool, &df, &sender, &consumer); - - // The sender sends the pool id along the new edge and exits. - let mut running = HashMap::from([(Uuid::nil(), df)]); - reclaim_memory_pools_after_exit(&memory_pool, &running, Uuid::nil(), &sender); - assert_eq!( - memory_pool.table_size(), - 1, - "the newly connected consumer can still read the pool" - ); - - // ...and it is reclaimed once that consumer is gone too. - let df = running.get_mut(&Uuid::nil()).expect("dataflow is present"); - df.running_nodes.remove(&sender); - reclaim_memory_pools_after_exit(&memory_pool, &running, Uuid::nil(), &consumer); - assert_eq!(memory_pool.table_size(), 0); - } - - #[test] - fn finishing_a_dataflow_releases_every_pool_it_still_holds() { - let (_running, memory_pool) = dataflow_with_pool(&[("sender", "recv")], &["recv"]); - - // Owned nodes only scope the /dev/shm sweep; the table release is - // unconditional. `recv` was still running and could still have - // reached this pool, but the dataflow is over and no later event - // would ever reclaim it. - release_memory_pools_of_finished_dataflow(&memory_pool, Uuid::nil(), &BTreeSet::new()); - - assert_eq!(memory_pool.table_size(), 0); - } - - /// `finish_dataflow` is a *per-daemon* finish, but `/dev/shm` is - /// host-wide: with two daemons serving one dataflow on one machine, the - /// first to finish must not unlink the segments of the second daemon's - /// still-running nodes (dora-rs/dora#2881 review). `cleanup_orphans` has - /// its own test for which names it matches; this one pins the wiring — - /// that the finish path passes its own nodes and not something wider. - #[cfg(target_os = "linux")] - #[test] - fn finishing_a_dataflow_sweeps_only_this_daemons_segments() { - let dataflow_id = Uuid::new_v4(); - let segment = |node: &str| { - std::path::PathBuf::from(format!("/dev/shm/dora_pool_{dataflow_id}_{node}_0")) - }; - // `ours` never registered its pool — a crash between creating the - // segment and registering it is exactly what the orphan sweep is - // for. `theirs` belongs to the co-located daemon. - for node in ["ours", "theirs"] { - std::fs::write(segment(node), b"x").unwrap(); - } - - let owned_nodes = BTreeSet::from([NodeId::from("ours".to_string())]); - release_memory_pools_of_finished_dataflow( - &MemoryPoolManager::new(), - dataflow_id, - &owned_nodes, - ); - - let (ours, theirs) = (segment("ours").exists(), segment("theirs").exists()); - for node in ["ours", "theirs"] { - let _ = std::fs::remove_file(segment(node)); - } - assert!(!ours, "this daemon's orphaned segment must be swept"); - assert!(theirs, "the co-located daemon's live segment must survive"); - } } #[cfg(test)] diff --git a/binaries/daemon/src/node_communication/mod.rs b/binaries/daemon/src/node_communication/mod.rs index 1de613d932..70f07d413a 100644 --- a/binaries/daemon/src/node_communication/mod.rs +++ b/binaries/daemon/src/node_communication/mod.rs @@ -360,15 +360,17 @@ impl Listener { ) .await?; } - DaemonRequest::RegisterPinnedMemory { - shared_memory_id, - metadata, + DaemonRequest::ExtensionStore { + namespace, + key, + value, } => { let (reply_sender, reply) = oneshot::channel(); self.process_daemon_event( - DaemonNodeEvent::RegisterPinnedMemory { - shared_memory_id, - metadata, + DaemonNodeEvent::ExtensionStore { + namespace, + key, + value, reply_sender, }, Some(reply), @@ -376,15 +378,17 @@ impl Listener { ) .await?; } - DaemonRequest::ReadPinnedMemory { - shared_memory_id, - free, + DaemonRequest::ExtensionLoad { + namespace, + key, + remove, } => { let (reply_sender, reply) = oneshot::channel(); self.process_daemon_event( - DaemonNodeEvent::ReadPinnedMemory { - shared_memory_id, - free, + DaemonNodeEvent::ExtensionLoad { + namespace, + key, + remove, reply_sender, }, Some(reply), @@ -392,11 +396,12 @@ impl Listener { ) .await?; } - DaemonRequest::FreePinnedMemory { shared_memory_id } => { + DaemonRequest::ExtensionDrop { namespace, key } => { let (reply_sender, reply) = oneshot::channel(); self.process_daemon_event( - DaemonNodeEvent::FreePinnedMemory { - shared_memory_id, + DaemonNodeEvent::ExtensionDrop { + namespace, + key, reply_sender, }, Some(reply), diff --git a/binaries/daemon/src/running_dataflow.rs b/binaries/daemon/src/running_dataflow.rs index 4c4642e7a6..034900c081 100644 --- a/binaries/daemon/src/running_dataflow.rs +++ b/binaries/daemon/src/running_dataflow.rs @@ -30,7 +30,7 @@ use crossbeam::queue::ArrayQueue; use eyre::eyre; use futures::FutureExt; use std::{ - collections::{BTreeMap, BTreeSet, HashMap, HashSet}, + collections::{BTreeMap, BTreeSet, HashMap}, sync::{ Arc, atomic::{self, AtomicBool, AtomicU32, AtomicU64}, @@ -292,11 +292,6 @@ pub struct RunningDataflow { /// Per-node pending message counters (incremented on send, decremented on recv) pub(crate) pending_messages: HashMap>, pub(crate) mappings: HashMap>, - /// Edges whose receiver lives on another daemon, which `mappings` (this - /// daemon's delivery table) does not record. Only `downstream_closure` - /// reads them, to follow a chain that leaves this daemon and returns to - /// it (dora-rs/dora#2881). - pub(crate) remote_edges: HashMap>, pub(crate) timers: BTreeMap>, /// Nodes subscribing to `dora/logs` virtual input. pub(crate) log_subscribers: Vec, @@ -304,12 +299,6 @@ pub struct RunningDataflow { pub(crate) input_deadlines: HashMap<(NodeId, DataId), InputDeadline>, pub(crate) broken_inputs: HashMap<(NodeId, DataId), Duration>, pub(crate) running_nodes: BTreeMap, - /// Every node of this dataflow that runs on *this* daemon, whether or - /// not it is still running. Only grows, so it still answers "was this - /// node ours?" at teardown, when `running_nodes` has long since dropped - /// the exited ones. Scopes the `/dev/shm` orphan sweep — see - /// `MemoryPoolManager::cleanup_orphans` (dora-rs/dora#2881). - pub(crate) owned_nodes: BTreeSet, pub(crate) dynamic_nodes: BTreeSet, pub(crate) open_external_mappings: BTreeSet, pub(crate) _timer_handles: BTreeMap>, @@ -407,14 +396,12 @@ impl RunningDataflow { subscribe_channels: HashMap::new(), pending_messages: HashMap::new(), mappings: HashMap::new(), - remote_edges: HashMap::new(), timers: BTreeMap::new(), log_subscribers: Vec::new(), open_inputs: BTreeMap::new(), input_deadlines: HashMap::new(), broken_inputs: HashMap::new(), running_nodes: BTreeMap::new(), - owned_nodes: BTreeSet::new(), dynamic_nodes: BTreeSet::new(), open_external_mappings: Default::default(), _timer_handles: BTreeMap::new(), @@ -1020,48 +1007,6 @@ impl RunningDataflow { node_output_ids(&self.mappings, &self.open_external_mappings, node_id) } - /// Every node reachable from `source` by following output→input edges, - /// transitively, plus `source` itself. - /// - /// Used to snapshot who may still learn a memory-pool id registered by - /// `source`: the id travels the dataflow as ordinary message data, so any - /// node downstream of the registrar may still receive it — including - /// nodes further along that a direct consumer forwards it to. The daemon - /// cannot see pool ids inside payloads, so the whole downstream closure - /// counts as a potential reader (dora-rs/dora#2881). - /// - /// Remote receivers are followed too, via `remote_edges`: `mappings` - /// holds only edges this daemon delivers, so a chain that leaves the - /// daemon and comes back — `local -> remote -> local` — would otherwise - /// stop at the first hop and lose the local node at its end. A remote - /// node never counts as live (pools are host-local), so including it - /// costs nothing; reaching *past* it is the point. - /// - /// Returns node ids as `String` because that is how the pool table keys - /// them. - pub(crate) fn downstream_closure(&self, source: &NodeId) -> HashSet { - let mut reachable = HashSet::from([source.to_string()]); - let mut queue = vec![source.clone()]; - while let Some(node) = queue.pop() { - let local = self - .mappings - .iter() - .filter(|(output_id, _)| output_id.0 == node) - .flat_map(|(_, receivers)| receivers.iter().map(|(receiver, _input)| receiver)); - let remote = self - .remote_edges - .iter() - .filter(|(output_id, _)| output_id.0 == node) - .flat_map(|(_, receivers)| receivers.iter()); - for receiver in local.chain(remote) { - if reachable.insert(receiver.to_string()) { - queue.push(receiver.clone()); - } - } - } - reachable - } - /// Nodes blocking an otherwise-finished dataflow (dora-rs/dora#2152). /// /// Returns nodes that should be force-stopped because the dataflow is @@ -1679,70 +1624,6 @@ mod tests { } } - // ---- dora-rs/dora#2881: downstream closure for memory-pool reachability ---- - - /// Wire `from/out -> to/in`. - fn edge(df: &mut RunningDataflow, from: &str, to: &str) { - df.add_mapping(node_id(from), data_id("out"), node_id(to), data_id("in")); - } - - #[test] - fn downstream_closure_follows_edges_transitively() { - let mut df = - RunningDataflow::new(uuid::Uuid::nil(), DaemonId::new(None), empty_descriptor()); - edge(&mut df, "sender", "middle"); - edge(&mut df, "middle", "sink"); - edge(&mut df, "unrelated", "other"); - - // A direct consumer may forward a pool id further down, so the whole - // downstream chain counts — and nothing outside it does. - assert_eq!( - df.downstream_closure(&node_id("sender")), - HashSet::from([ - "sender".to_string(), - "middle".to_string(), - "sink".to_string(), - ]), - ); - } - - #[test] - fn downstream_closure_follows_a_chain_through_another_daemon() { - let mut df = - RunningDataflow::new(uuid::Uuid::nil(), DaemonId::new(None), empty_descriptor()); - // `sender -> middle` leaves this daemon, so it is not in `mappings`; - // `middle -> sink` comes back to a local node. A pool id can travel - // the whole chain, so `sink` must count as a potential reader even - // though the only path to it runs through the other daemon. - df.remote_edges - .entry(OutputId(node_id("sender"), data_id("out"))) - .or_default() - .insert(node_id("middle")); - edge(&mut df, "middle", "sink"); - - assert_eq!( - df.downstream_closure(&node_id("sender")), - HashSet::from([ - "sender".to_string(), - "middle".to_string(), - "sink".to_string(), - ]), - ); - } - - #[test] - fn downstream_closure_terminates_on_a_cycle() { - let mut df = - RunningDataflow::new(uuid::Uuid::nil(), DaemonId::new(None), empty_descriptor()); - edge(&mut df, "sender", "peer"); - edge(&mut df, "peer", "sender"); - - assert_eq!( - df.downstream_closure(&node_id("sender")), - HashSet::from(["sender".to_string(), "peer".to_string()]), - ); - } - #[test] fn propagate_node_failed_keeps_idle_receiver_counter_consistent() { let mut df = diff --git a/docs/extensions.md b/docs/extensions.md new file mode 100644 index 0000000000..d22e5fd896 --- /dev/null +++ b/docs/extensions.md @@ -0,0 +1,98 @@ +# Extensions: the transport seam + +Some transports do not belong in dora's stable core — whether or not they ship +in this repository. A CUDA tensor-pool needs `libcudart`, +hand-parsed shared-memory headers, a seqlock and GPU-specific transport +selection — none of which the framework should carry, and none of which it can +usefully test without GPU runners. But such a transport still needs one thing +only the daemon can provide: **someone to clean up after a node that crashed.** + +The extension seam is that one thing and nothing more. It lets a transport ship +in-tree and be built (behind a feature flag) without any part of it entering +the 1.0 compatibility surface. + +## What dora provides + +A dataflow-scoped table of opaque byte values, keyed by `(namespace, key)`. +dora never interprets any of the three — it brokers lifetime only. + +```python +# Producer +node.extension_store("my-transport", "frame-7", descriptor_bytes) +node.send_output("handle", pa.array(["frame-7"])) + +# Consumer +descriptor = node.extension_load("my-transport", "frame-7") +... +node.extension_drop("my-transport", "frame-7") + +# Anywhere the extension runs: release what dropped keys referred to +for key in node.drain_dropped_extension_keys("my-transport"): + my_cache.pop(key, None) +``` + +Rust nodes get the same four operations on `DoraNode` +(`extension_store` / `extension_load` / `extension_drop`) plus +`dora_node_api::event_stream::extensions::drain_dropped_keys`. + +### Guarantees + +| | | +|---|---| +| **Reclamation on exit** | When the storing node exits — cleanly or by crashing — its entries are dropped and its readers notified. This is the reason the table is daemon-side; a node cannot do it for itself. | +| **Reclamation on finish** | Anything left when the dataflow finishes is released, so a long-lived `dora up` daemon does not accumulate entries. | +| **Notification** | Every node that stored *or read* a key is told when it goes away, so per-process resources keyed on it can be freed promptly instead of at process exit. | +| **Ownership** | Only the storing node may overwrite a key. A second node cannot redirect an existing key's readers. | +| **Scoping** | Entries are per-dataflow and per-namespace. Two dataflows, or two extensions in one process, cannot see or clobber each other. | +| **Idempotent drop** | Dropping an absent key succeeds and broadcasts nothing, so a retry after a lost reply is safe. | + +### Limits + +- **8192 entries per dataflow.** A store-per-frame loop that never drops will + hit the cap and get an error rather than exhausting the daemon. +- **4096 pending drop notifications per process**, oldest evicted. Missing one + means the extension frees that resource on its own teardown rather than + promptly — degraded, not incorrect. +- **Values are copied** through the daemon. This is a control-plane channel for + descriptors, not a data plane. Put the bulk payload in shared memory or a + regular dora output and keep the descriptor here. +- **A dropped notification is best-effort.** If the target's channel is full or + closed the daemon logs it rather than blocking the drop. + +## What dora deliberately does not provide + +No shared-memory helpers, no CUDA, no wire vocabulary for any particular +transport. The protocol variants are `ExtensionStore` / `ExtensionLoad` / +`ExtensionDrop` — deliberately generic, so that adding a second extension needs +no change to dora at all. + +If you find yourself wanting dora to grow a variant named after your transport, +that is the signal the seam is being used wrong. + +## Why not just send a dataflow message? + +You should, for the descriptor's *delivery* — that is what `send_output` is +for, and the example above does exactly that. What a message cannot do is +survive the sender crashing: the receiver keeps a mapping to memory nobody +owns, and nothing ever tells it otherwise. dora-rs/dora#2881 is that failure +mode with a real transport attached. + +Use dataflow messages to hand the key around; use the extension table for the +descriptor whose lifetime has to outlive a crash. + +## The first consumer + +The pinned/CUDA tensor-pool transport lives at +[`libraries/extensions/tensor-pool`](../libraries/extensions/tensor-pool) and +uses exactly these four operations. It is opt-in (`--features tensor-pool`, off +by default) and **outside the 1.0 compatibility guarantees** — that combination +is the point: a transport can ship in-tree, be built and used, and still not +freeze anything into dora's stable surface. + +It was previously integrated directly: ~3,000 lines inside the Python binding +with 64 `unsafe` sites, 950 lines of daemon lifecycle logic, and pool-specific +wire-protocol variants. Reworking it onto this channel removed all of that from +dora proper while keeping the feature usable. + +That is the shape this seam exists to produce: an extension is a package that +uses these four operations, not a fork of the framework. diff --git a/libraries/extensions/memory-pool/Cargo.toml b/libraries/extensions/memory-pool/Cargo.toml deleted file mode 100644 index 9360446798..0000000000 --- a/libraries/extensions/memory-pool/Cargo.toml +++ /dev/null @@ -1,13 +0,0 @@ -[package] -name = "dora-memory-pool" -version.workspace = true -edition.workspace = true -rust-version.workspace = true -documentation.workspace = true -readme.workspace = true -description.workspace = true -license.workspace = true -repository.workspace = true - -[dependencies] -tracing = "0.1.44" diff --git a/libraries/extensions/tensor-pool/Cargo.toml b/libraries/extensions/tensor-pool/Cargo.toml new file mode 100644 index 0000000000..112ed7466c --- /dev/null +++ b/libraries/extensions/tensor-pool/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "dora-tensor-pool" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +documentation.workspace = true +readme.workspace = true +description = "Daemon-side helper for the tensor-pool extension: reclaims orphaned /dev/shm segments. Not covered by dora's 1.0 compatibility guarantees — see libraries/extensions/tensor-pool/README.md." +license.workspace = true +repository.workspace = true +publish = false + +[dependencies] +tracing = { workspace = true } diff --git a/libraries/extensions/tensor-pool/README.md b/libraries/extensions/tensor-pool/README.md new file mode 100644 index 0000000000..9e778a8c7c --- /dev/null +++ b/libraries/extensions/tensor-pool/README.md @@ -0,0 +1,124 @@ +# tensor-pool — pinned-host / CUDA tensor transport + +Zero-copy handoff of large tensors between dora nodes: the producer registers a +shared-tensor pool once and overwrites it per frame, so a repeated 8 MB camera +frame does not travel through the normal message path every tick. CUDA sources +and sinks get DMA and IPC paths on top; a CPU-only machine gets plain shared +memory. + +## ⚠️ Not covered by dora's 1.0 compatibility guarantees + +**This is an opt-in extension. It is not part of the 1.0 API surface.** + +- **It may break in any release**, including a patch. The four Python methods, + the segment layout and the descriptor format are all free to change. +- **It has known open defects** — see below. They are real, reproducible, and + not scheduled against 1.0. +- **Its GPU paths have no automated coverage.** dora's CI has no GPU runners. + Only the CPU path is exercised, and only in a nightly job. + +Everything else in dora — the node APIs, the CLI, the descriptor format, the +wire protocol — carries the normal 1.0 stability promise. This does not. + +Use it if the throughput matters more to you than the stability, and pin your +dora version if you do. + +## Building with it + +Off by default, so a standard build neither compiles nor exposes it. + +```bash +# Python wheel with the transport +maturin develop -m apis/python/node/Cargo.toml --features tensor-pool + +# Daemon with orphaned-segment reclamation +cargo build -p dora-daemon --features tensor-pool +``` + +Both flags are independent. Without the daemon feature the transport still +works; shared-memory segments left by a node that crashed are reclaimed when +the dataflow finishes rather than at the next start. + +```python +from dora import Node +from dora.cuda import get_tensor_info, tensor_from_info # helpers, torch-only + +node = Node() +pool_id = node.register_tensor_pool(get_tensor_info(tensor), device="cuda:0") +node.send_output("frame", pool_id) # hand the id along as normal data +... +node.write_tensor_pool(pool_id, get_tensor_info(next_tensor)) +node.free_tensor_pool(pool_id) +``` + +Runnable dataflows are in [`examples/`](examples/), covering CPU↔CPU, CPU↔CUDA, +CUDA↔CUDA and the negative lifecycle cases. + +## Known open defects + +| Issue | Symptom | +|---|---| +| [#3015](https://github.com/dora-rs/dora/issues/3015) | Pool ids collide across node restarts — a restarted node cannot re-register its pool | +| [#2935](https://github.com/dora-rs/dora/issues/2935) | Cross-process free cleanup can be silently skipped | +| [#2890](https://github.com/dora-rs/dora/issues/2890) | The seqlock overflow fix (#2866) is incomplete: two inline end-write paths still use non-wrapping `old_gen + 1` | + +[#2881](https://github.com/dora-rs/dora/issues/2881) (pools not released when a +node crashes) is **fixed**, by #3014 plus the daemon-side reclamation described +below. + +## How it reaches dora + +Through the public [extension channel](../../../docs/extensions.md) and nothing +else. dora has no knowledge of pools, CUDA or the segment layout — it brokers +the lifetime of an opaque descriptor: + +| This transport | dora | +|---|---| +| publish a pool descriptor | `extension_store("dora-tensor-pool", id, bytes)` | +| look one up | `extension_load("dora-tensor-pool", id, remove=…)` | +| withdraw one | `extension_drop("dora-tensor-pool", id)` | +| learn what went away | `drain_dropped_extension_keys("dora-tensor-pool")` | + +The descriptor is JSON in `python/src/seam.rs`; dora never parses it. That +boundary is deliberate — it is what lets this extension change its own metadata +without touching dora's wire protocol, and what keeps `unsafe` pointer +arithmetic, the seqlock and the embedded `libcudart` bindings on this side of +the line. + +**Keep it that way.** If a future change wants dora to grow a request named +after this transport, that is the signal something is being done wrong: widen +the generic channel instead. + +## Layout + +| Path | What | +|---|---| +| `src/` | `dora-tensor-pool` — daemon-side reclamation of orphaned `/dev/shm` segments | +| `python/src/transport.rs` | the transport: segment layout, seqlock, CUDA helpers, the four operations | +| `python/src/seam.rs` | descriptor encode/decode over the extension channel | +| `python/tensor_info_helpers.py` | `get_tensor_info` / `tensor_from_info` for torch tensors | +| `examples/` | dataflows, CPU and CUDA | +| `tests/` | smoke tests (need `torch`; not wired to the default suite) | + +## History + +Built in-tree (#2168, #2386, #2619), extracted before 1.0 because it had grown +to ~3,000 lines inside the Python binding with 64 `unsafe` sites and 950 lines +of daemon lifecycle logic, then brought back here — behind a feature flag, +reaching dora only through the extension channel, and explicitly outside the +1.0 guarantees. + +The design questions its origin issue +([#1872](https://github.com/dora-rs/dora/issues/1872)) posed and that were +never answered still stand, and are the right starting point for anyone +reworking this: + +1. Why a new transport rather than extending zenoh-shm or making the Arrow IPC + path CUDA-aware? +2. What is the cross-platform story? (`/dev/shm` reclamation is Linux-only + today.) +3. Lifecycle: bounded pool size, and what happens when a producer crashes + mid-write. +4. API shape: four methods versus a flag on `send_output`. +5. Measurement against zenoh-shm **plus caller-side pinning**, on a real + workload — not against non-pinned zenoh-shm. diff --git a/examples/memory-pool/README.md b/libraries/extensions/tensor-pool/examples/README.md similarity index 76% rename from examples/memory-pool/README.md rename to libraries/extensions/tensor-pool/examples/README.md index bcd0d055cd..0b4255450a 100644 --- a/examples/memory-pool/README.md +++ b/libraries/extensions/tensor-pool/examples/README.md @@ -1,8 +1,8 @@ -# Memory Pool Example +# Tensor Pool Example ## Overview -This example exercises Dora's pinned memory-pool transport for repeated tensor transfer between a sender node and a receiver node. The positive scenarios keep the existing throughput-oriented behavior, and the negative scenarios verify that lifecycle errors are surfaced as warnings instead of crashing the nodes. +This example exercises Dora's pinned tensor-pool transport for repeated tensor transfer between a sender node and a receiver node. The positive scenarios keep the existing throughput-oriented behavior, and the negative scenarios verify that lifecycle errors are surfaced as warnings instead of crashing the nodes. ## Install @@ -24,14 +24,14 @@ python -c "import torch; assert torch.cuda.is_available()" ## Files -- `sender.py` — registers and updates a memory pool from the sender side. -- `receiver.py` — reads from the memory pool, measures throughput, and triggers lifecycle scenarios. +- `sender.py` — registers and updates a tensor pool from the sender side. +- `receiver.py` — reads from the tensor pool, measures throughput, and triggers lifecycle scenarios. - `cpu2cpu.yml` — positive throughput test for CPU sender → CPU receiver (GPU-less CI safe). - `cpu2cuda.yml` — positive throughput test for CPU sender → CUDA receiver. - `cuda2cpu.yml` — positive throughput test for CUDA sender → CPU receiver. -- `duplicate_free.yml` — receiver frees the same memory pool twice (CPU receiver). -- `read_after_free.yml` — receiver frees, then reads the same memory pool again (CPU receiver). -- `write_after_free.yml` — sender frees, then writes the same memory pool again (CPU receiver). +- `duplicate_free.yml` — receiver frees the same tensor pool twice (CPU receiver). +- `read_after_free.yml` — receiver frees, then reads the same tensor pool again (CPU receiver). +- `write_after_free.yml` — sender frees, then writes the same tensor pool again (CPU receiver). - `auto_cleanup.yml` — receiver does not free; daemon cleanup is expected on shutdown (CPU receiver). ## Run @@ -39,9 +39,9 @@ python -c "import torch; assert torch.cuda.is_available()" ### Positive throughput scenarios ```bash -dora run examples/memory-pool/cpu2cpu.yml -dora run examples/memory-pool/cpu2cuda.yml -dora run examples/memory-pool/cuda2cpu.yml +dora run examples/cpu2cpu.yml +dora run examples/cpu2cuda.yml +dora run examples/cuda2cpu.yml ``` Expected behavior: @@ -53,21 +53,21 @@ Expected behavior: ### Negative-path scenarios ```bash -dora run examples/memory-pool/duplicate_free.yml -dora run examples/memory-pool/read_after_free.yml -dora run examples/memory-pool/write_after_free.yml -dora run examples/memory-pool/auto_cleanup.yml +dora run examples/duplicate_free.yml +dora run examples/read_after_free.yml +dora run examples/write_after_free.yml +dora run examples/auto_cleanup.yml ``` Expected warnings/info: - duplicate free: - - `Attempt to release memory pool [memory_pool_id] failed - reason: pool does not exist. Operation aborted.` + - `Attempt to release tensor pool [tensor_pool_id] failed - reason: pool does not exist. Operation aborted.` - read after free: - - `Attempt to read memory pool [memory_pool_id] failed - reason: pool does not exist. Operation aborted.` + - `Attempt to read tensor pool [tensor_pool_id] failed - reason: pool does not exist. Operation aborted.` - write after free: - - `Attempt to write memory pool [memory_pool_id] failed - reason: pool does not exist. Operation aborted.` + - `Attempt to write tensor pool [tensor_pool_id] failed - reason: pool does not exist. Operation aborted.` - auto cleanup: - - `Detected xx unreleased memory pool of finished dataflow , releasing...` + - `Detected xx unreleased tensor pool of finished dataflow , releasing...` ## Pool lifetime @@ -86,7 +86,7 @@ remove` no longer strands its pools for the rest of the dataflow ## Notes -- The scenario is controlled through the `memory_pool_scenario` environment variable in each YAML file. +- The scenario is controlled through the `tensor_pool_scenario` environment variable in each YAML file. - `cpu2cpu.yml` and the four negative-lifecycle YAMLs use CPU-only receiver (`receiver_device: cpu`) and are safe for GPU-less CI runners. - The CUDA receiver scenarios (`cpu2cuda.yml`, `cuda2cpu.yml`) require a working CUDA runtime. - The negative scenarios use a reduced message count to keep lifecycle validation short and focused. diff --git a/examples/memory-pool/auto_cleanup.yml b/libraries/extensions/tensor-pool/examples/auto_cleanup.yml similarity index 83% rename from examples/memory-pool/auto_cleanup.yml rename to libraries/extensions/tensor-pool/examples/auto_cleanup.yml index 4f3fc754ed..6851219530 100644 --- a/examples/memory-pool/auto_cleanup.yml +++ b/libraries/extensions/tensor-pool/examples/auto_cleanup.yml @@ -1,9 +1,9 @@ -# Negative-path test: daemon releases an unreleased memory pool on shutdown. +# Negative-path test: daemon releases an unreleased tensor pool on shutdown. env: sender_device: cpu receiver_device: cpu message_num: 2 - memory_pool_scenario: auto_cleanup + tensor_pool_scenario: auto_cleanup nodes: - id: sender_node build: pip install torch --extra-index-url https://download.pytorch.org/whl/cpu numpy diff --git a/examples/memory-pool/cpu2cpu.yml b/libraries/extensions/tensor-pool/examples/cpu2cpu.yml similarity index 95% rename from examples/memory-pool/cpu2cpu.yml rename to libraries/extensions/tensor-pool/examples/cpu2cpu.yml index 65361fb2a1..d9a55ed2be 100644 --- a/examples/memory-pool/cpu2cpu.yml +++ b/libraries/extensions/tensor-pool/examples/cpu2cpu.yml @@ -4,7 +4,7 @@ env: sender_device: cpu receiver_device: cpu message_num: 100 - memory_pool_scenario: throughput + tensor_pool_scenario: throughput nodes: - id: sender_node build: pip install torch --extra-index-url https://download.pytorch.org/whl/cpu numpy diff --git a/examples/memory-pool/cpu2cuda.yml b/libraries/extensions/tensor-pool/examples/cpu2cuda.yml similarity index 91% rename from examples/memory-pool/cpu2cuda.yml rename to libraries/extensions/tensor-pool/examples/cpu2cuda.yml index 8d15ec5ccd..7b5aebccfa 100644 --- a/examples/memory-pool/cpu2cuda.yml +++ b/libraries/extensions/tensor-pool/examples/cpu2cuda.yml @@ -3,7 +3,7 @@ env: sender_device: cpu receiver_device: cuda message_num: 100 - memory_pool_scenario: throughput + tensor_pool_scenario: throughput nodes: - id: sender_node path: sender.py diff --git a/examples/memory-pool/cuda2cpu.yml b/libraries/extensions/tensor-pool/examples/cuda2cpu.yml similarity index 91% rename from examples/memory-pool/cuda2cpu.yml rename to libraries/extensions/tensor-pool/examples/cuda2cpu.yml index 8227e3c41b..d5d776fbb7 100644 --- a/examples/memory-pool/cuda2cpu.yml +++ b/libraries/extensions/tensor-pool/examples/cuda2cpu.yml @@ -3,7 +3,7 @@ env: sender_device: cuda receiver_device: cpu message_num: 100 - memory_pool_scenario: throughput + tensor_pool_scenario: throughput nodes: - id: sender_node path: sender.py diff --git a/examples/memory-pool/cuda2cuda.yml b/libraries/extensions/tensor-pool/examples/cuda2cuda.yml similarity index 100% rename from examples/memory-pool/cuda2cuda.yml rename to libraries/extensions/tensor-pool/examples/cuda2cuda.yml diff --git a/examples/memory-pool/cuda_inner.yml b/libraries/extensions/tensor-pool/examples/cuda_inner.yml similarity index 100% rename from examples/memory-pool/cuda_inner.yml rename to libraries/extensions/tensor-pool/examples/cuda_inner.yml diff --git a/examples/memory-pool/duplicate_free.yml b/libraries/extensions/tensor-pool/examples/duplicate_free.yml similarity index 83% rename from examples/memory-pool/duplicate_free.yml rename to libraries/extensions/tensor-pool/examples/duplicate_free.yml index 0d49c932e8..d986e50576 100644 --- a/examples/memory-pool/duplicate_free.yml +++ b/libraries/extensions/tensor-pool/examples/duplicate_free.yml @@ -1,9 +1,9 @@ -# Negative-path test: receiver frees the same memory pool twice. +# Negative-path test: receiver frees the same tensor pool twice. env: sender_device: cpu receiver_device: cpu message_num: 2 - memory_pool_scenario: duplicate_free + tensor_pool_scenario: duplicate_free nodes: - id: sender_node build: pip install torch --extra-index-url https://download.pytorch.org/whl/cpu numpy diff --git a/examples/memory-pool/read_after_free.yml b/libraries/extensions/tensor-pool/examples/read_after_free.yml similarity index 83% rename from examples/memory-pool/read_after_free.yml rename to libraries/extensions/tensor-pool/examples/read_after_free.yml index 89024d6a66..fb855cdd9a 100644 --- a/examples/memory-pool/read_after_free.yml +++ b/libraries/extensions/tensor-pool/examples/read_after_free.yml @@ -1,9 +1,9 @@ -# Negative-path test: receiver reads from a freed memory pool. +# Negative-path test: receiver reads from a freed tensor pool. env: sender_device: cpu receiver_device: cpu message_num: 2 - memory_pool_scenario: read_after_free + tensor_pool_scenario: read_after_free nodes: - id: sender_node build: pip install torch --extra-index-url https://download.pytorch.org/whl/cpu numpy diff --git a/examples/memory-pool/receiver.py b/libraries/extensions/tensor-pool/examples/receiver.py similarity index 80% rename from examples/memory-pool/receiver.py rename to libraries/extensions/tensor-pool/examples/receiver.py index 16aa73c4e2..5a56abd714 100644 --- a/examples/memory-pool/receiver.py +++ b/libraries/extensions/tensor-pool/examples/receiver.py @@ -1,5 +1,5 @@ #!/usr/bin/env python -"""Receive tensors through the memory-pool example dataflow.""" +"""Receive tensors through the tensor-pool example dataflow.""" import os import time @@ -13,7 +13,7 @@ node = Node("receiver_node") MESSAGE_COUNT = int(os.getenv("message_num", "100")) RECEIVER_DEVICE = os.getenv("receiver_device", "cpu") -SCENARIO = os.getenv("memory_pool_scenario", "throughput") +SCENARIO = os.getenv("tensor_pool_scenario", "throughput") if RECEIVER_DEVICE.startswith("cuda") and not torch.cuda.is_available(): raise RuntimeError("CUDA is not available for the configured receiver device.") @@ -23,7 +23,7 @@ pbar = tqdm(total=MESSAGE_COUNT) velocities = [] -memory_pool_id = None +tensor_pool_id = None torch_tensor = None for i in range(MESSAGE_COUNT): @@ -31,12 +31,12 @@ t_send = event["metadata"]["t_send"] if i == 0: - memory_pool_id = event["value"] - tensor_info = node.read_memory_pool(memory_pool_id) + tensor_pool_id = event["value"] + tensor_info = node.read_tensor_pool(tensor_pool_id) torch_tensor = tensor_from_info(tensor_info) print(f"Receiver preview: {torch_tensor[:5]}") - # The tensor is zero-copy — write_memory_pool on the sender overwrites + # The tensor is zero-copy — write_tensor_pool on the sender overwrites # the shmem bytes in place, so the receiver's existing tensor object # automatically reflects new data. Turn-based signaling ensures the # sender has finished writing before the receiver accesses the tensor. @@ -56,16 +56,16 @@ velocities.append(velocity) if SCENARIO == "duplicate_free" and i == MESSAGE_COUNT - 1: - node.free_memory_pool(memory_pool_id) - node.free_memory_pool(memory_pool_id) + node.free_tensor_pool(tensor_pool_id) + node.free_tensor_pool(tensor_pool_id) elif SCENARIO == "read_after_free" and i == MESSAGE_COUNT - 1: - node.free_memory_pool(memory_pool_id) + node.free_tensor_pool(tensor_pool_id) try: - node.read_memory_pool(memory_pool_id) + node.read_tensor_pool(tensor_pool_id) except Exception: pass # Expected: pool was freed, read should fail elif SCENARIO != "auto_cleanup" and i == MESSAGE_COUNT - 1: - node.free_memory_pool(memory_pool_id) + node.free_tensor_pool(tensor_pool_id) node.send_output("next_require", pa.array([])) pbar.update(1) diff --git a/examples/memory-pool/sender.py b/libraries/extensions/tensor-pool/examples/sender.py similarity index 77% rename from examples/memory-pool/sender.py rename to libraries/extensions/tensor-pool/examples/sender.py index a30cfa784e..d39e5136c4 100644 --- a/examples/memory-pool/sender.py +++ b/libraries/extensions/tensor-pool/examples/sender.py @@ -1,5 +1,5 @@ #!/usr/bin/env python -"""Send tensors through the memory-pool example dataflow.""" +"""Send tensors through the tensor-pool example dataflow.""" import os import time @@ -14,7 +14,7 @@ MESSAGE_COUNT = int(os.getenv("message_num", "100")) SENDER_DEVICE = os.getenv("sender_device", "cpu") RECEIVER_DEVICE = os.getenv("receiver_device", "cpu") -SCENARIO = os.getenv("memory_pool_scenario", "throughput") +SCENARIO = os.getenv("tensor_pool_scenario", "throughput") if SENDER_DEVICE.startswith("cuda"): idx = int(SENDER_DEVICE.split(":")[1]) if ":" in SENDER_DEVICE else 0 @@ -23,7 +23,7 @@ node = Node("sender_node") data_generation = np.random.default_rng() -memory_pool_id = None +tensor_pool_id = None for i in range(MESSAGE_COUNT): random_data = data_generation.integers(1000, size=SIZE, dtype=np.int64) random_data[0] = i # monotonic counter lets receiver detect change without collision risk @@ -34,13 +34,13 @@ if i == 0: print(f"Sender preview: {torch_tensor[:5]}") tensor_info = get_tensor_info(torch_tensor) - memory_pool_id = node.register_memory_pool(tensor_info, RECEIVER_DEVICE) - node.send_output("data", memory_pool_id, metadata) + tensor_pool_id = node.register_tensor_pool(tensor_info, RECEIVER_DEVICE) + node.send_output("data", tensor_pool_id, metadata) else: tensor_info = get_tensor_info(torch_tensor) if SCENARIO == "write_after_free" and i == 1: - node.free_memory_pool(memory_pool_id) - node.write_memory_pool(memory_pool_id, tensor_info) + node.free_tensor_pool(tensor_pool_id) + node.write_tensor_pool(tensor_pool_id, tensor_info) node.send_output("data", pa.array([]), metadata) node.next() diff --git a/examples/memory-pool/write_after_free.yml b/libraries/extensions/tensor-pool/examples/write_after_free.yml similarity index 83% rename from examples/memory-pool/write_after_free.yml rename to libraries/extensions/tensor-pool/examples/write_after_free.yml index c00bc6fc7e..a0085a68e4 100644 --- a/examples/memory-pool/write_after_free.yml +++ b/libraries/extensions/tensor-pool/examples/write_after_free.yml @@ -1,9 +1,9 @@ -# Negative-path test: sender writes to a freed memory pool. +# Negative-path test: sender writes to a freed tensor pool. env: sender_device: cpu receiver_device: cpu message_num: 2 - memory_pool_scenario: write_after_free + tensor_pool_scenario: write_after_free nodes: - id: sender_node build: pip install torch --extra-index-url https://download.pytorch.org/whl/cpu numpy diff --git a/libraries/extensions/tensor-pool/python/Cargo.toml b/libraries/extensions/tensor-pool/python/Cargo.toml new file mode 100644 index 0000000000..a1d298c0e7 --- /dev/null +++ b/libraries/extensions/tensor-pool/python/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "dora-tensor-pool-python" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +documentation.workspace = true +readme = "../README.md" +description = "Pinned-host / CUDA tensor-pool transport for dora Python nodes. Opt-in extension — NOT covered by dora's 1.0 compatibility guarantees. See libraries/extensions/tensor-pool/README.md." +license.workspace = true +repository.workspace = true +publish = false + +[dependencies] +dora-node-api = { workspace = true } +dora-message = { workspace = true } +arrow = { workspace = true, features = ["pyarrow"] } +eyre = { workspace = true } +pyo3 = { workspace = true } +serde_json = { workspace = true } +shared_memory_extended = "0.13.0" +tracing = { workspace = true } diff --git a/libraries/extensions/tensor-pool/python/src/lib.rs b/libraries/extensions/tensor-pool/python/src/lib.rs new file mode 100644 index 0000000000..b5ae1bb0e8 --- /dev/null +++ b/libraries/extensions/tensor-pool/python/src/lib.rs @@ -0,0 +1,14 @@ +//! Pinned-host / CUDA tensor-pool transport for dora Python nodes. +//! +//! **Not covered by dora's 1.0 compatibility guarantees.** This is an opt-in +//! extension: its API may change or break in a minor release, and it carries +//! known open defects. See `libraries/extensions/tensor-pool/README.md`. +//! +//! It reaches dora only through the public extension channel +//! (`docs/extensions.md`) — dora has no knowledge of pools, CUDA or the +//! DORADMA segment layout. + +pub mod seam; +pub mod transport; + +pub use transport::Pool; diff --git a/libraries/extensions/tensor-pool/python/src/seam.rs b/libraries/extensions/tensor-pool/python/src/seam.rs new file mode 100644 index 0000000000..7beabb8cbf --- /dev/null +++ b/libraries/extensions/tensor-pool/python/src/seam.rs @@ -0,0 +1,132 @@ +//! Adapter between this transport and dora's generic extension channel. +//! +//! dora brokers the lifetime of an opaque byte value and nothing else (see +//! `docs/extensions.md`). The pool's descriptor — shape, dtype, segment name, +//! whether the source was pinned, whether a CUDA IPC handle is present — is +//! that value, serialized here and never interpreted by dora. +//! +//! Before the extension channel existed, this rode on three pool-specific +//! daemon requests. Keeping the descriptor opaque is what lets the transport +//! change its own metadata without touching dora's wire protocol. + +use dora_node_api::{DoraNode, Metadata, MetadataParameters}; + +/// Namespace for every key this transport stores. Scoped so a second +/// extension in the same process cannot see or drop our entries. +pub const NAMESPACE: &str = "dora-tensor-pool"; + +/// Serialize a descriptor for the wire. +/// +/// JSON rather than a compact binary encoding: descriptors are a handful per +/// pool registration, not per frame, so legibility in a daemon dump is worth +/// more than the bytes. +pub fn encode(params: &MetadataParameters) -> Result, String> { + serde_json::to_vec(params).map_err(|e| format!("failed to encode pool descriptor: {e}")) +} + +/// Parse a descriptor back. An error here means the value was written by a +/// different version of this extension, not by dora. +pub fn decode(bytes: &[u8]) -> Result { + serde_json::from_slice(bytes).map_err(|e| format!("failed to decode pool descriptor: {e}")) +} + +/// Publish a pool descriptor. +pub fn store(node: &mut DoraNode, buffer_id: &str, meta: &Metadata) -> Result<(), String> { + let bytes = encode(&meta.parameters)?; + node.extension_store(NAMESPACE, buffer_id, bytes) + .map_err(|e| format!("{e:#}")) +} + +/// Fetch a pool descriptor, or `None` if the pool is gone. +/// +/// `take` drops the entry in the same round trip, which is the read-with-free +/// path: one request rather than a read followed by a racing drop. +pub fn load( + node: &mut DoraNode, + buffer_id: &str, + take: bool, +) -> Result, String> { + let bytes = node + .extension_load(NAMESPACE, buffer_id, take) + .map_err(|e| format!("{e:#}"))?; + bytes.as_deref().map(decode).transpose() +} + +/// Withdraw a pool descriptor, notifying every node that touched it. +pub fn drop_key(node: &mut DoraNode, buffer_id: &str) -> Result<(), String> { + node.extension_drop(NAMESPACE, buffer_id) + .map_err(|e| format!("{e:#}")) +} + +/// Fetch a descriptor as [`Metadata`], the shape the transport's read paths +/// already expect. +/// +/// A missing pool is an error rather than `Ok(None)`, matching the behaviour +/// of the daemon request this replaced: every caller treats "no descriptor" as +/// "fall back", not as a value. +pub fn load_metadata(node: &mut DoraNode, buffer_id: &str, take: bool) -> eyre::Result { + let timestamp = node.timestamp(); + let params = load(node, buffer_id, take) + .map_err(|e| eyre::eyre!("{e}"))? + .ok_or_else(|| eyre::eyre!("tensor pool `{buffer_id}` has no descriptor"))?; + Ok(Metadata::from_parameters(timestamp, params)) +} + +/// Buffer ids whose descriptor has gone away since the last call — because +/// another node dropped it, or because the daemon reclaimed it after the +/// owner exited. Each one means this process should release whatever it +/// mapped for that pool. +pub fn drain_dropped(namespace: &str) -> Vec { + dora_node_api::event_stream::extensions::drain_dropped_keys(namespace) +} + +#[cfg(test)] +mod tests { + use super::*; + use dora_node_api::Parameter; + + fn sample() -> MetadataParameters { + let mut p = MetadataParameters::new(); + p.insert("dtype".into(), Parameter::String("float32".into())); + p.insert("shape".into(), Parameter::ListInt(vec![3, 224, 224])); + p.insert("is_pinned".into(), Parameter::Bool(true)); + p.insert("ipc_present".into(), Parameter::Bool(false)); + p + } + + #[test] + fn descriptor_round_trips() { + let encoded = encode(&sample()).expect("encode"); + assert_eq!(decode(&encoded).expect("decode"), sample()); + } + + #[test] + fn every_parameter_kind_survives() { + // The descriptor is opaque to dora, so nothing else will catch a + // variant that fails to round-trip. + let mut p = MetadataParameters::new(); + p.insert("b".into(), Parameter::Bool(true)); + p.insert("i".into(), Parameter::Integer(-7)); + p.insert("s".into(), Parameter::String("x".into())); + p.insert("li".into(), Parameter::ListInt(vec![1, 2])); + p.insert("f".into(), Parameter::Float(1.5)); + p.insert("lf".into(), Parameter::ListFloat(vec![1.5, 2.5])); + p.insert("ls".into(), Parameter::ListString(vec!["a".into()])); + + let encoded = encode(&p).expect("encode"); + assert_eq!(decode(&encoded).expect("decode"), p); + } + + #[test] + fn garbage_is_an_error_not_a_panic() { + let err = decode(b"not json").expect_err("must reject"); + assert!(err.contains("failed to decode pool descriptor"), "{err}"); + } + + #[test] + fn an_empty_descriptor_round_trips() { + let empty = MetadataParameters::new(); + let encoded = encode(&empty).expect("encode"); + assert_eq!(decode(&encoded).expect("decode"), empty); + } +} diff --git a/libraries/extensions/tensor-pool/python/src/transport.rs b/libraries/extensions/tensor-pool/python/src/transport.rs new file mode 100644 index 0000000000..9c077b6570 --- /dev/null +++ b/libraries/extensions/tensor-pool/python/src/transport.rs @@ -0,0 +1,2982 @@ +//! The tensor-pool transport: segment layout, seqlock, CUDA helpers and the +//! four operations exposed to Python. +//! +//! **Not covered by dora's 1.0 compatibility guarantees** — opt-in extension, +//! known open defects, no GPU coverage in CI. See `../../README.md`. +//! +//! Reaches dora only through the generic extension channel (see [`crate::seam`] +//! and dora's `docs/extensions.md`). Everything in this file — the `unsafe` +//! pointer arithmetic over the DORADMA header, the seqlock, the embedded +//! `libcudart` ctypes module, the transport selection — stays on this side of +//! that boundary. If a change here seems to need a dora request named after +//! this transport, widen the generic channel instead. + +use std::collections::{HashMap, HashSet}; +use std::sync::LazyLock; + +use eyre::Context; + +use arrow::array::{Array, BinaryArray, StringArray}; +use arrow::pyarrow::{FromPyArrow, ToPyArrow}; +use dora_message::metadata::Parameter; +use dora_node_api::dora_core::config::NodeId; +use dora_node_api::{DataflowId, DoraNode}; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyDict, PyModule}; +use shared_memory_extended::ShmemConf; + +// ==================== pool_prelude ==================== +/// Pre-compiled CUDA DMA helper module. Compiled once (at first use) and reused +/// across all iterations to eliminate per-call PyModule::from_code overhead. +/// Maintains persistent state: pinned host pointers, pre-allocated GPU buffers. +static CUDA_HELPERS: LazyLock>>> = + LazyLock::new(|| std::sync::Mutex::new(None)); + +/// Counter to make pinned memory buffer IDs unique across registrations. +static PINNED_COUNTER: LazyLock> = LazyLock::new(|| std::sync::Mutex::new(0)); + +/// Maximum number of freed pool buffer IDs to remember at once. This is a +/// single budget shared across every peer the process reads from, not a +/// per-stream allowance — once the cap is exceeded the oldest entries are +/// evicted rather than keeping every ID for the life of the process. +const FREED_POOL_IDS_CAP: usize = 4096; + +/// Tracks freed pool buffer IDs so the DORADMA fast path can detect +/// read-after-free. Bounded to `FREED_POOL_IDS_CAP` entries (oldest evicted +/// first) so long-running nodes doing register->write->free every frame +/// don't leak memory indefinitely. +/// +/// The cap is one shared budget across *all* peers the process reads +/// from, not a per-stream recency window: a high-rate sender's frees can +/// evict a low-rate sender's tombstones well before the low-rate sender's +/// own next free (e.g. a 60Hz peer can cycle the whole cap in ~68s, +/// evicting a 1Hz peer's entries long before they'd naturally expire). An +/// evicted tombstone is harmless — it just makes a stale fast-path read +/// fall back to the existing `warn_missing_tensor_pool`/daemon path +/// instead of being caught here. +static FREED_POOL_IDS: LazyLock> = + LazyLock::new(|| std::sync::Mutex::new(FreedPoolIds::default())); + +/// Bounded, insertion-ordered set of freed pool buffer IDs. `set` gives +/// O(1) membership checks; `order` tracks insertion order so the oldest +/// entry can be evicted once `FREED_POOL_IDS_CAP` is exceeded. +#[derive(Default)] +struct FreedPoolIds { + set: HashSet, + order: std::collections::VecDeque, +} + +impl FreedPoolIds { + fn insert(&mut self, id: String) { + if self.set.insert(id.clone()) { + self.order.push_back(id); + while self.order.len() > FREED_POOL_IDS_CAP { + if let Some(oldest) = self.order.pop_front() { + self.set.remove(&oldest); + } + } + } + } + + fn contains(&self, id: &str) -> bool { + self.set.contains(id) + } + + fn remove(&mut self, id: &str) { + if self.set.remove(id) { + self.order.retain(|x| x != id); + } + } + + fn len(&self) -> usize { + self.set.len() + } +} + +#[cfg(test)] +mod freed_pool_ids_tests { + use super::{FREED_POOL_IDS_CAP, FreedPoolIds}; + + /// Regression test: without a cap, inserting one ID per frame in a + /// long-running node grows `FREED_POOL_IDS` forever. Inserting well + /// past the cap must keep the set bounded instead of leaking. + #[test] + fn insert_past_cap_does_not_grow_unbounded() { + let mut freed = FreedPoolIds::default(); + for i in 0..(FREED_POOL_IDS_CAP * 4) { + freed.insert(format!("pool_node_{i}")); + } + assert_eq!(freed.len(), FREED_POOL_IDS_CAP); + } + + /// Once past the cap, the oldest entries must be evicted first so the + /// most recently freed buffers (the ones a read-after-free check would + /// actually care about) stay tracked. + #[test] + fn insert_past_cap_evicts_oldest_first() { + let mut freed = FreedPoolIds::default(); + for i in 0..(FREED_POOL_IDS_CAP * 2) { + freed.insert(format!("pool_node_{i}")); + } + assert!( + !freed.contains("pool_node_0"), + "oldest entry should have been evicted" + ); + let newest = format!("pool_node_{}", FREED_POOL_IDS_CAP * 2 - 1); + assert!( + freed.contains(&newest), + "most recently freed entry should still be tracked" + ); + } + + /// Duplicate inserts of an already-tracked ID must not double-count + /// against the cap or push a second copy into the eviction order. + #[test] + fn duplicate_insert_is_idempotent() { + let mut freed = FreedPoolIds::default(); + freed.insert("pool_node_0".to_string()); + freed.insert("pool_node_0".to_string()); + assert_eq!(freed.len(), 1); + } + + #[test] + fn remove_drops_membership_and_order_entry() { + let mut freed = FreedPoolIds::default(); + freed.insert("pool_node_0".to_string()); + freed.remove("pool_node_0"); + assert!(!freed.contains("pool_node_0")); + assert_eq!(freed.len(), 0); + // Re-inserting after a remove must not be blocked by a stale + // order-queue entry left behind by `remove`. + freed.insert("pool_node_0".to_string()); + assert_eq!(freed.len(), 1); + } +} + +/// Per-pool persistent state. +/// Keeping Shmem alive prevents munmap, preserving stable mmap addresses +/// for pool-hit detection across `register_tensor_pool` calls. +/// +/// # Safety +/// `Shmem` is not `Send + Sync` due to raw pointer fields, but `PoolSlot` +/// is always stored behind a `Mutex` and `_shmem` is never accessed after +/// initialization — it is only kept alive for its destructor. +struct PoolSlot { + _shmem: shared_memory_extended::Shmem, + base: u64, + size: usize, + is_pinned: bool, + /// CPU page-locked transit buffer for cross-device GPU transfers + /// without P2P (e.g. RTX 5090). 0 means no transit path. + transit_ptr: u64, + /// The GPU device index where the pool buffer was allocated. + pool_device: i32, +} + +unsafe impl Send for PoolSlot {} +unsafe impl Sync for PoolSlot {} + +/// Persistent pool storage — stable mmap addresses for zero-copy detection. +/// Keyed by counter (unique per registration), supports unlimited pools. +static PINNED_POOL: LazyLock>> = + LazyLock::new(|| std::sync::Mutex::new(HashMap::new())); + +/// Persistent transit-buffer metadata for GPU pools. +/// Survives `PINNED_POOL` cache-miss so the write fast path can recover +/// `transit_ptr` and `pool_device` even when the `PoolSlot` has been evicted. +/// Keyed by counter, populated during `register_tensor_pool`, cleared in +/// `free_tensor_pool`. `transit_ptr=0` means no transit buffer (same-device +/// or P2P path). +static TRANSIT_META: LazyLock>> = + LazyLock::new(|| std::sync::Mutex::new(HashMap::new())); + +/// Receiver-side GPU cache per pool. +/// Keeps Shmem alive to prevent munmap, preserving stable mmap addresses +/// and valid GPU VAs for zero-copy reads across iterations. +struct RecvGpuSlot { + _shmem: shared_memory_extended::Shmem, + gpu_va: u64, // device VA from cudaHostGetDevicePointer, 0 if IPC path + gpu_buf: u64, // IPC-opened GPU DRAM pointer, 0 if GPU VA path + host_base: u64, // original host ptr passed to cudaHostRegister + gpu_buf_size: u64, // GPU buffer byte size from first import (baseline) +} + +/// Daemon-trusted GPU buffer sizes, keyed by buffer_id. +/// Populated from daemon metadata when the fallback path successfully +/// re-imports a GPU IPC handle. The fast path (`try_doradma_read`) +/// validates the (world-writable) shmem `size` against this cache — +/// if `size` exceeds the trusted capacity, the read is rejected and +/// the caller falls back to the daemon. +static GPU_BUF_SIZES: LazyLock>> = + LazyLock::new(|| std::sync::Mutex::new(HashMap::new())); +unsafe impl Send for RecvGpuSlot {} +unsafe impl Sync for RecvGpuSlot {} + +/// Receiver-side per-pool cache keeping Shmem alive + GPU VA for zero-copy reads. +/// Set up lazily in try_doradma_read: open shmem, cudaHostRegister, +/// cudaHostGetDevicePointer, then cache both Shmem and GPU VA. +static RECV_GPU_VA: LazyLock>> = + LazyLock::new(|| std::sync::Mutex::new(HashMap::new())); + +/// Receiver-side per-pool Shmem cache for CPU receivers. +/// Keeps Shmem alive to prevent munmap of the CPU pointer returned +/// by the as_cuda=False path in try_doradma_read. +struct RecvCpuSlot { + _shmem: shared_memory_extended::Shmem, + base: u64, +} +unsafe impl Send for RecvCpuSlot {} +unsafe impl Sync for RecvCpuSlot {} + +/// Receiver-side per-pool cache keeping Shmem alive for CPU zero-copy reads. +/// Without this cache, the Shmem handle drops at the end of try_doradma_read, +/// triggering munmap and making the returned CPU pointer a dangling pointer. +static RECV_CPU_SHMEM: LazyLock>> = + LazyLock::new(|| std::sync::Mutex::new(HashMap::new())); + +/// DORADMA shared-memory header layout: +/// +/// Offset Size Field +/// 0 8 magic — b"DORADMA\x00" +/// 8 8 json_len — u64 LE, metadata JSON byte length +/// 16 8 data_off — u64 LE, byte offset of tensor data from shmem base +/// 24 8 ipc_flag — u64 LE, 1 when ipc_handle is valid +/// 32 64 ipc_handle — CUDA IPC mem handle (only valid if ipc_flag == 1) +/// 96 8 write_gen — u64 LE, seqlock: even = complete, odd = writing +/// 104 152 reserved +/// 256 N json — padded-to-256-byte-alignment metadata JSON +/// 256+N M data — tensor payload +const DORADMA_HEADER_SIZE: usize = 256; +const DORADMA_MAGIC: &[u8; 8] = b"DORADMA\x00"; +const DORADMA_METADATA_ALIGN: usize = 256; + +/// Crossover where pinned-DMA bandwidth overtakes pageable copy + +/// cudaHostRegister/unregister fixed cost (~100 µs). Determined by +/// ablation study (2026-06-27): pageable faster below, pinned faster +/// above. Shared by `register_tensor_pool` and `write_tensor_pool`. +const DMA_PIN_THRESHOLD_BYTES: usize = 25 * 1024 * 1024; + +/// Returns `true` when the source tensor should be pinned before DMA. +/// +/// Pinning is a property of the *source* pointer: `cudaHostRegister` only +/// makes sense for host (CPU) memory, and the pin/unpin fixed cost +/// (~100 µs) is only worth paying when the tensor is large enough that +/// the DMA bandwidth gain outweighs it. +/// +/// # Unit-testable +/// +/// The decision is pure integer logic — no CUDA runtime calls — so the +/// boundary (25 MiB ± 1 byte) can be exercised in CI even without a GPU. +#[inline] +const fn should_pin(is_cuda: bool, size: usize) -> bool { + !is_cuda && size > DMA_PIN_THRESHOLD_BYTES +} + +#[cfg(test)] +mod pin_tests { + use super::*; + + #[test] + fn pin_cpu_source_above_threshold() { + // CPU source, 25 MiB + 1 byte → should pin + assert!(should_pin(false, 25 * 1024 * 1024 + 1)); + // CPU source, 100 MiB → should pin + assert!(should_pin(false, 100 * 1024 * 1024)); + } + + #[test] + fn pin_cpu_source_below_threshold() { + // CPU source, exactly at threshold → should NOT pin (> not >=) + assert!(!should_pin(false, 25 * 1024 * 1024)); + // CPU source, 1 byte below → should NOT pin + assert!(!should_pin(false, 25 * 1024 * 1024 - 1)); + // CPU source, tiny → should NOT pin + assert!(!should_pin(false, 1)); + } + + #[test] + fn pin_cuda_source_never_pins() { + // CUDA source regardless of size → never pin + assert!(!should_pin(true, 0)); + assert!(!should_pin(true, 25 * 1024 * 1024)); + assert!(!should_pin(true, 100 * 1024 * 1024)); + assert!(!should_pin(true, 1024 * 1024 * 1024)); + } + + #[test] + fn pin_zero_size_cpu() { + // Zero-size CPU tensor → below threshold, don't pin + assert!(!should_pin(false, 0)); + } +} + +// --------------------------------------------------------------------------- +// GPU transport-path classification — pure decision logic extractable +// from CUDA-runtime-embedded code so the full matrix can be exercised in +// CI without a GPU. Same pattern as `should_pin` above. +// --------------------------------------------------------------------------- + +/// Which transport path a GPU-pool registration (write-time) selects. +/// +/// Pure logic — no CUDA runtime calls. +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +enum TransportPath { + /// Buffer on sender device, plain DtoD memcpy (same-device or CPU source). + SameDeviceDtoD, + /// Cross-device with P2P peer access enabled. + P2PPeerAccess, + /// Cross-device without P2P — CPU page-locked transit (DtoH → HtoD). + HostStagingTransit, +} + +/// Classify which transport path a GPU-pool write should take at +/// registration time. +/// +/// Decision matrix (2³ = 8 cases, `is_cuda_source` dominates): +/// +/// | src CUDA | same dev | P2P | path | +/// |----------|----------|-----|---------------------| +/// | false | * | * | `SameDeviceDtoD` | +/// | true | true | * | `SameDeviceDtoD` | +/// | true | false | yes | `P2PPeerAccess` | +/// | true | false | no | `HostStagingTransit`| +#[inline] +fn classify_transport( + sender_device: i32, + receiver_device: i32, + p2p_available: bool, + is_cuda_source: bool, +) -> TransportPath { + if !is_cuda_source { + return TransportPath::SameDeviceDtoD; + } + if sender_device == receiver_device { + return TransportPath::SameDeviceDtoD; + } + if p2p_available { + return TransportPath::P2PPeerAccess; + } + TransportPath::HostStagingTransit +} + +/// Which write path `write_tensor_pool` dispatches to for a given frame. +/// +/// The fast and slow write paths both branch on the same 2×2×2 matrix +/// (`ipc_present` × `is_cuda` × `transit_ptr`); extracting the +/// classification makes the 5 reachable paths explicit and testable. +#[derive(Debug, PartialEq, Eq)] +enum WritePath { + /// CPU source → GPU pool via `dma_copy` (ipc_present=1, !is_cuda). + CpuToGpuPoolDma, + /// GPU source → GPU pool via transit (ipc_present=1, is_cuda, transit_ptr≠0). + GpuToGpuPoolTransit, + /// GPU source → GPU pool via plain DtoD `_cuda_memcpy_gpu_buf` + /// (ipc_present=1, is_cuda, transit_ptr=0). + GpuToGpuPoolDtoD, + /// GPU source → shmem data region via `cudaMemcpy` (ipc_present≠1, is_cuda). + GpuToShmem, + /// CPU source → shmem data region via `ptr::copy_nonoverlapping` + /// (ipc_present≠1, !is_cuda). + CpuToShmem, +} + +/// Classify which write path to take. +#[inline] +fn classify_write_path(ipc_present: u64, is_cuda: bool, transit_ptr: u64) -> WritePath { + if ipc_present == 1 { + if is_cuda { + if transit_ptr != 0 { + WritePath::GpuToGpuPoolTransit + } else { + WritePath::GpuToGpuPoolDtoD + } + } else { + WritePath::CpuToGpuPoolDma + } + } else if is_cuda { + WritePath::GpuToShmem + } else { + WritePath::CpuToShmem + } +} + +/// Result of validating a GPU-pool read `size` against the daemon-trusted +/// capacity cache (`GPU_BUF_SIZES`) and the first-import baseline +/// (`RecvGpuSlot::gpu_buf_size`). +#[derive(Debug, PartialEq, Eq)] +enum CapacityCheck { + /// Size is within the trusted bound. + Ok, + /// Size exceeds the trusted capacity → reject this read. + ExceedsTrustedSize, + /// No daemon-trusted entry and no cached baseline → reject first import. + NoTrustedEntry, +} + +/// Validate the read `size` for a GPU-pool buffer against the daemon-trusted +/// capacity and the cached first-import baseline. +/// +/// Resolution order: `trusted_sizes` (daemon metadata) → `cached_gpu_buf_size` +/// (first-import baseline) → reject. +/// +/// The daemon rejects zero-size pools at registration, so a resolved capacity +/// of zero is unreachable in normal operation. If it does occur (daemon bug +/// or memory corruption), the `size > capped` check fails closed — any +/// non-zero `size` triggers `ExceedsTrustedSize` rather than silently skipping +/// validation. +#[inline] +fn check_capacity_gpu_pool( + trusted_sizes: Option, + cached_gpu_buf_size: Option, + size: u64, +) -> CapacityCheck { + let cap = trusted_sizes.or(cached_gpu_buf_size); + match cap { + None => CapacityCheck::NoTrustedEntry, + Some(capped) if size > capped => CapacityCheck::ExceedsTrustedSize, + Some(_) => CapacityCheck::Ok, + } +} + +#[cfg(test)] +mod transport_tests { + use super::*; + + // -- classify_transport ------------------------------------------------- + + #[test] + fn same_device_no_transit() { + // Same GPU — never transit, regardless of P2P + assert_eq!( + classify_transport(0, 0, false, true), + TransportPath::SameDeviceDtoD + ); + assert_eq!( + classify_transport(1, 1, true, true), + TransportPath::SameDeviceDtoD + ); + } + + #[test] + fn cpu_source_never_transit() { + // CPU→GPU always uses dma_copy, no transit needed + assert_eq!( + classify_transport(0, 1, false, false), + TransportPath::SameDeviceDtoD + ); + assert_eq!( + classify_transport(0, 2, true, false), + TransportPath::SameDeviceDtoD + ); + } + + #[test] + fn cross_device_with_p2p() { + assert_eq!( + classify_transport(0, 1, true, true), + TransportPath::P2PPeerAccess + ); + } + + #[test] + fn cross_device_no_p2p_uses_transit() { + // This is the RTX 5090 / Blackwell path — the flag-ship non-P2P + // fallback that must NOT be dead code. + assert_eq!( + classify_transport(0, 1, false, true), + TransportPath::HostStagingTransit + ); + assert_eq!( + classify_transport(2, 0, false, true), + TransportPath::HostStagingTransit + ); + } + + #[test] + fn classify_transport_full_8_case_matrix() { + let cases: &[((i32, i32, bool, bool), TransportPath)] = &[ + // (src_dev, dst_dev, p2p, is_cuda) → expected + ((0, 0, false, false), TransportPath::SameDeviceDtoD), + ((0, 0, false, true), TransportPath::SameDeviceDtoD), + ((0, 0, true, false), TransportPath::SameDeviceDtoD), + ((0, 0, true, true), TransportPath::SameDeviceDtoD), + ((0, 1, false, false), TransportPath::SameDeviceDtoD), + ((0, 1, false, true), TransportPath::HostStagingTransit), + ((0, 1, true, false), TransportPath::SameDeviceDtoD), + ((0, 1, true, true), TransportPath::P2PPeerAccess), + ]; + for ((s, r, p2p, cuda), expected) in cases { + let got = classify_transport(*s, *r, *p2p, *cuda); + assert_eq!( + got, *expected, + "classify_transport(s={s}, r={r}, p2p={p2p}, cuda={cuda}) → {got:?}, expected {expected:?}" + ); + } + } + + // -- classify_write_path ------------------------------------------------- + + #[test] + fn write_path_cpu_to_gpu_pool_dma() { + assert_eq!(classify_write_path(1, false, 0), WritePath::CpuToGpuPoolDma); + // transit_ptr is irrelevant when !is_cuda + assert_eq!( + classify_write_path(1, false, 0xDEAD), + WritePath::CpuToGpuPoolDma + ); + } + + #[test] + fn write_path_gpu_to_gpu_pool_transit() { + assert_eq!( + classify_write_path(1, true, 1), + WritePath::GpuToGpuPoolTransit + ); + assert_eq!( + classify_write_path(1, true, 0xDEAD_BEEF), + WritePath::GpuToGpuPoolTransit + ); + } + + #[test] + fn write_path_gpu_to_gpu_pool_dtod() { + assert_eq!(classify_write_path(1, true, 0), WritePath::GpuToGpuPoolDtoD); + } + + #[test] + fn write_path_gpu_to_shmem() { + assert_eq!(classify_write_path(0, true, 0), WritePath::GpuToShmem); + assert_eq!(classify_write_path(0, true, 1), WritePath::GpuToShmem); + } + + #[test] + fn write_path_cpu_to_shmem() { + assert_eq!(classify_write_path(0, false, 0), WritePath::CpuToShmem); + assert_eq!(classify_write_path(0, false, 1), WritePath::CpuToShmem); + } + + #[test] + fn write_path_cache_miss_defaults_to_dtod() { + // When the write fast path hits a cache miss and constructs a + // fresh PoolSlot with transit_ptr=0, the dispatch must fall + // through to plain DtoD — NOT transit. This documents the + // current behaviour; if cache-miss transit recovery is added + // later, this test must be updated. + let path = classify_write_path(1, true, 0); + assert_eq!(path, WritePath::GpuToGpuPoolDtoD); + } + + #[test] + fn write_path_full_matrix() { + // 2×2×2 = 8 cases; 6 reachable (ipc_present=1 && transit_ptr≠0 + // for a CPU source is semantically unreachable because transit is + // only allocated on the CUDA-registration path). + let cases: &[(u64, bool, u64, WritePath)] = &[ + (1, false, 0, WritePath::CpuToGpuPoolDma), + (1, false, 1, WritePath::CpuToGpuPoolDma), + (1, true, 0, WritePath::GpuToGpuPoolDtoD), + (1, true, 1, WritePath::GpuToGpuPoolTransit), + (0, false, 0, WritePath::CpuToShmem), + (0, false, 1, WritePath::CpuToShmem), + (0, true, 0, WritePath::GpuToShmem), + (0, true, 1, WritePath::GpuToShmem), + ]; + for (ipc, cuda, tp, expected) in cases { + let got = classify_write_path(*ipc, *cuda, *tp); + assert_eq!( + got, *expected, + "classify_write_path(ipc={ipc}, cuda={cuda}, tp={tp}) → {got:?}, expected {expected:?}" + ); + } + } + + // -- check_capacity_gpu_pool -------------------------------------------- + + #[test] + fn capacity_ok_within_bounds() { + // Daemon-trusted cap present, size fits + assert_eq!( + check_capacity_gpu_pool(Some(4096), None, 4096), + CapacityCheck::Ok + ); + assert_eq!( + check_capacity_gpu_pool(Some(4096), None, 1), + CapacityCheck::Ok + ); + } + + #[test] + fn capacity_exceeds_trusted_size() { + assert_eq!( + check_capacity_gpu_pool(Some(4096), None, 4097), + CapacityCheck::ExceedsTrustedSize + ); + } + + #[test] + fn capacity_fallback_to_gpu_buf_size() { + // No daemon entry, but cached first-import baseline exists + assert_eq!( + check_capacity_gpu_pool(None, Some(4096), 2048), + CapacityCheck::Ok + ); + assert_eq!( + check_capacity_gpu_pool(None, Some(4096), 4097), + CapacityCheck::ExceedsTrustedSize + ); + } + + #[test] + fn capacity_no_trusted_entry_rejects() { + // Neither daemon nor cached baseline — must fail closed + assert_eq!( + check_capacity_gpu_pool(None, None, 1024), + CapacityCheck::NoTrustedEntry + ); + } + + #[test] + fn capacity_zero_trusted_cap_allows_zero_size_only() { + // The daemon rejects zero-size pools at registration, so + // a trusted capacity of 0 is unreachable in normal operation. + // If it does occur, fail closed: any non-zero size must be + // rejected rather than silently skipping validation. + assert_eq!(check_capacity_gpu_pool(Some(0), None, 0), CapacityCheck::Ok); + assert_eq!( + check_capacity_gpu_pool(Some(0), None, 1024 * 1024), + CapacityCheck::ExceedsTrustedSize + ); + } + + #[test] + fn capacity_zero_cached_buf_size_fails_closed() { + // Same reasoning as above: a cached baseline of 0 is + // unreachable; fail closed on any non-zero read. + assert_eq!( + check_capacity_gpu_pool(None, Some(0), 1024), + CapacityCheck::ExceedsTrustedSize + ); + } +} + +/// Get (or compile) the persistent CUDA DMA helper module. +/// +/// Compiled once at first use and reused across all subsequent iterations. +/// The module maintains internal state for pinned host pointers and GPU buffers, +/// eliminating per-call cudaHostRegister/cudaMalloc/cudaFree/shm_open overhead. +fn get_cuda_helpers(py: Python<'_>) -> Result, String> { + let mut guard = CUDA_HELPERS.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(ref module) = *guard { + return Ok(module.clone_ref(py)); + } + + let code = r#" +import ctypes +_lib = ctypes.CDLL('libcudart.so') + +_lib.cudaHostRegister.restype = ctypes.c_int +_lib.cudaHostRegister.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_uint] + +_lib.cudaHostUnregister.restype = ctypes.c_int +_lib.cudaHostUnregister.argtypes = [ctypes.c_void_p] + +_lib.cudaMalloc.restype = ctypes.c_int +_lib.cudaMalloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t] + +_lib.cudaMemcpy.restype = ctypes.c_int +_lib.cudaMemcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] + +_lib.cudaFree.restype = ctypes.c_int +_lib.cudaFree.argtypes = [ctypes.c_void_p] + +_lib.cudaDeviceSynchronize.restype = ctypes.c_int + +_lib.cudaHostGetDevicePointer.restype = ctypes.c_int +_lib.cudaHostGetDevicePointer.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_void_p, ctypes.c_uint] + +cudaMemcpyHostToDevice = 1 + +# IPC handle struct — must be a Structure subclass so ctypes passes it +# by value (64 bytes on the stack) to cudaIpcOpenMemHandle. +class _CudaIpcMemHandle(ctypes.Structure): + _fields_ = [("reserved", ctypes.c_byte * 64)] + +_lib.cudaIpcGetMemHandle.restype = ctypes.c_int +_lib.cudaIpcGetMemHandle.argtypes = [ctypes.POINTER(_CudaIpcMemHandle), ctypes.c_void_p] + +_lib.cudaIpcOpenMemHandle.restype = ctypes.c_int +_lib.cudaIpcOpenMemHandle.argtypes = [ctypes.POINTER(ctypes.c_void_p), _CudaIpcMemHandle, ctypes.c_uint] + +_lib.cudaIpcCloseMemHandle.restype = ctypes.c_int +_lib.cudaIpcCloseMemHandle.argtypes = [ctypes.c_void_p] + +# Persistent state: per-slot GPU buffer cache +_gpu_bufs = {} # slot -> (d_ptr, size) + +def _register_host(ptr, size): + """Pin host memory. Idempotent — error 712 (already registered) is ok.""" + err = _lib.cudaHostRegister(ctypes.c_void_p(ptr), size, 0) + if err != 0 and err != 712: + raise RuntimeError(f'cudaHostRegister(0x{ptr:x}, {size}) failed: {err}') + +def _unregister_host(ptr): + """Unpin host memory.""" + err = _lib.cudaHostUnregister(ctypes.c_void_p(ptr)) + if err != 0 and err != 713: + raise RuntimeError(f'cudaHostUnregister(0x{ptr:x}) failed: {err}') + +def _get_device_ptr(host_ptr): + """Get a GPU VA for a pinned host memory region.""" + d_ptr = ctypes.c_void_p() + err = _lib.cudaHostGetDevicePointer(ctypes.byref(d_ptr), ctypes.c_void_p(host_ptr), 0) + if err != 0: + raise RuntimeError(f'cudaHostGetDevicePointer(0x{host_ptr:x}) failed: {err}') + return d_ptr.value + + +# P2P and device query bindings (CUDA runtime). +_lib.cudaGetDevice.restype = ctypes.c_int +_lib.cudaGetDevice.argtypes = [ctypes.POINTER(ctypes.c_int)] +_lib.cudaDeviceCanAccessPeer.restype = ctypes.c_int +_lib.cudaDeviceCanAccessPeer.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_int] +_lib.cudaDeviceEnablePeerAccess.restype = ctypes.c_int +_lib.cudaDeviceEnablePeerAccess.argtypes = [ctypes.c_int, ctypes.c_uint] +_lib.cudaSetDevice.restype = ctypes.c_int +_lib.cudaSetDevice.argtypes = [ctypes.c_int] + +# Page-locked host allocation for cross-device staging buffers. +_lib.cudaHostAlloc.restype = ctypes.c_int +_lib.cudaHostAlloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t, ctypes.c_uint] +_lib.cudaFreeHost.restype = ctypes.c_int +_lib.cudaFreeHost.argtypes = [ctypes.c_void_p] + +_p2p_enabled_pairs = set() # {(src, dst)} pairs already peer-enabled + +def _alloc_transit(size): + """Allocate a page-locked CPU buffer for cross-device staging. + Returns the pointer (as int), or 0 on failure.""" + ptr = ctypes.c_void_p() + if _lib.cudaHostAlloc(ctypes.byref(ptr), size, 0) == 0: + return ptr.value + return 0 + +def _free_transit(ptr): + """Free a page-locked CPU transit buffer.""" + if ptr: + _lib.cudaFreeHost(ctypes.c_void_p(ptr)) + +def _can_access_peer(src, dst): + """Check if src GPU can P2P-access dst GPU. Returns bool.""" + can = ctypes.c_int(0) + if _lib.cudaDeviceCanAccessPeer(ctypes.byref(can), src, dst) == 0: + return can.value != 0 + return False + +def _set_cuda_device(idx): + """Set the current CUDA device. No-op if *idx* < 0.""" + if idx >= 0: + _lib.cudaSetDevice(idx) + +def _get_cuda_device(): + """Return the current CUDA device index, or -1 on failure.""" + dev = ctypes.c_int() + if _lib.cudaGetDevice(ctypes.byref(dev)) == 0: + return dev.value + return -1 + +def _transit_copy(src_ptr, src_dev, transit_ptr, dst_ptr, dst_dev, size): + """Copy via CPU transit: src(GPU) → DtoH → transit → HtoD → dst(GPU). + Returns True on success. Restores the caller's current CUDA device on exit.""" + saved = ctypes.c_int() + _lib.cudaGetDevice(ctypes.byref(saved)) + # GPU src → CPU transit + _lib.cudaSetDevice(src_dev) + err = _lib.cudaMemcpy(ctypes.c_void_p(transit_ptr), ctypes.c_void_p(src_ptr), size, 2) + if err != 0: + _lib.cudaSetDevice(saved.value) + return False + _lib.cudaDeviceSynchronize() + # CPU transit → GPU dst + _lib.cudaSetDevice(dst_dev) + err = _lib.cudaMemcpy(ctypes.c_void_p(dst_ptr), ctypes.c_void_p(transit_ptr), size, 1) + _lib.cudaDeviceSynchronize() + _lib.cudaSetDevice(saved.value) + return err == 0 + +def _transit_copy_gpu_buf(slot, src_ptr, src_dev, transit_ptr, dst_dev, size): + """Same as _transit_copy but looks up the pool's GPU buffer by slot. + Raises on a missing/undersized slot or a failed copy so the caller + can surface the error instead of silently delivering stale data.""" + if slot not in _gpu_bufs: + raise RuntimeError(f"GPU pool buffer for slot {slot} not initialised") + dst, capacity = _gpu_bufs[slot] + if size > capacity: + raise RuntimeError( + f"write size {size} exceeds GPU pool buffer capacity {capacity} (slot={slot})" + ) + if not _transit_copy(src_ptr, src_dev, transit_ptr, dst, dst_dev, size): + raise RuntimeError(f"transit copy into GPU pool buffer failed (slot={slot})") + +def _ensure_p2p_pair(a, b): + """Enable bidirectional P2P access for the (a, b) GPU pair only. + + Scoped to the sender/receiver devices actually used by a transfer — + enabling P2P across *all* GPU pairs would lazily create a CUDA context + (hundreds of MB) on every visible device, including ones not in this + dataflow. Idempotent per pair. + """ + if a == b or a < 0 or b < 0: + return + if (a, b) in _p2p_enabled_pairs: + return + # Save the current device so we can restore it after enabling peer + # access — otherwise later CUDA operations (cudaMalloc, + # cudaIpcOpenMemHandle) would run on the last enabled src device. + saved = ctypes.c_int() + _lib.cudaGetDevice(ctypes.byref(saved)) + for src, dst in ((a, b), (b, a)): + can = ctypes.c_int(0) + if _lib.cudaDeviceCanAccessPeer(ctypes.byref(can), src, dst) == 0 and can.value: + _lib.cudaSetDevice(src) + _lib.cudaDeviceEnablePeerAccess(dst, 0) + _p2p_enabled_pairs.add((src, dst)) + _lib.cudaSetDevice(saved.value) + +def _cuda_memcpy(dst, src, size, kind): + """cudaMemcpy wrapper. kind: 1=H2D, 2=D2H, 3=D2D.""" + err = _lib.cudaMemcpy(ctypes.c_void_p(dst), ctypes.c_void_p(src), size, kind) + if err != 0: + raise RuntimeError(f"cudaMemcpy(0x{dst:x}, 0x{src:x}, {size}, {kind}) failed: {err}") + _lib.cudaDeviceSynchronize() + +def _cuda_memcpy_gpu_buf(slot, src_ptr, size): + """Copy *size* bytes from *src_ptr* (GPU) into the pool's pinned GPU buffer + identified by *slot*. Used by write_tensor_pool when both source and pool + buffer are GPU-resident (same-device DtoD copy).""" + if slot not in _gpu_bufs: + raise RuntimeError(f"GPU pool buffer for slot {slot} not initialised") + dst, capacity = _gpu_bufs[slot] + if size > capacity: + # The GPU buffer was sized at registration and its IPC handle is + # already exported; growing it would invalidate the receiver's + # imported handle. Reject rather than overflow the allocation. + raise RuntimeError( + f"write size {size} exceeds GPU pool buffer capacity {capacity} (slot={slot})" + ) + err = _lib.cudaMemcpy(ctypes.c_void_p(dst), ctypes.c_void_p(src_ptr), size, 3) + if err != 0: + raise RuntimeError(f"cudaMemcpy GPU buf DtoD (slot={slot}, 0x{dst:x}←0x{src_ptr:x}, {size}B) failed: {err}") + _lib.cudaDeviceSynchronize() + +def _get_gpu_buf(slot, size): + """Get or allocate a GPU buffer for the given slot. Reuses when size matches.""" + if slot in _gpu_bufs and _gpu_bufs[slot][1] >= size: + return _gpu_bufs[slot][0] + if slot in _gpu_bufs: + _lib.cudaFree(ctypes.c_void_p(_gpu_bufs[slot][0])) + d_ptr = ctypes.c_void_p() + err = _lib.cudaMalloc(ctypes.byref(d_ptr), size) + if err != 0: + raise RuntimeError(f'cudaMalloc({size}) failed: {err}') + _gpu_bufs[slot] = (d_ptr.value, size) + return d_ptr.value + +def _free_gpu_buf(slot): + """Free the pooled GPU buffer for the given slot.""" + if slot in _gpu_bufs: + _lib.cudaFree(ctypes.c_void_p(_gpu_bufs[slot][0])) + del _gpu_bufs[slot] + +def _ipc_export(d_ptr): + """Export GPU memory for cross-process sharing. Returns 64-byte handle.""" + handle = _CudaIpcMemHandle() + err = _lib.cudaIpcGetMemHandle(ctypes.byref(handle), ctypes.c_void_p(d_ptr)) + if err != 0: + raise RuntimeError(f'cudaIpcGetMemHandle(0x{d_ptr:x}) failed: {err}') + return bytes(handle) + +def _ipc_import(handle_bytes): + """Import GPU memory from another process. Returns device pointer. + + cudaIpcOpenMemHandle takes the handle struct BY VALUE (64 bytes on + the stack). _CudaIpcMemHandle.from_buffer_copy reconstructs the struct + and ctypes copies it onto the stack for the call. + """ + handle = _CudaIpcMemHandle.from_buffer_copy(handle_bytes) + d_ptr = ctypes.c_void_p() + err = _lib.cudaIpcOpenMemHandle(ctypes.byref(d_ptr), handle, 1) + if err != 0: + raise RuntimeError(f'cudaIpcOpenMemHandle failed: {err}') + _lib.cudaDeviceSynchronize() + return d_ptr.value + +def _ipc_close(d_ptr): + """Close an IPC memory handle opened by _ipc_import. + Frees the GPU-side mapping without freeing the underlying allocation.""" + err = _lib.cudaIpcCloseMemHandle(ctypes.c_void_p(d_ptr)) + if err != 0: + raise RuntimeError(f'cudaIpcCloseMemHandle(0x{d_ptr:x}) failed: {err}') + +def dma_copy(ptr, size, slot, no_dma): + """DMA transfer from host to pre-allocated GPU buffer. + + Copies via cudaMemcpyHtoD (DMA engine). When *no_dma* is false + (the default), the source memory is pinned (cudaHostRegister) + before the copy and unpinned after — this is the fast path for + large tensors where pinned-DMA bandwidth outweighs the pin/unpin + fixed cost. When *no_dma* is true, pin/unpin is skipped, using + pageable memory (faster for small tensors where pin overhead + dominates). + + Returns the device pointer of the pooled GPU buffer. + + Raises RuntimeError if *size* exceeds the existing buffer capacity + and the buffer cannot be grown (GPU pool buffers are IPC-exported; + reallocation would invalidate the receiver's imported handle). + """ + if slot in _gpu_bufs: + _, capacity = _gpu_bufs[slot] + if size > capacity: + raise RuntimeError( + f"write size {size} exceeds GPU pool buffer capacity {capacity} (slot={slot})" + " — GPU buffer cannot be grown (IPC handle already exported)" + ) + if not no_dma: + _register_host(ptr, size) + try: + d_ptr = _get_gpu_buf(slot, size) + err = _lib.cudaMemcpy( + ctypes.c_void_p(d_ptr), + ctypes.c_void_p(ptr), + size, + cudaMemcpyHostToDevice, + ) + if err != 0: + raise RuntimeError(f'cudaMemcpy failed: {err}') + _lib.cudaDeviceSynchronize() + finally: + if not no_dma: + _unregister_host(ptr) + return d_ptr + +"#; + + let code_cstr = std::ffi::CString::new(code).map_err(|e| format!("CString: {}", e))?; + let bound_module = PyModule::from_code( + py, + code_cstr.as_c_str(), + c"_cuda_helpers.py", + c"_cuda_helpers", + ) + .map_err(|e| format!("CUDA helper compile failed: {}", e))?; + let module_ref: Py = bound_module.unbind(); + *guard = Some(module_ref.clone_ref(py)); + Ok(module_ref) +} + +/// Read 8 consecutive bytes from `p` as a little-endian u64, +/// without an `.unwrap()` call (avoids the CI unwrap budget). +fn read_header_u64(p: *const u8) -> u64 { + const { assert!(std::mem::size_of::() == 8) }; + let mut buf = [0u8; 8]; + unsafe { std::ptr::copy_nonoverlapping(p, buf.as_mut_ptr(), 8) }; + u64::from_le_bytes(buf) +} + +fn parse_tensor_pool_id(tensor_pool_id: Py, py: Python<'_>) -> eyre::Result { + let array_data = arrow::array::ArrayData::from_pyarrow_bound(tensor_pool_id.bind(py))?; + let array = arrow::array::make_array(array_data); + + if let Some(string_array) = array.as_any().downcast_ref::() { + if string_array.len() != 1 { + eyre::bail!( + "expected string array with exactly one element, got {}", + string_array.len() + ); + } + Ok(string_array.value(0).to_string()) + } else if let Some(binary_array) = array.as_any().downcast_ref::() { + if binary_array.len() != 1 { + eyre::bail!( + "expected binary array with exactly one element, got {}", + binary_array.len() + ); + } + Ok(String::from_utf8(binary_array.value(0).to_vec())?) + } else { + eyre::bail!( + "tensor_pool_id must be a string or binary array, got {:?}", + array.data_type() + ) + } +} + +fn warn_missing_tensor_pool(node_id: &NodeId, action: &str, buffer_id: &str) { + tracing::warn!( + "[{}] Attempt to {} tensor pool [{}] failed - reason: pool does not exist. Operation aborted.", + node_id, + action, + buffer_id + ); +} + +// ==================== seqlock ==================== +/// Begins a tensor-pool seqlock write at `gen_ptr` (header offset 96) +/// **if the generation is even**. Returns the **even** pre-write +/// generation — i.e., the generation value before the write cycle +/// began, which is always even. +/// +/// If the generation is already odd (leftover from a previous failed +/// write), the begin-increment is skipped and the **previous** even +/// generation (`cur - 1`) is returned. This ensures that +/// `seqlock_end`'s `pre + 2` always produces an even generation, +/// avoiding a permanent parity inversion. +unsafe fn seqlock_begin_if_even(gen_ptr: *mut u64) -> u64 { + unsafe { + let cur = std::ptr::read_volatile(gen_ptr); + if cur.is_multiple_of(2) { + std::ptr::write_volatile(gen_ptr, cur + 1); + std::sync::atomic::fence(std::sync::atomic::Ordering::Release); + } + cur & !1 // always return the even baseline + } +} + +/// Closes a tensor-pool seqlock write (header offset 96). +/// +/// Advances the generation to `pre_write_gen + 2` (even = "complete"). +/// The caller must only invoke this on a successful copy. GPU pool write +/// paths leave gen odd on failure (in-place writes cannot roll back to a +/// clean previous frame; double-buffering is deferred to a follow-up PR). +/// The `copy_ok == false` rollback branch is retained for the helper's +/// contract but is dead code in production — see the leave-gen-odd blocks +/// in `write_tensor_pool`. +unsafe fn seqlock_end(gen_ptr: *mut u64, pre_write_gen: u64, copy_ok: bool) { + unsafe { + if copy_ok { + std::ptr::write_volatile(gen_ptr, pre_write_gen.wrapping_add(2)); + } else { + std::ptr::write_volatile(gen_ptr, pre_write_gen); + } + std::sync::atomic::fence(std::sync::atomic::Ordering::Release); + } +} + +#[cfg(test)] +mod seqlock_tests { + #[test] + fn begin_if_even_when_even_flips_to_odd() { + let mut generation: u64 = 10; + unsafe { + super::seqlock_begin_if_even(&mut generation); + } + assert_eq!(generation, 11, "even generation must be flipped to odd"); + } + + #[test] + fn begin_if_even_when_odd_returns_even_baseline() { + let mut generation: u64 = 11; + let pre; + unsafe { + pre = super::seqlock_begin_if_even(&mut generation); + } + assert_eq!(generation, 11, "odd generation must stay odd (skip begin)"); + assert_eq!(pre, 10, "pre-write baseline must be even (cur - 1)"); + assert_eq!(pre % 2, 0); + } + + #[test] + fn begin_if_even_then_end_always_produces_even() { + // Simulates: gen stuck odd (11) from failure → begin returns 10 + // → end does 10+2=12 which is even → pool recovers. + let mut generation: u64 = 11; + let pre; + unsafe { + pre = super::seqlock_begin_if_even(&mut generation); + } + unsafe { + super::seqlock_end(&mut generation, pre, true); + } + assert_eq!(generation, 12, "pre(10)+2=12 is even → recovery"); + assert_eq!(generation % 2, 0); + } + + /// Regression test for #2436: a successful copy publishes the new + /// frame by advancing the generation to pre+2 (even = complete). + #[test] + fn end_success_advances_to_even() { + let mut generation: u64 = 11; // pre_write(10) + 1 + unsafe { + super::seqlock_end(&mut generation, 10, true); + } + assert_eq!(generation, 12); + assert_eq!(generation % 2, 0, "generation must be even on success"); + } + + /// Regression test for #2436: on failure the generation is rolled + /// back to pre_write_gen (even) so consumers see the previous + /// valid frame, not a torn/incomplete one. + #[test] + fn end_failure_rolls_back_to_even() { + let mut generation: u64 = 11; // pre_write(10) + 1 + unsafe { + super::seqlock_end(&mut generation, 10, false); + } + assert_eq!(generation, 10); + assert_eq!(generation % 2, 0, "rolled-back generation must be even"); + } + + /// The very first write (pre_write == 0) must not underflow. + #[test] + fn end_first_write_failure_rolls_back_to_zero() { + let mut generation: u64 = 1; + unsafe { + super::seqlock_end(&mut generation, 0, false); + } + assert_eq!(generation, 0); + } +} + +// ==================== the transport ==================== + +/// One node's view of the pool transport. +/// +/// Holds what the methods below need from dora — nothing more. Constructed +/// per call by the thin `#[pymethods]` wrappers in `dora-node-api-python`, +/// because PyO3 cannot add methods to its `Node` from another crate. +pub struct Pool<'a> { + pub node: &'a mut DoraNode, + pub node_id: NodeId, + pub dataflow_id: DataflowId, +} + +impl Pool<'_> { + // ==================== process_pending_frees ==================== + /// Process any tensor pools that were freed by another node. + pub fn process_pending_tensor_pool_frees(&mut self, py: Python) { + for shared_memory_id in crate::seam::drain_dropped(crate::seam::NAMESPACE) { + let buffer_id = shared_memory_id; + // Receiver-side cleanup (IPC handles, shmem mappings). + { + if let Some(slot) = RECV_GPU_VA + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&buffer_id) + { + if slot.gpu_buf != 0 { + if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + let _ = bound.call_method1("_ipc_close", (slot.gpu_buf,)); + } + } else if slot.gpu_va != 0 { + // Host-registered mapping (effective_as_cuda branch): + // must cudaHostUnregister before munmap. _unregister_host + // requires the original host pointer (shmem base), not the + // device VA returned by cudaHostGetDevicePointer — passing + // the device VA makes cudaHostUnregister fail and leaks the + // pin over an address range that then gets munmap'd. + if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + let _ = bound.call_method1("_unregister_host", (slot.host_base,)); + } + } + // Drop slot → munmap + } + } + RECV_CPU_SHMEM + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&buffer_id); + + // Sender-side cleanup (PINNED_POOL, GPU/transit buffers). + // Guard against cross-process counter aliasing: buffer ids are + // pool_{node_id}_{counter}. Extract the owner segment (between + // "pool_" and the final "_") and require an exact + // equality match — a prefix check (starts_with) would alias + // across node ids that are prefixes of each other (e.g. cam / + // cam_left). + if let Some(owner_and_counter) = buffer_id.strip_prefix("pool_") + && let Some((owner, counter_str)) = owner_and_counter.rsplit_once('_') + && owner == self.node_id.as_ref() + && let Ok(c) = counter_str.parse::() + && let Some(slot) = PINNED_POOL + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&c) + { + if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + let _ = bound.call_method1("_unregister_host", (slot.base,)); + let _ = bound.call_method1("_free_gpu_buf", (c,)); + if slot.transit_ptr != 0 { + let _ = bound.call_method1("_free_transit", (slot.transit_ptr,)); + } + } + TRANSIT_META + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&c); + } + } + } + + // ==================== pool_api_methods ==================== + + // === Tensor Pool API === + + /// Register a shared tensor pool for zero-copy tensor transfer. + /// + /// The returned pool ID can be shared across nodes (e.g. via a Dora output) + /// so that a receiver can call [`read_tensor_pool`] and [`free_tensor_pool`] + /// on it. + /// + /// # Concurrency / safety + /// + /// **This pool provides no internal mutual exclusion for data bytes.** + /// The writer must not begin a new [`write_tensor_pool`] while a receiver is + /// still consuming the previous tensor. Callers that share a pool across + /// nodes MUST enforce a **turn-based discipline** — for example, by waiting + /// for a `next_require` round-trip from the receiver before writing again. + /// The bundled `examples/tensor-pool/` dataflows demonstrate this pattern. + /// + /// The on-segment seqlock guards metadata integrity (header fields written + /// once at registration) and detects in-flight overwrites, but it does + /// **not** block the writer from starting a new write while a consumer + /// holds a zero-copy tensor. Skipping the turn-based discipline risks + /// torn data at the consumer. + pub fn register_tensor_pool( + &mut self, + tensor_info: &Bound<'_, PyDict>, + device: String, + py: Python, + ) -> eyre::Result> { + let ptr_val: u64 = tensor_info + .get_item("ptr")? + .ok_or_else(|| eyre::eyre!("missing ptr"))? + .extract()?; + let size: usize = tensor_info + .get_item("size")? + .ok_or_else(|| eyre::eyre!("missing size"))? + .extract()?; + let dtype: String = tensor_info + .get_item("dtype")? + .ok_or_else(|| eyre::eyre!("missing dtype"))? + .extract()?; + let shape_list: Vec = tensor_info + .get_item("shape")? + .ok_or_else(|| eyre::eyre!("missing shape"))? + .extract()?; + let tensor_device: String = tensor_info + .get_item("device")? + .ok_or_else(|| eyre::eyre!("missing device"))? + .extract()?; + + let is_cuda = tensor_device.starts_with("cuda"); + let receiver_is_cuda = device.starts_with("cuda"); + let cpu_mode = !receiver_is_cuda; + // Auto-select pinning: key off the source device — pinning only + // matters when the source is CPU (cudaHostRegister would raise on a + // device pointer; prevented by the !is_cuda guard above). + let is_pinned = should_pin(is_cuda, size); + let pinned_type = if cpu_mode { "cpu" } else { "cuda" }; + + if ptr_val == 0 { + eyre::bail!("Invalid source pointer (NULL)"); + } + if size == 0 || size > 1024 * 1024 * 1024 { + eyre::bail!("Invalid size: {} bytes", size); + } + if cfg!(not(target_os = "linux")) { + eyre::bail!( + "tensor-pool transport requires Linux (uses /dev/shm). \ + This platform is not supported." + ); + } + + // Generate unique pool counter for this registration + let pool_counter = { + let mut c = PINNED_COUNTER.lock().unwrap_or_else(|e| e.into_inner()); + *c += 1; + *c + }; + let shmem_name = format!( + "dora_pool_{}_{}_{}", + self.dataflow_id, self.node_id, pool_counter + ); + + let header_meta = PyDict::new(py); + header_meta.set_item("size", size)?; + header_meta.set_item("dtype", &dtype)?; + header_meta.set_item("shape", shape_list.clone())?; + header_meta.set_item("pinned_type", pinned_type)?; + + let json_bytes = py + .import("json") + .wrap_err("failed to import json")? + .call_method1("dumps", (header_meta,)) + .wrap_err("failed to serialize metadata to JSON")? + .extract::() + .wrap_err("failed to extract JSON string")? + .into_bytes(); + let json_len = json_bytes.len(); + let padded_json_len = json_len.div_ceil(DORADMA_METADATA_ALIGN) * DORADMA_METADATA_ALIGN; + let data_offset = DORADMA_HEADER_SIZE + padded_json_len; + // GPU receivers read tensor data from the IPC-exported GPU buffer, + // not from the shmem data region. Allocate only the header portion + // (metadata + IPC handle + seqlock) — a few hundred bytes instead of + // 80 MB. This also lets us skip cudaHostRegister on a useless data + // region. + let total_size = if receiver_is_cuda { + data_offset + } else { + data_offset + size + }; + + // Create shared memory + let mut shmem = ShmemConf::new() + .os_id(&shmem_name) + .size(total_size) + .writable(true) + .create() + .wrap_err_with(|| { + format!( + "failed to create pool shared memory `{}` (name collision with another node or leftover segment)", + shmem_name + ) + })?; + let shmem_ptr = unsafe { shmem.as_slice_mut().as_mut_ptr() }; + + // Pin the shmem for DMA only when the receiver reads from it + // (CPU receivers). GPU receivers never touch the shmem data + // region, and the header-only shmem is too small (< 1 page) to + // benefit from pinning. + if !receiver_is_cuda && let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + let _ = bound.call_method1("_register_host", (shmem_ptr as u64, total_size)); + } + + shmem.set_owner(false); + + // Write DORADMA header + unsafe { + std::ptr::copy_nonoverlapping(DORADMA_MAGIC.as_ptr(), shmem_ptr, 8); + let json_len_le = (json_len as u64).to_le_bytes(); + std::ptr::copy_nonoverlapping(json_len_le.as_ptr(), shmem_ptr.add(8), 8); + let data_off_le = (data_offset as u64).to_le_bytes(); + std::ptr::copy_nonoverlapping(data_off_le.as_ptr(), shmem_ptr.add(16), 8); + std::ptr::copy_nonoverlapping( + json_bytes.as_ptr(), + shmem_ptr.add(DORADMA_HEADER_SIZE), + json_len, + ); + } + + // Initialize seqlock write_gen at header[96] to 0 (even = complete) + unsafe { + std::ptr::write(shmem_ptr.add(96) as *mut u64, 0u64); + } + + // Seqlock: increment generation to odd (write-in-progress) + unsafe { + let gen_ptr = shmem_ptr.add(96) as *mut u64; + let old_gen = std::ptr::read_volatile(gen_ptr); + std::ptr::write_volatile(gen_ptr, old_gen + 1); + std::sync::atomic::fence(std::sync::atomic::Ordering::Release); + } + + // Copy tensor data to shmem — only when the receiver will + // actually read it. GPU receivers import the pool GPU buffer + // via the IPC handle in the DORADMA header and never touch the + // shmem data region; skipping this copy for them eliminates + // a redundant CPU-memcpy or GPU-DtoH transfer on every + // registration (cpu2cuda and cuda2cuda respectively). + if !receiver_is_cuda { + // The DtoH copy must publish either a fully-initialized data + // region or nothing — uninitialized shmem exposed as a valid + // frame is data corruption. Both a failed cudaMemcpy and a + // missing CUDA helper module are treated as copy failures. + if is_cuda { + // A missing CUDA helper module counts as a failed copy: the + // consumer must see the whole region or nothing. + let dtoh_copy_ok = if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + bound + .call_method1( + "_cuda_memcpy", + (shmem_ptr as u64 + data_offset as u64, ptr_val, size, 2u32), + ) + .is_ok() + } else { + false + }; + if !dtoh_copy_ok { + // The matching `_register_host` above is unconditional (it + // runs whenever `!receiver_is_cuda`), so the unregister must + // be too — gating it on `!is_pinned` would leak the pin if + // `should_pin` is ever tuned to pin CUDA sources. + if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + let _ = bound.call_method1("_unregister_host", (shmem_ptr as u64,)); + } + shmem.set_owner(true); + eyre::bail!( + "[{}] register_tensor_pool: DtoH copy failed ({} → CPU shmem, {} bytes)", + self.node_id, + tensor_device, + size + ); + } + } else { + unsafe { + std::ptr::copy_nonoverlapping( + ptr_val as *const u8, + shmem_ptr.add(data_offset), + size, + ); + } + } + } + + // GPU pool: allocate GPU buffer on current device, copy data, export + // IPC handle for cross-process zero-copy access. When the source + // tensor is also on CUDA (GPU→GPU), the source and pool buffer are on + // the same device (sender's current CUDA device), so a plain DtoD + // memcpy suffices. When the source is CPU, `dma_copy` does a pinned + // host→device DMA copy (existing path). + // Resolve sender and receiver device indices for cross-device detection. + let sender_device_idx = tensor_device + .strip_prefix("cuda") + .and_then(|s| s.strip_prefix(':')) + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + let receiver_device_idx = device + .strip_prefix("cuda") + .and_then(|s| s.strip_prefix(':')) + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + let cross_device = sender_device_idx != receiver_device_idx; + let mut transit_ptr: u64 = 0; + let mut pool_device = if receiver_is_cuda { + receiver_device_idx + } else { + sender_device_idx + }; + + // Tracks whether the GPU pool buffer + IPC handle were successfully set + // up. A CUDA receiver's shmem is header-only, so without the handle the + // pool is unusable — we fail registration rather than hand back a + // permanently-broken pool. + let mut ipc_written = false; + + if receiver_is_cuda && let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + + // Enable P2P for the sender/receiver pair before any IPC operations. + // Gate on a CUDA source: for a CPU source sender_device_idx defaults + // to 0, and enabling a spurious GPU0↔receiver P2P pair creates an + // unnecessary CUDA context. CPU-source registration snapshots the + // ambient device in the else branch below — running _set_cuda_device + // first would clobber the original value saved there. + if is_cuda { + let _ = bound + .call_method1("_ensure_p2p_pair", (sender_device_idx, receiver_device_idx)); + let _ = bound.call_method1("_set_cuda_device", (sender_device_idx,)); + } + + // Resolve transport path. classify_transport encodes the full + // 2³ decision matrix (pure, CI-tested); here we only need the + // single GPU-runtime-dependent input (p2p_available). + let p2p_available: bool = cross_device + && bound + .call_method1("_can_access_peer", (sender_device_idx, receiver_device_idx)) + .and_then(|r| r.extract::()) + .unwrap_or(false); + let transport_path = classify_transport( + sender_device_idx, + receiver_device_idx, + p2p_available, + is_cuda, + ); + let use_transit = transport_path == TransportPath::HostStagingTransit; + + let gpu_ptr: Option = if is_cuda { + if use_transit { + // Allocate pool buffer on receiver's GPU so the + // receiver can import the IPC handle on its own device. + // _transit_copy internally saves/restores the caller's + // device, so an explicit restore is unnecessary here. + let _ = bound.call_method1("_set_cuda_device", (receiver_device_idx,)); + let dst: u64 = bound + .call_method1("_get_gpu_buf", (pool_counter, size)) + .and_then(|r| r.extract::()) + .unwrap_or(0); + // Switch back to sender device. + let _ = bound.call_method1("_set_cuda_device", (sender_device_idx,)); + if dst != 0 { + // Allocate CPU page-locked transit buffer. + let tp: u64 = bound + .call_method1("_alloc_transit", (size,)) + .and_then(|r| r.extract::()) + .unwrap_or(0); + if tp != 0 { + // Copy via transit: sender GPU → CPU → receiver GPU. + let ok: bool = bound + .call_method1( + "_transit_copy", + ( + ptr_val, + sender_device_idx, + tp, + dst, + receiver_device_idx, + size, + ), + ) + .and_then(|r| r.extract::()) + .unwrap_or(false); + if ok { + transit_ptr = tp; + pool_device = receiver_device_idx; + Some(dst) + } else { + // _transit_copy failed: free both buffers + let _ = bound.call_method1("_free_transit", (tp,)); + let _ = bound.call_method1("_free_gpu_buf", (pool_counter,)); + None + } + } else { + // transit alloc failed: free GPU buffer + let _ = bound.call_method1("_free_gpu_buf", (pool_counter,)); + None + } + } else { + None + } + // Both branches start from sender_device_idx (restored + // after _ensure_p2p_pair above). _transit_copy does its + // own internal save/restore; the same-device branch below + // saves/restores explicitly so later cudaMalloc calls land + // on the right GPU. + } else { + // Same-device or P2P available: allocate on sender device. + // Save the current device and restore before returning + // so later cudaMalloc calls land on the right GPU. + let saved_dev: i32 = bound + .call_method0("_get_cuda_device") + .and_then(|r| r.extract::()) + .unwrap_or(0); + let _ = bound.call_method1("_set_cuda_device", (sender_device_idx,)); + let result = { + let dst: u64 = bound + .call_method1("_get_gpu_buf", (pool_counter, size)) + .and_then(|r| r.extract::()) + .unwrap_or(0); + if dst != 0 { + // Only export a handle if the DtoD copy succeeded — + // otherwise the receiver would import uninitialised memory. + if bound + .call_method1("_cuda_memcpy", (dst, ptr_val, size, 3u32)) + .is_ok() + { + Some(dst) + } else { + let _ = bound.call_method1("_free_gpu_buf", (pool_counter,)); + None + } + } else { + None + } + }; + let _ = bound.call_method1("_set_cuda_device", (saved_dev,)); + result + } + } else { + // CPU source → GPU pool: switch to the receiver's device + // so cudaMalloc in dma_copy → _get_gpu_buf lands on the + // correct GPU (matching the GPU-source branches). + let saved_dev: i32 = bound + .call_method0("_get_cuda_device") + .and_then(|r| r.extract::()) + .unwrap_or(0); + let _ = bound.call_method1("_set_cuda_device", (receiver_device_idx,)); + let result = bound + .call_method1("dma_copy", (ptr_val, size, pool_counter, !is_pinned)) + .and_then(|r| r.extract::()) + .ok(); + let _ = bound.call_method1("_set_cuda_device", (saved_dev,)); + result + }; + + if let Some(gpu_ptr) = gpu_ptr + && let Ok(handle) = bound + .call_method1("_ipc_export", (gpu_ptr,)) + .and_then(|r| r.extract::>()) + && handle.len() == 64 + { + unsafe { + // Write IPC handle into DORADMA reserved area [32..96) + std::ptr::copy_nonoverlapping(handle.as_ptr(), shmem_ptr.add(32), 64); + // ipc_present flag at byte 24 + std::ptr::write(shmem_ptr.add(24) as *mut u64, 1u64); + } + ipc_written = true; + } + } + + // A CUDA receiver's shmem is header-only and useless without the IPC + // handle. If GPU-pool setup failed (or the CUDA helpers were + // unavailable), fail registration instead of returning a pool that + // every later write/read would silently reject. Reclaim the shmem + // segment on the way out (it was created with owner=false). + if receiver_is_cuda && !ipc_written { + // The GPU pool buffer (and, on the transit path, the page-locked + // host transit buffer) were allocated before the IPC export, which + // failed. Free them before bailing — otherwise they leak for the + // life of the process since no PoolSlot was stored to track them. + if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + let _ = bound.call_method1("_free_gpu_buf", (pool_counter,)); + if transit_ptr != 0 { + let _ = bound.call_method1("_free_transit", (transit_ptr,)); + } + } + shmem.set_owner(true); + eyre::bail!( + "[{}] register_tensor_pool: failed to set up GPU pool buffer / IPC handle for CUDA receiver `{}`", + self.node_id, + tensor_device + ); + } + + // Seqlock: increment generation to even (write-complete) + unsafe { + let gen_ptr = shmem_ptr.add(96) as *mut u64; + let old_gen = std::ptr::read_volatile(gen_ptr); + std::ptr::write_volatile(gen_ptr, old_gen + 1); + std::sync::atomic::fence(std::sync::atomic::Ordering::Release); + } + + // Store shmem in pool (keep alive) + { + let mut pool = PINNED_POOL.lock().unwrap_or_else(|e| e.into_inner()); + pool.insert( + pool_counter, + PoolSlot { + _shmem: shmem, + base: shmem_ptr as u64, + size: total_size, + is_pinned, + transit_ptr, + pool_device, + }, + ); + } + + // Persist transit metadata so the write fast path can recover + // transit_ptr / pool_device on a PINNED_POOL cache-miss. + if transit_ptr != 0 || pool_device != 0 { + TRANSIT_META + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(pool_counter, (transit_ptr, pool_device)); + } + + let buffer_id = format!("pool_{}_{}", self.node_id, pool_counter); + + // Register with daemon for lifecycle tracking + { + let hlc = dora_node_api::dora_core::uhlc::HLC::default(); + let ts = hlc.new_timestamp(); + let mut params = dora_node_api::MetadataParameters::new(); + params.insert( + "ptr".to_string(), + dora_node_api::Parameter::Integer(ptr_val as i64), + ); + params.insert( + "size".to_string(), + dora_node_api::Parameter::Integer(size as i64), + ); + params.insert("dtype".to_string(), dora_node_api::Parameter::String(dtype)); + params.insert( + "shape".to_string(), + dora_node_api::Parameter::ListInt(shape_list), + ); + params.insert( + "shared_memory_name".to_string(), + dora_node_api::Parameter::String(shmem_name), + ); + params.insert( + "is_pinned".to_string(), + dora_node_api::Parameter::Bool(is_pinned), + ); + params.insert( + "pinned_type".to_string(), + dora_node_api::Parameter::String(pinned_type.to_string()), + ); + params.insert( + "ipc_present".to_string(), + dora_node_api::Parameter::Bool(ipc_written), + ); + params.insert( + "buffer_id".to_string(), + dora_node_api::Parameter::String(buffer_id.clone()), + ); + + let meta = dora_node_api::Metadata::from_parameters(ts, params); + if let Err(e) = crate::seam::store(self.node, &buffer_id, &meta) { + tracing::warn!("[{}] failed to register tensor pool: {:#}", self.node_id, e); + } + } + + let buffer_id_array = arrow::array::StringArray::from(vec![buffer_id]); + let buf_py: Py = buffer_id_array.to_data().to_pyarrow(py)?.unbind(); + Ok(buf_py) + } + + /// Write tensor data to an existing tensor pool. + /// + /// Overwrites the data region of a previously-registered pool without + /// re-registering, enabling memory reuse across iterations. + /// + /// # Concurrency / safety + /// + /// **This is a non-blocking overwrite.** The writer must not call this + /// while a receiver is consuming the previous tensor. Data-byte + /// consistency relies entirely on callers honoring a **turn-based + /// discipline**: wait for the receiver to signal completion (e.g. via + /// a `next_require` round-trip) before writing the next frame. + /// + /// The seqlock at header offset 96 detects in-flight overwrites so the + /// reader can retry, but it does **not** prevent the overwrite itself. + /// A writer that ignores the turn-based contract will produce torn + /// (partially updated) data at the consumer. + /// + /// The bundled `examples/tensor-pool/` dataflows demonstrate correct + /// turn-based usage: the sender writes, outputs the pool ID, and waits + /// for the next input event before writing again. + pub fn write_tensor_pool( + &mut self, + tensor_pool_id: Py, + tensor_info: &Bound<'_, PyDict>, + py: Python, + ) -> eyre::Result<()> { + let buffer_id = parse_tensor_pool_id(tensor_pool_id, py)?; + + let ptr_val: u64 = tensor_info + .get_item("ptr")? + .ok_or_else(|| eyre::eyre!("missing ptr"))? + .extract()?; + let size: usize = tensor_info + .get_item("size")? + .ok_or_else(|| eyre::eyre!("missing size"))? + .extract()?; + let tensor_device: String = tensor_info + .get_item("device")? + .ok_or_else(|| eyre::eyre!("missing device"))? + .extract()?; + let is_cuda = tensor_device.starts_with("cuda"); + + { + let freed = FREED_POOL_IDS.lock().unwrap_or_else(|e| e.into_inner()); + if freed.contains(&buffer_id) { + warn_missing_tensor_pool(&self.node_id, "write", &buffer_id); + return Ok(()); + } + } + + // Auto-select pinning based on tensor size (25 MiB threshold). + // Shared by cache-miss PoolSlot construction and slow-path dma_copy; + // cache-hit reuses the slot's stored is_pinned. + let auto_pin = should_pin(is_cuda, size); + + // Fast path: pool_ format -> DORADMA + if buffer_id.starts_with("pool_") { + // Extract counter from the last underscore segment — node_id + // may legitimately contain underscores. + if let Some((_, counter_str)) = buffer_id.rsplit_once('_') + && let Ok(counter) = counter_str.parse::() + { + // Try PINNED_POOL cache first to avoid per-iteration mmap/munmap. + // register_tensor_pool already stored the Shmem here; taking it + // prevents munmap, and storing it back keeps the mapping alive. + let pool_slot = { + PINNED_POOL + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&counter) + }; + + // Both cache-hit and cache-miss produce a PoolSlot that is + // stored back into PINNED_POOL after the write — this keeps + // the shmem mapping alive for the duration of the data copy. + let (shmem_ptr, shmem_capacity, mut store_back, is_pinned) = + if let Some(mut slot_data) = pool_slot { + // Cache hit: reuse the persistent mapping (no mmap). + // Recompute is_pinned from the current tensor size + // so the auto-selection reflects each write's payload. + let cap = slot_data.size; + slot_data.is_pinned = auto_pin; + let pinned = auto_pin; + (slot_data.base as *mut u8, cap, Some(slot_data), pinned) + } else { + // Cache miss: open via ShmemConf, wrap immediately + // so the mapping stays alive until post-write re-insert. + let shmem_name = format!( + "dora_pool_{}_{}_{}", + self.dataflow_id, self.node_id, counter + ); + match ShmemConf::new().os_id(&shmem_name).open() { + Ok(shmem) => { + let cap = shmem.len(); + let base = shmem.as_ptr() as u64; + let slot = PoolSlot { + _shmem: shmem, + base, + size: cap, + is_pinned: auto_pin, + transit_ptr: 0, + pool_device: 0, + }; + (base as *mut u8, cap, Some(slot), auto_pin) + } + Err(_) => (std::ptr::null_mut(), 0, None, false), + } + }; + + if !shmem_ptr.is_null() { + // Guard against truncated segments before any + // pointer arithmetic (mirrors slow-path + read guards). + if shmem_capacity < DORADMA_HEADER_SIZE { + if let Some(slot_data) = store_back { + PINNED_POOL + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(counter, slot_data); + } + return Ok(()); + } + let magic = unsafe { std::slice::from_raw_parts(shmem_ptr, 8) }; + if magic == DORADMA_MAGIC { + let data_offset = unsafe { read_header_u64(shmem_ptr.add(16)) as usize }; + + // Check if this pool has GPU DMA path enabled + let ipc_present = + unsafe { std::ptr::read(shmem_ptr.add(24) as *const u64) }; + + // Validate write size against pool capacity. The + // `size == 0` guard always applies. For GPU-buffer pools + // (ipc_present == 1) the shmem data region is unused, so + // its capacity is irrelevant — the GPU-buffer helpers + // validate `size` against the actual GPU allocation. + if size == 0 + || (ipc_present != 1 + && size > shmem_capacity.saturating_sub(data_offset)) + { + tracing::warn!( + "[{}] write_tensor_pool: size {} exceeds available pool capacity (data_offset={}, total={}), operation aborted", + self.node_id, + size, + data_offset, + shmem_capacity + ); + // Store back to PINNED_POOL to keep shmem alive + if let Some(slot_data) = store_back { + PINNED_POOL + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(counter, slot_data); + } + return Ok(()); + } + + if ipc_present == 1 && !is_cuda { + // Seqlock: begin (noop if still odd from a + // previous failed copy). + let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; + let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; + let mut copy_ok = true; + if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + if let Err(e) = bound + .call_method1("dma_copy", (ptr_val, size, counter, !is_pinned)) + { + copy_ok = false; + tracing::error!( + "[{}] write_tensor_pool: DMA copy failed: {}", + self.node_id, + e + ); + } + } else { + copy_ok = false; + } + if copy_ok { + // Publish: gen was odd (in-progress), flip to even. + unsafe { + seqlock_end(gen_ptr, pre_write_gen, true); + } + } else { + // GPU in-place write to the live IPC buffer: a + // partial copy cannot be rolled back to a clean + // previous frame, so leave gen odd. The reader + // retries until the next successful write. + } + if !copy_ok { + // Re-insert the slot so free_tensor_pool + // can clean up the GPU buffer and transit + // allocation (mirrors the is_cuda branch). + if let Some(slot_data) = store_back.take() { + PINNED_POOL + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(counter, slot_data); + } + return Err(eyre::eyre!( + "[{}] write_tensor_pool: DMA copy failed", + self.node_id + )); + } + } else if is_cuda { + // Seqlock: begin (noop if still odd from a + // previous failed copy). + let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; + let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; + let mut copy_ok = true; + if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + // Resolve transit metadata: cache-hit from PoolSlot, + // cache-miss from TRANSIT_META (populated during registration). + let mut transit_ptr = + store_back.as_ref().map_or(0, |s| s.transit_ptr); + let transit_from_cache; + if transit_ptr == 0 { + // Cache-miss fallback: TRANSIT_META survives + // PINNED_POOL eviction so the write fast path + // always knows whether a transit buffer exists. + let meta = + TRANSIT_META.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(&(tp, _pd)) = meta.get(&counter) { + transit_ptr = tp; + transit_from_cache = true; + } else { + transit_from_cache = false; + } + } else { + transit_from_cache = false; + } + let pool_dev = store_back.as_ref().map_or(0, |s| s.pool_device); + let write_path = classify_write_path( + ipc_present, + /*is_cuda=*/ true, + transit_ptr, + ); + let res = match write_path { + WritePath::GpuToGpuPoolTransit => { + // Recover pool_device from TRANSIT_META on + // cache-miss; otherwise use the PoolSlot value. + let pool_dev = if transit_from_cache { + TRANSIT_META + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(&counter) + .copied() + .map(|(_tp, pd)| pd) + .unwrap_or(pool_dev) + } else { + pool_dev + }; + let sender_dev = tensor_device + .strip_prefix("cuda") + .and_then(|d| d.strip_prefix(':')) + .and_then(|d| d.parse::().ok()) + .unwrap_or(0); + bound + .call_method1( + "_transit_copy_gpu_buf", + ( + counter, + ptr_val, + sender_dev, + transit_ptr, + pool_dev, + size, + ), + ) + .map(|_| ()) + } + WritePath::GpuToGpuPoolDtoD => bound + .call_method1( + "_cuda_memcpy_gpu_buf", + (counter, ptr_val, size), + ) + .map(|_| ()), + _ => { + // GpuToShmem: ipc_present ≠ 1, + // copy to shared-memory data region. + bound + .call_method1( + "_cuda_memcpy", + ( + shmem_ptr as u64 + data_offset as u64, + ptr_val, + size, + 2u32, + ), + ) + .map(|_| ()) + } + }; + if let Err(e) = res { + copy_ok = false; + tracing::error!( + "[{}] write_tensor_pool: GPU pool copy failed: {}", + self.node_id, + e + ); + } + } else { + copy_ok = false; + } + if copy_ok { + // Publish: gen was odd (in-progress), flip to even. + unsafe { + seqlock_end(gen_ptr, pre_write_gen, true); + } + } else { + // GPU in-place write to the live IPC buffer: a + // partial copy cannot be rolled back to a clean + // previous frame, so leave gen odd. The reader + // retries until the next successful write. + } + if !copy_ok { + if let Some(slot_data) = store_back.take() { + PINNED_POOL + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(counter, slot_data); + } + return Err(eyre::eyre!( + "[{}] write_tensor_pool: GPU pool copy failed", + self.node_id + )); + } + } else { + // Seqlock: begin (noop if still odd from a + // previous failed GPU-to-shmem copy). + let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; + let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; + unsafe { + std::ptr::copy_nonoverlapping( + ptr_val as *const u8, + shmem_ptr.add(data_offset), + size, + ); + } + // Publish: gen was odd (in-progress), flip to even. + // The shmem data region is not shared with the reader + // through IPC, so a copy failure means segfault — the + // process is dead before reaching here. No rollback needed. + unsafe { + seqlock_end(gen_ptr, pre_write_gen, true); + } + } + + // Store back to PINNED_POOL to keep shmem alive + if let Some(slot_data) = store_back { + PINNED_POOL + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(counter, slot_data); + } + + return Ok(()); + } + } + } + } + + // Slow path: query daemon for pool metadata + match crate::seam::load_metadata(self.node, &buffer_id, false) { + Ok(metadata) => { + let shmem_name = metadata.parameters.get("shared_memory_name").and_then(|p| { + if let Parameter::String(s) = p { + Some(s.clone()) + } else { + None + } + }); + + if let Some(ref name) = shmem_name + && let Ok(shmem) = ShmemConf::new().os_id(name).open() + { + // Mirror fast-path guard: reject segments smaller + // than the header before any pointer arithmetic. + if shmem.len() < DORADMA_HEADER_SIZE { + return Ok(()); + } + let shmem_ptr = shmem.as_ptr(); + + let magic = unsafe { std::slice::from_raw_parts(shmem_ptr, 8) }; + if magic == DORADMA_MAGIC { + let data_offset = unsafe { read_header_u64(shmem_ptr.add(16)) as usize }; + + // Check if this pool has GPU DMA path enabled + let ipc_present = + unsafe { std::ptr::read(shmem_ptr.add(24) as *const u64) }; + + // Validate write size against pool capacity. The + // `size == 0` guard always applies; the shmem-region + // capacity check is skipped for GPU-buffer pools + // (ipc_present == 1), which the GPU helpers validate. + let shmem_len = shmem.len(); + if size == 0 + || (ipc_present != 1 && size > shmem_len.saturating_sub(data_offset)) + { + tracing::warn!( + "[{}] write_tensor_pool (slow path): size {} exceeds available pool capacity (data_offset={}, total={}), operation aborted", + self.node_id, + size, + data_offset, + shmem_len + ); + return Ok(()); + } + + if ipc_present == 1 && !is_cuda { + // Extract counter for the DMA slot from buffer_id. + let slow_counter = buffer_id + .rsplit_once('_') + .and_then(|(_, c)| c.parse::().ok()); + let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; + let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; + let mut copy_ok = true; + if let (Ok(helpers), Some(c)) = (get_cuda_helpers(py), slow_counter) { + let bound = helpers.bind(py); + let slow_no_dma = !auto_pin; + if let Err(e) = + bound.call_method1("dma_copy", (ptr_val, size, c, slow_no_dma)) + { + copy_ok = false; + tracing::error!( + "[{}] write_tensor_pool (slow path): DMA copy failed: {}", + self.node_id, + e + ); + } + } else { + copy_ok = false; + } + if copy_ok { + // Publish: gen was odd (in-progress), flip to even. + unsafe { + seqlock_end(gen_ptr, pre_write_gen, true); + } + } else { + // GPU in-place write to the live IPC buffer: a + // partial copy cannot be rolled back to a clean + // previous frame, so leave gen odd. The reader + // retries until the next successful write. + } + if !copy_ok { + return Err(eyre::eyre!( + "[{}] write_tensor_pool (slow path): DMA copy failed", + self.node_id + )); + } + } else if is_cuda { + let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; + let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; + let mut copy_ok = true; + if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + // Slow path transit look-up: PINNED_POOL + // (contrast fast path which uses store_back). + let (transit_ptr, pool_device) = if let Some((_, counter_str)) = + buffer_id.rsplit_once('_') + && let Ok(c) = counter_str.parse::() + { + PINNED_POOL + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(&c) + .map(|s| (s.transit_ptr, s.pool_device)) + .unwrap_or((0, 0)) + } else { + (0, 0) + }; + let write_path = classify_write_path( + ipc_present, + /*is_cuda=*/ true, + transit_ptr, + ); + let res = match write_path { + WritePath::GpuToGpuPoolTransit => { + let sender_dev = tensor_device + .strip_prefix("cuda") + .and_then(|d| d.strip_prefix(':')) + .and_then(|d| d.parse::().ok()) + .unwrap_or(0); + bound + .call_method1( + "_transit_copy_gpu_buf", + ( + buffer_id + .rsplit_once('_') + .and_then(|(_, cs)| cs.parse::().ok()) + .unwrap_or(0), + ptr_val, + sender_dev, + transit_ptr, + pool_device, + size, + ), + ) + .map(|_| ()) + } + WritePath::GpuToGpuPoolDtoD => bound + .call_method1( + "_cuda_memcpy_gpu_buf", + ( + buffer_id + .rsplit_once('_') + .and_then(|(_, cs)| cs.parse::().ok()) + .unwrap_or(0), + ptr_val, + size, + ), + ) + .map(|_| ()), + _ => { + // GpuToShmem: copy to shmem data region. + bound + .call_method1( + "_cuda_memcpy", + ( + shmem_ptr as u64 + data_offset as u64, + ptr_val, + size, + 2u32, + ), + ) + .map(|_| ()) + } + }; + if let Err(e) = res { + copy_ok = false; + tracing::error!( + "[{}] write_tensor_pool (slow path): GPU pool copy failed: {}", + self.node_id, + e + ); + } + } else { + copy_ok = false; + } + if copy_ok { + // Publish: gen was odd (in-progress), flip to even. + unsafe { + seqlock_end(gen_ptr, pre_write_gen, true); + } + } else { + // GPU in-place write to the live IPC buffer: a + // partial copy cannot be rolled back to a clean + // previous frame, so leave gen odd. The reader + // retries until the next successful write. + } + if !copy_ok { + return Err(eyre::eyre!( + "[{}] write_tensor_pool (slow path): GPU pool copy failed", + self.node_id + )); + } + } else { + // Seqlock: begin (noop if still odd from a + // previous failed GPU-to-shmem copy). + let gen_ptr = unsafe { shmem_ptr.add(96) as *mut u64 }; + let pre_write_gen = unsafe { seqlock_begin_if_even(gen_ptr) }; + unsafe { + std::ptr::copy_nonoverlapping( + ptr_val as *const u8, + shmem_ptr.add(data_offset), + size, + ); + } + // Publish: gen was odd (in-progress), flip to even. + // The shmem data region is not shared through IPC — + // copy_nonoverlapping failure means segfault, + // so the process is dead before reaching here. + unsafe { + seqlock_end(gen_ptr, pre_write_gen, true); + } + } + } + } + } + Err(_) => { + warn_missing_tensor_pool(&self.node_id, "write", &buffer_id); + } + } + + Ok(()) + } + + /// Read tensor info from an existing tensor pool (zero-copy). + /// + /// Returns a `tensor_info` dict compatible with `tensor_from_info`. + /// The returned tensor shares the underlying shared-memory mapping — + /// no copy is made, so data bytes reflect whatever the writer has most + /// recently stored. + /// + /// # Concurrency / safety + /// + /// **The returned tensor is a zero-copy view into shared memory.** + /// Its data bytes can be overwritten at any time by a concurrent (or + /// subsequent) [`write_tensor_pool`] on the sender. The seqlock + /// re-check at end-of-read detects whether an overwrite occurred + /// mid-consumption, but it is the **caller's responsibility** to + /// ensure the tensor is not used after the writer is allowed to write + /// again. + /// + /// Correct consumers follow a **turn-based discipline**: read the + /// pool, consume the tensor fully, then signal the sender (e.g. via + /// the dataflow graph's `next_require` round-trip) that it is safe to + /// write the next frame. The bundled `examples/tensor-pool/` dataflows + /// demonstrate this pattern. + pub fn read_tensor_pool( + &mut self, + tensor_pool_id: Py, + py: Python, + ) -> eyre::Result> { + let buffer_id = parse_tensor_pool_id(tensor_pool_id, py)?; + + // Populate the trusted GPU buffer size from daemon metadata + // on the first read. Subsequent reads (if any) reuse the + // cached entry — the daemon query runs at most once per pool. + { + let trusted = GPU_BUF_SIZES.lock().unwrap_or_else(|e| e.into_inner()); + if !trusted.contains_key(&buffer_id) { + drop(trusted); + if let Ok(metadata) = crate::seam::load_metadata(self.node, &buffer_id, false) + && let Some(size) = metadata.parameters.get("size").and_then(|p| { + if let Parameter::Integer(v) = p { + Some(*v) + } else { + None + } + }) + { + GPU_BUF_SIZES + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(buffer_id.clone(), size as u64); + } + } + } + + // Fast path: DORADMA header read with daemon-trusted size validation. + if buffer_id.starts_with("pool_") { + // Retry on transient failures (odd seqlock, shmem not yet + // mapped) so a concurrent writer doesn't cause a hard error. + // Time-bounded: a GPU copy (cudaMemcpy + synchronize) takes + // milliseconds, so we wait up to 500ms total with 1ms sleeps + // between attempts. + let deadline = std::time::Instant::now() + .checked_add(std::time::Duration::from_millis(500)) + .unwrap_or(std::time::Instant::now()); + loop { + match self.try_doradma_read(&buffer_id, py) { + Ok(Some(result)) => return Ok(result), + Ok(None) if std::time::Instant::now() < deadline => { + // Transient — yield the GIL and sleep so the + // writer can complete its copy+sync. + py.detach(|| { + std::thread::sleep(std::time::Duration::from_millis(1)); + }); + continue; + } + Ok(None) => break, + Err(e) => { + warn_missing_tensor_pool(&self.node_id, "read", &buffer_id); + eyre::bail!("tensor pool {}: fast path failed: {}", buffer_id, e); + } + } + } + // Retries exhausted — fall back to the daemon for CPU pools. + if let Ok(metadata) = crate::seam::load_metadata(self.node, &buffer_id, false) { + let size = metadata + .parameters + .get("size") + .and_then(|p| { + if let Parameter::Integer(v) = p { + Some(*v) + } else { + None + } + }) + .unwrap_or(0); + let dtype = metadata + .parameters + .get("dtype") + .and_then(|p| { + if let Parameter::String(s) = p { + Some(s.clone()) + } else { + None + } + }) + .unwrap_or_default(); + let shape = metadata + .parameters + .get("shape") + .and_then(|p| { + if let Parameter::ListInt(v) = p { + Some(v.clone()) + } else { + None + } + }) + .unwrap_or_default(); + // Only for CPU pools — GPU pools need IPC import. + let ipc_present = metadata + .parameters + .get("ipc_present") + .and_then(|p| { + if let Parameter::Bool(v) = p { + Some(*v) + } else { + None + } + }) + .unwrap_or(false); + if ipc_present { + warn_missing_tensor_pool(&self.node_id, "read", &buffer_id); + eyre::bail!( + "tensor pool {}: fast path retries exhausted for GPU pool \ + (daemon fallback cannot provide a GPU pointer)", + buffer_id + ); + } + let shmem_name = metadata.parameters.get("shared_memory_name").and_then(|p| { + if let Parameter::String(s) = p { + Some(s.clone()) + } else { + None + } + }); + if let Some(ref name) = shmem_name + && let Ok(shmem) = ShmemConf::new().os_id(name).open() + && shmem.len() >= DORADMA_HEADER_SIZE + { + let shmem_ptr = shmem.as_ptr(); + let magic = unsafe { std::slice::from_raw_parts(shmem_ptr, 8) }; + if magic == DORADMA_MAGIC { + let data_offset = unsafe { read_header_u64(shmem_ptr.add(16)) as usize }; + // Mirror fast-path bounds check. + if data_offset > shmem.len() + || (size as usize) > shmem.len().saturating_sub(data_offset) + { + warn_missing_tensor_pool(&self.node_id, "read", &buffer_id); + eyre::bail!( + "tensor pool {}: header bounds exceeded: \ + data_offset {} + size {} > shmem_len {}", + buffer_id, + data_offset, + size, + shmem.len() + ); + } + // Seqlock: reject a torn mid-write frame. The + // fallback is reached when the fast path retries + // are exhausted, typically because the generation + // is stuck odd (crashed writer). + let read_gen = + unsafe { std::ptr::read_volatile(shmem_ptr.add(96) as *const u64) }; + if read_gen % 2 != 0 { + warn_missing_tensor_pool(&self.node_id, "read", &buffer_id); + eyre::bail!( + "tensor pool {}: daemon fallback: seqlock write in progress \ + (generation={}, odd)", + buffer_id, + read_gen + ); + } + // Cache-hit: use the stored mapping's base so + // the pointer stays valid after the fresh shmem + // is dropped. Cache-miss: insert the fresh + // mapping and drop the old one (if any). + let read_ptr; + { + let mut cpu_cache = + RECV_CPU_SHMEM.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(cached) = cpu_cache.get(&buffer_id) { + read_ptr = (cached.base + data_offset as u64) as i64; + } else { + let base = shmem_ptr as u64; + read_ptr = (base + data_offset as u64) as i64; + cpu_cache.entry(buffer_id.clone()).or_insert(RecvCpuSlot { + _shmem: shmem, + base, + }); + } + } + let dict = PyDict::new(py); + dict.set_item("ptr", read_ptr)?; + dict.set_item("size", size)?; + dict.set_item("dtype", dtype)?; + dict.set_item("shape", shape)?; + dict.set_item("device", "cpu")?; + return Ok(dict.into()); + } + } + } + warn_missing_tensor_pool(&self.node_id, "read", &buffer_id); + eyre::bail!( + "tensor pool {}: fast path retries exhausted — pool not ready after 500ms", + buffer_id + ); + } + + warn_missing_tensor_pool(&self.node_id, "read", &buffer_id); + eyre::bail!("tensor pool {} not found", buffer_id); + } + + /// Free a tensor pool. + pub fn free_tensor_pool(&mut self, tensor_pool_id: Py, py: Python) -> eyre::Result<()> { + let buffer_id = parse_tensor_pool_id(tensor_pool_id, py)?; + + if crate::seam::drop_key(self.node, &buffer_id).is_err() { + warn_missing_tensor_pool(&self.node_id, "release", &buffer_id); + } + + // Clean up sender-side pinned pool mapping (Shmem + CUDA host register). + // Without this, each register->write->free cycle leaks one Shmem mapping + // and one cudaHostRegister pinned region for the process lifetime. + // PINNED_POOL is sender-side (per-process), so bare counter is sufficient. + { + let counter = buffer_id + .strip_prefix("pool_") + .and_then(|s| s.rsplit_once('_').map(|(_, c)| c)) + .and_then(|c| c.parse::().ok()); + if let Some(c) = counter + && let Some(slot) = PINNED_POOL + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&c) + { + if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + let _ = bound.call_method1("_unregister_host", (slot.base,)); + let _ = bound.call_method1("_free_gpu_buf", (c,)); + if slot.transit_ptr != 0 { + let _ = bound.call_method1("_free_transit", (slot.transit_ptr,)); + } + } + // Remove transit metadata regardless of whether CUDA helpers + // are available — a missing _free_transit is a leak, but a stale + // TRANSIT_META entry is a correctness bug on re-registration. + TRANSIT_META + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&c); + } + // PoolSlot dropped here -> Shmem unmapped + } + + // Clean up receiver-side caches so the shmem mappings are released. + // Keyed by full buffer_id (namespaced) to correctly handle + // multiple sender nodes with the same per-process counter. + { + // Close GPU IPC handle before dropping the cache entry. + if let Some(slot) = RECV_GPU_VA + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&buffer_id) + { + if slot.gpu_buf != 0 { + if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + let _ = bound.call_method1("_ipc_close", (slot.gpu_buf,)); + } + } else if slot.gpu_va != 0 { + // Host-registered mapping (effective_as_cuda branch): + // must cudaHostUnregister before munmap. _unregister_host + // requires the original host pointer (shmem base), not the + // device VA returned by cudaHostGetDevicePointer. + if let Ok(helpers) = get_cuda_helpers(py) { + let bound = helpers.bind(py); + let _ = bound.call_method1("_unregister_host", (slot.host_base,)); + } + } + // slot._shmem drops here -> munmap + } + } + RECV_CPU_SHMEM + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&buffer_id); + GPU_BUF_SIZES + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&buffer_id); + + { + let mut freed = FREED_POOL_IDS.lock().unwrap_or_else(|e| e.into_inner()); + tracing::debug!( + "[{}] free_tensor_pool: adding {} to FREED_POOL_IDS (set size={})", + self.node_id, + buffer_id, + freed.len() + ); + freed.insert(buffer_id); + } + + Ok(()) + } + + // ==================== try_doradma_read ==================== + + /// DORADMA fast path for read_tensor_pool: reads metadata directly from + /// the shmem header, bypassing the daemon for zero-copy metadata retrieval. + /// + /// Buffer ID format: `"pool_{node_id}_{counter}"` → + /// shmem name: `"dora_pool_{dataflow_id}_{node_id}_{counter}"`. + /// + /// # Synchronization model + /// + /// The seqlock (write_gen at header offset 96) guards **data-byte** + /// consistency across `write_tensor_pool` overwrites — the end-of-read + /// generation re-check detects if a write occurred mid-read. Header + /// fields (json_len, data_offset) are written once at registration and + /// never change, so they are not subject to torn-read risk. + /// + /// The seqlock does NOT protect the tensor data bytes from being + /// overwritten while a consumer is iterating the zero-copy tensor. + /// Callers must enforce a **turn-based** discipline: the writer must + /// not begin a new `write_tensor_pool` until the receiver has finished + /// consuming the previous tensor. The example dataflow enforces this + /// via `next_require` round-trip signaling. + /// + /// Returns `Ok(Some(tensor_info_dict))` on success, `Ok(None)` to fall back to daemon. + fn try_doradma_read( + &mut self, + buffer_id: &str, + py: Python<'_>, + ) -> eyre::Result>> { + // Format: "pool_{node_id}_{counter}". + // Use rsplit to extract the counter from the end — the node_id + // portion may itself contain underscores (legal in dora node ids). + let counter: u64 = match buffer_id.rsplit_once('_') { + Some((_, c)) => match c.parse() { + Ok(c) => c, + Err(_) => return Ok(None), + }, + None => return Ok(None), + }; + let pool_node_id = buffer_id + .strip_prefix("pool_") + .and_then(|s| s.strip_suffix(&format!("_{counter}"))) + .unwrap_or(""); + + // Check freed tracking -> if this buffer was freed, fall back to daemon + { + let freed = FREED_POOL_IDS.lock().unwrap_or_else(|e| e.into_inner()); + if freed.contains(buffer_id) { + tracing::debug!( + "[{}] try_doradma_read: buffer {} is freed, fallback to daemon", + self.node_id, + buffer_id + ); + return Ok(None); + } + } + + let shmem_name = format!( + "dora_pool_{}_{}_{}", + self.dataflow_id, pool_node_id, counter + ); + + // Open shared memory + let shmem = match ShmemConf::new().os_id(&shmem_name).open() { + Ok(s) => s, + Err(_) => return Ok(None), + }; + + let shmem_ptr = shmem.as_ptr(); + let shmem_size = shmem.len(); + + // Reject truncated segments: need at least DORADMA_HEADER_SIZE bytes for the header + if shmem_size < DORADMA_HEADER_SIZE { + tracing::warn!( + "[{}] try_doradma_read: shmem size {} < DORADMA_HEADER_SIZE {}, rejecting", + self.node_id, + shmem_size, + DORADMA_HEADER_SIZE, + ); + return Ok(None); + } + + // Verify DORADMA magic header + + unsafe { + let magic = std::slice::from_raw_parts(shmem_ptr, 8); + if magic != DORADMA_MAGIC { + return Ok(None); + } + } + + // Read header: [magic:8][json_len:8][data_offset:8][reserved:232] + let json_len = unsafe { read_header_u64(shmem_ptr.add(8)) as usize }; + let data_offset = unsafe { read_header_u64(shmem_ptr.add(16)) as usize }; + + // Validate JSON length fits within the segment + if json_len > shmem_size.saturating_sub(DORADMA_HEADER_SIZE) { + tracing::warn!( + "[{}] try_doradma_read: json_len {} exceeds shmem bounds (size={}, header={})", + self.node_id, + json_len, + shmem_size, + DORADMA_HEADER_SIZE, + ); + return Ok(None); + } + + // Read JSON metadata from header + let json_slice = + unsafe { std::slice::from_raw_parts(shmem_ptr.add(DORADMA_HEADER_SIZE), json_len) }; + let json_str = match std::str::from_utf8(json_slice) { + Ok(s) => s, + Err(_) => return Ok(None), + }; + + // Parse JSON to Python dict + let metadata_dict: Bound<'_, PyDict> = match py.import("json") { + Ok(m) => match m.call_method1("loads", (json_str,)) { + Ok(v) => match v.cast_into::() { + Ok(d) => d, + Err(_) => return Ok(None), + }, + Err(_) => return Ok(None), + }, + Err(_) => return Ok(None), + }; + + let size: usize = match metadata_dict.get_item("size") { + Ok(Some(v)) => v.extract().unwrap_or(0), + _ => 0, + }; + if size == 0 { + return Ok(None); + } + + // Check if this pool uses GPU DMA (IPC handle in header). + // Reading ipc_present early lets us skip the shmem data-region + // size check for GPU-buffer pools (receiver_is_cuda registrations + // allocate header-only shmem). + let ipc_present = unsafe { std::ptr::read(shmem_ptr.add(24) as *const u64) }; + + // Verify data_offset + size fits within shared memory segment. + // GPU-buffer reads (ipc_present == 1) don't access the shmem data + // region, so the size check is only required for CPU-receiver paths. + if ipc_present != 1 { + // Use saturating operations to guard against corrupted/hostile + // headers with a near-usize::MAX data_offset (overflow-safe, + // matching the write-path checks). + if data_offset > shmem_size || size > shmem_size.saturating_sub(data_offset) { + tracing::warn!( + "[{}] try_doradma_read: data_offset {} + size {} exceeds shmem_size {}", + self.node_id, + data_offset, + size, + shmem_size, + ); + return Ok(None); + } + } + + // Auto-detect read path from pinned_type + let pinned_type: Option = metadata_dict + .get_item("pinned_type") + .ok() + .flatten() + .and_then(|v| v.extract::().ok()); + + let effective_as_cuda = ipc_present == 1 || pinned_type.as_deref() != Some("cpu"); + + // Seqlock: read generation, validate stable after data access + let read_gen = unsafe { std::ptr::read_volatile(shmem_ptr.add(96) as *const u64) }; + if read_gen % 2 != 0 { + // Writer is in progress — retry (caller falls back to daemon) + return Ok(None); + } + std::sync::atomic::fence(std::sync::atomic::Ordering::Acquire); + + let read_ptr: u64; + + if ipc_present == 1 { + // GPU DMA pool: import IPC handle once, cache GPU buffer ptr + read_ptr = { + let cache = RECV_GPU_VA.lock().unwrap_or_else(|e| e.into_inner()); + match cache.get(buffer_id) { + Some(slot_data) if slot_data.gpu_buf != 0 => { + // Validate size against the GPU buffer's registered + // capacity. GPU_BUF_SIZES (populated from daemon + // metadata) is authoritative; gpu_buf_size (populated + // from shmem at first import) is the baseline. + let trusted_sizes = { + let trusted = GPU_BUF_SIZES.lock().unwrap_or_else(|e| e.into_inner()); + trusted.get(buffer_id).copied() + }; + if check_capacity_gpu_pool( + trusted_sizes, + Some(slot_data.gpu_buf_size), + size as u64, + ) == CapacityCheck::ExceedsTrustedSize + { + return Ok(None); + } + slot_data.gpu_buf + } + _ => { + // First IPC import: validate size against + // daemon-trusted capacity. Fail closed — + // NoTrustedEntry rejects the import rather than + // trusting the world-writable shmem size. + let trusted_sizes = { + let trusted = GPU_BUF_SIZES.lock().unwrap_or_else(|e| e.into_inner()); + trusted.get(buffer_id).copied() + }; + match check_capacity_gpu_pool(trusted_sizes, None, size as u64) { + CapacityCheck::Ok => {} + CapacityCheck::ExceedsTrustedSize | CapacityCheck::NoTrustedEntry => { + return Ok(None); + } + } + drop(cache); + let handle_bytes = + unsafe { std::slice::from_raw_parts(shmem_ptr.add(32), 64) }; + let helpers = get_cuda_helpers(py) + .map_err(|e| eyre::eyre!("get_cuda_helpers: {}", e))?; + let bound = helpers.bind(py); + let handle_py = PyBytes::new(py, handle_bytes); + let gpu_ptr: u64 = bound + .call_method1("_ipc_import", (handle_py,)) + .map_err(|e| eyre::eyre!("_ipc_import: {}", e))? + .extract() + .map_err(|e| eyre::eyre!("extract gpu_ptr: {}", e))?; + let mut cache = RECV_GPU_VA.lock().unwrap_or_else(|e| e.into_inner()); + cache.insert( + buffer_id.to_string(), + RecvGpuSlot { + _shmem: shmem, + gpu_va: 0, + gpu_buf: gpu_ptr, + host_base: shmem_ptr as u64, + // Baseline from shmem — daemon metadata + // already populated GPU_BUF_SIZES above. + gpu_buf_size: size as u64, + }, + ); + gpu_ptr + } + } + }; + } else if effective_as_cuda { + read_ptr = { + let cache = RECV_GPU_VA.lock().unwrap_or_else(|e| e.into_inner()); + match cache.get(buffer_id) { + Some(slot_data) => { + // GPU VA is stable across data overwrites; + // cache is keyed by full buffer_id (namespaced). + slot_data.gpu_va + data_offset as u64 + } + None => { + drop(cache); + let helpers = get_cuda_helpers(py) + .map_err(|e| eyre::eyre!("get_cuda_helpers: {}", e))?; + let bound = helpers.bind(py); + bound + .call_method1("_register_host", (shmem_ptr as u64, shmem_size)) + .map_err(|e| eyre::eyre!("_register_host: {}", e))?; + let va: u64 = bound + .call_method1("_get_device_ptr", (shmem_ptr as u64,)) + .map_err(|e| eyre::eyre!("_get_device_ptr: {}", e))? + .extract() + .map_err(|e| eyre::eyre!("extract gpu_va: {}", e))?; + let mut cache = RECV_GPU_VA.lock().unwrap_or_else(|e| e.into_inner()); + cache.insert( + buffer_id.to_string(), + RecvGpuSlot { + _shmem: shmem, + gpu_va: va, + gpu_buf: 0, + host_base: shmem_ptr as u64, + gpu_buf_size: 0, // CPU memory, no GPU buffer + }, + ); + va + data_offset as u64 + } + } + }; + } else { + // On the first read the fresh mapping is cached; on subsequent + // reads the fresh mapping is dropped and the returned pointer + // must use the cached mapping's base (a different mmap address). + let mut cpu_cache = RECV_CPU_SHMEM.lock().unwrap_or_else(|e| e.into_inner()); + read_ptr = match cpu_cache.get(buffer_id) { + Some(cached) => cached.base + data_offset as u64, + None => { + let base = shmem_ptr as u64; + cpu_cache.insert( + buffer_id.to_string(), + RecvCpuSlot { + _shmem: shmem, + base, + }, + ); + base + data_offset as u64 + } + }; + } + + // Seqlock: re-read generation — mismatch means data changed during read + let read_gen2 = unsafe { std::ptr::read_volatile(shmem_ptr.add(96) as *const u64) }; + if read_gen2 != read_gen { + return Ok(None); + } + + // Build tensor_info dict + let dtype: String = metadata_dict + .get_item("dtype") + .ok() + .flatten() + .and_then(|v| v.extract::().ok()) + .unwrap_or_default(); + let shape: Vec = metadata_dict + .get_item("shape") + .ok() + .flatten() + .and_then(|v| v.extract::>().ok()) + .unwrap_or_default(); + let device = if effective_as_cuda { "cuda" } else { "cpu" }; + + let dict = PyDict::new(py); + dict.set_item("ptr", read_ptr as i64)?; + dict.set_item("size", size)?; + dict.set_item("dtype", dtype)?; + dict.set_item("shape", shape)?; + dict.set_item("device", device)?; + + // A successful fast-path read means the pool is alive. Clear + // any stale tombstone so that the fast path re-engages after a + // sender restart re-creates the same buffer_id (FREED_POOL_IDS + // is per-process and free_tensor_pool inserts there). + { + FREED_POOL_IDS + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(buffer_id); + } + + Ok(Some(dict.into())) + } +} diff --git a/libraries/extensions/tensor-pool/python/tensor_info_helpers.py b/libraries/extensions/tensor-pool/python/tensor_info_helpers.py new file mode 100644 index 0000000000..b967a6be1b --- /dev/null +++ b/libraries/extensions/tensor-pool/python/tensor_info_helpers.py @@ -0,0 +1,147 @@ +"""Pool-side tensor helpers, parked out of `dora/cuda.py`. + +`get_tensor_info` / `tensor_from_info` were the documented way to feed the +removed `register_tensor_pool` / `write_tensor_pool` / `read_tensor_pool` +methods, and the two dtype maps existed only to serve them. They have no +other caller in dora, so they travelled with the transport. + +`_CudaArrayInterface` stays in `dora/cuda.py` — `open_ipc_handle` still uses +it. Reinstating these belongs on the dora-pool side of the seam, not in the +`dora` package. See ../README.md. +""" + +_DTYPE_MAP = { + " dict: + """Serialize a tensor into a ``tensor_info`` dict containing pointer, + size, dtype, shape, and device. + + This is the canonical way to pass tensor metadata to tensor-pool + operations such as ``register_tensor_pool`` and ``write_tensor_pool``. + """ + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + return { + "ptr": tensor.data_ptr(), + "size": tensor.nbytes, + "dtype": str(tensor.dtype), + "shape": list(tensor.shape), + "device": str(tensor.device), + } + + +def tensor_from_info(tensor_info: dict) -> torch.Tensor: + """Reconstruct a PyTorch tensor from a ``tensor_info`` dict (zero-copy). + + The returned tensor shares the same underlying memory as the original + tensor that produced the ``tensor_info``. Used by consumers that read + a tensor pool via ``read_tensor_pool``. + """ + ptr = tensor_info.get('ptr', 0) + if ptr == 0: + raise ValueError("tensor_info has null pointer (ptr=0); pool may not exist or has been freed") + dtype_str = tensor_info["dtype"] + shape = tensor_info["shape"] + device = tensor_info.get("device", "cpu") + size = tensor_info.get("size", 0) + + dtype = _DTYPE_MAP.get(dtype_str, torch.int64) + + if device.startswith("cuda"): + # CUDA tensor — zero-copy via __cuda_array_interface__ + np_dtype = _TORCH_TO_NUMPY_DTYPE_MAP.get(dtype) + if np_dtype is None: + raise ValueError(f"Unsupported dtype: {dtype}") + + # Validate that product(shape) * itemsize(dtype) does not + # exceed the registered size — a peer-controlled header that + # claims a large shape over a small buffer would produce an + # out-of-bounds GPU tensor (the CPU path is saved by numpy + # reshape, but the GPU path has no equivalent backstop). + expected_bytes = np.dtype(np_dtype).itemsize + for dim in shape: + expected_bytes *= dim + if expected_bytes > size: + raise ValueError( + f"tensor shape {shape} * {np_dtype} itemsize = {expected_bytes} bytes " + f"exceeds registered size {size} bytes — header may be corrupted" + ) + + typestr = np.dtype(np_dtype).str + wrapper = _CudaArrayInterface(ptr, shape, None, typestr) + return torch.as_tensor(wrapper, device="cuda") + else: + # CPU tensor — zero-copy via numpy / torch.frombuffer + np_dtype = _TORCH_TO_NUMPY_DTYPE_MAP.get(dtype) + if np_dtype is None and dtype != torch.bfloat16: + raise ValueError(f"Unsupported dtype: {dtype}") + + c_array = (ctypes.c_byte * size).from_address(ptr) + + if dtype == torch.bfloat16: + byte_tensor = torch.frombuffer(c_array, dtype=torch.uint8) + return byte_tensor.view(dtype=torch.bfloat16).reshape(shape) + + np_array = np.frombuffer(c_array, dtype=np_dtype).reshape(shape) + return torch.from_numpy(np_array) diff --git a/libraries/extensions/tensor-pool/smoke-tests.rs.example b/libraries/extensions/tensor-pool/smoke-tests.rs.example new file mode 100644 index 0000000000..71098cb9d0 --- /dev/null +++ b/libraries/extensions/tensor-pool/smoke-tests.rs.example @@ -0,0 +1,113 @@ +// Smoke tests for the tensor-pool extension — NOT a cargo test target. +// +// They need two things the default build does not provide: a Python wheel +// built with `--features tensor-pool`, and `torch` + `tqdm` in the +// environment. Wiring them into `tests/example-smoke.rs` unconditionally +// would give a confusing failure to anyone running the standard suite +// without a feature-built wheel, so they live here as a `.example` and are +// copied in deliberately. +// +// To run them: install the wheel with the feature, then paste this into +// `tests/example-smoke.rs` (which provides `run_smoke_test` / +// `run_smoke_test_local`) and +// +// cargo test -p dora-examples --test example-smoke -- --ignored tensor_pool +// +// This is a real coverage gap while the extension is opt-in — see the README. + +// --------------------------------------------------------------------------- +// Tensor-pool CPU transport (#2168) +// +// Requires `torch` and `tqdm`, so these are `#[ignore]`-gated and skipped by +// the main `smoke-suite` nightly job. They run in their own nightly job — +// `tensor-pool-smoke` in .github/workflows/nightly.yml — which executes them +// with `--ignored`; the per-node `build:` steps pip-install CPU torch, so no +// GPU is needed. Run all six locally with the same filter CI uses: +// cargo test --test example-smoke -- --ignored tensor_pool +// (`tensor_pool` matches both `smoke_tensor_pool_*` and +// `smoke_local_tensor_pool_*`; `smoke_tensor_pool` would miss the latter). +// Or via `scripts/smoke-all.sh` which gates on `python3 -c "import torch"` +// and skips gracefully when download.pytorch.org is unreachable. +// --------------------------------------------------------------------------- + +#[test] +#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] +fn smoke_tensor_pool_cpu2cpu() { + run_smoke_test( + "tensor-pool-cpu2cpu", + "examples/cpu2cpu.yml", + Duration::from_secs(60), + ); +} + +#[test] +#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] +fn smoke_local_tensor_pool_cpu2cpu() { + run_smoke_test_local("local-tensor-pool-cpu2cpu", "examples/cpu2cpu.yml", 60); +} + +// Negative-lifecycle scenarios validate the "warn, don't crash" contract. +#[test] +#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] +fn smoke_local_tensor_pool_auto_cleanup() { + run_smoke_test_local( + "local-tensor-pool-auto-cleanup", + "examples/auto_cleanup.yml", + 10, + ); +} + +#[test] +#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] +fn smoke_local_tensor_pool_duplicate_free() { + run_smoke_test_local( + "local-tensor-pool-duplicate-free", + "examples/duplicate_free.yml", + 10, + ); +} + +#[test] +#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] +fn smoke_local_tensor_pool_read_after_free() { + run_smoke_test_local( + "local-tensor-pool-read-after-free", + "examples/read_after_free.yml", + 10, + ); +} + +#[test] +#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] +fn smoke_local_tensor_pool_write_after_free() { + run_smoke_test_local( + "local-tensor-pool-write-after-free", + "examples/write_after_free.yml", + 10, + ); +} + +// GPU tensor-pool tests: require CUDA-capable GPU(s). +// cuda_inner needs at least 1 GPU; cuda2cuda needs ≥2 distinct GPUs. +// Both are `#[ignore]`-gated because standard CI runners lack GPUs. +// Run locally: +// cargo test --test example-smoke -- --ignored cuda +#[test] +#[ignore = "requires CUDA GPU(s)"] +fn smoke_tensor_pool_cuda_inner() { + run_smoke_test( + "tensor-pool-cuda-inner", + "examples/cuda_inner.yml", + Duration::from_secs(60), + ); +} + +#[test] +#[ignore = "requires CUDA GPU(s) — ≥2 GPUs"] +fn smoke_tensor_pool_cuda2cuda() { + run_smoke_test( + "tensor-pool-cuda2cuda", + "examples/cuda2cuda.yml", + Duration::from_secs(60), + ); +} diff --git a/libraries/extensions/memory-pool/src/lib.rs b/libraries/extensions/tensor-pool/src/lib.rs similarity index 83% rename from libraries/extensions/memory-pool/src/lib.rs rename to libraries/extensions/tensor-pool/src/lib.rs index 673feff545..a0dbbbb03c 100644 --- a/libraries/extensions/memory-pool/src/lib.rs +++ b/libraries/extensions/tensor-pool/src/lib.rs @@ -1,16 +1,16 @@ use std::collections::{HashMap, HashSet}; use std::sync::{Arc, Mutex}; -/// Identifier for a memory pool buffer, scoped by dataflow. +/// Identifier for a tensor pool buffer, scoped by dataflow. #[derive(Debug, Clone, Hash, PartialEq, Eq)] -pub struct MemoryPoolId { +pub struct TensorPoolId { /// The dataflow that owns this pool. pub dataflow_id: String, /// The per-node buffer identifier. pub id: String, } -/// Metadata for a memory pool tensor. +/// Metadata for a tensor pool tensor. /// /// # Cross-process safety /// @@ -19,7 +19,7 @@ pub struct MemoryPoolId { /// retrieve the data pointer via `shared_memory_name` (opening the shmem /// file and reading the DORADMA header for `data_offset`), not via `ptr`. #[derive(Debug, Clone, Default)] -pub struct MemoryPoolMetadata { +pub struct TensorPoolMetadata { /// Raw pointer to tensor data in the registering process's address space. /// Only valid in the registering process; cross-process consumers must /// use `shared_memory_name` instead. @@ -44,11 +44,11 @@ pub struct MemoryPoolMetadata { pub pinned_type: Option, } -/// Entry in the memory pool table. +/// Entry in the tensor pool table. #[derive(Debug, Clone)] -pub struct MemoryPoolEntry { +pub struct TensorPoolEntry { /// Metadata about the tensor. - pub metadata: MemoryPoolMetadata, + pub metadata: TensorPoolMetadata, /// Node that registered this memory. pub registered_by: String, /// All nodes that have accessed this pool (registered or read). @@ -58,13 +58,13 @@ pub struct MemoryPoolEntry { /// id from a message in flight — the registering node's downstream /// consumers. Together with `touched_by` this is the set of nodes that /// may still reference the pool; see - /// [`MemoryPoolManager::reclaim_unreachable`]. Rewiring a running + /// [`TensorPoolManager::reclaim_unreachable`]. Rewiring a running /// dataflow adds to it, via - /// [`MemoryPoolManager::extend_potential_readers`]. + /// [`TensorPoolManager::extend_potential_readers`]. pub potential_readers: HashSet, } -impl MemoryPoolEntry { +impl TensorPoolEntry { /// Whether any node that could still reference this pool is alive. fn reachable(&self, is_live: &impl Fn(&str) -> bool) -> bool { self.touched_by @@ -86,17 +86,17 @@ pub struct CleanupSummary { pub released_count: usize, } -/// Manager for memory pool allocations. +/// Manager for tensor pool allocations. #[derive(Clone)] -pub struct MemoryPoolManager { - /// Table mapping memory pool IDs to their entries. - memory_pool_table: Arc>>, +pub struct TensorPoolManager { + /// Table mapping tensor pool IDs to their entries. + tensor_pool_table: Arc>>, } -impl MemoryPoolManager { +impl TensorPoolManager { pub fn new() -> Self { Self { - memory_pool_table: Arc::new(Mutex::new(HashMap::new())), + tensor_pool_table: Arc::new(Mutex::new(HashMap::new())), } } @@ -105,22 +105,22 @@ impl MemoryPoolManager { /// Poisoning should never happen in practice (no panics inside lock /// guards), but degrading gracefully is preferable to crashing the /// daemon on an edge case. - fn lock_table(&self) -> std::sync::MutexGuard<'_, HashMap> { - self.memory_pool_table + fn lock_table(&self) -> std::sync::MutexGuard<'_, HashMap> { + self.tensor_pool_table .lock() .unwrap_or_else(|poison| poison.into_inner()) } - /// Register a memory pool with the given ID and metadata. + /// Register a tensor pool with the given ID and metadata. /// /// `potential_readers` lists the nodes that may still ask for this pool /// without having opened it yet (the registrar's downstream consumers). /// It keeps the pool alive until none of them can use it any more — see /// [`Self::reclaim_unreachable`]. - pub fn register_memory_pool( + pub fn register_tensor_pool( &self, - id: MemoryPoolId, - metadata: MemoryPoolMetadata, + id: TensorPoolId, + metadata: TensorPoolMetadata, registered_by: String, potential_readers: HashSet, ) -> Result<(), String> { @@ -134,7 +134,7 @@ impl MemoryPoolManager { touched.insert(registered_by.clone()); table.insert( id, - MemoryPoolEntry { + TensorPoolEntry { metadata, registered_by, touched_by: touched, @@ -145,27 +145,27 @@ impl MemoryPoolManager { Ok(()) } - /// Get the current number of entries in the memory pool table. + /// Get the current number of entries in the tensor pool table. pub fn table_size(&self) -> usize { let table = self.lock_table(); table.len() } - /// Read memory pool metadata by ID. + /// Read tensor pool metadata by ID. /// /// `requested_by` is the node ID of the caller, used for audit logging. /// Cross-node reads are allowed (receivers must read senders' pools) /// but logged at debug level for diagnostics. - pub fn read_memory_pool( + pub fn read_tensor_pool( &self, - id: &MemoryPoolId, + id: &TensorPoolId, requested_by: &str, - ) -> Option { + ) -> Option { let mut table = self.lock_table(); table.get_mut(id).map(|entry| { if entry.registered_by != requested_by { tracing::debug!( - "memory pool {} (registered by {}) read by {}", + "tensor pool {} (registered by {}) read by {}", id.id, entry.registered_by, requested_by, @@ -176,7 +176,7 @@ impl MemoryPoolManager { }) } - /// Free memory pool by ID. + /// Free tensor pool by ID. /// /// Any node may free a pool — the normal lifecycle is that the /// registering (sender) node creates the pool and a reading @@ -184,11 +184,11 @@ impl MemoryPoolManager { /// /// Removes the entry from the table and attempts to clean up the /// underlying shared memory. - pub fn free_memory_pool( + pub fn free_tensor_pool( &self, - id: &MemoryPoolId, + id: &TensorPoolId, requested_by: &str, - ) -> Result<(MemoryPoolMetadata, HashSet), String> { + ) -> Result<(TensorPoolMetadata, HashSet), String> { // Only the table removal needs the lock. Releasing it before the // shared-memory unlink below keeps the `remove_file` syscall out of the // critical section, so concurrent `register`/`read`/`free` calls on @@ -197,12 +197,12 @@ impl MemoryPoolManager { let mut table = self.lock_table(); table .remove(id) - .ok_or_else(|| "memory pool not found".to_string())? + .ok_or_else(|| "tensor pool not found".to_string())? }; if entry.registered_by != requested_by { tracing::debug!( - "memory pool {} (registered by {}) freed by {}", + "tensor pool {} (registered by {}) freed by {}", id.id, entry.registered_by, requested_by, @@ -215,7 +215,7 @@ impl MemoryPoolManager { } /// Unlink an entry's backing segment, if it has one. - fn release_segment(&self, metadata: &MemoryPoolMetadata) -> Result<(), String> { + fn release_segment(&self, metadata: &TensorPoolMetadata) -> Result<(), String> { match &metadata.shared_memory_name { Some(name) if !name.is_empty() => self.free_shared_memory(name), _ => Ok(()), @@ -253,7 +253,7 @@ impl MemoryPoolManager { #[cfg(not(target_os = "linux"))] { Err(format!( - "memory-pool transport is unavailable on this platform; cannot clean up shared memory `{}`", + "tensor-pool transport is unavailable on this platform; cannot clean up shared memory `{}`", shm_name )) } @@ -275,8 +275,8 @@ impl MemoryPoolManager { /// is gone instead of only at dataflow teardown (dora-rs/dora#2881). /// /// Because a released pool has no live node in `touched_by`, there is - /// nobody left to send a `FreeMemoryPool` notification to — unlike - /// `free_memory_pool`, which must tell the other holders to drop their + /// nobody left to send a `FreeTensorPool` notification to — unlike + /// `free_tensor_pool`, which must tell the other holders to drop their /// per-process buffers. /// /// `is_live` reports whether a node is still running **on this daemon**. @@ -289,7 +289,7 @@ impl MemoryPoolManager { &self, dataflow_id: &str, is_live: impl Fn(&str) -> bool, - ) -> Vec { + ) -> Vec { self.release_matching(dataflow_id, |entry| !entry.reachable(&is_live)) } @@ -327,11 +327,11 @@ impl MemoryPoolManager { /// the dataflow is over, so none of its nodes can ask for a pool any /// more. Without it a daemon that outlives the dataflow (`dora up`) /// would hold every unfreed pool until it exits (dora-rs/dora#2881). - pub fn cleanup_dataflow(&self, dataflow_id: &str) -> Vec { + pub fn cleanup_dataflow(&self, dataflow_id: &str) -> Vec { let released = self.release_matching(dataflow_id, |_| true); if !released.is_empty() { tracing::info!( - "Detected {} unreleased memory pool of finished dataflow {dataflow_id}, releasing...", + "Detected {} unreleased tensor pool of finished dataflow {dataflow_id}, releasing...", released.len(), ); } @@ -342,14 +342,14 @@ impl MemoryPoolManager { /// unlink their segments, returning the released ids. /// /// The table lock is released before the unlinks, matching - /// `free_memory_pool`: the `remove_file` syscalls must not serialize + /// `free_tensor_pool`: the `remove_file` syscalls must not serialize /// unrelated pool operations. fn release_matching( &self, dataflow_id: &str, - mut should_release: impl FnMut(&MemoryPoolEntry) -> bool, - ) -> Vec { - let released: Vec<(MemoryPoolId, MemoryPoolEntry)> = { + mut should_release: impl FnMut(&TensorPoolEntry) -> bool, + ) -> Vec { + let released: Vec<(TensorPoolId, TensorPoolEntry)> = { let mut table = self.lock_table(); table .extract_if(|id, entry| id.dataflow_id == dataflow_id && should_release(entry)) @@ -360,7 +360,7 @@ impl MemoryPoolManager { .into_iter() .map(|(id, entry)| { if let Err(err) = self.release_segment(&entry.metadata) { - tracing::warn!("failed to release memory pool {}: {err}", id.id); + tracing::warn!("failed to release tensor pool {}: {err}", id.id); } id }) @@ -371,7 +371,7 @@ impl MemoryPoolManager { /// one of *this daemon's* nodes. /// /// Catches what the pool table cannot: a segment whose creator died - /// between `shm_open` and `register_memory_pool` has no entry to release. + /// between `shm_open` and `register_tensor_pool` has no entry to release. /// /// Ownership matters because `/dev/shm` is host-wide while the segment /// name carries only the *dataflow* id: sweeping by dataflow id alone @@ -430,7 +430,7 @@ impl MemoryPoolManager { } } - /// Cleanup all memory pools on shutdown. + /// Cleanup all tensor pools on shutdown. pub fn cleanup_all(&self) -> Result> { // Drain the table in one move instead of cloning every key into a // `Vec` only to look each one back up and remove it, and release the @@ -441,7 +441,7 @@ impl MemoryPoolManager { if unreleased_count > 0 { tracing::info!( - "Detected {} unreleased memory pool, releasing...", + "Detected {} unreleased tensor pool, releasing...", unreleased_count ); } @@ -457,7 +457,7 @@ impl MemoryPoolManager { if errors.is_empty() { if unreleased_count > 0 { tracing::info!( - "Successfully released {} unreleased memory pools!", + "Successfully released {} unreleased tensor pools!", released_count ); } @@ -467,7 +467,7 @@ impl MemoryPoolManager { }) } else { tracing::warn!( - "Released {} of {} unreleased memory pools; {} failed", + "Released {} of {} unreleased tensor pools; {} failed", released_count, unreleased_count, errors.len() @@ -477,7 +477,7 @@ impl MemoryPoolManager { } } -impl Default for MemoryPoolManager { +impl Default for TensorPoolManager { fn default() -> Self { Self::new() } @@ -487,8 +487,8 @@ impl Default for MemoryPoolManager { mod tests { use super::*; - fn make_metadata() -> MemoryPoolMetadata { - MemoryPoolMetadata { + fn make_metadata() -> TensorPoolMetadata { + TensorPoolMetadata { ptr: 0, size: 1024, dtype: "float32".into(), @@ -501,8 +501,8 @@ mod tests { } } - fn make_id(name: &str) -> MemoryPoolId { - MemoryPoolId { + fn make_id(name: &str) -> TensorPoolId { + TensorPoolId { dataflow_id: "test_df".to_string(), id: name.to_string(), } @@ -515,14 +515,14 @@ mod tests { #[test] fn register_and_read() { - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); let id = make_id("pool-1"); let meta = make_metadata(); - mgr.register_memory_pool(id.clone(), meta.clone(), "node_a".into(), readers(&[])) + mgr.register_tensor_pool(id.clone(), meta.clone(), "node_a".into(), readers(&[])) .unwrap(); let read = mgr - .read_memory_pool(&id, "node_a") + .read_tensor_pool(&id, "node_a") .expect("pool should exist"); assert_eq!(read.ptr, meta.ptr); @@ -533,14 +533,14 @@ mod tests { #[test] fn double_register_fails() { - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); let id = make_id("pool-1"); let meta = make_metadata(); - mgr.register_memory_pool(id.clone(), meta.clone(), "node_a".into(), readers(&[])) + mgr.register_tensor_pool(id.clone(), meta.clone(), "node_a".into(), readers(&[])) .unwrap(); let err = mgr - .register_memory_pool(id, meta, "node_a".into(), readers(&[])) + .register_tensor_pool(id, meta, "node_a".into(), readers(&[])) .unwrap_err(); assert!(err.contains("already registered")); @@ -550,43 +550,43 @@ mod tests { fn cross_owner_free_succeeds() { // Cross-node free is the normal lifecycle: sender registers, // receiver frees. Both operations must succeed. - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); let id = make_id("pool-1"); let meta = make_metadata(); - mgr.register_memory_pool(id.clone(), meta, "node_a".into(), readers(&[])) + mgr.register_tensor_pool(id.clone(), meta, "node_a".into(), readers(&[])) .unwrap(); // A different node frees the pool — this must succeed. - mgr.free_memory_pool(&id, "node_b").unwrap(); + mgr.free_tensor_pool(&id, "node_b").unwrap(); // Pool should be gone after successful free. - assert!(mgr.read_memory_pool(&id, "node_b").is_none()); + assert!(mgr.read_tensor_pool(&id, "node_b").is_none()); } #[test] fn double_free_second_fails() { - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); let id = make_id("pool-1"); let meta = make_metadata(); - mgr.register_memory_pool(id.clone(), meta, "node_a".into(), readers(&[])) + mgr.register_tensor_pool(id.clone(), meta, "node_a".into(), readers(&[])) .unwrap(); - mgr.free_memory_pool(&id, "node_a").unwrap(); + mgr.free_tensor_pool(&id, "node_a").unwrap(); - let err = mgr.free_memory_pool(&id, "node_a").unwrap_err(); - assert!(err.contains("memory pool not found")); + let err = mgr.free_tensor_pool(&id, "node_a").unwrap_err(); + assert!(err.contains("tensor pool not found")); } #[test] fn concurrent_frees_across_pools_all_succeed() { - // `free_memory_pool` only holds the table lock for the removal, not + // `free_tensor_pool` only holds the table lock for the removal, not // across the shared-memory unlink. Frees on distinct pools must proceed // independently and leave the table empty. Metadata carries no backing // segment (`shared_memory_name: None`) so the test stays cross-platform. - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); let n = 16; for i in 0..n { - mgr.register_memory_pool( + mgr.register_tensor_pool( make_id(&format!("pool-{i}")), make_metadata(), "node_a".into(), @@ -599,7 +599,7 @@ mod tests { .map(|i| { let mgr = mgr.clone(); std::thread::spawn(move || { - mgr.free_memory_pool(&make_id(&format!("pool-{i}")), "node_b") + mgr.free_tensor_pool(&make_id(&format!("pool-{i}")), "node_b") }) }) .collect(); @@ -611,10 +611,10 @@ mod tests { #[test] fn cleanup_all_tracks_counts() { - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); for i in 0..3 { - mgr.register_memory_pool( + mgr.register_tensor_pool( make_id(&format!("pool-{}", i)), make_metadata(), "node_a".into(), @@ -631,10 +631,10 @@ mod tests { #[test] fn cleanup_all_reports_partial_release_on_failure() { - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); // One entry frees cleanly (no backing shmem name)... - mgr.register_memory_pool( + mgr.register_tensor_pool( make_id("ok"), make_metadata(), "node_a".into(), @@ -645,7 +645,7 @@ mod tests { // in `free_shared_memory`, so its release errors. let mut bad_meta = make_metadata(); bad_meta.shared_memory_name = Some("invalid_name".to_string()); - mgr.register_memory_pool(make_id("bad"), bad_meta, "node_a".into(), readers(&[])) + mgr.register_tensor_pool(make_id("bad"), bad_meta, "node_a".into(), readers(&[])) .unwrap(); // cleanup_all must surface the failure (rather than silently claiming @@ -657,21 +657,21 @@ mod tests { #[test] fn table_size_tracks_entries() { - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); assert_eq!(mgr.table_size(), 0); let id = make_id("pool-1"); - mgr.register_memory_pool(id.clone(), make_metadata(), "node_a".into(), readers(&[])) + mgr.register_tensor_pool(id.clone(), make_metadata(), "node_a".into(), readers(&[])) .unwrap(); assert_eq!(mgr.table_size(), 1); - mgr.free_memory_pool(&id, "node_a").unwrap(); + mgr.free_tensor_pool(&id, "node_a").unwrap(); assert_eq!(mgr.table_size(), 0); } #[test] fn poison_recovery_lock_table() { - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); // lock_table is private but accessible from a child test module. // Verify it returns a guard, and that operations work after release. { @@ -684,7 +684,7 @@ mod tests { #[test] fn cleanup_orphans_runs_without_panic() { // Sweep should run cleanly without panicking regardless of platform. - MemoryPoolManager::cleanup_orphans("test-dataflow-uuid", |_| true); + TensorPoolManager::cleanup_orphans("test-dataflow-uuid", |_| true); } /// Two daemons on one host serve one dataflow, so both see the other's @@ -718,7 +718,7 @@ mod tests { std::fs::write(segment(file), b"x").unwrap(); } - MemoryPoolManager::cleanup_orphans(&dataflow_id, |node| { + TensorPoolManager::cleanup_orphans(&dataflow_id, |node| { matches!(node, "local" | "local_with_underscore") }); @@ -742,9 +742,9 @@ mod tests { // downstream and exits. The receiver has not opened the pool yet, // so `touched_by` is still just the sender — freeing here would // break the in-flight transfer. - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); let id = make_id("pool-1"); - mgr.register_memory_pool( + mgr.register_tensor_pool( id.clone(), make_metadata(), "sender".into(), @@ -758,14 +758,14 @@ mod tests { released.is_empty(), "pool must survive: `recv` can still read it" ); - assert!(mgr.read_memory_pool(&id, "recv").is_some()); + assert!(mgr.read_tensor_pool(&id, "recv").is_some()); } #[test] fn pool_is_released_once_no_live_node_can_reach_it() { - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); let id = make_id("pool-1"); - mgr.register_memory_pool( + mgr.register_tensor_pool( id.clone(), make_metadata(), "sender".into(), @@ -785,11 +785,11 @@ mod tests { // `potential_readers` is a snapshot of the topology at registration // time; a node that actually opened the pool must keep it alive on // its own, even if it is not in that snapshot. - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); let id = make_id("pool-1"); - mgr.register_memory_pool(id.clone(), make_metadata(), "sender".into(), readers(&[])) + mgr.register_tensor_pool(id.clone(), make_metadata(), "sender".into(), readers(&[])) .unwrap(); - mgr.read_memory_pool(&id, "late_reader").unwrap(); + mgr.read_tensor_pool(&id, "late_reader").unwrap(); let released = mgr.reclaim_unreachable("test_df", |node| node == "late_reader"); @@ -799,15 +799,15 @@ mod tests { #[test] fn reclaim_is_scoped_to_one_dataflow() { - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); let mine = make_id("pool-1"); - let other = MemoryPoolId { + let other = TensorPoolId { dataflow_id: "other_df".to_string(), id: "pool-1".to_string(), }; - mgr.register_memory_pool(mine.clone(), make_metadata(), "sender".into(), readers(&[])) + mgr.register_tensor_pool(mine.clone(), make_metadata(), "sender".into(), readers(&[])) .unwrap(); - mgr.register_memory_pool( + mgr.register_tensor_pool( other.clone(), make_metadata(), "sender".into(), @@ -819,19 +819,19 @@ mod tests { assert_eq!(released, vec![mine]); assert!( - mgr.read_memory_pool(&other, "sender").is_some(), + mgr.read_tensor_pool(&other, "sender").is_some(), "a same-named pool of another dataflow must not be touched" ); } #[test] fn extending_the_reader_set_only_touches_pools_the_source_can_reach() { - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); let mine = make_id("pool-1"); - mgr.register_memory_pool(mine.clone(), make_metadata(), "sender".into(), readers(&[])) + mgr.register_tensor_pool(mine.clone(), make_metadata(), "sender".into(), readers(&[])) .unwrap(); let other = make_id("pool-2"); - mgr.register_memory_pool( + mgr.register_tensor_pool( other.clone(), make_metadata(), "stranger".into(), @@ -849,16 +849,16 @@ mod tests { "only the pool `sender` can reach may gain the new reader" ); assert!( - mgr.read_memory_pool(&mine, "late_consumer").is_some(), + mgr.read_tensor_pool(&mine, "late_consumer").is_some(), "the consumer wired to `sender` must keep `pool-1` alive" ); } #[test] fn cleanup_dataflow_releases_every_pool_of_that_dataflow() { - let mgr = MemoryPoolManager::new(); + let mgr = TensorPoolManager::new(); for i in 0..3 { - mgr.register_memory_pool( + mgr.register_tensor_pool( make_id(&format!("pool-{i}")), make_metadata(), "sender".into(), @@ -866,11 +866,11 @@ mod tests { ) .unwrap(); } - let other = MemoryPoolId { + let other = TensorPoolId { dataflow_id: "other_df".to_string(), id: "pool-1".to_string(), }; - mgr.register_memory_pool( + mgr.register_tensor_pool( other.clone(), make_metadata(), "sender".into(), diff --git a/libraries/message/src/daemon_to_node.rs b/libraries/message/src/daemon_to_node.rs index f251d7b167..c690e0ccf6 100644 --- a/libraries/message/src/daemon_to_node.rs +++ b/libraries/message/src/daemon_to_node.rs @@ -92,8 +92,14 @@ pub enum DaemonCommunication { pub enum DaemonReply { Result(Result<(), String>), NextEvents(Vec>), - NodeConfig { result: Result }, - PinnedMemoryMetadata { metadata: Metadata }, + NodeConfig { + result: Result, + }, + /// Reply to [`DaemonRequest::ExtensionLoad`]. `None` means the key is not + /// in the table — either never stored, or already dropped. + ExtensionValue { + value: Option>, + }, Empty, } @@ -111,10 +117,10 @@ impl DaemonReply { DaemonReply::NextEvents(events) => { events.iter().map(|e| e.inner.encode_size_hint()).sum() } - DaemonReply::Result(_) - | DaemonReply::NodeConfig { .. } - | DaemonReply::PinnedMemoryMetadata { .. } - | DaemonReply::Empty => 0, + // The extension value is opaque bytes handed straight back, so + // its own length is the whole hint. + DaemonReply::ExtensionValue { value } => value.as_ref().map_or(0, |bytes| bytes.len()), + DaemonReply::Result(_) | DaemonReply::NodeConfig { .. } | DaemonReply::Empty => 0, } } } @@ -177,14 +183,15 @@ pub enum NodeEvent { error: String, source_node_id: NodeId, }, - /// A memory pool has been freed by another node in the dataflow. + /// An extension key this node stored or loaded has been dropped, by + /// another node or by the daemon reclaiming it. /// - /// When any node calls `free_memory_pool`, the daemon broadcasts this - /// event to every node so that per-process GPU/transit buffers and - /// shmem mappings are released regardless of which node initiated the - /// free. - FreeMemoryPool { - shared_memory_id: String, + /// Delivered out of band: the event stream consumes it rather than + /// surfacing it to user code, so a language binding polls + /// `drain_dropped_extension_keys()` instead. + ExtensionDropped { + namespace: String, + key: String, }, } @@ -211,8 +218,8 @@ impl NodeEvent { | NodeEvent::AllInputsClosed | NodeEvent::ParamUpdate { .. } | NodeEvent::ParamDeleted { .. } - | NodeEvent::NodeFailed { .. } - | NodeEvent::FreeMemoryPool { .. } => 0, + | NodeEvent::NodeFailed { .. } => 0, + NodeEvent::ExtensionDropped { namespace, key } => namespace.len() + key.len(), }; payload.saturating_add(PER_EVENT_ENVELOPE) } diff --git a/libraries/message/src/node_to_daemon.rs b/libraries/message/src/node_to_daemon.rs index 8103dad992..eba9a8e22e 100644 --- a/libraries/message/src/node_to_daemon.rs +++ b/libraries/message/src/node_to_daemon.rs @@ -29,16 +29,31 @@ pub enum DaemonRequest { NodeConfig { node_id: NodeId, }, - RegisterPinnedMemory { - shared_memory_id: String, - metadata: Metadata, + /// Store an opaque value in the daemon's dataflow-scoped extension table. + /// + /// dora attaches no meaning to `namespace`, `key` or `value`: this is the + /// out-of-band channel for transports that live outside the tree (see + /// `docs/extensions.md`). The daemon tracks who stored a key and who has + /// read it, so it can notify readers on removal and reclaim the entry when + /// the dataflow ends. + ExtensionStore { + namespace: String, + key: String, + value: Vec, }, - ReadPinnedMemory { - shared_memory_id: String, - free: bool, + /// Read an opaque value back. `remove: true` drops it in the same round + /// trip, so a consume-once handoff needs one request rather than two. + ExtensionLoad { + namespace: String, + key: String, + remove: bool, }, - FreePinnedMemory { - shared_memory_id: String, + /// Drop an opaque value. Every node that stored or loaded the key is sent + /// [`NodeEvent::ExtensionDropped`] so it can release whatever the value + /// referred to. + ExtensionDrop { + namespace: String, + key: String, }, } @@ -60,10 +75,10 @@ impl DaemonRequest { | DaemonRequest::OutputsDone | DaemonRequest::NextEvent | DaemonRequest::EventStreamDropped - | DaemonRequest::NodeConfig { .. } - | DaemonRequest::RegisterPinnedMemory { .. } - | DaemonRequest::ReadPinnedMemory { .. } - | DaemonRequest::FreePinnedMemory { .. } => 0, + | DaemonRequest::NodeConfig { .. } => 0, + // The stored value dominates; the namespace and key are short. + DaemonRequest::ExtensionStore { value, .. } => value.len(), + DaemonRequest::ExtensionLoad { .. } | DaemonRequest::ExtensionDrop { .. } => 0, } } @@ -79,9 +94,9 @@ impl DaemonRequest { | DaemonRequest::OutputsDone | DaemonRequest::NextEvent | DaemonRequest::EventStreamDropped - | DaemonRequest::RegisterPinnedMemory { .. } - | DaemonRequest::ReadPinnedMemory { .. } - | DaemonRequest::FreePinnedMemory { .. } => true, + | DaemonRequest::ExtensionStore { .. } + | DaemonRequest::ExtensionLoad { .. } + | DaemonRequest::ExtensionDrop { .. } => true, } } @@ -97,9 +112,9 @@ impl DaemonRequest { | DaemonRequest::SendMessage { .. } | DaemonRequest::OutputSent { .. } | DaemonRequest::EventStreamDropped - | DaemonRequest::RegisterPinnedMemory { .. } - | DaemonRequest::ReadPinnedMemory { .. } - | DaemonRequest::FreePinnedMemory { .. } => false, + | DaemonRequest::ExtensionStore { .. } + | DaemonRequest::ExtensionLoad { .. } + | DaemonRequest::ExtensionDrop { .. } => false, } } } diff --git a/scripts/qa/ci-nightly-jobs.sh b/scripts/qa/ci-nightly-jobs.sh index ec249a42b2..dcca517781 100755 --- a/scripts/qa/ci-nightly-jobs.sh +++ b/scripts/qa/ci-nightly-jobs.sh @@ -1,14 +1,14 @@ #!/usr/bin/env bash # scripts/qa/ci-nightly-jobs.sh -- local driver for the GHA nightly jobs. # -# The GHA nightly workflow (.github/workflows/nightly.yml) has 27 test jobs +# The GHA nightly workflow (.github/workflows/nightly.yml) has 26 test jobs # (re-counted in #2999: the previous 23 counted neither hub-smoke nor the two # ros2-zenoh-* jobs, and multi-daemon-late-subscriber is new) # (post-#1716, plus cluster-record-replay from #2013 and kani-proofs). # `cargo test -p dora-examples --test example-smoke` (run by qa-nightly's # example-smoke step) # covers 4 of them (smoke-suite + log-sinks + service-action + streaming). -# This script covers 20 of the rest -- memory-pool-smoke and hub-smoke have +# This script covers 20 of the rest -- hub-smoke has # no entry here, and one local cli-tests run covers both cli-tests halves -- with # platform-aware dispatch -- on macOS dev machines it runs the macOS subset, # on Linux it runs the Linux subset, etc. (#1716). diff --git a/scripts/smoke-all.sh b/scripts/smoke-all.sh index f42441c2b1..009b9775a8 100755 --- a/scripts/smoke-all.sh +++ b/scripts/smoke-all.sh @@ -539,28 +539,6 @@ if [ "$RUN_PYTHON" = true ]; then run_local "local-queue-size-and-timeout" "tests/queue_size_and_timeout_python/dataflow.yaml" 20 run_local "local-queue-size-latest-data-python" "tests/queue_size_latest_data_python/dataflow.yaml" 20 - echo "" - echo "=== Memory-pool CPU transport ===" - # Dependencies (torch, numpy, tqdm) are provisioned by per-node `build:` - # steps that pip-install from download.pytorch.org/whl/cpu. Skip - # gracefully on air-gapped / network-restricted machines where that index - # is unreachable; the gate is a lightweight TCP probe, not an import check. - if python3 -c " -import urllib.request, sys -try: - urllib.request.urlopen('https://download.pytorch.org/whl/cpu/', timeout=5) -except Exception: - sys.exit(1) -" 2>/dev/null; then - run_networked "memory-pool-cpu2cpu" "examples/memory-pool/cpu2cpu.yml" 60 - run_local "local-memory-pool-cpu2cpu" "examples/memory-pool/cpu2cpu.yml" 60 - run_local "local-memory-pool-auto-cleanup" "examples/memory-pool/auto_cleanup.yml" 10 - run_local "local-memory-pool-duplicate-free" "examples/memory-pool/duplicate_free.yml" 10 - run_local "local-memory-pool-read-after-free" "examples/memory-pool/read_after_free.yml" 10 - run_local "local-memory-pool-write-after-free" "examples/memory-pool/write_after_free.yml" 10 - else - log_skip "memory-pool" "download.pytorch.org unreachable (run on a machine with PyPI access to exercise this suite)" - fi fi # --------------------------------------------------------------------------- diff --git a/tests/example-smoke.rs b/tests/example-smoke.rs index a1d3d77d39..b3f5e5efb8 100644 --- a/tests/example-smoke.rs +++ b/tests/example-smoke.rs @@ -1914,107 +1914,6 @@ fn smoke_shell_node_blocked_without_flag() { let _ = std::fs::remove_file(&yaml); } -// --------------------------------------------------------------------------- -// Memory-pool CPU transport (#2168) -// -// Requires `torch` and `tqdm`, so these are `#[ignore]`-gated and skipped by -// the main `smoke-suite` nightly job. They run in their own nightly job — -// `memory-pool-smoke` in .github/workflows/nightly.yml — which executes them -// with `--ignored`; the per-node `build:` steps pip-install CPU torch, so no -// GPU is needed. Run all six locally with the same filter CI uses: -// cargo test --test example-smoke -- --ignored memory_pool -// (`memory_pool` matches both `smoke_memory_pool_*` and -// `smoke_local_memory_pool_*`; `smoke_memory_pool` would miss the latter). -// Or via `scripts/smoke-all.sh` which gates on `python3 -c "import torch"` -// and skips gracefully when download.pytorch.org is unreachable. -// --------------------------------------------------------------------------- - -#[test] -#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] -fn smoke_memory_pool_cpu2cpu() { - run_smoke_test( - "memory-pool-cpu2cpu", - "examples/memory-pool/cpu2cpu.yml", - Duration::from_secs(60), - ); -} - -#[test] -#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] -fn smoke_local_memory_pool_cpu2cpu() { - run_smoke_test_local( - "local-memory-pool-cpu2cpu", - "examples/memory-pool/cpu2cpu.yml", - 60, - ); -} - -// Negative-lifecycle scenarios validate the "warn, don't crash" contract. -#[test] -#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] -fn smoke_local_memory_pool_auto_cleanup() { - run_smoke_test_local( - "local-memory-pool-auto-cleanup", - "examples/memory-pool/auto_cleanup.yml", - 10, - ); -} - -#[test] -#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] -fn smoke_local_memory_pool_duplicate_free() { - run_smoke_test_local( - "local-memory-pool-duplicate-free", - "examples/memory-pool/duplicate_free.yml", - 10, - ); -} - -#[test] -#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] -fn smoke_local_memory_pool_read_after_free() { - run_smoke_test_local( - "local-memory-pool-read-after-free", - "examples/memory-pool/read_after_free.yml", - 10, - ); -} - -#[test] -#[ignore = "requires `torch` and `tqdm` (not in standard CI)"] -fn smoke_local_memory_pool_write_after_free() { - run_smoke_test_local( - "local-memory-pool-write-after-free", - "examples/memory-pool/write_after_free.yml", - 10, - ); -} - -// GPU memory-pool tests: require CUDA-capable GPU(s). -// cuda_inner needs at least 1 GPU; cuda2cuda needs ≥2 distinct GPUs. -// Both are `#[ignore]`-gated because standard CI runners lack GPUs. -// Run locally: -// cargo test --test example-smoke -- --ignored cuda -#[test] -#[ignore = "requires CUDA GPU(s)"] -fn smoke_memory_pool_cuda_inner() { - run_smoke_test( - "memory-pool-cuda-inner", - "examples/memory-pool/cuda_inner.yml", - Duration::from_secs(60), - ); -} - -#[test] -#[ignore = "requires CUDA GPU(s) — ≥2 GPUs"] -fn smoke_memory_pool_cuda2cuda() { - run_smoke_test( - "memory-pool-cuda2cuda", - "examples/memory-pool/cuda2cuda.yml", - Duration::from_secs(60), - ); -} - // --------------------------------------------------------------------------- // Examples under `examples/` that do NOT have a corresponding `smoke_*` or // `contract_*` test in this file. Some are blocked (filed issue or external @@ -2035,13 +1934,6 @@ fn smoke_memory_pool_cuda2cuda() { // // | Example | Where it's tested / blocker | Tracking | // |---------------------------|------------------------------------------------------|----------| -// | memory-pool | covered: smoke_memory_pool_cpu2cpu / | #2264 | -// | | smoke_local_memory_pool_{cpu2cpu, auto_cleanup, | | -// | | duplicate_free, read_after_free, write_after_free} | | -// | | (#[ignore]); nightly memory-pool-smoke job; | | -// | | smoke-all.sh gates on `import torch`, skips | | -// | | gracefully when download.pytorch.org unreachable. | | -// | | cuda2cpu/cpu2cuda/etc blocked: needs NVIDIA CUDA. | | // | cuda-benchmark | blocker: needs NVIDIA CUDA toolkit | — | // | dynamic-add-remove | blocker: `dora node add` times out + | #1682 | // | | corrupts dataflow state | |