Repository navigation
Conversation
## Problem
A storage unit binds its put/get ROUTER to a TCP port. Anything that
dials that
port is forwarded straight to the worker, and the worker then dies
decoding it:
`ZMQMessage.deserialize` sat outside the worker's `try`, so one stray
payload
terminated the request loop. The actor process, its bound socket and the
proxy
thread all stayed up, so requests kept being accepted with nobody left
to answer
them and every client blocked until its own timeout.
## Fix
1. **Filter at the proxy.** `zmq.proxy()` offers no interception point,
so the
proxy now forwards by hand and keeps only the two identity prefixes real
peers
use (storage managers and the metrics collector).
2. **Survive bad frames anyway.** The filter cannot cover everything: an
allowed
peer can still send frames that fail to decode. Decoding moves inside
the
`try`, mirroring the controller's request loop, and the peer gets an
error
reply instead of hanging on `recv`.
3. **One source for the prefixes.** They now live in `zmq_utils` and the
identities are built from them, so a rename cannot silently make every
request
unroutable.
The controller already tolerates undecodable messages, so it needs no
change.
## Tests
`tests/test_simple_storage_unit.py`: 17 passed.
- `test_foreign_payload_dropped_and_unit_stays_usable` — feeds a real
TLS
ClientHello and asserts it draws *no reply at all* (dropped before the
worker,
unlike a decode failure), then that the unit still serves.
- `test_undecodable_request_answered_and_worker_survives` — an allowed
identity
sending undecodable frames gets `PUT_GET_ERROR` and the worker keeps
serving.
Each was confirmed to fail when its corresponding fix is reverted.
Signed-off-by: huniu20 <huniumail@gmail.com>
As title Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>
### Problem Every RPC built a socket from scratch: create DEALER → connect → send → recv → close. That pays a TCP and ZMTP handshake per request, and makes the peer's ROUTER accrete a fresh identity on every call. ### Approach `ZMQSocketPool` lends connected DEALER sockets and takes them back after a clean send/recv. The decorator `with_zmq_socket` now leases from a pool instead of building a socket per call. Four properties carry the correctness argument: **A lease is exclusive.** Responses do not echo their request's `request_id` (see `ZMQMessage.create`), so a reply is matched to its request only by arrival order. Two concurrent users of one socket would read each other's replies. **A failed lease discards its socket.** Any exception — including timeout and `CancelledError`, which `asyncio.gather` raises into siblings on the first failure — closes the socket rather than returning it. The request may already be on the wire, so its reply could still arrive and the next lessee would read it as its own. The notify path closes its socket explicitly for the same reason, without raising, and the pool refuses to park an already-closed socket. **Sockets are keyed by lease owner, then by address.** The owner is the running event loop, or the thread outside one. pyzmq silently rebinds an async socket to whatever loop it next sees, and one bound to a *closed* loop is the `Bad file descriptor` / SIGABRT failure that per-call context churn used to cause; a finished owner's sockets are swept on the next access. Keying by address rather than peer id means a peer restarted under the same id at a new address is never handed a socket still connected to the old one. **One pool per role.** Each role owns the pools for the requests it makes, and sockets are never shared across roles. The pool borrows its context and never terminates it. **Concurrency waits rather than opening more.** An async lease takes a permit per (owner, address) first, so sockets in flight follow `TQ_SOCKET_POOL_SIZE` instead of peak concurrency — the context's budget becomes a property of configuration. Waiting also costs less than the handshake it replaces: locally, 300-way concurrency to one peer ran 48ms against 71ms, with peak sockets 4 against 300. ### One pool per role The split is deliberate rather than incidental: a pool is scoped to one role and one kind of request, so every socket it holds is interchangeable but for the address it is connected to. | Role | Pool | Dials | Owner | |---|---|---|---| | `CLIENT` | `controller_rpc_pool` | controller `request_handle_socket` | the loop `TransferQueueClient` builds in `__init__`, or the caller's for the async client | | `STORAGE` | `storage_rpc_pool` | storage units' `put_get_socket` | the caller's loop | | `STORAGE` | `notify_pool` | controller `request_handle_socket` | its own dedicated loop and thread | | `CONTROLLER` | metrics pool | storage units' `put_get_socket` | the collector thread, with no loop at all | Sharing one pool across these would reuse nothing in the first place: they differ by owning loop, socket name, or timeout, and a socket can only be reused when all three match. A ZMQ socket also cannot safely cross event loops, so `notify_pool` running on its own loop could never have shared with another scenario regardless of how the pools were arranged. Keeping them separate buys isolation on top of that. One role's timeouts and cancellations only ever poison its own sockets, a role's pool is closed with the component that owns it, and the eviction of a finished loop's sockets cannot reach into another role's buckets. Reuse itself does not depend on the split — every entry point above sits on a long-lived owner — but attributing a failure does. ### Sizing `TQ_SOCKET_POOL_SIZE` (default 4, must be >= 1) is the cap for every pool: idle sockets per (owner, address) bucket, and with an async lease the sockets in flight to that address. Concurrency past it waits for one to come back. The cap multiplies by peer count while a context's ceiling does not, which is what decides the default. A storage manager addressing two thousand units can hold `4 x 2000 = 8004` sockets against the client context's 8192 — inside the budget, where 8 would have put it at 16000. That worst case is reachable rather than theoretical: a single peer under 50-way concurrency does saturate its cap. Raise it if one peer routinely sees more overlap than that, and lower it if descriptors are tighter than latency; a single `put` or `get` sends one merged request per unit, so per-peer concurrency comes only from operations overlapping. `AsyncSimpleStorageManager` warns at construction when the product exceeds its context's `ZMQ_MAX_SOCKETS`, naming both knobs. Past that ceiling a lease fails outright with `EMFILE`, so it is worth saying before the first request rather than during one. Ceilings, for reference: the client context allows 8192 (`TQ_CLIENT_ZMQ_MAX_SOCKETS`), the controller's 4096 (`TQ_CONTROLLER_ZMQ_MAX_SOCKETS`), and a storage unit keeps libzmq's 1023 — it binds three sockets regardless of fleet size. Raising any of them needs file descriptors to match (`ulimit -n`). ### Fixes found along the way - `AsyncTransferQueueClient.__init__` leaked the context and its native I/O threads when it raised after allocating it (an out-of-range `max_sockets`). Everything checkable without a live context is now checked first, and the context is destroyed if the remaining check fails. - `StorageManager.close()` raised `AttributeError` when a subclass rejected its config before calling `super().__init__()`, as the KV managers do — burying the constructor's real error under a teardown failure. Pre-existing on `main`. - `AsyncSimpleStorageManager.close()` had the same problem for `storage_rpc_pool` when the base constructor raised, e.g. on a controller handshake timeout. - `_notify_and_wait` decremented its ACK budget by the poll interval, so unrelated traffic on the socket could stretch the wait. It now uses one deadline for the whole wait. - `TQMetricsExporter` lazily minted its own `zmq.Context` with nothing to close it. It now borrows the controller's. - Pooling those queries then made the collector hold a socket per storage unit for the life of the process: buckets are per address and only swept when their owner ends, and the collector is a permanent daemon thread. Each query now closes its socket, as the notify path does, taking resident sockets from one per unit to zero. Collection walks every unit once per cycle, so a kept socket would be reused only a cycle later while occupying budget for the whole walk. ### Behaviour and API changes - `with_zmq_socket(...)` drops `socket_name`, `get_identity`, `get_context` and `timeout` in favour of `get_pool`; socket name and timeout live on the pool. - `TQMetricsExporter(role=..., zmq_context=...)`. Without a context, querying storage units raises `RuntimeError` instead of silently building a second context. The controller always passes one. - New `TQ_SOCKET_POOL_SIZE` (default 4, must be >= 1) sizes every role's pool. Reuse cannot be disabled: below 1 nothing is ever parked, so every request would pay a fresh connect while still looking pooled. - `ZMQSocketPool.alease` is the async entry point and enforces the cap by waiting; `lease` stays for the synchronous collector, which issues one request at a time and cannot await. Nesting one lease inside another raises, including across pools — a permit lives until its body ends, so nesting can wait on one the task or a sibling already holds, and separate semaphores do not break that cycle. An RPC's reply is what frees its permit, so a second RPC belongs after the first returns. - The controller's context now sets `MAX_SOCKETS` (4096, clamped to `ZMQ_SOCKET_LIMIT`) instead of keeping libzmq's default 1023. Its own two ROUTERs are nowhere near that, but the metrics exporter borrows this context, and a controller that cannot open a socket stops answering requests at all. - `StorageManager` gains `notify_pool`; `AsyncSimpleStorageManager` gains `storage_rpc_pool`. Both are closed before the context they live on. ### Tests Reuse is asserted from the peer's side rather than from pool state. Every socket dials with its own ZMQ identity, so one identity across many requests means the connection was reused and a new one means the old socket was discarded — the same property observed from outside, with no access to the pool's internals. Each role's own reuse is covered where that role is tested, since a pool belongs to a role. `test_client.py` drives it through the public client: repeated `get_partition_list()` calls must reach the controller over one connection, which is the user-visible payoff of pooling. `test_metrics.py` covers the collector, whose pool keys by thread because it runs with no event loop at all. Both were checked by breaking the pool's take path so no socket is ever handed back — each then fails with one identity per request, so neither passes vacuously. What remains in `tests/test_zmq_socket_pool.py` needs control over reply timing and loop lifetime that a caller-level test cannot reach: the timed-out socket that must not answer the next request; poisoned leases from both an exception and a cancellation; no reuse across event loops; no stale socket after a peer re-registers at a new address; and a burst above the cap. The cap and its permits are covered there too: 40-way concurrency against a cap of 4 must peak at 4 sockets and drop nothing; a busy peer must not stall requests to another, since permits are keyed like the buckets; nesting raises, both on one pool and across two; and consecutive leases are not mistaken for nesting. The rest cover context sharing and pool wiring (`test_zmq_shared_context.py`), the exporter borrowing its owner's context (`test_metrics.py`), and quiet teardown after a rejected config (`test_kv_storage_manager.py`). ### Verification `python -m compileall -q transfer_queue tutorial tests` — clean. `ruff check`, `ruff format --check` and `mypy` — clean. `python -m pytest -q`: | | passed | skipped | errors | |---|---|---|---| | this branch | 607 | 10 | 8 | | `main` | 584 | 10 | 8 | The 8 errors are the same on both: `tests/test_yuanrong_storage_client_e2e.py` cannot import the optional Yuanrong backend in this environment. --------- Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
…scend#171) ## Problem A storage unit stops answering and every client routed to it fails when its own recv timeout expires. Raising the timeout (400s to 1800s in our deployment) changed nothing. Evidence from one such failure: the unit's node was alive and serving other traffic throughout, a TCP connect to its `put_get_socket` succeeded, and the unit's own counters showed it fully healthy (1025 GET_DATA served, 9.6ms p99, 1.33GB RSS). It had served exactly one GET_DATA fewer than its cohort. So the unit never saw the request that timed out; it was lost between the two ends, not queued behind slow work. Ascend#168 protected the worker thread from dying, which is a different cause of the same symptom; here the thread was intact. That loss is invisible by construction. ZMQ connect is asynchronous and SNDHWM is 0, so a DEALER accepts `send()` into an unbounded local queue for a peer it has not reached yet. The caller only learns anything when its own RCVTIMEO expires, which is why no timeout value can distinguish a lost request from a slow one. Lowering SNDHWM would not help either, it only trades silent queuing for silent dropping. ## Fix 1. **Retry on a new connection.** Up to `TQ_SIMPLE_STORAGE_MAX_ATTEMPTS` (default 3, floored at 1). A fresh socket and TCP connection is the part that matters. Only a missing answer is retried: `zmq.error.Again` now raises `StorageUnitTimeout`, while an error the unit reported surfaces on the first attempt. Replaying is safe where it is used: get is read-only and put overwrites the same global indexes. A put carrying a `data_parser` is the exception and is sent once, because a retry would re-run the parser on the unit and the public API does not constrain its side effects. 2. **Self-diagnosing residual failure.** After the last attempt, ask the unit for its own counters over a fresh socket with a short timeout. That probe is served by the same worker thread as put and get, so an answer proves the request was lost in flight and silence means the unit itself stopped. The failure log carries that verdict plus tcp reachability, the unit's op counts and RSS, and the wire size of the request that failed (measured from the serialized ZMQ frames on put). 3. **Log volume unchanged in the steady state.** A recovered request logs one line and skips the diagnosis entirely. Heavy requests are logged only above `TQ_STORAGE_SLOW_REQUEST_SECONDS` (5s) or `TQ_STORAGE_LARGE_PAYLOAD_MB` (256MB), via shared `log_heavy_operation`: puts on the manager (serialized request frames + end-to-end RTT), gets on the unit (serialized response frames + local handling time). No tensor-`nbytes` estimate remains. ## Tests `tests/test_storage_request_retry.py`: 14 passed. Covers recovery on retry, the bounded attempt count, that a nonpositive count still issues one request, that errors reported by the unit are not retried, that a parser-backed put is not replayed while the ordinary put still is, that retry logs carry wire-size detail from the failed attempt, shared heavy-request logging thresholds, and each diagnosis verdict. Full non-e2e run: 623 passed, 10 skipped, plus `python -m compileall -q transfer_queue tutorial tests`. `test_yuanrong_storage_client_e2e.py` is excluded locally, it needs the optional `openyuanrong-datasystem` dependency. --------- Signed-off-by: jathonzhang <jathonzhang@tencent.com> Co-authored-by: jathonzhang <jathonzhang@tencent.com>
As title. --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>
## Summary A storage-unit request can be lost between the manager and the unit without either end raising an error. Ascend#171 recovers such a request by retrying it; this PR closes one of the paths it can be lost on, and adds the instrumentation needed to tell the loss apart from a slow unit. It also corrects two defects in the diagnosis that Ascend#171 introduced — both are on `main` today. ## What's here ### 1. Raise the client-facing ROUTER's accept queue ZMQ defaults `BACKLOG` to 100. That was simply the unset default: a full accept queue is drained **without an RST** (`tcp_abort_on_overflow=0`), so the client stays in `ESTABLISHED` waiting for a reply nobody will ever send, and the server never learns the request existed. Neither end raises. Default is now 4096, tunable via `TQ_STORAGE_ZMQ_BACKLOG` (set it to 100 to restore ZMQ's default for an A/B run). This only raises a queue ceiling — it cannot reject a connection that previously succeeded. ### 2. Make the loss measurable rather than inferred - **Arrival counters** (`requests_arrived`, `arrivals_by_op` in `get_metrics`). The per-op counters advance inside `monitor.measure()`, so they only move once a request *completes* — a request that arrived and then stalled is indistinguishable from one that never arrived. Counting at decode time separates the two, and a shortfall against the caller's send count localizes the loss to one side of the connection. - **Accept-queue probe** (`transfer_queue/utils/accept_probe.py`). Samples `Recv-Q` / `sk_drops` / `ListenOverflows` from the kernel's own view. **Opt-in and off by default** (`TQ_ACCEPT_PROBE_INTERVAL=0`): it shells out to `ss` on a timer, and depth must be sampled sub-second because the queue drains in milliseconds — which is why inspecting a unit *after* it hung always read zero. ### 3. Report an accept-queue drop once per drop `sk_drops` is a cumulative kernel counter, but the initial check compared it against the probe's *first* sample. Once a socket had ever dropped a connection the condition stayed true for the life of the process — at a 0.1 s interval that is ten errors per second per storage unit. The flood also destroyed the signal: repeating the same cumulative total says nothing about *when* drops happened. Now compares against the previous sample and reports both deltas. ### 4. Base the timeout diagnosis on evidence (fixes two defects on `main`) `_diagnose_storage_unit` reported `verdict=request_lost_in_flight` **unconditionally** whenever the post-failure probe answered. But a unit that resumed inside the 10 s diagnostic window answers that probe too, having merely finished the original request late — so the verdict asserted *where* the request was lost on evidence that only shows the unit is serving *now*. It also printed `ops=` from `op_stats`, which `_handle_get_metrics` populates only when Prometheus is enabled, so the line read `ops={}` as if the unit had served nothing. The verdict now follows the unit's own arrival counter for the failed operation: | Probe answered, and… | Verdict | |---|---| | no arrival recorded for that op | `request_lost_in_flight` | | an arrival recorded for that op | `arrived_but_unfinished` | | counters unavailable (e.g. older unit) | `unit_serving_again` | `op_stats` is now reported as `unavailable(prometheus_disabled)` rather than as an empty dict. **The counters could not actually be correlated as intended.** The worker keyed `arrivals_by_op` by `str(operation)`, which for a `(str, Enum)` member renders as `"ZMQRequestType.GET_DATA"`, while `op_stats` is keyed `"GET_DATA"` — the two dicts shared no keys at all. Fixed by keying on `operation.name`. Note `.value` would *also* miss: the enum values are short wire tokens (`GET_DATA = "GET"`). ### 5. Stop the accept probe on shutdown `AcceptQueueProbe.start()` was called but `stop()` had no caller. The finalizer tore down only the ZMQ resources, so a unit finalized without its process exiting left the daemon thread spawning `ss` on a timer indefinitely, and the window summary was never logged. Now stopped in `_shutdown_resources`, before the ZMQ teardown. ## Behavior and compatibility - **No public API change.** `transfer_queue/utils/zmq_utils.py` is untouched. - **One default changes:** ROUTER `BACKLOG` 100 → 4096. - **The probe is off by default** and imported lazily, so a run that does not enable it never loads the module or spawns the thread. - **Steady-state log volume is unchanged**; the near-full-queue warning fires once per probe, not per sample. ## New configuration | Variable | Default | Purpose | |---|---|---| | `TQ_STORAGE_ZMQ_BACKLOG` | `4096` | Accept-queue depth for the client-facing ROUTER | | `TQ_ACCEPT_PROBE_INTERVAL` | `0` (off) | Accept-queue sampling period, seconds; sub-second when enabled | ## Note on the rebase This branch previously carried Ascend#171's commits as its base and included a per-call socket-timeout change. Since then Ascend#171 merged and Ascend#167 replaced per-call sockets with `ZMQSocketPool`, which resolves timeouts from the pool instead of the decorator. That change therefore had nowhere left to live and has been dropped — `main` already routes the two timeouts through separate `storage_rpc_pool` / `storage_probe_pool` instances, which satisfies the same intent. Its test file went with it; pool timeout behaviour is covered by `tests/test_zmq_socket_pool.py::test_timed_out_socket_is_not_reused`. The branch is now four commits on top of `ba8880c`, all mine. ## Testing - `tests/test_accept_probe.py` — utilization and delta arithmetic, peak tracking across a drained burst, drop reported once per drop, each new drop still reported, warn-once behaviour, zero-backlog not dividing by zero, and the probe being stopped on shutdown. - `tests/test_storage_request_retry.py` — the diagnosis table above, plus the empty-`op_stats` case and the arrival-key spelling. - Full suite: **644 passed, 10 skipped**, against a **623**-passing baseline measured on `ba8880c` itself — the +21 are the new tests, no regressions. - `check_license.py`, `check_docstrings.py`, `ruff check` and `ruff format --check` all clean. `tests/test_yuanrong_storage_client_e2e.py` is excluded from those runs: it errors identically on the unmodified base (optional Yuanrong backend not installed locally). --------- Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
CLA Signature PassGinray, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
1 of 9 tasks
Ginray
force-pushed
the
hixl-ucx-payload-zero-copy-3
branch
from
September 22, 2026 11:44
420233f to
5d32fc4
Compare
CLA Signature PassGinray, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
## Motivation Large training samples can make the long-lived host-memory footprint of `SimpleStorage` a capacity bottleneck as training scale and sequence length increase. ## Design This PR adds an optional hybrid storage layer inside `SimpleStorageUnit` without changing the public API, controller, or client behavior. Each `(field, global index)` sample is routed independently according to its actual stored payload size: samples at or above the internal 1 MiB threshold are written to node-local SSD, while smaller samples remain in host memory. The hybrid store preserves existing PUT, GET, CLEAR, overwrite, capacity, and checkpoint semantics; it reuses decoded contiguous buffers when possible and falls back to pickle for composite values. SSD files are temporary, isolated by run and storage unit, and removed during clear or graceful shutdown. The feature is disabled by default and configured through `ssd_offload.enabled` and `ssd_offload.path`. ## E2E Performance The following results are averages over four iterations with a batch size of 512 and eight fields. Throughput is reported in Gbit/s. **Test environment:** 2 × Intel Xeon Platinum 8480C CPUs (224 cores), 10 × NVIDIA/Mellanox ConnectX-7 NDR 400 Gb/s InfiniBand adapters plus an Intel E810-C 100 GbE adapter, and 8 × KIOXIA CM6-R 3.84 TB NVMe SSDs (`KCM6DRUL3T84`). 2 Node. | Per-field sample size | Total payload | Memory PUT | SSD PUT | Memory GET | SSD GET | Memory E2E | SSD E2E | E2E change | |---:|---:|---:|---:|---:|---:|---:|---:|---:| | 1 MiB | 4 GiB | 55.35 | 51.33 | 29.58 | 28.42 | 38.34 | 36.44 | -5.0% | | 2 MiB | 8 GiB | 57.17 | 53.53 | 25.37 | 24.65 | 35.04 | 33.73 | -3.7% | | 4 MiB | 16 GiB | 58.18 | 50.80 | 27.60 | 26.06 | 37.39 | 34.41 | -8.0% | --------- Signed-off-by: pinjie <pinjiex@nvidia.com>
As title Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>
As title Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>
…te… (Ascend#179) ## What Bind the Ascend ACL device context on every thread that calls Mooncake `register_buffer`/`unregister_buffer`, i.e. the worker threads of the per-call `ThreadPoolExecutor` in `put()`/`get()`. ## Why `MooncakeStoreClient.put()`/`get()` create a fresh `ThreadPoolExecutor` per call and register/unregister transient buffers inside those worker threads. ACL contexts are thread-local, so the workers had no context: - `register_buffer` silently fell back to host memory; - `unregister_buffer` failed with `INVALID_PARAMS`, leaking the registration and flooding logs with `aclrtGetCurrentContext failed`. ## Changes `transfer_queue/storage/clients/mooncake_client.py` - add `_ensure_accelerator_context()`: no-op when `torch.npu` is unavailable (CUDA/CPU/generic unaffected), otherwise `npu.set_device(...)` for the calling thread with a one-time context warmup; - call it once before `store.setup(...)` (the client thread needs a context too); - pass it as `initializer=_ensure_accelerator_context` to both `ThreadPoolExecutor`s (`put()` and `get()`). `_register_all_buffers`/`_unregister_all_buffers` run only inside those pools, so no other call sites are needed. ## Tests `tests/test_mooncake_client_thread_context.py` (hardware-free): - `_ensure_accelerator_context` sets the device and warms up only once; - it is a no-op when no NPU is available; - `put()` and `get()` build their executor with a non-None `initializer` that binds the context (executor faked, store faked). ``` python -m pytest -q tests/test_mooncake_client_thread_context.py # 4 passed python -m compileall -q transfer_queue tutorial tests ruff check / ruff format --check # clean ``` ## Verification on hardware (Ascend 910B + mooncake-transfer-engine 0.3.11.post1) Before: `Unregister buffer failed` / `aclrtGetCurrentContext failed` 1500+ times per short run. After: **0** occurrences; register/unregister succeed; no leak. Signed-off-by: ys2025-AI <yuanshuaien@163.com>
…d path (Ascend#175) Refs: Ascend#169 ## What this changes Tensor reads registered and unregistered a receive buffer around every transfer. `register_buffer` is a kernel operation (page pinning + MR setup) whose throughput is an order of magnitude below the RDMA transfer it enables, so it dominated read time and left the RDMA backend slower than the TCP one. This leases the receive buffer from the local buffer that `store.setup()` already registered, then copies into the caller's tensors so the returned tensors keep owning their memory. `split_by_bytes()` bounds each lease by one reader thread's share of the pool, so a read no longer needs `local_buffer_size` to exceed the batch size. Three cases still register their own receive buffers: builds without `mooncake.store.BufferPool` (mooncake-transfer-engine < v0.3.12), a group whose aligned total exceeds the per-thread share, and a pool that cannot serve the request. ### How the pool is sized `max_bytes=local_buffer_size` keeps in-flight leases inside the region `setup()` registered. On the default `max_bytes=0` mooncake budgets *twice* the local buffer and serves the excess through `allocate_overflow_unlocked()` — `posix_memalign` plus one `register_buffer` per lease — which is exactly the cost this path removes, and it would happen silently. `max_regions=MAX_BATCH_WORKER_THREADS` bounds concurrent leases, because each `get()` starts its own reader pool and nothing else limits how many run at once. The per-thread share is `local_buffer_size // (2 * MAX_BATCH_WORKER_THREADS)` — half of an equal split. The headroom is on our side rather than the pool's because `BufferPoolNative` clamps any `max_bytes` below `local_buffer_capacity_` back up to it, so `max_bytes` can only bound the overflow arena, never reserve space inside the registered region. mooncake allocates from that same region for its own read-side staging: remote DISK/DFS replica reads in `batch_get_into_multi_buffers()`, and the host staging `execute_ranged_read()` needs when the destination is device memory. Splitting the whole buffer across the reader threads would make those allocations fail — and inside one `batch_get_into` call, a failure returns `NO_AVAILABLE_HANDLE` while the lease is still held, so the 3 retries cannot help. ## Benchmarks Sections 1-3 were measured with the sizing as first submitted (share = 1/4 of the local buffer, `max_bytes=0`). Section 4 re-measures the sizing change on a different machine. ### 1. Tuning `BATCH_SIZE_LIMIT` does not substitute for the pool Register-per-read, threads=4, both reads swept over the same `BATCH_SIZE_LIMIT` grid (read time, ms, median; per-key size fixed at 1MB, only the key count differs): | BATCH_SIZE_LIMIT | 50 | 100 | 200 | 400 | 800 | 1600 | 3200 | |---|---|---|---|---|---|---|---| | 1024×1MB (1GB) | **109** | 138 | 155 | 160 | 253 | 384 | 329 | | 4096×1MB (4GB) | 528 | 461 | **455** | 496 | 666 | 723 | 1049 | | _pool @400 (ref)_ | | | | 1GB **103** / 4GB **434** | | | | The optimum sits at *small* values and grows worse as the batch grows: `register_buffer` cost is dominated by per-byte page pinning, not per-call overhead, so a larger batch does not reduce the pinning work — it only reduces how many worker threads share it. ### 2. The pool removes registration entirely Breakdown at threads=1, 1024×1MB (1GB), segment sums ≈ wall: | stage | register-per-read | pool | |---|---|---| | register_buffer | 279 ms (~79%) | 0 | | batch_get_into (RDMA) + copy/meta | ~75 ms | 104 ms | | **total** | **354 ms** | **104 ms (3.4×)** | <img width="869" height="685" alt="image" src="https://github.com/user-attachments/assets/1a93e94b-cfff-4371-ba4f-8d6ef898499d" /> ### 3. Where the win is big, and where it is honestly marginal baseline vs pool, threads=4, `BATCH_SIZE_LIMIT=400` (read time, ms, median): | workload | baseline | pool | speedup | |---|---|---|---| | 64×16MB (1GB) | 367 | 94 | **3.9×** | | 1024×1MB (1GB) | 178 | 109 | **1.6×** | | 4096×1MB (4GB) | 419 (best bsl) | 387 | 1.08× | | 16384×64KB (1GB) | 389 | 480 | **0.8× (slower)** | Two caveats: - **The advantage shrinks with more threads and larger reads.** At threads=1 the pool is ~3.4×; at threads=4 the register cost is spread across workers and the copy-out scales with bytes, so a 4GB read narrows the gain to ~1.08×. The pool also lowers step-to-step jitter (1024×1MB threads=1 stdev 19.6 → 7.8 ms). - **Many tiny keys can regress.** For 16384×64KB the mandatory copy-out does not pay off against that shape's cheap merged registration. Vectorizing the copy-out (uniform groups copied in one strided pass instead of a per-tensor loop) cut this case from ~2× slower to ~0.8×, but it is still a net loss at extreme fragmentation. A `num_keys`/per-key-size heuristic could fall back to register-per-read here; real trajectory reads (≈1024×1MB) do not hit it. The copy is what keeps returned tensors caller-owned so the buffer returns immediately. A zero-copy view mode would restore the large-read win but requires a "no cross-read retention" contract, so the copy path is the safe default. ### 4. What the reserved headroom costs protocol=rdma (mlx5), single node, `local_buffer_size`=1GB, threads=4, `BATCH_SIZE_LIMIT`=400, median of 10 reads ± stdev. All three modes are timed round-robin within each rep so machine drift is shared: | workload | register-per-read | share = 1/4 | share = 1/8 (this PR) | |---|---|---|---| | 64×16MB | 323.0 ± 7.3 | 96.6 ± 13.1 (3.34×) | 92.1 ± 3.7 (**3.51×**) | | 1024×1MB | 209.5 ± 40.9 | 115.6 ± 5.7 (1.81×) | 118.4 ± 5.0 (1.77×) | | 4096×1MB | 598.8 ± 56.4 | 417.7 ± 36.8 (1.43×) | 375.7 ± 70.1 (**1.59×**) | | 16384×64KB | 384.8 ± 39.0 | 466.7 ± 69.1 (0.82×) | 425.2 ± 48.0 (**0.90×**) | Leases taken per read confirm the extra rounds (4→8, 5→10, 21→41; 16384×64KB stays at 41 because `BATCH_SIZE_LIMIT` already caps each batch below both shares), but the extra rounds do not show up as time: three of four shapes are slightly faster with the smaller share, and 1024×1MB is 0.04× slower — well inside stdev. The headroom is effectively free. This run also reproduces two results from section 3 on different hardware: the large win on 64×16MB (3.9× there, 3.34-3.51× here) and the regression on 16384×64KB (0.8× there, 0.82× here, softened to 0.90× by the smaller share). Caveats: single node, so the transfer itself may be served by a local copy rather than RDMA — the eliminated registration and the round counts are real, but absolute numbers are not comparable to a cross-node run. Timing calls `MooncakeStoreClient.get` directly rather than `scripts/performance_test`, and the machine had residual load. Treat the three modes against each other as the result, not the absolute milliseconds. ## Compatibility The register-per-read path is preserved as a fallback: the client gates on `mooncake.store.BufferPool` availability, so builds without it (mooncake-transfer-engine < v0.3.12) transparently keep registering their own receive buffers. The "upgrade" warning is emitted once, and only where CPU RDMA reads would pay for registration — TCP and GDR reads stay quiet. Requests the pool cannot serve also fall back, warning once rather than per group per read. The GDR read/write paths and the non-tensor paths are untouched. ## Verification `tests/test_mooncake_buffer_pool.py` fakes the store and the pool, so it needs neither mooncake nor RDMA. The fake pool keeps the properties this path depends on: buffers are carved out of one arena and reused, exhaustion raises (which is what `acquire(block=False)` actually does — it never returns `None`), a buffer cannot be returned while an exported view of it is alive, and the pool cannot close while a buffer is out on loan. Covered: pooled reads with mixed dtypes and scalars, the strided copy-out for uniform groups, multi-round reads when the batch exceeds the share, register fallback for an oversized tensor / an exhausted pool / a build without `BufferPool`, buffer reuse not disturbing tensors returned by an earlier read, a failed transfer and a failed copy-out both returning the buffer and preserving the original exception, concurrent readers staying within the budget and the region cap, warn-once, the warning's scope, and `close()` ordering. Separately verified against a real `BufferPool` (TCP, in-process): the pool accepts these arguments, an over-capacity `acquire` raises `requested buffer size exceeds pool capacity`, the region cap raises `buffer pool is exhausted`, and `close()` is refused while a buffer is leased. --------- Signed-off-by: weiziyang <weiziyang@baidu.com> Co-authored-by: weiziyang <weiziyang@baidu.com>
Signed-off-by: Ginray <ginray0215@gmail.com>
Ginray
force-pushed
the
hixl-ucx-payload-zero-copy-3
branch
from
October 10, 2026 17:31
5d32fc4 to
ce3daea
Compare
Signed-off-by: Yinlei Sun <ginray0215@gmail.com>
Ginray
force-pushed
the
hixl-ucx-payload-zero-copy-3
branch
from
October 10, 2026 17:48
ce3daea to
90ad2c4
Compare
Ginray
marked this pull request as ready for review
October 12, 2026 12:33
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
Collaborator
[AI Review]Summary. The direction is good: replacing the packed format with frame-native scatter WRITE is a net simplification. In its current form, though, I'd suggest addressing two issues before merging: a performance regression on the default ZMQ backend, and a failure mode where one uncertain transfer permanently disables a client ↔ storage-unit path with no recovery owner. A few reliability risks also need validation under realistic RL patterns, and the change would be much easier to review if split. Locally, Scope note. Against Should fix before merge1. Default ZMQ backend: GET responses are serialized twice. TransferQueue/transfer_queue/storage/simple_storage.py Lines 1187 to 1198 in 90ad2c4 TransferQueue/transfer_queue/storage/simple_storage.py Lines 1049 to 1050 in 90ad2c4 For GET this walks the whole payload twice on the single storage worker thread (non-contiguous tensors are copied by 2. One uncertain failure permanently disables a path, with no recovery. TransferQueue/transfer_queue/storage/payload_transfer/nixl.py Lines 343 to 347 in 90ad2c4
For memory safety, quarantining the one receive buffer is already sufficient: it is never reused, and later GETs use a fresh transfer_id and buffer, so a late WRITE can only land in the quarantined memory. GET is read-only, so it can also be retried after quarantine. Suggest dropping the per-target permanent failure in favor of a byte budget on quarantined memory (fail fast only when exceeded). On the sender side, either reset the remote agent or fail the job explicitly rather than leaving a half-working client. Reliability risks to validate3. On storage units, all stored data becomes registered (pinned) memory, retained per whole PUT. TransferQueue/transfer_queue/storage/simple_storage.py Lines 276 to 277 in 90ad2c4
The 1000-cycle validation runs PUT/GET/CLEAR in lockstep, which is the one pattern that hides this. (The previous path copied into one per-PUT bytearray, so retention granularity was the same, but the memory was not pinned.) Consider keeping zero-copy on the client GET path only and copying out on storage-side PUT commit, which also makes the storage cache effective. At minimum, document this and validate with many outstanding PUTs. 4. The sender may remove and re-add the remote agent on almost every PUT. TransferQueue/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py Lines 262 to 264 in 90ad2c4 TransferQueue/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py Lines 646 to 652 in 90ad2c4 With storage-side buffers held long-term (point 3), nearly every PUT presents new metadata. Whether this rebuilds UCX endpoints/QPs depends on NIXL, but "QP count stayed at 2" was measured with one buffer reused 999 times. Please check QP/FD counts with many un-cleared PUTs. If NIXL supports loading partial metadata incrementally for a known agent, the 5. O(n²) overlap scan under the global runtime lock. TransferQueue/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py Lines 516 to 527 in 90ad2c4 Jagged nested tensors are unbound into one frame per sample on PUT. Timing this scan alone: 0.22 s for 4096 frames and 0.86 s for 8192, all while holding the runtime 6. Pending PUT/GET entries never expire. Simplicity and maintainability
Interfaces
Test gaps
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
9 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
PR #163 introduced the
PayloadTransferabstraction and the optional NIXL-UCX Host-to-Host payload backend forSimpleStorage.This PR further optimizes and hardens the NIXL-UCX path for RL workloads. The main goals are to reduce avoidable Host-side copies, support large frame-native payloads, avoid blocking
SimpleStorageworkers during GET transfers, and make buffer / transfer lifecycle handling safer under concurrency and failure.This work follows the optimization plan outlined in the Payload Transfer roadmap.
Summary
receive_buffer_cache_mbper runtime.SimpleStorageworker to continue processing requests while the NIXL transfer completes.Design
ZMQ remains the control plane, while non-empty Host payloads use NIXL-UCX.
The implementation intentionally keeps the scope limited to the
SimpleStorageHost payload fast path rather than introducing a general-purpose NIXL transport or MR management subsystem.Compatibility
SimpleStoragepublic APIs remain unchanged.Validation
Validation was performed on the current PR head with both two-node NIXL-UCX tests and local RL integration tests.
rma(rc_verbs/hns_0:1), including PUT / GET / CLEAR, datatype and nested-object handling, receive-buffer lease reuse, same-peer / different-peer concurrency, and controlled peer failure.data_ptr()addresses, with matching GET contents and SHA-256 hashes. Existing receive MRs were reused as send sources without another registration. No payload-sized staging copy was observed; small metadata frames were still copied.Follow-up
Continue reviewing whether the current lifecycle and failure-handling design can be simplified without weakening correctness.
Further validate performance and memory behavior in larger-scale distributed RL workloads.
Consider adding a resource limit to the idle registered receive pool. This may depend on the results of integration with RL frameworks, and I have not yet decided on the right approach.Current status: implemented a bounded idle receive-buffer cache with
receive_buffer_cache_mb(default: 256 MiB per runtime); oversized transient buffers are released after their leases end.Connect the runtime diagnostics to the existing TransferQueue metrics / Prometheus path.Current status: payload runtime diagnostics are exported through the existing metrics path as
tq_storage_payload_*metrics.