From 827aa033b85df37a4f2dc24a65e8d5e442ab2c69 Mon Sep 17 00:00:00 2001 From: Philipp Oppermann Date: Thu, 13 Aug 2026 13:07:28 +0000 Subject: [PATCH 1/3] refactor: park the memory-pool transport out of the tree MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extracts the pinned/CUDA memory-pool transport (#2168, #2386, #2619) into external/dora-pool, staged for lifting into its own repository. dora 1.0 ships no pool API: the four Python methods, the three DaemonRequest/two DaemonReply variants, the daemon-side registry and the node-api plumbing are all removed. No seam is left behind — #1872 declined to commit to this architecture, so dora should not ship a socket moulded to its shape. The parked copy is the post-#3014 code (pool reclamation, #2881), so it carries that fix rather than the pre-fix state. external/dora-pool/README.md leads with a seam contract: a budget for any future reinstatement (<200 lines of dora, no new unsafe), a table of what may never return in-tree, and re-entry criteria — #1872's five unanswered design questions, the open correctness bugs, and a GPU CI story. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/cargo-release.yml | 1 - .github/workflows/nightly.yml | 66 - .github/workflows/release.yml | 1 - CLAUDE.md | 2 +- Cargo.lock | 9 - Cargo.toml | 6 +- Makefile | 10 +- apis/python/node/Cargo.toml | 1 - apis/python/node/dora/__init__.pyi | 53 - apis/python/node/dora/cuda.py | 136 - apis/python/node/src/lib.rs | 3048 ---------------- .../node/src/daemon_connection/interactive.rs | 3 - .../node_integration_testing.rs | 3 - .../rust/node/src/event_stream/memory_pool.rs | 41 - apis/rust/node/src/event_stream/mod.rs | 2 - apis/rust/node/src/event_stream/thread.rs | 8 - apis/rust/node/src/lib.rs | 2 +- apis/rust/node/src/node/control_channel.rs | 62 - apis/rust/node/src/node/mod.rs | 30 - binaries/daemon/Cargo.toml | 1 - binaries/daemon/src/event_types.rs | 14 - binaries/daemon/src/lib.rs | 977 +----- binaries/daemon/src/node_communication/mod.rs | 44 - binaries/daemon/src/running_dataflow.rs | 121 +- external/dora-pool/.gitignore | 2 + external/dora-pool/README.md | 164 + external/dora-pool/daemon-side/Cargo.toml | 18 + .../dora-pool/daemon-side}/src/lib.rs | 0 .../dora-pool/examples}/README.md | 14 +- .../dora-pool/examples}/auto_cleanup.yml | 0 .../dora-pool/examples}/cpu2cpu.yml | 0 .../dora-pool/examples}/cpu2cuda.yml | 0 .../dora-pool/examples}/cuda2cpu.yml | 0 .../dora-pool/examples}/cuda2cuda.yml | 0 .../dora-pool/examples}/cuda_inner.yml | 0 .../dora-pool/examples}/duplicate_free.yml | 0 .../dora-pool/examples}/read_after_free.yml | 0 .../dora-pool/examples}/receiver.py | 0 .../dora-pool/examples}/sender.py | 0 .../dora-pool/examples}/write_after_free.yml | 0 .../dora-pool/python-binding/node_binding.rs | 3086 +++++++++++++++++ .../python-binding/tensor_info_helpers.py | 147 + external/dora-pool/tests/smoke-tests.rs | 111 + libraries/extensions/memory-pool/Cargo.toml | 13 - libraries/message/src/daemon_to_node.rs | 10 - libraries/message/src/node_to_daemon.rs | 21 +- scripts/qa/ci-nightly-jobs.sh | 4 +- scripts/smoke-all.sh | 22 - tests/example-smoke.rs | 108 - 49 files changed, 3553 insertions(+), 4808 deletions(-) delete mode 100644 apis/rust/node/src/event_stream/memory_pool.rs create mode 100644 external/dora-pool/.gitignore create mode 100644 external/dora-pool/README.md create mode 100644 external/dora-pool/daemon-side/Cargo.toml rename {libraries/extensions/memory-pool => external/dora-pool/daemon-side}/src/lib.rs (100%) rename {examples/memory-pool => external/dora-pool/examples}/README.md (92%) rename {examples/memory-pool => external/dora-pool/examples}/auto_cleanup.yml (100%) rename {examples/memory-pool => external/dora-pool/examples}/cpu2cpu.yml (100%) rename {examples/memory-pool => external/dora-pool/examples}/cpu2cuda.yml (100%) rename {examples/memory-pool => external/dora-pool/examples}/cuda2cpu.yml (100%) rename {examples/memory-pool => external/dora-pool/examples}/cuda2cuda.yml (100%) rename {examples/memory-pool => external/dora-pool/examples}/cuda_inner.yml (100%) rename {examples/memory-pool => external/dora-pool/examples}/duplicate_free.yml (100%) rename {examples/memory-pool => external/dora-pool/examples}/read_after_free.yml (100%) rename {examples/memory-pool => external/dora-pool/examples}/receiver.py (100%) rename {examples/memory-pool => external/dora-pool/examples}/sender.py (100%) rename {examples/memory-pool => external/dora-pool/examples}/write_after_free.yml (100%) create mode 100644 external/dora-pool/python-binding/node_binding.rs create mode 100644 external/dora-pool/python-binding/tensor_info_helpers.py create mode 100644 external/dora-pool/tests/smoke-tests.rs delete mode 100644 libraries/extensions/memory-pool/Cargo.toml 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..17d43f839b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2093,7 +2093,6 @@ dependencies = [ "dora-arrow-convert", "dora-core", "dora-download", - "dora-memory-pool", "dora-message", "dora-node-api", "dora-tracing", @@ -2227,13 +2226,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" @@ -2360,7 +2352,6 @@ dependencies = [ "pythonize", "serde_json", "serde_yaml", - "shared_memory_extended", "tokio", "tracing", ] diff --git a/Cargo.toml b/Cargo.toml index 7c34a2a79a..41f4d72b3f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,7 +67,6 @@ members = [ "libraries/extensions/download", "libraries/extensions/telemetry/*", "libraries/extensions/mavlink2-bridge", - "libraries/extensions/memory-pool", "libraries/extensions/ros2-bridge", "libraries/extensions/ros2-bridge/msg-gen", "libraries/extensions/ros2-bridge/python", @@ -92,6 +91,10 @@ members = [ "xtask", ] +# Parked, not built: the memory-pool transport extracted from the tree. +# See external/dora-pool/README.md. +exclude = ["external/dora-pool/daemon-side"] + [workspace.package] edition = "2024" # Raising this does not raise `dora-message`, which pins its own lower @@ -134,7 +137,6 @@ 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-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..434ed97069 100644 --- a/apis/python/node/Cargo.toml +++ b/apis/python/node/Cargo.toml @@ -35,7 +35,6 @@ pythonize = { workspace = true } futures = { workspace = true } dora-ros2-bridge-python = { workspace = true } dora-download = { workspace = true } -shared_memory_extended = "0.13.0" 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..1b51105ea2 100644 --- a/apis/python/node/dora/__init__.pyi +++ b/apis/python/node/dora/__init__.pyi @@ -215,59 +215,6 @@ class Node: node.send_output("string", b"string", {"open_telemetry_context": "7632e76"}) ```""" - def register_memory_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). - - 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. - - Args: - tensor_info: dict with keys "ptr", "size", "dtype", "shape", "device" - device: "cpu" or "cuda" — the receiver's device (controls DMA path) - - Returns: - pyarrow string array containing the buffer_id - """ - - def write_memory_pool( - self, memory_pool_id: pyarrow.Array, tensor_info: dict - ) -> None: - """Write new tensor data into an existing memory pool. - - Fast path: DORADMA header read from shmem, then memcpy data. - Slow path: falls back to daemon RPC. - - Args: - memory_pool_id: pyarrow string array from register_memory_pool - tensor_info: dict with keys "ptr", "size", "dtype", "shape", "device" - """ - - def read_memory_pool(self, memory_pool_id: pyarrow.Array) -> dict: - """Read tensor metadata from a memory pool. - - Fast path: DORADMA header read from shmem (zero-copy, no daemon). - Slow path: falls back to daemon RPC. - - 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" - """ - - def free_memory_pool(self, memory_pool_id: pyarrow.Array) -> None: - """Free a memory pool by removing it from the daemon's tracking table. - - Args: - memory_pool_id: pyarrow string array from register_memory_pool - """ - 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..40f58c1c4d 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,6 @@ 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); - let timeout = timeout_to_duration(timeout)?; let event = py.detach(|| self.events.recv(timeout)); if let Some(event) = event { @@ -1385,74 +251,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 +263,6 @@ impl Node { /// :rtype: list[dict] #[allow(clippy::should_implement_trait)] pub fn drain(&self, py: Python) -> PyResult>> { - self.process_pending_memory_pool_frees(py); let events = self .events .drain() @@ -1495,7 +292,6 @@ 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); match self.events.try_recv() { Ok(event) => event.to_py_dict(py).ok(), Err(_) => None, @@ -1530,13 +326,6 @@ 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)); - let timeout = timeout_to_duration(timeout)?; let event = self.events.recv_async_timeout(timeout).await; if let Some(event) = event { @@ -2011,1437 +800,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 +825,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) @@ -3647,347 +940,6 @@ impl Node { pub fn id(&self) -> String { 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, - }; - 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); - } - - Ok(Some(dict.into())) - } } /// 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/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..d417669379 100644 --- a/apis/rust/node/src/event_stream/mod.rs +++ b/apis/rust/node/src/event_stream/mod.rs @@ -40,8 +40,6 @@ mod data_conversion; mod event; /// 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..5cc024706f 100644 --- a/apis/rust/node/src/event_stream/thread.rs +++ b/apis/rust/node/src/event_stream/thread.rs @@ -144,14 +144,6 @@ 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()); - continue; - } - if let Some(tx) = tx.as_ref() { // `blocking_send` is used because this function runs on a // dedicated `std::thread` (not a tokio worker). Using 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..82def924ec 100644 --- a/apis/rust/node/src/node/control_channel.rs +++ b/apis/rust/node/src/node/control_channel.rs @@ -137,66 +137,4 @@ impl ControlChannel { other => bail!("unexpected OutputSent reply: {other:?}"), } } - - pub fn register_pinned_memory( - &mut self, - shared_memory_id: String, - metadata: Metadata, - ) -> eyre::Result<()> { - let request = DaemonRequest::RegisterPinnedMemory { - shared_memory_id, - metadata, - }; - let reply = self - .channel - .request(&Timestamped { - inner: request, - timestamp: self.clock.new_timestamp(), - }) - .wrap_err("failed to send RegisterPinnedMemory request to dora-daemon")?; - match reply { - DaemonReply::Result(Ok(())) => Ok(()), - DaemonReply::Result(Err(e)) => bail!("{e}"), - other => bail!("unexpected RegisterPinnedMemory reply: {other:?}"), - } - } - - pub fn read_pinned_memory( - &mut self, - shared_memory_id: String, - free: bool, - ) -> eyre::Result { - let request = DaemonRequest::ReadPinnedMemory { - shared_memory_id, - free, - }; - let reply = self - .channel - .request(&Timestamped { - inner: request, - timestamp: self.clock.new_timestamp(), - }) - .wrap_err("failed to send ReadPinnedMemory request to dora-daemon")?; - match reply { - DaemonReply::PinnedMemoryMetadata { metadata } => Ok(metadata), - DaemonReply::Result(Err(e)) => bail!("{e}"), - other => bail!("unexpected ReadPinnedMemory reply: {other:?}"), - } - } - - pub fn free_pinned_memory(&mut self, shared_memory_id: String) -> eyre::Result<()> { - let request = DaemonRequest::FreePinnedMemory { shared_memory_id }; - let reply = self - .channel - .request(&Timestamped { - inner: request, - timestamp: self.clock.new_timestamp(), - }) - .wrap_err("failed to send FreePinnedMemory request to dora-daemon")?; - match reply { - DaemonReply::Result(Ok(())) => Ok(()), - DaemonReply::Result(Err(e)) => bail!("{e}"), - other => bail!("unexpected FreePinnedMemory reply: {other:?}"), - } - } } diff --git a/apis/rust/node/src/node/mod.rs b/apis/rust/node/src/node/mod.rs index 1b06bbed86..fc2cdeaff6 100644 --- a/apis/rust/node/src/node/mod.rs +++ b/apis/rust/node/src/node/mod.rs @@ -2340,36 +2340,6 @@ impl DoraNode { ))), } } - - /// Register a pinned memory pool with the daemon for lifecycle tracking. - /// - /// 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( - &mut self, - shared_memory_id: String, - metadata: Metadata, - ) -> Result<(), eyre::Error> { - self.control_channel - .register_pinned_memory(shared_memory_id, metadata) - } - - /// Read pinned memory metadata from the daemon. - /// - /// When `free` is true, the daemon also frees the pool after reading. - pub fn read_pinned_memory( - &mut self, - shared_memory_id: String, - free: bool, - ) -> Result { - self.control_channel - .read_pinned_memory(shared_memory_id, free) - } - - /// 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) - } } /// Return the serialized log `fields` object when it fits `limit`, else `None`. diff --git a/binaries/daemon/Cargo.toml b/binaries/daemon/Cargo.toml index 252b0c2fcc..0e0eb078af 100644 --- a/binaries/daemon/Cargo.toml +++ b/binaries/daemon/Cargo.toml @@ -35,7 +35,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..91f310995a 100644 --- a/binaries/daemon/src/event_types.rs +++ b/binaries/daemon/src/event_types.rs @@ -147,20 +147,6 @@ pub enum DaemonNodeEvent { EventStreamDropped { reply_sender: oneshot::Sender, }, - RegisterPinnedMemory { - shared_memory_id: String, - metadata: metadata::Metadata, - reply_sender: oneshot::Sender, - }, - ReadPinnedMemory { - shared_memory_id: String, - free: bool, - reply_sender: oneshot::Sender, - }, - FreePinnedMemory { - shared_memory_id: String, - reply_sender: oneshot::Sender, - }, } #[derive(Debug)] diff --git a/binaries/daemon/src/lib.rs b/binaries/daemon/src/lib.rs index e12732010a..f9bb0adfbc 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::{ @@ -351,7 +350,6 @@ pub struct Daemon { pub(crate) builds: BTreeMap, pub(crate) git_manager: GitManager, pub(crate) metrics_system: Arc>, - pub(crate) memory_pool: MemoryPoolManager, /// 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 +649,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 +1493,6 @@ impl Daemon { zenoh_publish_tx, remote_daemon_events_tx, git_manager: Default::default(), - memory_pool: MemoryPoolManager::new(), builds, sessions: Default::default(), metrics_system: Arc::new(std::sync::Mutex::new(sysinfo::System::new())), @@ -2174,13 +1870,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 +2560,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 +2579,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 +2744,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 +3127,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 +3149,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,14 +3726,6 @@ 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| { - spawn_nodes.iter().any(|id| id.as_ref() == node) - }); - let mut logger = self .logger .for_dataflow(dataflow_id) @@ -4106,7 +3737,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 +3854,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)); } } } @@ -4943,147 +4565,6 @@ impl Daemon { let reply = inner.await.map_err(|err| format!("{err:?}")); let _ = reply_sender.send(DaemonReply::Result(reply)); } - DaemonNodeEvent::RegisterPinnedMemory { - shared_memory_id, - metadata, - 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 _ = reply_sender.send(DaemonReply::Result(result)); - } - DaemonNodeEvent::ReadPinnedMemory { - shared_memory_id, - free, - 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))); - } - } - } - DaemonNodeEvent::FreePinnedMemory { - shared_memory_id, - reply_sender, - } => { - let id = MemoryPoolId { - dataflow_id: dataflow_id.to_string(), - id: shared_memory_id, - }; - let result = free_pool_and_notify( - &self.memory_pool, - self.running.get(&dataflow_id), - &id, - node_id.as_ref(), - &self.clock, - ); - let _ = reply_sender.send(DaemonReply::Result(result)); - } } Ok(()) } @@ -5801,14 +5282,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,15 +5704,6 @@ 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, - dataflow_id, - &node_id, - ); - logger .log( if node_result.is_ok() { @@ -8154,238 +7619,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 +9040,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..026dac4efc 100644 --- a/binaries/daemon/src/node_communication/mod.rs +++ b/binaries/daemon/src/node_communication/mod.rs @@ -360,50 +360,6 @@ impl Listener { ) .await?; } - DaemonRequest::RegisterPinnedMemory { - shared_memory_id, - metadata, - } => { - let (reply_sender, reply) = oneshot::channel(); - self.process_daemon_event( - DaemonNodeEvent::RegisterPinnedMemory { - shared_memory_id, - metadata, - reply_sender, - }, - Some(reply), - connection, - ) - .await?; - } - DaemonRequest::ReadPinnedMemory { - shared_memory_id, - free, - } => { - let (reply_sender, reply) = oneshot::channel(); - self.process_daemon_event( - DaemonNodeEvent::ReadPinnedMemory { - shared_memory_id, - free, - reply_sender, - }, - Some(reply), - connection, - ) - .await?; - } - DaemonRequest::FreePinnedMemory { shared_memory_id } => { - let (reply_sender, reply) = oneshot::channel(); - self.process_daemon_event( - DaemonNodeEvent::FreePinnedMemory { - shared_memory_id, - reply_sender, - }, - Some(reply), - connection, - ) - .await?; - } // `DaemonRequest` is `#[non_exhaustive]`: a node built against a newer // dora-node-api may send a request this daemon predates. Answer with an // explicit error so the node fails loudly instead of hanging on a 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/external/dora-pool/.gitignore b/external/dora-pool/.gitignore new file mode 100644 index 0000000000..2c96eb1b65 --- /dev/null +++ b/external/dora-pool/.gitignore @@ -0,0 +1,2 @@ +target/ +Cargo.lock diff --git a/external/dora-pool/README.md b/external/dora-pool/README.md new file mode 100644 index 0000000000..14c39ee05e --- /dev/null +++ b/external/dora-pool/README.md @@ -0,0 +1,164 @@ +# dora-pool — parked memory-pool transport + +Pinned-host / CUDA memory-pool transport for zero-copy tensor handoff between +dora nodes. **Extracted from `dora-rs/dora` before the 1.0 release and parked +here.** It does not build against dora today, and dora 1.0 ships no pool API. + +This directory is staged for lifting into its own repository. Nothing here is +compiled by a dora build. + +--- + +## ⚠️ Read this before writing any code — the seam contract + +If you are reinstating this transport (human or agent), the single most +important constraint is: + +> **Add a seam to dora. Do not re-integrate the transport into dora.** + +The version parked here was *woven through* dora: 3,043 lines inside +`apis/python/node/src/lib.rs` (70% of that file), 952 lines of handlers and +reclamation logic in `binaries/daemon/src/lib.rs`, five wire-protocol variants, +and plumbing in six more files. That is the shape to avoid, and it is why this +was parked rather than maintained. + +### What dora may gain + +A seam, in this order of preference: + +1. **Nothing.** Check first whether the transport can be built entirely on + dora's existing public API (`DoraNode`, `send_output`, shared memory). + #1872 Q1 asks exactly this and was never answered — see below. +2. **One accessor**, if step 1 fails. Something shaped like + `node._pool_handle() -> PoolHandle`, where `PoolHandle` is an opaque + capability object the external package drives. Target: **under 100 lines in + dora, zero `unsafe`, no CUDA symbols, no knowledge of the DORADMA layout.** +3. **A generic extension point**, if a second consumer ever justifies it — + e.g. a registered side-channel over the existing daemon connection, with + dora carrying no pool-specific vocabulary at all. + +### What must NOT go back into dora + +Non-negotiable, because each of these is what made the parked version +unmaintainable: + +| Never in dora | Why | +|---|---| +| `unsafe` pointer arithmetic against the DORADMA header | ~40 sites reading hardcoded offsets (8/16/24/32/96) out of memory another process writes | +| The seqlock (`seqlock_begin_write` / `_begin_if_even` / `_end`) | Concurrency primitive owned by the transport, not the framework; still has an open correctness bug (#2890) | +| `unsafe impl Send`/`Sync` on raw-pointer slots | Six of them, asserting thread-safety on process-wide statics | +| Embedded Python that `ctypes`-loads `libcudart.so` | A ~300-line Python program living as a Rust string literal; it belongs in a Python package | +| CUDA transport selection (P2P, IPC handles, transit buffers, host staging) | Entirely the external package's concern | +| Pool-specific `DaemonRequest`/`DaemonReply` variants | Freezes pool vocabulary into dora's wire protocol; see the generic side-channel option above | +| Pool lifecycle logic in daemon node-exit / dataflow-finish paths | The parked version threaded reachability tracking through five separate exit paths | + +**Litmus test:** if `grep -ri 'cuda\|doradma\|pinned\|seqlock' ` over the dora +tree returns anything outside a docs file, the seam is wrong. + +**Budget:** a correct reinstatement should touch **under 200 lines of dora**, +across no more than two or three files, with no new `unsafe`. For scale: the +integration this replaces touched 4,356 lines across 17 files. + +No patch file is shipped here, deliberately — a ready-to-apply re-integration +sitting next to this contract would invite exactly the outcome it argues +against. If you need to see what the old integration touched, it is the reverse +of the commit that removed it: + +```bash +git -C /path/to/dora show # dora-rs/dora#3152 +``` + +--- + +## Why this was parked + +Not a judgement on the idea — the gap it addresses is real. It was parked +because the implementation never cleared the bar its own design issue set. + +[dora-rs/dora#1872](https://github.com/dora-rs/dora/issues/1872) ("Opt-in +pinned-host memory pool for high-throughput CPU→GPU transfers") was filed on +2026-05-19, the same day PR #1623 (+2,717 lines) was closed for being "7× +larger than the feature needed". #1872 says, verbatim: + +> This issue documents the gap and invites proposals. **It does NOT commit to +> an architecture.** dora has historically moved AWAY from custom shared-memory +> infrastructure (see #1745…). Any new transport has to clear a high bar. +> +> Priority: **not currently scheduled.** Filed for visibility. + +PR #2168 landed the same architecture 24 days later, merged 5 days after +opening, answering none of the five questions #1872 said a proposal must answer +"before writing code". + +### Re-entry criteria + +**A. Answer #1872's five questions** — in an issue or RFC, before code: + +1. **Why a new transport, not extending an existing one?** Can zenoh-shm take + an optionally-pinned host-memory provider backend? Can the Arrow IPC path be + made CUDA-aware? Why are these worse? +2. **What is the cross-platform story?** POSIX shm is Linux/macOS; Windows + needs `CreateFileMapping`. What is the graceful fallback when CUDA is absent + entirely — the feature must be invisible to the majority of users who never + touch a GPU. +3. **Lifecycle: who owns pinned memory and when is it freed?** Pinned memory is + a finite system resource. Bounded pool size? Freed on producer drop, daemon + shutdown, or explicit free? What happens when the producer crashes + mid-write? *This is the question the four bugs below all descend from.* +4. **API shape.** `node.send_output(..., pinned=True)` versus four new methods? + How does a consumer learn whether it can DMA directly or must fall back? +5. **Measurement.** Benchmark against **zenoh-shm plus caller-side pinning** — + the workaround that exists today — not against non-pinned zenoh-shm. Real + workload (1080p RGBA at 30 Hz ≈ 8 MB/frame), not a microbenchmark. #2168 + reported "4.5 GB/s locally" with no baseline at all. + +**B. Close the four open correctness bugs.** All were open against dora's 1.0 +milestone when this was parked: + +| Bug | Symptom | State when parked | +|---|---|---| +| [#3015](https://github.com/dora-rs/dora/issues/3015) | Pool ids collide across node restarts — a restarted node cannot re-register | PR #3056 open | +| [#2881](https://github.com/dora-rs/dora/issues/2881) | Pools not released when a node crashes or is dynamically removed | **fixed** — #3014 landed on dora `main` before this extraction, so the parked copy is the post-fix code | +| [#2935](https://github.com/dora-rs/dora/issues/2935) | Cross-process `FreeMemoryPool` cleanup silently skipped | **partial fix included** — `free_pool_and_notify` / `notify_memory_pool_freed` plus four regression tests; issue still open, verify before relying on it | +| [#2890](https://github.com/dora-rs/dora/issues/2890) | Seqlock overflow fix (#2866) incomplete — two inline end-write paths still use non-wrapping `old_gen + 1` | PR #3149 draft | + +**C. Have a GPU CI story.** The CUDA paths — IPC handles, P2P selection, +transit staging — never had automated coverage in dora. The only tested path +was CPU (`smoke_memory_pool_cpu2cpu` and five `smoke_local_memory_pool_*`, all +`#[ignore]`-gated, run in a dedicated nightly job). The unit tests that exist +are pure decision-matrix logic (`should_pin`, `classify_transport`, +`classify_write_path`) — they test which branch is chosen, never what the +branch does. + +--- + +## What is in here + +| Path | What it is | Builds? | +|---|---|---| +| `daemon-side/` | The `dora-memory-pool` crate: pool registry, metadata, reachability, orphan cleanup. Only depends on `tracing`. | **Yes**, standalone | +| `python-binding/node_binding.rs` | The 3,043 lines lifted out of `apis/python/node/src/lib.rs`, in original order with section markers. Statics, DORADMA header handling, seqlock, CUDA ctypes helpers, the four `#[pymethods]`, the receive-path free drain, and the `try_doradma_read` fast path. | **No** — needs the seam | +| `examples/` | Nine dataflow YAMLs (`cpu2cpu`, `cpu2cuda`, `cuda2cpu`, `cuda2cuda`, `cuda_inner`, plus four negative-lifecycle scenarios) with `sender.py` / `receiver.py`. | n/a | +| `python-binding/tensor_info_helpers.py` | `get_tensor_info` / `tensor_from_info` and their dtype maps, moved out of `dora/cuda.py` — they only ever fed the pool methods. | n/a | +| `tests/smoke-tests.rs` | The eight smoke tests removed from `tests/example-smoke.rs`. | **No** — needs a harness | + +`python-binding/node_binding.rs` does not compile on its own by design. It +needs from dora only: `&mut DoraNode` (for the three pinned-memory calls), +`node_id`, `dataflow_id`, and a free-queue drain — four calls and two fields. +That narrowness is the argument that a small seam is achievable; it is not a +licence to restore the patch. + +[#3014](https://github.com/dora-rs/dora/pull/3014) (the #2881 fix) landed on +dora `main` before this extraction, so `daemon-side/src/lib.rs` here is +byte-identical to the post-fix version — the fix travelled with the transport +rather than being stranded. + +## Lifting this into its own repository + +```bash +git subtree split -P external/dora-pool -b dora-pool-split +# then push dora-pool-split to the new repo's main +``` + +After lifting, delete `external/dora-pool` from dora and keep only a pointer to +the new repository. diff --git a/external/dora-pool/daemon-side/Cargo.toml b/external/dora-pool/daemon-side/Cargo.toml new file mode 100644 index 0000000000..27f35fc5b2 --- /dev/null +++ b/external/dora-pool/daemon-side/Cargo.toml @@ -0,0 +1,18 @@ +# Standalone: deliberately its own workspace root so this directory can be +# lifted into `dora-rs/dora-pool` with `git subtree` / `cp -r` and build +# unchanged. It is also listed in the dora workspace's `exclude` so the +# parked code is never compiled as part of a dora build. +[workspace] + +[package] +name = "dora-memory-pool" +version = "1.0.0-rc.4" +edition = "2024" +rust-version = "1.88" # HashMap::extract_if (1.87) + let-chains (1.88) +description = "Daemon-side memory-pool registry: pool metadata, lifecycle, reachability and orphan cleanup. Parked out of dora-rs/dora — see ../README.md." +license = "Apache-2.0" +repository = "https://github.com/dora-rs/dora-pool" +publish = false + +[dependencies] +tracing = "0.1.44" diff --git a/libraries/extensions/memory-pool/src/lib.rs b/external/dora-pool/daemon-side/src/lib.rs similarity index 100% rename from libraries/extensions/memory-pool/src/lib.rs rename to external/dora-pool/daemon-side/src/lib.rs diff --git a/examples/memory-pool/README.md b/external/dora-pool/examples/README.md similarity index 92% rename from examples/memory-pool/README.md rename to external/dora-pool/examples/README.md index bcd0d055cd..2634634079 100644 --- a/examples/memory-pool/README.md +++ b/external/dora-pool/examples/README.md @@ -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,10 +53,10 @@ 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: diff --git a/examples/memory-pool/auto_cleanup.yml b/external/dora-pool/examples/auto_cleanup.yml similarity index 100% rename from examples/memory-pool/auto_cleanup.yml rename to external/dora-pool/examples/auto_cleanup.yml diff --git a/examples/memory-pool/cpu2cpu.yml b/external/dora-pool/examples/cpu2cpu.yml similarity index 100% rename from examples/memory-pool/cpu2cpu.yml rename to external/dora-pool/examples/cpu2cpu.yml diff --git a/examples/memory-pool/cpu2cuda.yml b/external/dora-pool/examples/cpu2cuda.yml similarity index 100% rename from examples/memory-pool/cpu2cuda.yml rename to external/dora-pool/examples/cpu2cuda.yml diff --git a/examples/memory-pool/cuda2cpu.yml b/external/dora-pool/examples/cuda2cpu.yml similarity index 100% rename from examples/memory-pool/cuda2cpu.yml rename to external/dora-pool/examples/cuda2cpu.yml diff --git a/examples/memory-pool/cuda2cuda.yml b/external/dora-pool/examples/cuda2cuda.yml similarity index 100% rename from examples/memory-pool/cuda2cuda.yml rename to external/dora-pool/examples/cuda2cuda.yml diff --git a/examples/memory-pool/cuda_inner.yml b/external/dora-pool/examples/cuda_inner.yml similarity index 100% rename from examples/memory-pool/cuda_inner.yml rename to external/dora-pool/examples/cuda_inner.yml diff --git a/examples/memory-pool/duplicate_free.yml b/external/dora-pool/examples/duplicate_free.yml similarity index 100% rename from examples/memory-pool/duplicate_free.yml rename to external/dora-pool/examples/duplicate_free.yml diff --git a/examples/memory-pool/read_after_free.yml b/external/dora-pool/examples/read_after_free.yml similarity index 100% rename from examples/memory-pool/read_after_free.yml rename to external/dora-pool/examples/read_after_free.yml diff --git a/examples/memory-pool/receiver.py b/external/dora-pool/examples/receiver.py similarity index 100% rename from examples/memory-pool/receiver.py rename to external/dora-pool/examples/receiver.py diff --git a/examples/memory-pool/sender.py b/external/dora-pool/examples/sender.py similarity index 100% rename from examples/memory-pool/sender.py rename to external/dora-pool/examples/sender.py diff --git a/examples/memory-pool/write_after_free.yml b/external/dora-pool/examples/write_after_free.yml similarity index 100% rename from examples/memory-pool/write_after_free.yml rename to external/dora-pool/examples/write_after_free.yml diff --git a/external/dora-pool/python-binding/node_binding.rs b/external/dora-pool/python-binding/node_binding.rs new file mode 100644 index 0000000000..38bde199ee --- /dev/null +++ b/external/dora-pool/python-binding/node_binding.rs @@ -0,0 +1,3086 @@ +// PARKED CODE — DOES NOT COMPILE STANDALONE. DO NOT ADD TO A CARGO TARGET. +// +// The memory-pool transport lifted verbatim out of +// `apis/python/node/src/lib.rs` in dora-rs/dora before the 1.0 release +// (3,043 lines, ~70% of that file). Sections appear in their original order, +// separated by `// ==== ====` markers. +// +// It references `self.node_id`, `self.dataflow_id` and `self.node` from the +// `#[pyclass] Node` it used to live inside. Reinstating it means building a +// seam, NOT restoring the original integration. +// +// >>> READ ../README.md "The seam contract" BEFORE TOUCHING THIS FILE. <<< +// +// Short version: dora may gain one opaque accessor (<100 lines, no `unsafe`, +// no CUDA symbols, no DORADMA knowledge). Everything below — the seqlock, the +// header pointer arithmetic, the `unsafe impl Send`/`Sync` slots, the embedded +// libcudart ctypes module, the transport selection — stays on this side of +// the seam. What the old integration touched is recoverable as the reverse of +// dora-rs/dora#3152 — no patch file is shipped here on purpose. +// +// Open correctness bugs against this code when it was parked: #3015, #2935, +// #2890. (#2881 is fixed here — PR #3014 is included.) + + +// ==================== 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_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 + ); +} + + +// ==================== seqlock ==================== +/// 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); + } +} + +// ==================== process_pending_frees ==================== + /// 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); + } + } + } + + +// ==================== pool_api_methods ==================== + + // === 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(()) + } + +// ==================== drain_test_module ==================== + +#[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}`"); + } +} + +// ==================== try_doradma_read ==================== + + /// 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, + }; + 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); + } + + Ok(Some(dict.into())) + } + +// ==================== next_drain ==================== + // 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); + + +// ==================== drain_drain ==================== + self.process_pending_memory_pool_frees(py); + +// ==================== try_recv_drain ==================== + self.process_pending_memory_pool_frees(py); + +// ==================== recv_async_drain ==================== + // 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)); + diff --git a/external/dora-pool/python-binding/tensor_info_helpers.py b/external/dora-pool/python-binding/tensor_info_helpers.py new file mode 100644 index 0000000000..3ceaeddce3 --- /dev/null +++ b/external/dora-pool/python-binding/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_memory_pool` / `write_memory_pool` / `read_memory_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 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/external/dora-pool/tests/smoke-tests.rs b/external/dora-pool/tests/smoke-tests.rs new file mode 100644 index 0000000000..e7c9499d65 --- /dev/null +++ b/external/dora-pool/tests/smoke-tests.rs @@ -0,0 +1,111 @@ +// PARKED — not wired to any test harness. +// +// The eight smoke tests removed from `tests/example-smoke.rs` in dora-rs/dora +// when the memory-pool transport was parked. Paths have been rewritten to this +// package's `examples/`; the `run_smoke_test` / `run_smoke_test_local` helpers +// they call live in dora's test harness and would need reproducing here. +// +// The `memory-pool-smoke` nightly job referenced below no longer exists in +// dora's nightly.yml — it was removed with the transport. + +// --------------------------------------------------------------------------- +// 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/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/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/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/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/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/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/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/cuda2cuda.yml", + Duration::from_secs(60), + ); +} + 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/message/src/daemon_to_node.rs b/libraries/message/src/daemon_to_node.rs index f251d7b167..ca72dbe227 100644 --- a/libraries/message/src/daemon_to_node.rs +++ b/libraries/message/src/daemon_to_node.rs @@ -93,7 +93,6 @@ pub enum DaemonReply { Result(Result<(), String>), NextEvents(Vec>), NodeConfig { result: Result }, - PinnedMemoryMetadata { metadata: Metadata }, Empty, } @@ -177,15 +176,6 @@ pub enum NodeEvent { error: String, source_node_id: NodeId, }, - /// A memory pool has been freed by another node in the dataflow. - /// - /// 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, - }, } impl NodeEvent { diff --git a/libraries/message/src/node_to_daemon.rs b/libraries/message/src/node_to_daemon.rs index 8103dad992..a687211596 100644 --- a/libraries/message/src/node_to_daemon.rs +++ b/libraries/message/src/node_to_daemon.rs @@ -29,17 +29,6 @@ pub enum DaemonRequest { NodeConfig { node_id: NodeId, }, - RegisterPinnedMemory { - shared_memory_id: String, - metadata: Metadata, - }, - ReadPinnedMemory { - shared_memory_id: String, - free: bool, - }, - FreePinnedMemory { - shared_memory_id: String, - }, } impl DaemonRequest { @@ -78,10 +67,7 @@ impl DaemonRequest { | DaemonRequest::CloseOutputs(_) | DaemonRequest::OutputsDone | DaemonRequest::NextEvent - | DaemonRequest::EventStreamDropped - | DaemonRequest::RegisterPinnedMemory { .. } - | DaemonRequest::ReadPinnedMemory { .. } - | DaemonRequest::FreePinnedMemory { .. } => true, + | DaemonRequest::EventStreamDropped => true, } } @@ -96,10 +82,7 @@ impl DaemonRequest { | DaemonRequest::NextEvent | DaemonRequest::SendMessage { .. } | DaemonRequest::OutputSent { .. } - | DaemonRequest::EventStreamDropped - | DaemonRequest::RegisterPinnedMemory { .. } - | DaemonRequest::ReadPinnedMemory { .. } - | DaemonRequest::FreePinnedMemory { .. } => false, + | DaemonRequest::EventStreamDropped => 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 | | From 989a098a2bba838e91959452fd5734c20e4e5cca Mon Sep 17 00:00:00 2001 From: Philipp Oppermann Date: Thu, 13 Aug 2026 13:58:35 +0000 Subject: [PATCH 2/3] feat(extensions): add a generic extension channel for out-of-tree transports MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A transport that lives outside dora still needs one thing only the daemon can provide: reclamation after a crash. A node that dies cannot withdraw the descriptor it published, so its readers keep mappings to memory nobody owns (dora-rs/dora#2881 is that failure mode with a real transport attached). Adds a dataflow-scoped table of opaque byte values whose lifetime the daemon brokers — store / load / drop, plus a drained notification when a key goes away. dora never interprets the namespace, key or value. Deliberately generic rather than shaped around any one transport: naming the protocol variants after the memory pool would freeze that architecture into dora, which #1872 explicitly declined to do. A second extension needs no change here at all. Guarantees: only the storing node may overwrite a key; entries are scoped per dataflow and per namespace; every node that stored or read a key is notified when it is dropped; a dropped key is reclaimed on owner exit and on dataflow finish; dropping an absent key is a no-op so retries are safe. Bounded at 8192 entries per dataflow and 4096 pending notifications per process. This is a control plane for descriptors, not a data plane — values are copied through the daemon. See docs/extensions.md. Co-Authored-By: Claude Opus 5 (1M context) --- apis/python/node/dora/__init__.pyi | 29 ++ apis/python/node/src/lib.rs | 73 ++++ apis/rust/node/src/event_stream/extensions.rs | 101 +++++ apis/rust/node/src/event_stream/mod.rs | 2 + apis/rust/node/src/event_stream/thread.rs | 8 + apis/rust/node/src/node/control_channel.rs | 66 ++++ apis/rust/node/src/node/mod.rs | 44 +++ binaries/daemon/src/event_types.rs | 17 + binaries/daemon/src/extension_table.rs | 361 ++++++++++++++++++ binaries/daemon/src/lib.rs | 177 ++++++++- binaries/daemon/src/node_communication/mod.rs | 49 +++ docs/extensions.md | 90 +++++ external/dora-pool/README.md | 56 ++- .../dora-pool/python-binding/node_binding.rs | 13 +- libraries/message/src/daemon_to_node.rs | 19 +- libraries/message/src/node_to_daemon.rs | 36 +- 16 files changed, 1117 insertions(+), 24 deletions(-) create mode 100644 apis/rust/node/src/event_stream/extensions.rs create mode 100644 binaries/daemon/src/extension_table.rs create mode 100644 docs/extensions.md diff --git a/apis/python/node/dora/__init__.pyi b/apis/python/node/dora/__init__.pyi index 1b51105ea2..1afe05ec96 100644 --- a/apis/python/node/dora/__init__.pyi +++ b/apis/python/node/dora/__init__.pyi @@ -215,6 +215,35 @@ class Node: node.send_output("string", b"string", {"open_telemetry_context": "7632e76"}) ```""" + def extension_store(self, namespace: str, key: str, value: bytes) -> None: + """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 (see drain_dropped_extension_keys). + """ + + 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. + + 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/src/lib.rs b/apis/python/node/src/lib.rs index 40f58c1c4d..de422dc964 100644 --- a/apis/python/node/src/lib.rs +++ b/apis/python/node/src/lib.rs @@ -696,6 +696,79 @@ impl Node { self.dataflow_id.to_string() } + /// 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. 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/mod.rs b/apis/rust/node/src/event_stream/mod.rs index d417669379..9f234117aa 100644 --- a/apis/rust/node/src/event_stream/mod.rs +++ b/apis/rust/node/src/event_stream/mod.rs @@ -38,6 +38,8 @@ 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; /// Merged event streams combining internal and external event sources. diff --git a/apis/rust/node/src/event_stream/thread.rs b/apis/rust/node/src/event_stream/thread.rs index 5cc024706f..4316d6ebee 100644 --- a/apis/rust/node/src/event_stream/thread.rs +++ b/apis/rust/node/src/event_stream/thread.rs @@ -144,6 +144,14 @@ fn event_stream_loop( close_tx = true; } + // 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; + } + if let Some(tx) = tx.as_ref() { // `blocking_send` is used because this function runs on a // dedicated `std::thread` (not a tokio worker). Using diff --git a/apis/rust/node/src/node/control_channel.rs b/apis/rust/node/src/node/control_channel.rs index 82def924ec..8b0ac1de44 100644 --- a/apis/rust/node/src/node/control_channel.rs +++ b/apis/rust/node/src/node/control_channel.rs @@ -137,4 +137,70 @@ impl ControlChannel { other => bail!("unexpected OutputSent reply: {other:?}"), } } + + pub fn extension_store( + &mut self, + namespace: String, + key: String, + value: Vec, + ) -> eyre::Result<()> { + let request = DaemonRequest::ExtensionStore { + namespace, + key, + value, + }; + let reply = self + .channel + .request(&Timestamped { + inner: request, + timestamp: self.clock.new_timestamp(), + }) + .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 ExtensionStore reply: {other:?}"), + } + } + + pub fn extension_load( + &mut self, + namespace: String, + key: String, + remove: bool, + ) -> eyre::Result>> { + let request = DaemonRequest::ExtensionLoad { + namespace, + key, + remove, + }; + let reply = self + .channel + .request(&Timestamped { + inner: request, + timestamp: self.clock.new_timestamp(), + }) + .wrap_err("failed to send ExtensionLoad request to dora-daemon")?; + match reply { + DaemonReply::ExtensionValue { value } => Ok(value), + DaemonReply::Result(Err(e)) => bail!("{e}"), + other => bail!("unexpected ExtensionLoad reply: {other:?}"), + } + } + + 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 ExtensionDrop request to dora-daemon")?; + match reply { + DaemonReply::Result(Ok(())) => Ok(()), + DaemonReply::Result(Err(e)) => bail!("{e}"), + 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 fc2cdeaff6..b68d9af8ad 100644 --- a/apis/rust/node/src/node/mod.rs +++ b/apis/rust/node/src/node/mod.rs @@ -2340,6 +2340,50 @@ impl DoraNode { ))), } } + + /// Store an opaque value in the daemon's dataflow-scoped extension table. + /// + /// 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, + namespace: impl Into, + key: impl Into, + value: Vec, + ) -> Result<(), eyre::Error> { + self.control_channel + .extension_store(namespace.into(), key.into(), value) + } + + /// Read an opaque value back, optionally removing it in the same round trip. + /// + /// Returns `None` if the key is not in the table — never stored, or + /// already dropped. + pub fn extension_load( + &mut self, + namespace: impl Into, + key: impl Into, + remove: bool, + ) -> Result>, eyre::Error> { + self.control_channel + .extension_load(namespace.into(), key.into(), remove) + } + + /// 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()) + } } /// Return the serialized log `fields` object when it fits `limit`, else `None`. diff --git a/binaries/daemon/src/event_types.rs b/binaries/daemon/src/event_types.rs index 91f310995a..85afe78de1 100644 --- a/binaries/daemon/src/event_types.rs +++ b/binaries/daemon/src/event_types.rs @@ -147,6 +147,23 @@ pub enum DaemonNodeEvent { EventStreamDropped { reply_sender: oneshot::Sender, }, + ExtensionStore { + namespace: String, + key: String, + value: Vec, + reply_sender: oneshot::Sender, + }, + ExtensionLoad { + namespace: String, + key: String, + remove: bool, + reply_sender: oneshot::Sender, + }, + ExtensionDrop { + namespace: String, + key: String, + reply_sender: oneshot::Sender, + }, } #[derive(Debug)] 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 f9bb0adfbc..be97b331a4 100644 --- a/binaries/daemon/src/lib.rs +++ b/binaries/daemon/src/lib.rs @@ -159,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; @@ -181,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); @@ -242,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, @@ -350,6 +444,9 @@ pub struct Daemon { pub(crate) builds: BTreeMap, pub(crate) git_manager: GitManager, pub(crate) metrics_system: Arc>, + /// 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`. @@ -1493,6 +1590,7 @@ impl Daemon { zenoh_publish_tx, remote_daemon_events_tx, git_manager: Default::default(), + extensions: ExtensionTable::new(), builds, sessions: Default::default(), metrics_system: Arc::new(std::sync::Mutex::new(sysinfo::System::new())), @@ -4552,6 +4650,64 @@ impl Daemon { } => self .output_sent(dataflow_id, node_id, output_id, metadata) .context("failed to mark output sent")?, + DaemonNodeEvent::ExtensionStore { + namespace, + key, + value, + reply_sender, + } => { + 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::ExtensionLoad { + namespace, + key, + remove, + reply_sender, + } => { + 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::ExtensionDrop { + namespace, + key, + reply_sender, + } => { + let ext_key = ExtensionKey { + dataflow_id: dataflow_id.to_string(), + namespace, + key, + }; + drop_extension_and_notify( + &mut self.extensions, + self.running.get(&dataflow_id), + &ext_key, + &self.clock, + ); + // 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 @@ -5253,6 +5409,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, @@ -5704,6 +5868,17 @@ impl Daemon { dataflow.connected_nodes.remove(&node_id); } + // 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 .log( if node_result.is_ok() { diff --git a/binaries/daemon/src/node_communication/mod.rs b/binaries/daemon/src/node_communication/mod.rs index 026dac4efc..70f07d413a 100644 --- a/binaries/daemon/src/node_communication/mod.rs +++ b/binaries/daemon/src/node_communication/mod.rs @@ -360,6 +360,55 @@ impl Listener { ) .await?; } + DaemonRequest::ExtensionStore { + namespace, + key, + value, + } => { + let (reply_sender, reply) = oneshot::channel(); + self.process_daemon_event( + DaemonNodeEvent::ExtensionStore { + namespace, + key, + value, + reply_sender, + }, + Some(reply), + connection, + ) + .await?; + } + DaemonRequest::ExtensionLoad { + namespace, + key, + remove, + } => { + let (reply_sender, reply) = oneshot::channel(); + self.process_daemon_event( + DaemonNodeEvent::ExtensionLoad { + namespace, + key, + remove, + reply_sender, + }, + Some(reply), + connection, + ) + .await?; + } + DaemonRequest::ExtensionDrop { namespace, key } => { + let (reply_sender, reply) = oneshot::channel(); + self.process_daemon_event( + DaemonNodeEvent::ExtensionDrop { + namespace, + key, + reply_sender, + }, + Some(reply), + connection, + ) + .await?; + } // `DaemonRequest` is `#[non_exhaustive]`: a node built against a newer // dora-node-api may send a request this daemon predates. Answer with an // explicit error so the node fails loudly instead of hanging on a reply diff --git a/docs/extensions.md b/docs/extensions.md new file mode 100644 index 0000000000..333ccad8d8 --- /dev/null +++ b/docs/extensions.md @@ -0,0 +1,90 @@ +# Extensions: the out-of-tree transport seam + +Some transports do not belong in dora. A CUDA memory-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. + +## 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. + +## Prior art in this repo + +The pinned/CUDA memory-pool transport was built *inside* dora and removed +before 1.0 — it had grown to 3,000 lines in the Python binding with 64 `unsafe` +sites, plus 950 lines of daemon-side lifecycle logic, and no CI could exercise +its GPU paths. It now lives in `external/dora-pool`, whose README records both +the design questions it never answered (dora-rs/dora#1872) and the constraints +on bringing it back. + +That is the shape this seam exists to prevent: a transport should be a package +that uses these four operations, not a fork of the framework. diff --git a/external/dora-pool/README.md b/external/dora-pool/README.md index 14c39ee05e..fdabb26315 100644 --- a/external/dora-pool/README.md +++ b/external/dora-pool/README.md @@ -22,20 +22,44 @@ reclamation logic in `binaries/daemon/src/lib.rs`, five wire-protocol variants, and plumbing in six more files. That is the shape to avoid, and it is why this was parked rather than maintained. -### What dora may gain +### The seam already exists — use it -A seam, in this order of preference: +dora ships a generic **extension channel**: a dataflow-scoped table of opaque +byte values that the daemon brokers the *lifetime* of and nothing else. See +[`docs/extensions.md`](../../docs/extensions.md) in dora. -1. **Nothing.** Check first whether the transport can be built entirely on - dora's existing public API (`DoraNode`, `send_output`, shared memory). - #1872 Q1 asks exactly this and was never answered — see below. -2. **One accessor**, if step 1 fails. Something shaped like - `node._pool_handle() -> PoolHandle`, where `PoolHandle` is an opaque - capability object the external package drives. Target: **under 100 lines in - dora, zero `unsafe`, no CUDA symbols, no knowledge of the DORADMA layout.** -3. **A generic extension point**, if a second consumer ever justifies it — - e.g. a registered side-channel over the existing daemon connection, with - dora carrying no pool-specific vocabulary at all. +```python +node.extension_store(namespace, key, value) # bytes in +node.extension_load(namespace, key, remove=False) # bytes out, or None +node.extension_drop(namespace, key) # withdraw + notify +node.drain_dropped_extension_keys(namespace) # what went away +``` + +Rust nodes get the same on `DoraNode`, plus +`event_stream::extensions::drain_dropped_keys`. + +It carries no pool vocabulary at all — no CUDA, no DORADMA, no shared-memory +helpers. That is deliberate: a transport-shaped API in dora would freeze this +package's architecture into the framework, which is exactly what #1872 +declined to do. + +**How this package maps onto it** + +| Old in-tree call | Now | +|---|---| +| `register_pinned_memory(id, metadata)` | `extension_store("dora-pool", id, serialized_metadata)` | +| `read_pinned_memory(id, free)` | `extension_load("dora-pool", id, remove=free)` | +| `free_pinned_memory(id)` | `extension_drop("dora-pool", id)` | +| `drain_freed_pools()` | `drain_dropped_extension_keys("dora-pool")` | +| daemon-side registry, reachability, orphan sweep | the daemon's own reclamation — a crashed owner's entries are dropped and its readers notified | + +Everything else — the shared-memory segment, the DORADMA header, the seqlock, +the CUDA transport selection — stays on this side. The metadata that used to +travel as `MetadataParameters` becomes whatever bytes this package chooses; +the daemon does not look at them. + +If the seam turns out to be insufficient, the fix is to make it *more generic*, +not to add a pool-shaped API to dora. ### What must NOT go back into dora @@ -55,9 +79,11 @@ unmaintainable: **Litmus test:** if `grep -ri 'cuda\|doradma\|pinned\|seqlock' ` over the dora tree returns anything outside a docs file, the seam is wrong. -**Budget:** a correct reinstatement should touch **under 200 lines of dora**, -across no more than two or three files, with no new `unsafe`. For scale: the -integration this replaces touched 4,356 lines across 17 files. +**Budget:** reinstating this package should touch **zero lines of dora** — the +extension channel above is already there. If it genuinely cannot carry some +requirement, widen the *generic* channel rather than adding anything +pool-shaped, and keep it under a couple of hundred lines with no new `unsafe`. +For scale: the integration this replaces touched 4,356 lines across 17 files. No patch file is shipped here, deliberately — a ready-to-apply re-integration sitting next to this contract would invite exactly the outcome it argues diff --git a/external/dora-pool/python-binding/node_binding.rs b/external/dora-pool/python-binding/node_binding.rs index 38bde199ee..6843bb35df 100644 --- a/external/dora-pool/python-binding/node_binding.rs +++ b/external/dora-pool/python-binding/node_binding.rs @@ -11,11 +11,14 @@ // // >>> READ ../README.md "The seam contract" BEFORE TOUCHING THIS FILE. <<< // -// Short version: dora may gain one opaque accessor (<100 lines, no `unsafe`, -// no CUDA symbols, no DORADMA knowledge). Everything below — the seqlock, the -// header pointer arithmetic, the `unsafe impl Send`/`Sync` slots, the embedded -// libcudart ctypes module, the transport selection — stays on this side of -// the seam. What the old integration touched is recoverable as the reverse of +// Short version: dora already ships the seam — a generic extension channel +// (`extension_store` / `extension_load` / `extension_drop` / +// `drain_dropped_extension_keys`, see dora's `docs/extensions.md`). The three +// `*_pinned_memory` calls and `drain_freed_pools` below map onto it directly; +// the README has the table. Everything else — the seqlock, the header pointer +// arithmetic, the `unsafe impl Send`/`Sync` slots, the embedded libcudart +// ctypes module, the transport selection — stays on this side of the seam. +// What the old integration touched is recoverable as the reverse of // dora-rs/dora#3152 — no patch file is shipped here on purpose. // // Open correctness bugs against this code when it was parked: #3015, #2935, diff --git a/libraries/message/src/daemon_to_node.rs b/libraries/message/src/daemon_to_node.rs index ca72dbe227..a4ccdf73f9 100644 --- a/libraries/message/src/daemon_to_node.rs +++ b/libraries/message/src/daemon_to_node.rs @@ -92,7 +92,14 @@ pub enum DaemonCommunication { pub enum DaemonReply { Result(Result<(), String>), NextEvents(Vec>), - NodeConfig { result: Result }, + 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, } @@ -176,6 +183,16 @@ pub enum NodeEvent { error: String, source_node_id: NodeId, }, + /// An extension key this node stored or loaded has been dropped, by + /// another node or by the daemon reclaiming it. + /// + /// 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, + }, } impl NodeEvent { diff --git a/libraries/message/src/node_to_daemon.rs b/libraries/message/src/node_to_daemon.rs index a687211596..c93e8873a9 100644 --- a/libraries/message/src/node_to_daemon.rs +++ b/libraries/message/src/node_to_daemon.rs @@ -29,6 +29,32 @@ pub enum DaemonRequest { NodeConfig { node_id: NodeId, }, + /// 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, + }, + /// 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, + }, + /// 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, + }, } impl DaemonRequest { @@ -67,7 +93,10 @@ impl DaemonRequest { | DaemonRequest::CloseOutputs(_) | DaemonRequest::OutputsDone | DaemonRequest::NextEvent - | DaemonRequest::EventStreamDropped => true, + | DaemonRequest::EventStreamDropped + | DaemonRequest::ExtensionStore { .. } + | DaemonRequest::ExtensionLoad { .. } + | DaemonRequest::ExtensionDrop { .. } => true, } } @@ -82,7 +111,10 @@ impl DaemonRequest { | DaemonRequest::NextEvent | DaemonRequest::SendMessage { .. } | DaemonRequest::OutputSent { .. } - | DaemonRequest::EventStreamDropped => false, + | DaemonRequest::EventStreamDropped + | DaemonRequest::ExtensionStore { .. } + | DaemonRequest::ExtensionLoad { .. } + | DaemonRequest::ExtensionDrop { .. } => false, } } } From 61a9876352c4d64915fe3d7b82c1837fd4750ed5 Mon Sep 17 00:00:00 2001 From: Philipp Oppermann Date: Thu, 13 Aug 2026 14:25:24 +0000 Subject: [PATCH 3/3] feat(tensor-pool): reinstate the transport as an opt-in extension MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keeps the feature and its Python API surface, but out of dora's core: it lives at libraries/extensions/tensor-pool (beside ros2-bridge, same python/ subcrate shape) and reaches dora only through the generic extension channel. Opt-in and off by default, two independent flags: maturin develop -m apis/python/node/Cargo.toml --features tensor-pool cargo build -p dora-daemon --features tensor-pool A default build neither compiles nor exposes it — the receive-path drain becomes an empty no-op and dora's core keeps zero pool vocabulary. NOT covered by the 1.0 compatibility guarantees, stated in the README, the crate description, the module header, every pymethod docstring and the .pyi stubs, along with the open defects (#3015, #2935, #2890). Named tensor-pool, not memory-pool: dora already has an unrelated shared_memory_pool_size descriptor key (and DORA_NODE_SHM_POOL_SIZE) for the Zenoh SHM buffer pool, and one name for both was a persistent source of confusion. 'gpu-' would have been the other obvious fix but is inaccurate — the CPU path works without CUDA and is the only one with CI coverage. This renames the four Python methods (register_memory_pool -> register_tensor_pool and likewise for write/read/free), the crates, the feature flags and the example env keys: a user-visible break, permissible because the feature sits outside the 1.0 guarantees, and cheaper now than later. The six former daemon calls now go through extension_store / extension_load / extension_drop / drain_dropped_extension_keys, with the descriptor encoded as JSON dora never parses (python/src/seam.rs). The unsafe pointer arithmetic, the seqlock and the embedded libcudart bindings stay on the extension's side. Returning the crate to the workspace put it under -D warnings for the first time, which surfaced three latent bugs: a deprecated downcast_into, a dead initializer in the device-to-host copy path, and seqlock_begin_write with no callers at all (every write path uses begin_if_even). All fixed. Known gap: the smoke tests need a feature-built wheel plus torch, so they ship as smoke-tests.rs.example rather than a cargo target that cannot compile. The extension has 53 unit tests but no in-tree end-to-end coverage. Co-Authored-By: Claude Opus 5 (1M context) --- Cargo.lock | 23 ++ Cargo.toml | 10 +- apis/python/node/Cargo.toml | 5 + apis/python/node/dora/__init__.pyi | 30 ++ apis/python/node/src/lib.rs | 105 +++++ binaries/daemon/Cargo.toml | 5 + binaries/daemon/src/lib.rs | 9 + docs/extensions.md | 32 +- external/dora-pool/.gitignore | 2 - external/dora-pool/README.md | 190 --------- external/dora-pool/daemon-side/Cargo.toml | 18 - libraries/extensions/tensor-pool/Cargo.toml | 14 + libraries/extensions/tensor-pool/README.md | 124 ++++++ .../tensor-pool}/examples/README.md | 24 +- .../tensor-pool}/examples/auto_cleanup.yml | 4 +- .../tensor-pool}/examples/cpu2cpu.yml | 2 +- .../tensor-pool}/examples/cpu2cuda.yml | 2 +- .../tensor-pool}/examples/cuda2cpu.yml | 2 +- .../tensor-pool}/examples/cuda2cuda.yml | 0 .../tensor-pool}/examples/cuda_inner.yml | 0 .../tensor-pool}/examples/duplicate_free.yml | 4 +- .../tensor-pool}/examples/read_after_free.yml | 4 +- .../tensor-pool}/examples/receiver.py | 22 +- .../tensor-pool}/examples/sender.py | 14 +- .../examples/write_after_free.yml | 4 +- .../extensions/tensor-pool/python/Cargo.toml | 21 + .../extensions/tensor-pool/python/src/lib.rs | 14 + .../extensions/tensor-pool/python/src/seam.rs | 132 ++++++ .../tensor-pool/python/src/transport.rs | 385 +++++++----------- .../python}/tensor_info_helpers.py | 10 +- .../tensor-pool/smoke-tests.rs.example | 70 ++-- .../extensions/tensor-pool}/src/lib.rs | 216 +++++----- libraries/message/src/daemon_to_node.rs | 12 +- libraries/message/src/node_to_daemon.rs | 8 +- 34 files changed, 847 insertions(+), 670 deletions(-) delete mode 100644 external/dora-pool/.gitignore delete mode 100644 external/dora-pool/README.md delete mode 100644 external/dora-pool/daemon-side/Cargo.toml create mode 100644 libraries/extensions/tensor-pool/Cargo.toml create mode 100644 libraries/extensions/tensor-pool/README.md rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/README.md (84%) rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/auto_cleanup.yml (83%) rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/cpu2cpu.yml (95%) rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/cpu2cuda.yml (91%) rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/cuda2cpu.yml (91%) rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/cuda2cuda.yml (100%) rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/cuda_inner.yml (100%) rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/duplicate_free.yml (83%) rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/read_after_free.yml (83%) rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/receiver.py (80%) rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/sender.py (77%) rename {external/dora-pool => libraries/extensions/tensor-pool}/examples/write_after_free.yml (83%) create mode 100644 libraries/extensions/tensor-pool/python/Cargo.toml create mode 100644 libraries/extensions/tensor-pool/python/src/lib.rs create mode 100644 libraries/extensions/tensor-pool/python/src/seam.rs rename external/dora-pool/python-binding/node_binding.rs => libraries/extensions/tensor-pool/python/src/transport.rs (91%) rename {external/dora-pool/python-binding => libraries/extensions/tensor-pool/python}/tensor_info_helpers.py (94%) rename external/dora-pool/tests/smoke-tests.rs => libraries/extensions/tensor-pool/smoke-tests.rs.example (53%) rename {external/dora-pool/daemon-side => libraries/extensions/tensor-pool}/src/lib.rs (83%) diff --git a/Cargo.lock b/Cargo.lock index 17d43f839b..48d5b03e0a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2095,6 +2095,7 @@ dependencies = [ "dora-download", "dora-message", "dora-node-api", + "dora-tensor-pool", "dora-tracing", "dunce", "eyre", @@ -2344,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", @@ -2621,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 41f4d72b3f..74ead52416 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,6 +67,10 @@ members = [ "libraries/extensions/download", "libraries/extensions/telemetry/*", "libraries/extensions/mavlink2-bridge", + # 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", @@ -91,10 +95,6 @@ members = [ "xtask", ] -# Parked, not built: the memory-pool transport extracted from the tree. -# See external/dora-pool/README.md. -exclude = ["external/dora-pool/daemon-side"] - [workspace.package] edition = "2024" # Raising this does not raise `dora-message`, which pins its own lower @@ -137,6 +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-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/apis/python/node/Cargo.toml b/apis/python/node/Cargo.toml index 434ed97069..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,6 +39,7 @@ pythonize = { workspace = true } futures = { workspace = true } dora-ros2-bridge-python = { workspace = true } dora-download = { workspace = true } +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 1afe05ec96..48957faa0b 100644 --- a/apis/python/node/dora/__init__.pyi +++ b/apis/python/node/dora/__init__.pyi @@ -215,6 +215,36 @@ class Node: node.send_output("string", b"string", {"open_telemetry_context": "7632e76"}) ```""" + def register_tensor_pool( + self, tensor_info: dict, device: str = "cpu" + ) -> pyarrow.Array: + """Register a tensor's memory as a shared pool for zero-copy transfer. + + 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. + """ + + def write_tensor_pool( + self, tensor_pool_id: pyarrow.Array, tensor_info: dict + ) -> None: + """Write tensor data into an existing tensor pool. + + Extension method; see register_tensor_pool. + """ + + def read_tensor_pool(self, tensor_pool_id: pyarrow.Array) -> dict: + """Read tensor metadata from a tensor pool (zero-copy). + + Extension method; see register_tensor_pool. + """ + + 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 extension_store(self, namespace: str, key: str, value: bytes) -> None: """Store an opaque value in the daemon's dataflow-scoped extension table. diff --git a/apis/python/node/src/lib.rs b/apis/python/node/src/lib.rs index de422dc964..42cdaa0bd7 100644 --- a/apis/python/node/src/lib.rs +++ b/apis/python/node/src/lib.rs @@ -239,6 +239,9 @@ impl Node { #[pyo3(signature = (timeout=None))] #[allow(clippy::should_implement_trait)] pub fn next(&self, py: Python, timeout: Option) -> PyResult>> { + // 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 { @@ -263,6 +266,9 @@ impl Node { /// :rtype: list[dict] #[allow(clippy::should_implement_trait)] pub fn drain(&self, py: Python) -> PyResult>> { + // 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() @@ -292,6 +298,9 @@ impl Node { /// :rtype: dict #[allow(clippy::should_implement_trait)] pub fn try_recv(&mut self, py: Python) -> Option> { + // 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, @@ -326,6 +335,9 @@ impl Node { #[pyo3(signature = (timeout=None))] #[allow(clippy::should_implement_trait)] pub async fn recv_async(&self, timeout: Option) -> PyResult>> { + // 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 { @@ -696,6 +708,72 @@ 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 @@ -1013,6 +1091,33 @@ impl Node { pub fn id(&self) -> String { self.node_id.to_string() } + + /// 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, + }; + f(&mut pool) + } + + /// 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/binaries/daemon/Cargo.toml b/binaries/daemon/Cargo.toml index 0e0eb078af..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"] } diff --git a/binaries/daemon/src/lib.rs b/binaries/daemon/src/lib.rs index be97b331a4..af12374f76 100644 --- a/binaries/daemon/src/lib.rs +++ b/binaries/daemon/src/lib.rs @@ -3824,6 +3824,15 @@ impl Daemon { uv: bool, write_events_to: Option, ) -> eyre::Result> + use<>> { + // 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) + }); + let mut logger = self .logger .for_dataflow(dataflow_id) diff --git a/docs/extensions.md b/docs/extensions.md index 333ccad8d8..d22e5fd896 100644 --- a/docs/extensions.md +++ b/docs/extensions.md @@ -1,12 +1,15 @@ -# Extensions: the out-of-tree transport seam +# Extensions: the transport seam -Some transports do not belong in dora. A CUDA memory-pool needs `libcudart`, +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. +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 @@ -77,14 +80,19 @@ 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. -## Prior art in this repo +## The first consumer -The pinned/CUDA memory-pool transport was built *inside* dora and removed -before 1.0 — it had grown to 3,000 lines in the Python binding with 64 `unsafe` -sites, plus 950 lines of daemon-side lifecycle logic, and no CI could exercise -its GPU paths. It now lives in `external/dora-pool`, whose README records both -the design questions it never answered (dora-rs/dora#1872) and the constraints -on bringing it back. +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. -That is the shape this seam exists to prevent: a transport should be a package -that uses these four operations, not a fork of the framework. +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/external/dora-pool/.gitignore b/external/dora-pool/.gitignore deleted file mode 100644 index 2c96eb1b65..0000000000 --- a/external/dora-pool/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -target/ -Cargo.lock diff --git a/external/dora-pool/README.md b/external/dora-pool/README.md deleted file mode 100644 index fdabb26315..0000000000 --- a/external/dora-pool/README.md +++ /dev/null @@ -1,190 +0,0 @@ -# dora-pool — parked memory-pool transport - -Pinned-host / CUDA memory-pool transport for zero-copy tensor handoff between -dora nodes. **Extracted from `dora-rs/dora` before the 1.0 release and parked -here.** It does not build against dora today, and dora 1.0 ships no pool API. - -This directory is staged for lifting into its own repository. Nothing here is -compiled by a dora build. - ---- - -## ⚠️ Read this before writing any code — the seam contract - -If you are reinstating this transport (human or agent), the single most -important constraint is: - -> **Add a seam to dora. Do not re-integrate the transport into dora.** - -The version parked here was *woven through* dora: 3,043 lines inside -`apis/python/node/src/lib.rs` (70% of that file), 952 lines of handlers and -reclamation logic in `binaries/daemon/src/lib.rs`, five wire-protocol variants, -and plumbing in six more files. That is the shape to avoid, and it is why this -was parked rather than maintained. - -### The seam already exists — use it - -dora ships a generic **extension channel**: a dataflow-scoped table of opaque -byte values that the daemon brokers the *lifetime* of and nothing else. See -[`docs/extensions.md`](../../docs/extensions.md) in dora. - -```python -node.extension_store(namespace, key, value) # bytes in -node.extension_load(namespace, key, remove=False) # bytes out, or None -node.extension_drop(namespace, key) # withdraw + notify -node.drain_dropped_extension_keys(namespace) # what went away -``` - -Rust nodes get the same on `DoraNode`, plus -`event_stream::extensions::drain_dropped_keys`. - -It carries no pool vocabulary at all — no CUDA, no DORADMA, no shared-memory -helpers. That is deliberate: a transport-shaped API in dora would freeze this -package's architecture into the framework, which is exactly what #1872 -declined to do. - -**How this package maps onto it** - -| Old in-tree call | Now | -|---|---| -| `register_pinned_memory(id, metadata)` | `extension_store("dora-pool", id, serialized_metadata)` | -| `read_pinned_memory(id, free)` | `extension_load("dora-pool", id, remove=free)` | -| `free_pinned_memory(id)` | `extension_drop("dora-pool", id)` | -| `drain_freed_pools()` | `drain_dropped_extension_keys("dora-pool")` | -| daemon-side registry, reachability, orphan sweep | the daemon's own reclamation — a crashed owner's entries are dropped and its readers notified | - -Everything else — the shared-memory segment, the DORADMA header, the seqlock, -the CUDA transport selection — stays on this side. The metadata that used to -travel as `MetadataParameters` becomes whatever bytes this package chooses; -the daemon does not look at them. - -If the seam turns out to be insufficient, the fix is to make it *more generic*, -not to add a pool-shaped API to dora. - -### What must NOT go back into dora - -Non-negotiable, because each of these is what made the parked version -unmaintainable: - -| Never in dora | Why | -|---|---| -| `unsafe` pointer arithmetic against the DORADMA header | ~40 sites reading hardcoded offsets (8/16/24/32/96) out of memory another process writes | -| The seqlock (`seqlock_begin_write` / `_begin_if_even` / `_end`) | Concurrency primitive owned by the transport, not the framework; still has an open correctness bug (#2890) | -| `unsafe impl Send`/`Sync` on raw-pointer slots | Six of them, asserting thread-safety on process-wide statics | -| Embedded Python that `ctypes`-loads `libcudart.so` | A ~300-line Python program living as a Rust string literal; it belongs in a Python package | -| CUDA transport selection (P2P, IPC handles, transit buffers, host staging) | Entirely the external package's concern | -| Pool-specific `DaemonRequest`/`DaemonReply` variants | Freezes pool vocabulary into dora's wire protocol; see the generic side-channel option above | -| Pool lifecycle logic in daemon node-exit / dataflow-finish paths | The parked version threaded reachability tracking through five separate exit paths | - -**Litmus test:** if `grep -ri 'cuda\|doradma\|pinned\|seqlock' ` over the dora -tree returns anything outside a docs file, the seam is wrong. - -**Budget:** reinstating this package should touch **zero lines of dora** — the -extension channel above is already there. If it genuinely cannot carry some -requirement, widen the *generic* channel rather than adding anything -pool-shaped, and keep it under a couple of hundred lines with no new `unsafe`. -For scale: the integration this replaces touched 4,356 lines across 17 files. - -No patch file is shipped here, deliberately — a ready-to-apply re-integration -sitting next to this contract would invite exactly the outcome it argues -against. If you need to see what the old integration touched, it is the reverse -of the commit that removed it: - -```bash -git -C /path/to/dora show # dora-rs/dora#3152 -``` - ---- - -## Why this was parked - -Not a judgement on the idea — the gap it addresses is real. It was parked -because the implementation never cleared the bar its own design issue set. - -[dora-rs/dora#1872](https://github.com/dora-rs/dora/issues/1872) ("Opt-in -pinned-host memory pool for high-throughput CPU→GPU transfers") was filed on -2026-05-19, the same day PR #1623 (+2,717 lines) was closed for being "7× -larger than the feature needed". #1872 says, verbatim: - -> This issue documents the gap and invites proposals. **It does NOT commit to -> an architecture.** dora has historically moved AWAY from custom shared-memory -> infrastructure (see #1745…). Any new transport has to clear a high bar. -> -> Priority: **not currently scheduled.** Filed for visibility. - -PR #2168 landed the same architecture 24 days later, merged 5 days after -opening, answering none of the five questions #1872 said a proposal must answer -"before writing code". - -### Re-entry criteria - -**A. Answer #1872's five questions** — in an issue or RFC, before code: - -1. **Why a new transport, not extending an existing one?** Can zenoh-shm take - an optionally-pinned host-memory provider backend? Can the Arrow IPC path be - made CUDA-aware? Why are these worse? -2. **What is the cross-platform story?** POSIX shm is Linux/macOS; Windows - needs `CreateFileMapping`. What is the graceful fallback when CUDA is absent - entirely — the feature must be invisible to the majority of users who never - touch a GPU. -3. **Lifecycle: who owns pinned memory and when is it freed?** Pinned memory is - a finite system resource. Bounded pool size? Freed on producer drop, daemon - shutdown, or explicit free? What happens when the producer crashes - mid-write? *This is the question the four bugs below all descend from.* -4. **API shape.** `node.send_output(..., pinned=True)` versus four new methods? - How does a consumer learn whether it can DMA directly or must fall back? -5. **Measurement.** Benchmark against **zenoh-shm plus caller-side pinning** — - the workaround that exists today — not against non-pinned zenoh-shm. Real - workload (1080p RGBA at 30 Hz ≈ 8 MB/frame), not a microbenchmark. #2168 - reported "4.5 GB/s locally" with no baseline at all. - -**B. Close the four open correctness bugs.** All were open against dora's 1.0 -milestone when this was parked: - -| Bug | Symptom | State when parked | -|---|---|---| -| [#3015](https://github.com/dora-rs/dora/issues/3015) | Pool ids collide across node restarts — a restarted node cannot re-register | PR #3056 open | -| [#2881](https://github.com/dora-rs/dora/issues/2881) | Pools not released when a node crashes or is dynamically removed | **fixed** — #3014 landed on dora `main` before this extraction, so the parked copy is the post-fix code | -| [#2935](https://github.com/dora-rs/dora/issues/2935) | Cross-process `FreeMemoryPool` cleanup silently skipped | **partial fix included** — `free_pool_and_notify` / `notify_memory_pool_freed` plus four regression tests; issue still open, verify before relying on it | -| [#2890](https://github.com/dora-rs/dora/issues/2890) | Seqlock overflow fix (#2866) incomplete — two inline end-write paths still use non-wrapping `old_gen + 1` | PR #3149 draft | - -**C. Have a GPU CI story.** The CUDA paths — IPC handles, P2P selection, -transit staging — never had automated coverage in dora. The only tested path -was CPU (`smoke_memory_pool_cpu2cpu` and five `smoke_local_memory_pool_*`, all -`#[ignore]`-gated, run in a dedicated nightly job). The unit tests that exist -are pure decision-matrix logic (`should_pin`, `classify_transport`, -`classify_write_path`) — they test which branch is chosen, never what the -branch does. - ---- - -## What is in here - -| Path | What it is | Builds? | -|---|---|---| -| `daemon-side/` | The `dora-memory-pool` crate: pool registry, metadata, reachability, orphan cleanup. Only depends on `tracing`. | **Yes**, standalone | -| `python-binding/node_binding.rs` | The 3,043 lines lifted out of `apis/python/node/src/lib.rs`, in original order with section markers. Statics, DORADMA header handling, seqlock, CUDA ctypes helpers, the four `#[pymethods]`, the receive-path free drain, and the `try_doradma_read` fast path. | **No** — needs the seam | -| `examples/` | Nine dataflow YAMLs (`cpu2cpu`, `cpu2cuda`, `cuda2cpu`, `cuda2cuda`, `cuda_inner`, plus four negative-lifecycle scenarios) with `sender.py` / `receiver.py`. | n/a | -| `python-binding/tensor_info_helpers.py` | `get_tensor_info` / `tensor_from_info` and their dtype maps, moved out of `dora/cuda.py` — they only ever fed the pool methods. | n/a | -| `tests/smoke-tests.rs` | The eight smoke tests removed from `tests/example-smoke.rs`. | **No** — needs a harness | - -`python-binding/node_binding.rs` does not compile on its own by design. It -needs from dora only: `&mut DoraNode` (for the three pinned-memory calls), -`node_id`, `dataflow_id`, and a free-queue drain — four calls and two fields. -That narrowness is the argument that a small seam is achievable; it is not a -licence to restore the patch. - -[#3014](https://github.com/dora-rs/dora/pull/3014) (the #2881 fix) landed on -dora `main` before this extraction, so `daemon-side/src/lib.rs` here is -byte-identical to the post-fix version — the fix travelled with the transport -rather than being stranded. - -## Lifting this into its own repository - -```bash -git subtree split -P external/dora-pool -b dora-pool-split -# then push dora-pool-split to the new repo's main -``` - -After lifting, delete `external/dora-pool` from dora and keep only a pointer to -the new repository. diff --git a/external/dora-pool/daemon-side/Cargo.toml b/external/dora-pool/daemon-side/Cargo.toml deleted file mode 100644 index 27f35fc5b2..0000000000 --- a/external/dora-pool/daemon-side/Cargo.toml +++ /dev/null @@ -1,18 +0,0 @@ -# Standalone: deliberately its own workspace root so this directory can be -# lifted into `dora-rs/dora-pool` with `git subtree` / `cp -r` and build -# unchanged. It is also listed in the dora workspace's `exclude` so the -# parked code is never compiled as part of a dora build. -[workspace] - -[package] -name = "dora-memory-pool" -version = "1.0.0-rc.4" -edition = "2024" -rust-version = "1.88" # HashMap::extract_if (1.87) + let-chains (1.88) -description = "Daemon-side memory-pool registry: pool metadata, lifecycle, reachability and orphan cleanup. Parked out of dora-rs/dora — see ../README.md." -license = "Apache-2.0" -repository = "https://github.com/dora-rs/dora-pool" -publish = false - -[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/external/dora-pool/examples/README.md b/libraries/extensions/tensor-pool/examples/README.md similarity index 84% rename from external/dora-pool/examples/README.md rename to libraries/extensions/tensor-pool/examples/README.md index 2634634079..0b4255450a 100644 --- a/external/dora-pool/examples/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 @@ -61,13 +61,13 @@ 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/external/dora-pool/examples/auto_cleanup.yml b/libraries/extensions/tensor-pool/examples/auto_cleanup.yml similarity index 83% rename from external/dora-pool/examples/auto_cleanup.yml rename to libraries/extensions/tensor-pool/examples/auto_cleanup.yml index 4f3fc754ed..6851219530 100644 --- a/external/dora-pool/examples/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/external/dora-pool/examples/cpu2cpu.yml b/libraries/extensions/tensor-pool/examples/cpu2cpu.yml similarity index 95% rename from external/dora-pool/examples/cpu2cpu.yml rename to libraries/extensions/tensor-pool/examples/cpu2cpu.yml index 65361fb2a1..d9a55ed2be 100644 --- a/external/dora-pool/examples/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/external/dora-pool/examples/cpu2cuda.yml b/libraries/extensions/tensor-pool/examples/cpu2cuda.yml similarity index 91% rename from external/dora-pool/examples/cpu2cuda.yml rename to libraries/extensions/tensor-pool/examples/cpu2cuda.yml index 8d15ec5ccd..7b5aebccfa 100644 --- a/external/dora-pool/examples/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/external/dora-pool/examples/cuda2cpu.yml b/libraries/extensions/tensor-pool/examples/cuda2cpu.yml similarity index 91% rename from external/dora-pool/examples/cuda2cpu.yml rename to libraries/extensions/tensor-pool/examples/cuda2cpu.yml index 8227e3c41b..d5d776fbb7 100644 --- a/external/dora-pool/examples/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/external/dora-pool/examples/cuda2cuda.yml b/libraries/extensions/tensor-pool/examples/cuda2cuda.yml similarity index 100% rename from external/dora-pool/examples/cuda2cuda.yml rename to libraries/extensions/tensor-pool/examples/cuda2cuda.yml diff --git a/external/dora-pool/examples/cuda_inner.yml b/libraries/extensions/tensor-pool/examples/cuda_inner.yml similarity index 100% rename from external/dora-pool/examples/cuda_inner.yml rename to libraries/extensions/tensor-pool/examples/cuda_inner.yml diff --git a/external/dora-pool/examples/duplicate_free.yml b/libraries/extensions/tensor-pool/examples/duplicate_free.yml similarity index 83% rename from external/dora-pool/examples/duplicate_free.yml rename to libraries/extensions/tensor-pool/examples/duplicate_free.yml index 0d49c932e8..d986e50576 100644 --- a/external/dora-pool/examples/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/external/dora-pool/examples/read_after_free.yml b/libraries/extensions/tensor-pool/examples/read_after_free.yml similarity index 83% rename from external/dora-pool/examples/read_after_free.yml rename to libraries/extensions/tensor-pool/examples/read_after_free.yml index 89024d6a66..fb855cdd9a 100644 --- a/external/dora-pool/examples/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/external/dora-pool/examples/receiver.py b/libraries/extensions/tensor-pool/examples/receiver.py similarity index 80% rename from external/dora-pool/examples/receiver.py rename to libraries/extensions/tensor-pool/examples/receiver.py index 16aa73c4e2..5a56abd714 100644 --- a/external/dora-pool/examples/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/external/dora-pool/examples/sender.py b/libraries/extensions/tensor-pool/examples/sender.py similarity index 77% rename from external/dora-pool/examples/sender.py rename to libraries/extensions/tensor-pool/examples/sender.py index a30cfa784e..d39e5136c4 100644 --- a/external/dora-pool/examples/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/external/dora-pool/examples/write_after_free.yml b/libraries/extensions/tensor-pool/examples/write_after_free.yml similarity index 83% rename from external/dora-pool/examples/write_after_free.yml rename to libraries/extensions/tensor-pool/examples/write_after_free.yml index c00bc6fc7e..a0085a68e4 100644 --- a/external/dora-pool/examples/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/external/dora-pool/python-binding/node_binding.rs b/libraries/extensions/tensor-pool/python/src/transport.rs similarity index 91% rename from external/dora-pool/python-binding/node_binding.rs rename to libraries/extensions/tensor-pool/python/src/transport.rs index 6843bb35df..9c077b6570 100644 --- a/external/dora-pool/python-binding/node_binding.rs +++ b/libraries/extensions/tensor-pool/python/src/transport.rs @@ -1,29 +1,29 @@ -// PARKED CODE — DOES NOT COMPILE STANDALONE. DO NOT ADD TO A CARGO TARGET. -// -// The memory-pool transport lifted verbatim out of -// `apis/python/node/src/lib.rs` in dora-rs/dora before the 1.0 release -// (3,043 lines, ~70% of that file). Sections appear in their original order, -// separated by `// ==== ====` markers. -// -// It references `self.node_id`, `self.dataflow_id` and `self.node` from the -// `#[pyclass] Node` it used to live inside. Reinstating it means building a -// seam, NOT restoring the original integration. -// -// >>> READ ../README.md "The seam contract" BEFORE TOUCHING THIS FILE. <<< -// -// Short version: dora already ships the seam — a generic extension channel -// (`extension_store` / `extension_load` / `extension_drop` / -// `drain_dropped_extension_keys`, see dora's `docs/extensions.md`). The three -// `*_pinned_memory` calls and `drain_freed_pools` below map onto it directly; -// the README has the table. Everything else — the seqlock, the header pointer -// arithmetic, the `unsafe impl Send`/`Sync` slots, the embedded libcudart -// ctypes module, the transport selection — stays on this side of the seam. -// What the old integration touched is recoverable as the reverse of -// dora-rs/dora#3152 — no patch file is shipped here on purpose. -// -// Open correctness bugs against this code when it was parked: #3015, #2935, -// #2890. (#2881 is fixed here — PR #3014 is included.) - +//! 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 @@ -52,7 +52,7 @@ const FREED_POOL_IDS_CAP: usize = 4096; /// 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 +/// 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())); @@ -155,7 +155,7 @@ mod freed_pool_ids_tests { /// Per-pool persistent state. /// Keeping Shmem alive prevents munmap, preserving stable mmap addresses -/// for pool-hit detection across `register_memory_pool` calls. +/// for pool-hit detection across `register_tensor_pool` calls. /// /// # Safety /// `Shmem` is not `Send + Sync` due to raw pointer fields, but `PoolSlot` @@ -184,8 +184,8 @@ static PINNED_POOL: LazyLock>> = /// 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 +/// 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())); @@ -253,7 +253,7 @@ 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`. +/// 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. @@ -359,7 +359,7 @@ fn classify_transport( TransportPath::HostStagingTransit } -/// Which write path `write_memory_pool` dispatches to for a given frame. +/// 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 @@ -851,7 +851,7 @@ def _cuda_memcpy(dst, src, size, kind): 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 + 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") @@ -983,8 +983,8 @@ fn read_header_u64(p: *const u8) -> u64 { 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))?; +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::() { @@ -1005,37 +1005,23 @@ fn parse_memory_pool_id(memory_pool_id: Py, py: Python<'_>) -> eyre::Resu Ok(String::from_utf8(binary_array.value(0).to_vec())?) } else { eyre::bail!( - "memory_pool_id must be a string or binary array, got {:?}", + "tensor_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) { +fn warn_missing_tensor_pool(node_id: &NodeId, action: &str, buffer_id: &str) { tracing::warn!( - "[{}] Attempt to {} memory pool [{}] failed - reason: pool does not exist. Operation aborted.", + "[{}] Attempt to {} tensor pool [{}] failed - reason: pool does not exist. Operation aborted.", node_id, action, buffer_id ); } - // ==================== seqlock ==================== -/// 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) +/// 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. @@ -1056,7 +1042,7 @@ unsafe fn seqlock_begin_if_even(gen_ptr: *mut u64) -> u64 { } } -/// Closes a memory-pool seqlock write (header offset 96). +/// 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 @@ -1064,7 +1050,7 @@ unsafe fn seqlock_begin_if_even(gen_ptr: *mut u64) -> u64 { /// 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`. +/// in `write_tensor_pool`. unsafe fn seqlock_end(gen_ptr: *mut u64, pre_write_gen: u64, copy_ok: bool) { unsafe { if copy_ok { @@ -1151,10 +1137,24 @@ mod seqlock_tests { } } -// ==================== process_pending_frees ==================== - /// 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() { +// ==================== 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). { @@ -1220,34 +1220,32 @@ mod seqlock_tests { } } + // ==================== pool_api_methods ==================== -// ==================== pool_api_methods ==================== - - // === Memory Pool API === + // === Tensor Pool API === - /// Register a shared memory pool for zero-copy tensor transfer. + /// 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_memory_pool`] and [`free_memory_pool`] + /// 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_memory_pool`] while a receiver is + /// 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/memory-pool/` dataflows demonstrate this pattern. + /// 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. - #[pyo3(signature = (tensor_info, device))] - pub fn register_memory_pool( - &self, + pub fn register_tensor_pool( + &mut self, tensor_info: &Bound<'_, PyDict>, device: String, py: Python, @@ -1290,7 +1288,7 @@ mod seqlock_tests { } if cfg!(not(target_os = "linux")) { eyre::bail!( - "memory-pool transport requires Linux (uses /dev/shm). \ + "tensor-pool transport requires Linux (uses /dev/shm). \ This platform is not supported." ); } @@ -1397,19 +1395,20 @@ mod seqlock_tests { // 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) { + // 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); - dtoh_copy_ok = bound + bound .call_method1( "_cuda_memcpy", (shmem_ptr as u64 + data_offset as u64, ptr_val, size, 2u32), ) - .is_ok(); + .is_ok() } else { - dtoh_copy_ok = false; - } + false + }; if !dtoh_copy_ok { // The matching `_register_host` above is unconditional (it // runs whenever `!receiver_is_cuda`), so the unregister must @@ -1421,7 +1420,7 @@ mod seqlock_tests { } shmem.set_owner(true); eyre::bail!( - "[{}] register_memory_pool: DtoH copy failed ({} → CPU shmem, {} bytes)", + "[{}] register_tensor_pool: DtoH copy failed ({} → CPU shmem, {} bytes)", self.node_id, tensor_device, size @@ -1643,7 +1642,7 @@ mod seqlock_tests { } shmem.set_owner(true); eyre::bail!( - "[{}] register_memory_pool: failed to set up GPU pool buffer / IPC handle for CUDA receiver `{}`", + "[{}] register_tensor_pool: failed to set up GPU pool buffer / IPC handle for CUDA receiver `{}`", self.node_id, tensor_device ); @@ -1724,12 +1723,8 @@ mod seqlock_tests { ); 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); + if let Err(e) = crate::seam::store(self.node, &buffer_id, &meta) { + tracing::warn!("[{}] failed to register tensor pool: {:#}", self.node_id, e); } } @@ -1738,7 +1733,7 @@ mod seqlock_tests { Ok(buf_py) } - /// Write tensor data to an existing memory pool. + /// 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. @@ -1756,17 +1751,16 @@ mod seqlock_tests { /// 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 + /// 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. - #[pyo3(signature = (memory_pool_id, tensor_info))] - pub fn write_memory_pool( - &self, - memory_pool_id: Py, + pub fn write_tensor_pool( + &mut self, + tensor_pool_id: Py, tensor_info: &Bound<'_, PyDict>, py: Python, ) -> eyre::Result<()> { - let buffer_id = parse_memory_pool_id(memory_pool_id, py)?; + let buffer_id = parse_tensor_pool_id(tensor_pool_id, py)?; let ptr_val: u64 = tensor_info .get_item("ptr")? @@ -1785,7 +1779,7 @@ mod seqlock_tests { { 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); + warn_missing_tensor_pool(&self.node_id, "write", &buffer_id); return Ok(()); } } @@ -1803,7 +1797,7 @@ mod seqlock_tests { && 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 + // 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 @@ -1879,7 +1873,7 @@ mod seqlock_tests { && size > shmem_capacity.saturating_sub(data_offset)) { tracing::warn!( - "[{}] write_memory_pool: size {} exceeds available pool capacity (data_offset={}, total={}), operation aborted", + "[{}] write_tensor_pool: size {} exceeds available pool capacity (data_offset={}, total={}), operation aborted", self.node_id, size, data_offset, @@ -1908,7 +1902,7 @@ mod seqlock_tests { { copy_ok = false; tracing::error!( - "[{}] write_memory_pool: DMA copy failed: {}", + "[{}] write_tensor_pool: DMA copy failed: {}", self.node_id, e ); @@ -1928,7 +1922,7 @@ mod seqlock_tests { // retries until the next successful write. } if !copy_ok { - // Re-insert the slot so free_memory_pool + // 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() { @@ -1938,7 +1932,7 @@ mod seqlock_tests { .insert(counter, slot_data); } return Err(eyre::eyre!( - "[{}] write_memory_pool: DMA copy failed", + "[{}] write_tensor_pool: DMA copy failed", self.node_id )); } @@ -2035,7 +2029,7 @@ mod seqlock_tests { if let Err(e) = res { copy_ok = false; tracing::error!( - "[{}] write_memory_pool: GPU pool copy failed: {}", + "[{}] write_tensor_pool: GPU pool copy failed: {}", self.node_id, e ); @@ -2062,7 +2056,7 @@ mod seqlock_tests { .insert(counter, slot_data); } return Err(eyre::eyre!( - "[{}] write_memory_pool: GPU pool copy failed", + "[{}] write_tensor_pool: GPU pool copy failed", self.node_id )); } @@ -2102,11 +2096,7 @@ mod seqlock_tests { } // Slow path: query daemon for pool metadata - match self - .node - .get_mut() - .read_pinned_memory(buffer_id.clone(), false) - { + 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 { @@ -2143,7 +2133,7 @@ mod seqlock_tests { || (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", + "[{}] write_tensor_pool (slow path): size {} exceeds available pool capacity (data_offset={}, total={}), operation aborted", self.node_id, size, data_offset, @@ -2168,7 +2158,7 @@ mod seqlock_tests { { copy_ok = false; tracing::error!( - "[{}] write_memory_pool (slow path): DMA copy failed: {}", + "[{}] write_tensor_pool (slow path): DMA copy failed: {}", self.node_id, e ); @@ -2189,7 +2179,7 @@ mod seqlock_tests { } if !copy_ok { return Err(eyre::eyre!( - "[{}] write_memory_pool (slow path): DMA copy failed", + "[{}] write_tensor_pool (slow path): DMA copy failed", self.node_id )); } @@ -2274,7 +2264,7 @@ mod seqlock_tests { if let Err(e) = res { copy_ok = false; tracing::error!( - "[{}] write_memory_pool (slow path): GPU pool copy failed: {}", + "[{}] write_tensor_pool (slow path): GPU pool copy failed: {}", self.node_id, e ); @@ -2295,7 +2285,7 @@ mod seqlock_tests { } if !copy_ok { return Err(eyre::eyre!( - "[{}] write_memory_pool (slow path): GPU pool copy failed", + "[{}] write_tensor_pool (slow path): GPU pool copy failed", self.node_id )); } @@ -2323,14 +2313,14 @@ mod seqlock_tests { } } Err(_) => { - warn_missing_memory_pool(&self.node_id, "write", &buffer_id); + warn_missing_tensor_pool(&self.node_id, "write", &buffer_id); } } Ok(()) } - /// Read tensor info from an existing memory pool (zero-copy). + /// 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 — @@ -2341,7 +2331,7 @@ mod seqlock_tests { /// /// **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 + /// 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 @@ -2350,14 +2340,14 @@ mod seqlock_tests { /// 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 + /// write the next frame. The bundled `examples/tensor-pool/` dataflows /// demonstrate this pattern. - pub fn read_memory_pool( - &self, - memory_pool_id: Py, + pub fn read_tensor_pool( + &mut self, + tensor_pool_id: Py, py: Python, ) -> eyre::Result> { - let buffer_id = parse_memory_pool_id(memory_pool_id, py)?; + 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 @@ -2366,10 +2356,7 @@ mod seqlock_tests { 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) + 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) @@ -2409,17 +2396,13 @@ mod seqlock_tests { } 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); + 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) = self - .node - .get_mut() - .read_pinned_memory(buffer_id.clone(), false) - { + if let Ok(metadata) = crate::seam::load_metadata(self.node, &buffer_id, false) { let size = metadata .parameters .get("size") @@ -2466,9 +2449,9 @@ mod seqlock_tests { }) .unwrap_or(false); if ipc_present { - warn_missing_memory_pool(&self.node_id, "read", &buffer_id); + warn_missing_tensor_pool(&self.node_id, "read", &buffer_id); eyre::bail!( - "memory pool {}: fast path retries exhausted for GPU pool \ + "tensor pool {}: fast path retries exhausted for GPU pool \ (daemon fallback cannot provide a GPU pointer)", buffer_id ); @@ -2492,9 +2475,9 @@ mod seqlock_tests { if data_offset > shmem.len() || (size as usize) > shmem.len().saturating_sub(data_offset) { - warn_missing_memory_pool(&self.node_id, "read", &buffer_id); + warn_missing_tensor_pool(&self.node_id, "read", &buffer_id); eyre::bail!( - "memory pool {}: header bounds exceeded: \ + "tensor pool {}: header bounds exceeded: \ data_offset {} + size {} > shmem_len {}", buffer_id, data_offset, @@ -2509,9 +2492,9 @@ mod seqlock_tests { 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); + warn_missing_tensor_pool(&self.node_id, "read", &buffer_id); eyre::bail!( - "memory pool {}: daemon fallback: seqlock write in progress \ + "tensor pool {}: daemon fallback: seqlock write in progress \ (generation={}, odd)", buffer_id, read_gen @@ -2546,27 +2529,23 @@ mod seqlock_tests { } } } - warn_missing_memory_pool(&self.node_id, "read", &buffer_id); + warn_missing_tensor_pool(&self.node_id, "read", &buffer_id); eyre::bail!( - "memory pool {}: fast path retries exhausted — pool not ready after 500ms", + "tensor 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); + warn_missing_tensor_pool(&self.node_id, "read", &buffer_id); + eyre::bail!("tensor 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)?; + /// 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)?; - match self.node.get_mut().free_pinned_memory(buffer_id.clone()) { - Ok(_) => {} - Err(_) => { - warn_missing_memory_pool(&self.node_id, "release", &buffer_id); - } + 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). @@ -2643,7 +2622,7 @@ mod seqlock_tests { { 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={})", + "[{}] free_tensor_pool: adding {} to FREED_POOL_IDS (set size={})", self.node_id, buffer_id, freed.len() @@ -2654,76 +2633,9 @@ mod seqlock_tests { Ok(()) } -// ==================== drain_test_module ==================== - -#[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}`"); - } -} + // ==================== try_doradma_read ==================== -// ==================== try_doradma_read ==================== - - /// DORADMA fast path for read_memory_pool: reads metadata directly from + /// 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}"` → @@ -2732,7 +2644,7 @@ mod memory_pool_free_drain_tests { /// # Synchronization model /// /// The seqlock (write_gen at header offset 96) guards **data-byte** - /// consistency across `write_memory_pool` overwrites — the end-of-read + /// 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. @@ -2740,12 +2652,16 @@ mod memory_pool_free_drain_tests { /// 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 + /// 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(&self, buffer_id: &str, py: Python<'_>) -> eyre::Result>> { + 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). @@ -2835,7 +2751,7 @@ mod memory_pool_free_drain_tests { // 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(v) => match v.cast_into::() { Ok(d) => d, Err(_) => return Ok(None), }, @@ -3053,7 +2969,7 @@ mod memory_pool_free_drain_tests { // 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). + // is per-process and free_tensor_pool inserts there). { FREED_POOL_IDS .lock() @@ -3063,27 +2979,4 @@ mod memory_pool_free_drain_tests { Ok(Some(dict.into())) } - -// ==================== next_drain ==================== - // 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); - - -// ==================== drain_drain ==================== - self.process_pending_memory_pool_frees(py); - -// ==================== try_recv_drain ==================== - self.process_pending_memory_pool_frees(py); - -// ==================== recv_async_drain ==================== - // 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)); - +} diff --git a/external/dora-pool/python-binding/tensor_info_helpers.py b/libraries/extensions/tensor-pool/python/tensor_info_helpers.py similarity index 94% rename from external/dora-pool/python-binding/tensor_info_helpers.py rename to libraries/extensions/tensor-pool/python/tensor_info_helpers.py index 3ceaeddce3..b967a6be1b 100644 --- a/external/dora-pool/python-binding/tensor_info_helpers.py +++ b/libraries/extensions/tensor-pool/python/tensor_info_helpers.py @@ -1,7 +1,7 @@ """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_memory_pool` / `write_memory_pool` / `read_memory_pool` +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. @@ -55,7 +55,7 @@ # --------------------------------------------------------------------------- -# Tensor info helpers for memory-pool operations +# Tensor info helpers for tensor-pool operations # --------------------------------------------------------------------------- @@ -77,8 +77,8 @@ 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``. + 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() @@ -96,7 +96,7 @@ def tensor_from_info(tensor_info: dict) -> torch.Tensor: 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``. + a tensor pool via ``read_tensor_pool``. """ ptr = tensor_info.get('ptr', 0) if ptr == 0: diff --git a/external/dora-pool/tests/smoke-tests.rs b/libraries/extensions/tensor-pool/smoke-tests.rs.example similarity index 53% rename from external/dora-pool/tests/smoke-tests.rs rename to libraries/extensions/tensor-pool/smoke-tests.rs.example index e7c9499d65..71098cb9d0 100644 --- a/external/dora-pool/tests/smoke-tests.rs +++ b/libraries/extensions/tensor-pool/smoke-tests.rs.example @@ -1,33 +1,40 @@ -// PARKED — not wired to any test harness. +// Smoke tests for the tensor-pool extension — NOT a cargo test target. // -// The eight smoke tests removed from `tests/example-smoke.rs` in dora-rs/dora -// when the memory-pool transport was parked. Paths have been rewritten to this -// package's `examples/`; the `run_smoke_test` / `run_smoke_test_local` helpers -// they call live in dora's test harness and would need reproducing here. +// 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. // -// The `memory-pool-smoke` nightly job referenced below no longer exists in -// dora's nightly.yml — it was removed with the transport. +// 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. // --------------------------------------------------------------------------- -// Memory-pool CPU transport (#2168) +// 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 — -// `memory-pool-smoke` in .github/workflows/nightly.yml — which executes them +// `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 memory_pool -// (`memory_pool` matches both `smoke_memory_pool_*` and -// `smoke_local_memory_pool_*`; `smoke_memory_pool` would miss the latter). +// 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_memory_pool_cpu2cpu() { +fn smoke_tensor_pool_cpu2cpu() { run_smoke_test( - "memory-pool-cpu2cpu", + "tensor-pool-cpu2cpu", "examples/cpu2cpu.yml", Duration::from_secs(60), ); @@ -35,20 +42,16 @@ fn smoke_memory_pool_cpu2cpu() { #[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/cpu2cpu.yml", - 60, - ); +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_memory_pool_auto_cleanup() { +fn smoke_local_tensor_pool_auto_cleanup() { run_smoke_test_local( - "local-memory-pool-auto-cleanup", + "local-tensor-pool-auto-cleanup", "examples/auto_cleanup.yml", 10, ); @@ -56,9 +59,9 @@ fn smoke_local_memory_pool_auto_cleanup() { #[test] #[ignore = "requires `torch` and `tqdm` (not in standard CI)"] -fn smoke_local_memory_pool_duplicate_free() { +fn smoke_local_tensor_pool_duplicate_free() { run_smoke_test_local( - "local-memory-pool-duplicate-free", + "local-tensor-pool-duplicate-free", "examples/duplicate_free.yml", 10, ); @@ -66,9 +69,9 @@ fn smoke_local_memory_pool_duplicate_free() { #[test] #[ignore = "requires `torch` and `tqdm` (not in standard CI)"] -fn smoke_local_memory_pool_read_after_free() { +fn smoke_local_tensor_pool_read_after_free() { run_smoke_test_local( - "local-memory-pool-read-after-free", + "local-tensor-pool-read-after-free", "examples/read_after_free.yml", 10, ); @@ -76,24 +79,24 @@ fn smoke_local_memory_pool_read_after_free() { #[test] #[ignore = "requires `torch` and `tqdm` (not in standard CI)"] -fn smoke_local_memory_pool_write_after_free() { +fn smoke_local_tensor_pool_write_after_free() { run_smoke_test_local( - "local-memory-pool-write-after-free", + "local-tensor-pool-write-after-free", "examples/write_after_free.yml", 10, ); } -// GPU memory-pool tests: require CUDA-capable GPU(s). +// 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_memory_pool_cuda_inner() { +fn smoke_tensor_pool_cuda_inner() { run_smoke_test( - "memory-pool-cuda-inner", + "tensor-pool-cuda-inner", "examples/cuda_inner.yml", Duration::from_secs(60), ); @@ -101,11 +104,10 @@ fn smoke_memory_pool_cuda_inner() { #[test] #[ignore = "requires CUDA GPU(s) — ≥2 GPUs"] -fn smoke_memory_pool_cuda2cuda() { +fn smoke_tensor_pool_cuda2cuda() { run_smoke_test( - "memory-pool-cuda2cuda", + "tensor-pool-cuda2cuda", "examples/cuda2cuda.yml", Duration::from_secs(60), ); } - diff --git a/external/dora-pool/daemon-side/src/lib.rs b/libraries/extensions/tensor-pool/src/lib.rs similarity index 83% rename from external/dora-pool/daemon-side/src/lib.rs rename to libraries/extensions/tensor-pool/src/lib.rs index 673feff545..a0dbbbb03c 100644 --- a/external/dora-pool/daemon-side/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 a4ccdf73f9..c690e0ccf6 100644 --- a/libraries/message/src/daemon_to_node.rs +++ b/libraries/message/src/daemon_to_node.rs @@ -117,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, } } } @@ -218,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 c93e8873a9..eba9a8e22e 100644 --- a/libraries/message/src/node_to_daemon.rs +++ b/libraries/message/src/node_to_daemon.rs @@ -75,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, } }