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>
## 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>
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. |
1727d70 to
d4ca943
Compare
|
NIXL retries only cover a lost control-plane reply, not a failed or stuck WRITE. Retried today: Not retried:
Mapping
Related: when the peer answers PREPARE with an explicit error, Suggested:
|
|
This brings back the caller-side ERROR log that #160 removed. After #160, a key that a concurrent In this PR, Test (runs unchanged on both refs; it could go in import logging
import pytest
import transfer_queue.storage.managers.simple_storage_manager as manager_module
from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager
from transfer_queue.storage.simple_storage import KEY_NOT_FOUND_MARKER, StorageKeyNotFoundError
from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType
class _GetErrorSocket:
async def send_multipart(self, *args, **kwargs):
pass
async def recv_multipart(self, *args, **kwargs):
return ZMQMessage.create(
request_type=ZMQRequestType.GET_ERROR,
sender_id="su0",
body={
"message": "Failed to get data from storage unit id #su0, "
f"detail error message: {KEY_NOT_FOUND_MARKER}: key 1 not found"
},
).serialize()
@pytest.mark.asyncio
async def test_key_not_found_reply_is_not_logged_as_error_by_caller(caplog):
# A key cleared mid-read is a normal race: the caller rebuilds the type without an ERROR log.
mgr = object.__new__(AsyncSimpleStorageManager)
mgr.storage_manager_id = "TQ_STORAGE_test"
mgr.storage_unit_infos = {}
if hasattr(manager_module, "create_payload_transfer"):
mgr.payload_transfer = manager_module.create_payload_transfer(None)
with caplog.at_level(logging.DEBUG, logger="transfer_queue"):
with pytest.raises(StorageKeyNotFoundError):
await AsyncSimpleStorageManager._get_from_single_storage_unit.__wrapped__(
mgr, [1], ["log_probs"], target_storage_unit="su0", socket=_GetErrorSocket()
)
errors = [r for r in caplog.records if r.name.startswith("transfer_queue") and r.levelno >= logging.ERROR]
assert errors == [], [r.getMessage() for r in errors]Results (
Suggested fix: raise @@ -80,6 +80,9 @@ class ZmqPayloadTransfer(PayloadTransfer):
global_indexes: list[int],
fields: list[str],
) -> dict[str, Any]:
+ # Imported here because simple_storage imports this package.
+ from transfer_queue.storage.simple_storage import StorageKeyNotFoundError
+
request = ZMQMessage.create(
request_type=ZMQRequestType.GET_DATA,
sender_id=sender_id,
@@ -90,12 +93,15 @@ class ZmqPayloadTransfer(PayloadTransfer):
await control_socket.send_multipart(request.serialize())
response = ZMQMessage.deserialize(await control_socket.recv_multipart(copy=False))
if response.request_type != ZMQRequestType.GET_DATA_RESPONSE:
- raise RuntimeError(
- f"Failed to get data from storage unit {target_id}: {response.body.get('message', 'Unknown error')}"
- )
+ message = response.body.get("message", "Unknown error")
+ error_type = StorageKeyNotFoundError if KEY_NOT_FOUND_MARKER in message else RuntimeError
+ raise error_type(f"Failed to get data from storage unit {target_id}: {message}")
return response.body["data"]
except zmq.error.Again as exc:
self._raise_timeout(sender_id, target_id, "get", exc)
+ except StorageKeyNotFoundError:
+ # A key cleared mid-read is expected; the storage unit already logged it at debug.
+ raise
except Exception as exc:
logger.error(f"[{sender_id}]: Unexpected error from storage unit {target_id}: {type(exc).__name__}: {exc}")
raise RuntimeError( |
|
TransferQueue/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py Lines 250 to 262 in d4ca943 NIXL registers and deregisters memory on every transfer. This is the same cost #175 removed from the Mooncake read path, which this PR also merges, but the NIXL path still pays it. Send side (client for put, storage unit for get): Receive side: For reference, #175 measured Suggested direction, mirroring #175:
Trade-off: packing into the arena costs one memcpy. #175 saw this lose on many tiny keys (16384 x 64 KiB, 0.8-0.9x), so the fallback threshold matters. Could you add a register-vs-WRITE time breakdown for the two-node NIXL run (for example 64 x 16 MiB, 1024 x 1 MiB, 16384 x 64 KiB) before and after, so we can see what this buys on RDMA? |
e081fa5 to
2033bd1
Compare
Signed-off-by: Ginray <ginray0215@gmail.com>
2033bd1 to
332cc69
Compare
Thanks a lot for the detailed review. This PR is mainly for syncing/merging Most of these concerns are already addressed in the follow-up PR #176, commit One intentional difference is that we take a more conservative approach once a WRITE may have started: we do not automatically retry a generic Pending-transfer expiry and skipping the redundant cancel after an explicit PREPARE failure are not included yet. We’ll evaluate whether they are worth adding based on the implementation complexity and whether they can be kept simple without over-designing the lifecycle handling. |
Thanks a lot for catching this. It has been fixed in the current PR, and we also added the caller-side regression test as suggested. |
Thanks a lot for the detailed analysis. The follow-up PR #176, commit For the send side, we currently prefer to keep frame-native scatter rather than pack every payload into a long-lived registered arena, since that essentially trades registration overhead for an additional payload-sized memcpy. We also have a previous 1 GiB cold-start profile: for PUT, receive-buffer registration was 6.15 ms, send-source registration was 2.51 ms, while the detach copy was 237.60 ms; for GET, they were 2.45 ms, 5.27 ms, and 245.64 ms respectively. In that profile, the detach-copy cost was roughly 30x higher than the combined receive- and send-side registration cost. For the RL + TQ workloads we are currently targeting, payloads are more likely to contain relatively large tensor frames rather than being dominated by thousands of very small frames, so avoiding the additional payload-sized copy currently looks more suitable for the workload we are targeting. The multi-frame benchmarks you suggested are very useful for understanding how registration overhead scales with frame count. I’ll add those experiments, and we can also continue evaluating and adjusting the design based on validation with real RL workloads. |
|
[AI Review] I reviewed the merge commit itself ( Should fix before merge1. With NIXL + SSD offload, offloaded payloads stay in memoryTransferQueue/transfer_queue/storage/payload_transfer/nixl.py Lines 348 to 352 in 332cc69 TransferQueue/transfer_queue/storage/simple_storage.py Lines 537 to 545 in 332cc69
Repro: 8 x 8 MiB samples are offloaded, plus one tiny field. The output shows the bytes still alive after Mixed puts like this are normal in RL. The "NIXL+SSD public API checks" in the description cover behavior, not memory. Suggested fix: copy the values that stay in memory. They are below the threshold, so the copy is cheap. With this change the NIXL case retains 0.0 MiB, and the SSD tests in else:
# Kept values are small; copy so a view cannot pin the whole receive buffer.
prepared_values.append(value.clone() if isinstance(value, torch.Tensor) else value)Repro scriptimport gc
import tempfile
import tracemalloc
import torch
from transfer_queue.storage.simple_storage import HybridStorageUnitData
from transfer_queue.utils.serial_utils import calc_packed_size, decode, encode, pack_into, unpack_from
MIB = 2**20
def zmq_frames(data): # ZMQ recv: one independent buffer per frame
return [bytearray(memoryview(f).cast("B")) for f in encode(data)]
def nixl_frames(data): # NixlRuntime.receive: one detached packed bytearray
frames = encode(data)
packed = bytearray(calc_packed_size(frames))
pack_into(packed, frames)
return unpack_from(packed[:])
for name, make_frames in [("zmq", zmq_frames), ("nixl", nixl_frames)]:
with tempfile.TemporaryDirectory() as tmp:
store = HybridStorageUnitData(None, tmp, "run", "unit", threshold_bytes=MIB)
data = {"log_probs": torch.randn(8, 2 * MIB), "reward": torch.randn(8)} # 8 MiB/sample + tiny field
tracemalloc.start()
frames = make_frames(data)
del data
store.put_data(decode(frames), list(range(8)))
del frames
gc.collect()
retained = tracemalloc.get_traced_memory()[0] / MIB
tracemalloc.stop()
print(f"{name}: ssd_active_values={store.ssd_active_values} retained={retained:.1f} MiB")2. Errors are handled in both
|
Signed-off-by: Ginray <ginray0215@gmail.com>
Thanks for the detailed review. We went through the points again and added another round of fixes. The current HEAD has been updated. In #189, we focused on regressions introduced by the merge:
We did not duplicate the NIXL lifecycle / retry work in #189 because that is already handled by the follow-up #176, including:
Automatic expiry of Main has also moved forward since this PR was created. New changes such as We also left optional cleanups such as simplifying Thanks again for the review. In particular, the SSD receive-buffer retention and ZMQ error-handling findings helped us catch real regressions introduced by the merge. |
8 participants
Motivation
Bring the rdma branch up to date with main and verify that the upstream changes remain compatible with the NIXL-UCX payload path.
Summary
Validation
The following checks were run against this merge to verify regression coverage and compatibility with the NIXL-UCX path.
kv_put,kv_batch_get, andkv_clear.