From 3ece2166b444513a3911dab9766e4f6d332a87a8 Mon Sep 17 00:00:00 2001 From: huniu20 <70004880+huniu20@users.noreply.github.com> Date: Mon, 7 Sep 2026 19:47:08 +0800 Subject: [PATCH 01/12] [fix] Reject foreign ZMQ peers at the storage proxy (#168) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 --- tests/test_simple_storage_unit.py | 49 ++++++++++++++++++++++-- transfer_queue/metrics.py | 3 +- transfer_queue/storage/managers/base.py | 10 ++++- transfer_queue/storage/simple_storage.py | 41 +++++++++++++++++++- transfer_queue/utils/zmq_utils.py | 9 +++++ 5 files changed, 104 insertions(+), 8 deletions(-) diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index e36063ee..1145d963 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -14,6 +14,7 @@ # limitations under the License. import time +from uuid import uuid4 import pytest import ray @@ -22,15 +23,24 @@ import zmq from transfer_queue.storage.simple_storage import SimpleStorageUnit -from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket +from transfer_queue.utils.zmq_utils import ( + STORAGE_MANAGER_IDENTITY_PREFIX, + ZMQMessage, + ZMQRequestType, + create_zmq_socket, +) class MockStorageClient: """Mock client for testing storage unit operations.""" - def __init__(self, storage_put_get_address, storage_ip): + def __init__(self, storage_put_get_address, storage_ip, identity=None): self.context = zmq.Context() - self.socket = create_zmq_socket(self.context, zmq.DEALER, storage_ip) + # The storage proxy drops identities outside STORAGE_CLIENT_IDENTITY_PREFIXES, so mirror + # the prefix a real storage manager uses instead of letting ZMQ auto-assign one. + if identity is None: + identity = f"{STORAGE_MANAGER_IDENTITY_PREFIX}{uuid4().hex[:8]}".encode() + self.socket = create_zmq_socket(self.context, zmq.DEALER, storage_ip, identity=identity) self.socket.setsockopt(zmq.RCVTIMEO, 5000) # 5 second timeout self.socket.connect(storage_put_get_address) @@ -731,3 +741,36 @@ def test_storage_unit_checkpoint_load_missing_file(storage_setup, tmp_path): assert "message" in response.body client.close() + + +def test_foreign_payload_dropped_and_unit_stays_usable(storage_setup): + """A non-TransferQueue peer (e.g. a TLS probe) must not stop the proxy.""" + _, put_get_address, storage_ip = storage_setup + + foreign = MockStorageClient(put_get_address, storage_ip, identity=b"foreign_probe") + foreign.socket.setsockopt(zmq.RCVTIMEO, 1000) + # A TLS ClientHello record header: what a port scanner sends to an unknown TCP port. + foreign.socket.send(b"\x16\x03\x01\x02\x00") + # Dropped before the worker, so unlike a decode failure it draws no reply at all. + with pytest.raises(zmq.Again): + foreign.socket.recv_multipart() + foreign.close() + + client = MockStorageClient(put_get_address, storage_ip) + response = client.send_put(0, [0], {"val": [torch.tensor([1.0])]}) + assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE + client.close() + + +def test_undecodable_request_answered_and_worker_survives(storage_setup): + """An allowed identity passes the proxy, so the worker itself must survive bad frames.""" + _, put_get_address, storage_ip = storage_setup + + client = MockStorageClient(put_get_address, storage_ip) + client.socket.send_multipart([b"\x16\x03\x01\x02\x00"]) + response = ZMQMessage.deserialize(client.socket.recv_multipart(copy=False)) + assert response.request_type == ZMQRequestType.PUT_GET_ERROR + + response = client.send_put(0, [0], {"val": [torch.tensor([1.0])]}) + assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE + client.close() diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index 8d7d7b51..e90e9c59 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -26,6 +26,7 @@ from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( + METRICS_COLLECTOR_IDENTITY_PREFIX, ZMQMessage, ZMQRequestType, ZMQServerInfo, @@ -380,7 +381,7 @@ def _get_or_create_socket(self, su_id: str, su_info: ZMQServerInfo) -> zmq.Socke if sock is not None and not sock.closed: return sock - identity = f"metrics_collector_{uuid4().hex[:8]}".encode() + identity = f"{METRICS_COLLECTOR_IDENTITY_PREFIX}{uuid4().hex[:8]}".encode() sock = create_zmq_socket(self._zmq_ctx, zmq.DEALER, su_info.ip, identity) timeout_ms = TQ_METRICS_STORAGE_TIMEOUT * 1000 address = format_zmq_address(su_info.ip, su_info.ports["put_get_socket"]) diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 9c32e213..14790523 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -36,7 +36,13 @@ from transfer_queue.metadata import BatchMeta, extract_field_schema from transfer_queue.storage.clients.base import StorageClientFactory from transfer_queue.utils.logging_utils import get_logger -from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo, create_zmq_socket +from transfer_queue.utils.zmq_utils import ( + STORAGE_MANAGER_IDENTITY_PREFIX, + ZMQMessage, + ZMQRequestType, + ZMQServerInfo, + create_zmq_socket, +) logger = get_logger(__name__) @@ -70,7 +76,7 @@ def __init__( config: DictConfig, zmq_context: zmq.asyncio.Context | None = None, ): - self.storage_manager_id = f"TQ_STORAGE_{uuid4().hex[:8]}" + self.storage_manager_id = f"{STORAGE_MANAGER_IDENTITY_PREFIX}{uuid4().hex[:8]}" self.config = config self.controller_info = controller_info diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index f80f905e..bf9d616f 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -30,6 +30,7 @@ from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.perf_utils import IntervalPerfMonitor from transfer_queue.utils.zmq_utils import ( + STORAGE_CLIENT_IDENTITY_PREFIXES, ZMQMessage, ZMQRequestType, ZMQServerInfo, @@ -262,8 +263,29 @@ def _start_process_put_get(self) -> None: def _proxy_routine(self) -> None: """ZMQ proxy for message forwarding between frontend ROUTER and backend DEALER.""" logger.info(f"[{self.storage_unit_id}]: start ZMQ proxy...") + assert self.put_get_socket is not None, "put_get_socket is not properly initialized" + front, back = self.put_get_socket, self.worker_socket + poller = zmq.Poller() + poller.register(front, zmq.POLLIN) + poller.register(back, zmq.POLLIN) try: - zmq.proxy(self.put_get_socket, self.worker_socket) + # Forwarding by hand rather than via zmq.proxy() so a non-TransferQueue peer on + # this exposed TCP port cannot reach the worker and terminate its request loop. + while not self._shutdown_event.is_set(): + events = dict(poller.poll(1000)) + if front in events: + messages = front.recv_multipart(copy=False) + identity = bytes(messages[0]) if messages else b"" + if not identity.startswith(STORAGE_CLIENT_IDENTITY_PREFIXES): + logger.warning( + "[%s]: dropping request with unrecognized ZMQ identity", + self.storage_unit_id, + ) + continue + back.send_multipart(messages, copy=False) + + if back in events: + front.send_multipart(back.recv_multipart(copy=False), copy=False) except zmq.ContextTerminated: logger.info(f"[{self.storage_unit_id}]: ZMQ Proxy stopped gracefully (Context Terminated)") except Exception as e: @@ -305,7 +327,22 @@ def _worker_routine(self) -> None: identity = messages[0] serialized_msg = messages[1:] - request_msg = ZMQMessage.deserialize(serialized_msg) + try: + request_msg = ZMQMessage.deserialize(serialized_msg) + except Exception as e: + # The identity filter cannot cover this: an allowed peer can still send + # frames that fail to decode, and decoding here used to kill the thread. + logger.error( + f"[{self.storage_unit_id}]: undecodable request from " + f"identity={bytes(identity)!r}: {type(e).__name__}: {e}" + ) + error_msg = ZMQMessage.create( + request_type=ZMQRequestType.PUT_GET_ERROR, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={"message": f"undecodable request: {type(e).__name__}: {e}"}, + ) + worker_socket.send_multipart([identity] + error_msg.serialize(), copy=False) + continue operation = request_msg.request_type try: diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 2ca63d7f..c4b04ce7 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -32,6 +32,15 @@ logger = get_logger(__name__) +# Identity prefixes of the peers allowed to reach a storage unit. The storage proxy drops +# anything else, so an identity built without these prefixes is silently unreachable. +STORAGE_MANAGER_IDENTITY_PREFIX = "TQ_STORAGE_" +METRICS_COLLECTOR_IDENTITY_PREFIX = "metrics_collector_" +STORAGE_CLIENT_IDENTITY_PREFIXES = ( + STORAGE_MANAGER_IDENTITY_PREFIX.encode(), + METRICS_COLLECTOR_IDENTITY_PREFIX.encode(), +) + bytestr: TypeAlias = bytes | bytearray | memoryview From d45bfe84c9940ee7650e624dba3d60c5dbf95c62 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Han=20Zhenyu=20=E9=9F=A9=E6=8C=AF=E5=AE=87?= Date: Tue, 8 Sep 2026 13:09:32 +0800 Subject: [PATCH 02/12] [doc] Update `README` (#170) As title Signed-off-by: 0oshowero0 --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 0e026e80..0a04641e 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,7 @@ TransferQueue offers **fine-grained, sub-sample-level** data management and **lo

🔄 Updates

+ - **September 8, 2026**: 🔥 TransferQueue has been adopted in [Meshy](https://github.com/OpenBMB/Meshy). Meshy models each role as a dedicated service, and TransferQueue serves as the data flow manager that connects each role. - **June 18, 2026**: 🔥 TransferQueue has been adopted in [ROLL](https://github.com/alibaba/ROLL/pull/463). This integration introduces a [`RemoteBatch`](https://github.com/alibaba/ROLL/blob/main/docs_roll/docs/User%20Guides/Advanced%20Features/remote_batch_transfer.md) abstraction, enabling seamless compatibility with existing `DataProto` design. - **June 9, 2026**: 🔥 TransferQueue has been adopted in [UniRL](https://github.com/Tencent-Hunyuan/UniRL), a unified RL framework for multimodal models developed by Tencent Hunyuan. - **April 15, 2026**: 🔥 TransferQueue has been adopted in [Relax](https://github.com/redai-infra/Relax)! By leveraging the `StreamingDataLoader` abstraction, it schedules training data across the cluster at micro-batch granularity, reducing synchronization barriers in a single-controller setup. From 57527957d504dc3bdaa32be8ca309d4f5d8fb1b0 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 10 Sep 2026 16:06:16 +0800 Subject: [PATCH 03/12] [perf, fix] Pool and reuse ZMQ request sockets (#167) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### 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 --- tests/test_client.py | 24 ++ tests/test_controller.py | 23 +- tests/test_kv_storage_manager.py | 22 + tests/test_metrics.py | 78 +++- tests/test_zmq_shared_context.py | 201 ++++++++- tests/test_zmq_socket_pool.py | 390 ++++++++++++++++++ transfer_queue/client.py | 238 ++++++----- transfer_queue/controller.py | 11 +- transfer_queue/metrics.py | 88 ++-- transfer_queue/storage/managers/base.py | 83 ++-- .../managers/simple_storage_manager.py | 43 +- transfer_queue/utils/zmq_utils.py | 293 +++++++++++-- 12 files changed, 1248 insertions(+), 246 deletions(-) create mode 100644 tests/test_zmq_socket_pool.py diff --git a/tests/test_client.py b/tests/test_client.py index d287108e..aa887656 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -68,6 +68,7 @@ def __init__(self, controller_id="controller_0"): ) self.running = True + self.caller_identities: set[bytes] = set() self.request_thread = Thread(target=self._handle_requests, daemon=True) self.request_thread.start() @@ -85,6 +86,9 @@ def _handle_requests(self): if self.request_socket in socks: messages = self.request_socket.recv_multipart(copy=False) identity = messages.pop(0) + # Each pooled socket dials with its own identity, so this records how + # many distinct connections the client opened. + self.caller_identities.add(bytes(identity)) serialized_msg = messages request_msg = ZMQMessage.deserialize(serialized_msg) @@ -1330,6 +1334,26 @@ def test_kv_retrieve_keys_type_validation(self, client_setup): ) +# ===================================================== +# Controller RPC Socket Pool Tests +# ===================================================== + + +def test_controller_rpc_reuses_one_socket(client_setup): + """Repeated calls must share a single connection to the controller. + + This is the client-side payoff of the socket pool: TransferQueueClient runs every call + on the one event loop it builds in __init__, so all of them land in the same pool + bucket. Counted from the controller, which sees one ZMQ identity per socket dialled. + """ + client, mock_controller, _ = client_setup + + for _ in range(5): + assert client.get_partition_list() == ["partition_0", "partition_1", "test_partition"] + + assert len(mock_controller.caller_identities) == 1, "each request opened its own socket" + + # ===================================================== # Checkpoint Interface Tests # ===================================================== diff --git a/tests/test_controller.py b/tests/test_controller.py index 8a735f6a..32c1603a 100644 --- a/tests/test_controller.py +++ b/tests/test_controller.py @@ -21,7 +21,7 @@ import torch import zmq -from transfer_queue.controller import TransferQueueController +from transfer_queue.controller import TQ_CONTROLLER_ZMQ_MAX_SOCKETS, TransferQueueController from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket # Set up logging @@ -1465,6 +1465,27 @@ def test_controller_checkpoint_load_nonexistent_file(self, ray_setup, tmp_path): print("✓ load_checkpoint raises on missing file") +class TestControllerSocketBudget: + """The controller's context must hold more sockets than libzmq's default 1023. + + The metrics exporter borrows this context to query storage units, so the budget has to + cover the fleet rather than just the controller's own two ROUTERs. Past the ceiling a + socket cannot be opened at all, which would take out the control plane too. + """ + + def test_context_ceiling_is_raised_above_the_libzmq_default(self, ray_setup): + controller = TransferQueueController.remote() + try: + # Read it from inside the actor: the ceiling is only meaningful on the context + # the controller actually built. + budget = ray.get(controller.__ray_call__.remote(lambda self: self.zmq_context.get(zmq.MAX_SOCKETS))) + finally: + ray.kill(controller) + + assert budget > 1023, "the controller kept libzmq's default socket ceiling" + assert budget == TQ_CONTROLLER_ZMQ_MAX_SOCKETS + + class TestTransferQueueControllerBadRequests: """The request loop must survive requests it cannot decode, does not handle, or fails on. diff --git a/tests/test_kv_storage_manager.py b/tests/test_kv_storage_manager.py index 7ac3744d..14980fab 100644 --- a/tests/test_kv_storage_manager.py +++ b/tests/test_kv_storage_manager.py @@ -23,6 +23,28 @@ from transfer_queue.storage.managers.base import KVStorageManager +def test_close_is_quiet_when_config_is_rejected_before_super_init(): + """A config rejected before super().__init__() leaves nothing for close() to release. + + Every KV manager validates its config first, so none of the base attributes close() + reaches for exist yet. __del__ calls close() regardless, and an AttributeError there + buries the real ValueError under a spurious teardown error. + """ + built = [] + + class Probe(KVStorageManager): + def __init__(self, *args, **kwargs): + built.append(self) + super().__init__(*args, **kwargs) + + with pytest.raises(ValueError, match="Missing client_name"): + Probe(None, {}) + + manager = built[0] + assert not hasattr(manager, "notify_pool"), "the test no longer exercises the pre-super path" + manager.close() + + def get_meta(data, global_indexes=None): if not global_indexes: global_indexes = list(range(data.batch_size[0])) diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 15c43882..a8ac3e2b 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -16,12 +16,17 @@ """Unit tests for the Prometheus metrics exporter (transfer_queue.metrics).""" import time -from unittest.mock import MagicMock +from threading import Thread +from unittest.mock import MagicMock, patch import pytest try: + import zmq + from transfer_queue.metrics import TQMetricsExporter + from transfer_queue.utils.enum_utils import Role + from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo _HAS_DEPS = True except (ImportError, OSError): @@ -207,6 +212,77 @@ def test_multiple_ops_tracked_independently(self): # --------------------------------------------------------------------------- +class TestStorageQuerySocketPool: + def test_pool_borrows_the_owner_context(self): + """The exporter must query storage units over the context it was handed. + + The controller already holds a long-lived synchronous context. A second one would + add another native I/O thread and leave a context nobody closes, since the exporter + lives for the whole life of its Ray actor. + """ + ctx = zmq.Context() + try: + exporter = TQMetricsExporter(zmq_context=ctx) + with patch("zmq.Context") as minted: + exporter._get_socket_pool() + minted.assert_not_called() + finally: + ctx.destroy(linger=0) + + def test_missing_context_is_reported(self): + """Without a context there is nothing to query over, so say so rather than mint one.""" + exporter = TQMetricsExporter() + with pytest.raises(RuntimeError, match="without a ZMQ context"): + exporter._get_socket_pool() + + def test_collector_parks_no_socket_between_queries(self): + """A queried unit must leave nothing in the pool. + + Collection walks every unit once per cycle, so a parked socket is reused only a + cycle later while holding a slot in the controller context's budget for the whole + walk. At a few thousand units that budget is what runs out first. + """ + identities: set[bytes] = set() + ctx_peer = zmq.Context() + router = ctx_peer.socket(zmq.ROUTER) + port = router.bind_to_random_port("tcp://127.0.0.1") + running = True + + def serve(): + poller = zmq.Poller() + poller.register(router, zmq.POLLIN) + while running: + if not dict(poller.poll(50)): + continue + identity, _ = router.recv_multipart() + identities.add(bytes(identity)) + response = ZMQMessage.create( + request_type=ZMQRequestType.METRICS_RESPONSE, + sender_id="storage_0", + body={}, + ) + router.send_multipart([identity, *response.serialize()]) + + server = Thread(target=serve, daemon=True) + server.start() + + su_info = ZMQServerInfo(role=Role.STORAGE, id="storage_0", ip="127.0.0.1", ports={"put_get_socket": port}) + ctx = zmq.Context() + try: + exporter = TQMetricsExporter(zmq_context=ctx) + for _ in range(3): + assert exporter._query_storage_unit(su_info, "storage_0") == {} + # A parked socket would be reused, so a fresh identity per query is the + # externally visible proof that nothing was kept. + assert len(identities) == 3, "a queried unit left its socket in the pool" + finally: + running = False + server.join(timeout=2.0) + ctx.destroy(linger=0) + router.close(linger=0) + ctx_peer.term() + + class TestStorageMetricsCollection: def test_collect_with_no_storage_units(self): """No storage units registered — collect should be a no-op.""" diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 2bd9a6c7..127e4d56 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -17,11 +17,13 @@ with_zmq_socket used to create and term() a context per RPC call, which churned libzmq signaler file descriptors and crashed the process under concurrency (Bad file descriptor --> SIGABRT). It now reuses the owner's context and only creates the socket per call, so +-> SIGABRT). It now reuses the owner's context and leases sockets from a shared pool, so these tests assert every concurrent call sees the SAME context, alive until close(). +Pool-specific behavior is covered in test_zmq_socket_pool.py. """ import asyncio +from contextlib import contextmanager from threading import Thread from unittest.mock import patch @@ -138,8 +140,10 @@ def _spy_create(ctx, *args, **kwargs): assert len(results) == num_calls assert all(isinstance(meta, BatchMeta) for meta in results) - # Every call must have used the SAME context, and it must be the client's context. - assert len(seen_contexts) == num_calls + # Every socket must come from the SAME context, and it must be the client's. The count + # is bounded by the pool's cap rather than by the call count, since concurrency past it + # waits for a socket instead of opening one. + assert seen_contexts, "no socket was created at all" assert all(ctx is client.zmq_context for ctx in seen_contexts) # The shared context must NOT have been terminated by any call. assert not client.zmq_context.closed @@ -168,6 +172,83 @@ def test_client_rejects_invalid_context_pool_size(echo_controller): ) +@contextmanager +def _no_context_left_open(): + """Assert every ZMQ context built inside the block is closed by the time it exits. + + __init__ arms the finalizer only on success, so an invalid argument must either be + rejected before the context is built or destroy it on the way out -- otherwise the + context and its native I/O threads leak for the process lifetime. + """ + created = [] + real_context = zmq.asyncio.Context + + def _spy(*args, **kwargs): + ctx = real_context(*args, **kwargs) + created.append(ctx) + return ctx + + with patch("zmq.asyncio.Context", side_effect=_spy): + try: + yield created + finally: + leaked = [ctx for ctx in created if not ctx.closed] + for ctx in leaked: + ctx.destroy(linger=0) + assert not leaked, f"{len(leaked)} ZMQ context(s) left open on the failure path" + + +def test_client_rejects_invalid_socket_pool_size(echo_controller): + """A bad TQ_SOCKET_POOL_SIZE must name the variable, not silently disable reuse. + + Below 1 nothing is ever parked, so every request pays a fresh connect while the client + still looks pooled. + """ + for bad in (-1, 0): + with patch("transfer_queue.client.TQ_SOCKET_POOL_SIZE", bad): + with _no_context_left_open() as created: + with pytest.raises(ValueError, match="TQ_SOCKET_POOL_SIZE must be at least 1"): + AsyncTransferQueueClient( + client_id="client_invalid_socket_pool", + controller_info=echo_controller.zmq_server_info, + ) + assert created == [], "checkable without a context, so none should be built" + + +def test_invalid_max_sockets_does_not_leak_a_context(echo_controller): + """No invalid max-sockets input may leave the context or its I/O threads behind.""" + # Rejectable without a live context. + with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "not-a-number"): + with _no_context_left_open() as created: + with pytest.raises(ValueError, match="TQ_CLIENT_ZMQ_MAX_SOCKETS must be an integer"): + AsyncTransferQueueClient( + client_id="client_garbage_max_sockets", + controller_info=echo_controller.zmq_server_info, + ) + assert created == [], "parsing needs no context, so none should be built" + + for bad in (0, -5): + with _no_context_left_open() as created: + with pytest.raises(ValueError, match="at least 1"): + AsyncTransferQueueClient( + client_id="client_low_max_sockets", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=bad, + ) + assert created == [], "the lower bound needs no context, so none should be built" + + # Above ZMQ_SOCKET_LIMIT: this one genuinely needs a live context, so it must be + # destroyed rather than hoisted. + with _no_context_left_open() as created: + with pytest.raises(ValueError, match="ZMQ_SOCKET_LIMIT"): + AsyncTransferQueueClient( + client_id="client_huge_max_sockets", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=10**9, + ) + assert len(created) == 1, "the limit check requires a built context" + + def test_simple_storage_borrows_client_context(echo_controller): client = AsyncTransferQueueClient( client_id="client_simple_storage_context", @@ -188,6 +269,105 @@ def test_simple_storage_borrows_client_context(echo_controller): client.close() +def test_each_scenario_gets_its_own_pool(echo_controller): + """Controller RPC, storage RPC and notify must each hold a separate pool. + + They share one context but never one pool: each dials a different peer or socket name, + so a shared pool reused nothing while letting one scenario's sockets be swept by + another's. Separate pools keep each scenario's failures to itself. + """ + client = AsyncTransferQueueClient( + client_id="client_scenario_pools", + controller_info=echo_controller.zmq_server_info, + ) + + with patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"): + manager = AsyncSimpleStorageManager( + echo_controller.zmq_server_info, + {"zmq_info": {"storage_0": echo_controller.zmq_server_info}}, + zmq_context=client.zmq_context, + ) + + pools = [client.controller_rpc_pool, manager.storage_rpc_pool, manager.notify_pool] + assert len({id(pool) for pool in pools}) == 3, "scenarios must not share a pool" + + manager.close() + client.close() + + +@pytest.mark.parametrize( + "pool_size, units, warns", + [ + (8, 4, False), # 32 idle sockets against a 64-socket budget + (8, 32, True), # 256 would exceed it + ], +) +def test_warns_when_pool_size_times_units_exceeds_the_context(echo_controller, caplog, pool_size, units, warns): + """The per-address cap multiplies by storage-unit count; the context ceiling does not. + + At a few thousand units the product passes ZMQ_MAX_SOCKETS, where opening a socket + fails outright, so the mismatch is worth naming at construction rather than at a lease. + """ + client = AsyncTransferQueueClient( + client_id="client_pool_budget", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=64, + ) + zmq_info = { + f"storage_{i}": ZMQServerInfo( + role=Role.STORAGE, + id=f"storage_{i}", + ip="127.0.0.1", + ports={"put_get_socket": 5600 + i}, + ) + for i in range(units) + } + + with ( + patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"), + patch("transfer_queue.storage.managers.simple_storage_manager.TQ_SOCKET_POOL_SIZE", pool_size), + caplog.at_level("WARNING"), + ): + manager = AsyncSimpleStorageManager( + echo_controller.zmq_server_info, + {"zmq_info": zmq_info}, + zmq_context=client.zmq_context, + ) + + assert ("above this context's ZMQ_MAX_SOCKETS" in caplog.text) is warns + + manager.close() + client.close() + + +def test_close_after_failed_handshake_still_releases_the_context(echo_controller): + """A manager whose base constructor raised must still tear down on close(). + + storage_rpc_pool is assigned after super().__init__(), so a handshake timeout leaves it + unset while the context and its native I/O threads are already allocated. __del__ calls + close() regardless, so dereferencing the pool there aborts the base teardown and leaks + the context. + """ + built = [] + + def _fail(self): + built.append(self) + raise TimeoutError("handshake failed") + + with patch.object(StorageManager, "_connect_to_controller", _fail): + with pytest.raises(TimeoutError): + AsyncSimpleStorageManager( + echo_controller.zmq_server_info, + {"zmq_info": {"storage_0": echo_controller.zmq_server_info}}, + ) + + manager = built[0] + assert not hasattr(manager, "storage_rpc_pool"), "the test no longer exercises the partial-build path" + + manager.close() + assert manager.zmq_context.closed, "the context outlived a failed construction" + + def test_simple_storage_does_not_destroy_borrowed_context(echo_controller): client = AsyncTransferQueueClient( client_id="client_borrowed_context_lifecycle", @@ -266,8 +446,11 @@ def test_close_skips_destroy_while_loop_thread_alive(echo_controller): context.destroy(linger=0) -def _make_borrowing_manager(zmq_context): - """A minimal manager that borrows a caller's context, like SimpleStorage does.""" +def _make_borrowing_manager(client=None): + """A minimal manager borrowing *client*'s context, like SimpleStorage does. + + With no client it creates its own context. Either way it builds its own notify pool. + """ class Borrower(StorageManager): def _connect_to_controller(self): @@ -282,7 +465,7 @@ async def get_data(self, *args, **kwargs): async def clear_data(self, *args, **kwargs): return None - return Borrower(None, {}, zmq_context=zmq_context) + return Borrower(None, {}) if client is None else Borrower(None, {}, zmq_context=client.zmq_context) def test_stuck_notify_thread_vetoes_destroy_of_borrowed_context(echo_controller): @@ -295,7 +478,7 @@ def test_stuck_notify_thread_vetoes_destroy_of_borrowed_context(echo_controller) client_id="client_notify_thread_veto", controller_info=echo_controller.zmq_server_info, ) - client.storage_manager = _make_borrowing_manager(client.zmq_context) + client.storage_manager = _make_borrowing_manager(client) context = client.zmq_context with patch.object(client.storage_manager._notify_thread, "is_alive", return_value=True): @@ -314,7 +497,7 @@ def test_healthy_notify_thread_does_not_block_destroy(echo_controller): client_id="client_notify_thread_clean", controller_info=echo_controller.zmq_server_info, ) - client.storage_manager = _make_borrowing_manager(client.zmq_context) + client.storage_manager = _make_borrowing_manager(client) client.close() assert client.zmq_context.closed @@ -326,7 +509,7 @@ def test_manager_with_own_context_does_not_veto(echo_controller): client_id="client_independent_manager", controller_info=echo_controller.zmq_server_info, ) - client.storage_manager = _make_borrowing_manager(None) # creates its own context + client.storage_manager = _make_borrowing_manager() # creates its own context assert client.storage_manager.zmq_context is not client.zmq_context # Even a stuck notify thread on an unrelated context must not block the client. diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py new file mode 100644 index 00000000..6af2ebcf --- /dev/null +++ b/tests/test_zmq_socket_pool.py @@ -0,0 +1,390 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for ZMQSocketPool's concurrency invariants. + +Sockets used to be created and closed per request. The pool reuses them, which is only +safe because a lease is exclusive and a socket that did not complete a clean send/recv is +discarded: replies carry no request id (see ZMQMessage.create), so a reply left in flight +by a timed-out or cancelled request would be read by the next user of that socket as its +own. test_timed_out_socket_is_not_reused pins exactly that. + +Each role's own reuse is asserted where that role is tested -- the client's controller RPC +in test_client.py, the metrics collector in test_metrics.py, pool wiring and context +lifecycle in test_zmq_shared_context.py. What is left here needs control over reply timing +and loop lifetime that a caller-level test cannot reach. + +Reuse is asserted from the peer's side rather than from the pool's internals: every socket +dials with its own ZMQ identity, so one identity across many requests means the connection +was reused, and a fresh identity means the old socket was discarded. +""" + +import asyncio +import threading +from unittest.mock import patch + +import pytest +import zmq +import zmq.asyncio + +import transfer_queue.utils.zmq_utils as zmq_utils +from transfer_queue.utils.enum_utils import Role +from transfer_queue.utils.zmq_utils import ZMQServerInfo, ZMQSocketPool + + +class _Peer: + """A ROUTER that echoes one reply per request, recording who dialled it. + + ``tag`` distinguishes which endpoint answered, for the re-registration test. + """ + + def __init__(self, delay_first_reply: float = 0.0, peer_id: str = "peer_0", tag: bytes = b"reply-to-"): + self.context = zmq.Context() + self.socket = self.context.socket(zmq.ROUTER) + port = self.socket.bind_to_random_port("tcp://127.0.0.1") + self.info = ZMQServerInfo(role=Role.STORAGE, id=peer_id, ip="127.0.0.1", ports={"put_get_socket": port}) + self.identities: list[bytes] = [] + self._tag = tag + self._delay_first_reply = delay_first_reply + self._replies = 0 + self.running = True + self.thread = threading.Thread(target=self._serve, daemon=True) + self.thread.start() + + @property + def callers(self) -> int: + """How many distinct sockets have dialled this peer.""" + return len(set(self.identities)) + + def _serve(self): + poller = zmq.Poller() + poller.register(self.socket, zmq.POLLIN) + while self.running: + if not dict(poller.poll(50)): + continue + identity, request = self.socket.recv_multipart() + self.identities.append(identity) + if self._replies == 0 and self._delay_first_reply: + # Reply late enough that the requester has already timed out, leaving this + # reply in flight -- the poisoned-socket scenario. + time_left = self._delay_first_reply + while time_left > 0 and self.running: + sleep = min(0.05, time_left) + threading.Event().wait(sleep) + time_left -= sleep + self._replies += 1 + self.socket.send_multipart([identity, self._tag + request]) + + def stop(self): + self.running = False + self.thread.join(timeout=2.0) + self.socket.close(linger=0) + self.context.term() + + +@pytest.fixture +def peer(): + p = _Peer() + yield p + p.stop() + + +async def _round_trip(pool, peer_info, payload=b"req"): + with pool.lease(peer_info) as sock: + await sock.send_multipart([payload]) + return (await sock.recv_multipart())[0] + + +@pytest.mark.asyncio +async def test_timed_out_socket_is_not_reused(): + """A timed-out request must not leave its socket -- or its late reply -- in the pool. + + Regression guard for the core hazard: with the socket pooled, the *next* request would + receive the previous request's reply, silently attributing one response to another. + """ + # Replies to the first request only after the pool's 1s timeout has expired, so that + # reply is still in flight when the socket would otherwise go to the next caller. The + # peer serves serially, so the delay stays well inside the second request's own timeout. + peer = _Peer(delay_first_reply=1.4) + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", timeout=1) + try: + with pytest.raises(zmq.error.Again): + await _round_trip(pool, peer.info, b"first") + + # The late reply to "first" must not surface as the answer to "second". + assert await _round_trip(pool, peer.info, b"second") == b"reply-to-second" + assert peer.callers == 2, "the timed-out socket was handed to the next request" + finally: + pool.close() + ctx.destroy(linger=0) + peer.stop() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["exception", "cancellation"]) +async def test_poisoned_lease_is_discarded(failure): + """A lease that did not complete cleanly must not return its socket to the pool. + + The request is on the wire before the failure and its reply is still in flight, so + parking the socket would hand the next caller a reply belonging to someone else. + Cancellation counts as well as a raise: asyncio.gather cancels its siblings on the + first failure, so this is the routine case rather than an exotic one. + """ + # Answers the first request only after it has been abandoned, which is what leaves the + # stale reply in flight. + peer = _Peer(delay_first_reply=0.6) + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") + + async def poisoned(): + with pool.lease(peer.info) as sock: + await sock.send_multipart([b"doomed"]) + # Fail only once the peer has the request, so its reply really is outstanding. + while not peer.identities: + await asyncio.sleep(0.01) + if failure == "exception": + raise RuntimeError("handler blew up") + await asyncio.sleep(60) # cancelled here, after the lease was handed out + + try: + if failure == "exception": + with pytest.raises(RuntimeError): + await poisoned() + else: + task = asyncio.create_task(poisoned()) + while not peer.identities: + await asyncio.sleep(0.01) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + abandoned = set(peer.identities) + # The reply to "doomed" must not surface as the answer to "second". + assert await _round_trip(pool, peer.info, b"second") == b"reply-to-second" + assert set(peer.identities) - abandoned, "the poisoned socket was reused" + finally: + pool.close() + ctx.destroy(linger=0) + peer.stop() + + +def test_sockets_are_not_reused_across_event_loops(peer): + """A socket bound to a finished loop must never be handed to another one. + + pyzmq rebinds an async socket to whatever loop it next sees; one bound to a *closed* + loop is the "Bad file descriptor / SIGABRT" failure this keying exists to prevent. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") + + async def lease_twice(tag): + # Twice per loop, so a socket IS reused within a loop -- which is what makes the + # cross-loop comparison meaningful rather than trivially true. + for i in range(2): + assert await _round_trip(pool, peer.info, f"{tag}{i}".encode()) == f"reply-to-{tag}{i}".encode() + + for tag in ("a", "b", "c"): + asyncio.run(lease_twice(tag)) + + assert peer.callers == 3, "one socket per loop, reused within it" + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_reregistered_peer_is_not_served_a_stale_socket(): + """A peer that moves to a new address must not be answered by its old endpoint. + + A storage unit or controller restarted or re-registered keeps its id but gets a fresh + port, so keying reuse on the id alone would keep leasing a socket wired to the address + it no longer answers on -- silently talking to a dead or reassigned endpoint. All calls + share one event loop, as a real client or storage manager does; a loop per call would + discard the socket via owner eviction and hide the bug. + """ + old = _Peer(peer_id="su0", tag=b"OLD:") + new = _Peer(peer_id="su0", tag=b"NEW:") + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") + try: + assert await _round_trip(pool, old.info) == b"OLD:req" + # Same id, different port -- exactly what re-registration produces. + assert await _round_trip(pool, new.info) == b"NEW:req" + # And a peer that moves back is still reachable. + assert await _round_trip(pool, old.info) == b"OLD:req" + finally: + pool.close() + ctx.destroy(linger=0) + old.stop() + new.stop() + + +@pytest.mark.asyncio +async def test_burst_beyond_pool_size_is_served(peer): + """Concurrency above maxsize must not be refused or blocked, only left unparked.""" + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=2) + + results = await asyncio.gather(*[_round_trip(pool, peer.info, f"c{i}".encode()) for i in range(8)]) + assert sorted(results) == sorted(f"reply-to-c{i}".encode() for i in range(8)) + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_alease_caps_sockets_in_flight(peer): + """alease waits for a socket rather than opening one past maxsize. + + This is what makes the context's socket budget a function of configuration: without it + the peak equals real concurrency, which at a few thousand peers exhausts ZMQ_MAX_SOCKETS + and fails a lease with EMFILE. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=4) + + async def one(i): + async with pool.alease(peer.info) as sock: + await sock.send_multipart([f"c{i}".encode()]) + return (await sock.recv_multipart())[0] + + results = await asyncio.gather(*[one(i) for i in range(40)]) + + assert sorted(results) == sorted(f"reply-to-c{i}".encode() for i in range(40)), "a request was dropped" + assert peer.callers == 4, "concurrency past maxsize opened extra sockets instead of waiting" + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_alease_permits_are_per_address(peer): + """A busy peer must not stall requests to a different one. + + Permits are keyed like the buckets, so a fan-out across N peers is not serialized by a + cap meant to bound one peer's concurrency. + """ + other = _Peer(peer_id="peer_1", tag=b"other-") + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=1) + + async def hold(): + async with pool.alease(peer.info) as sock: + await sock.send_multipart([b"slow"]) + await sock.recv_multipart() + await asyncio.sleep(0.3) # keeps peer's only permit + + async def other_peer(): + async with pool.alease(other.info) as sock: + await sock.send_multipart([b"req"]) + return (await sock.recv_multipart())[0] + + # Would time out if one peer's permit gated the other. + _, reply = await asyncio.wait_for(asyncio.gather(hold(), other_peer()), timeout=5) + assert reply == b"other-req" + + pool.close() + ctx.destroy(linger=0) + other.stop() + + +@pytest.mark.asyncio +async def test_nested_alease_raises_instead_of_hanging(peer): + """A lease inside another lease must fail loudly rather than wait on a held permit. + + A permit is released when its body ends, so nesting can wait on one the same task -- + or a sibling mid-cycle -- already holds. That hangs with no timeout and no error, the + failure mode hardest to diagnose in production. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=8) + + with pytest.raises(RuntimeError, match="already holding"): + async with pool.alease(peer.info): + async with pool.alease(peer.info): + pass + + # The outer permit is returned, so the pool still works. + async with pool.alease(peer.info) as sock: + await sock.send_multipart([b"after"]) + assert (await sock.recv_multipart())[0] == b"reply-to-after" + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_nesting_is_rejected_across_pools(peer): + """Two pools do not make nesting safe: a cycle between them deadlocks just as well. + + Each task holds what the other waits for, and separate semaphores do not break that. + """ + other = _Peer(peer_id="peer_1", tag=b"other-") + ctx = zmq.asyncio.Context() + first = ZMQSocketPool(ctx, "first", "put_get_socket", maxsize=1) + second = ZMQSocketPool(ctx, "second", "put_get_socket", maxsize=1) + + with pytest.raises(RuntimeError, match="already holding"): + async with first.alease(peer.info): + async with second.alease(other.info): + pass + + first.close() + second.close() + ctx.destroy(linger=0) + other.stop() + + +@pytest.mark.asyncio +async def test_consecutive_leases_are_not_nesting(peer): + """Leasing again after the previous body ended is the supported shape.""" + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=2) + + for i in range(3): + async with pool.alease(peer.info) as sock: + await sock.send_multipart([f"s{i}".encode()]) + assert (await sock.recv_multipart())[0] == f"reply-to-s{i}".encode() + + pool.close() + ctx.destroy(linger=0) + + +def test_pool_size_comes_from_the_env_var_at_construction(): + """Every role's pool takes TQ_SOCKET_POOL_SIZE, resolved per pool rather than at import. + + A default bound in the signature would freeze whatever the environment held when this + module first loaded, which is what makes such a knob look settable but do nothing. + """ + ctx = zmq.Context() + try: + with patch.object(zmq_utils, "TQ_SOCKET_POOL_SIZE", 7): + assert ZMQSocketPool(ctx, "owner", "put_get_socket")._maxsize == 7 + # An explicit argument still wins, which is what lets a caller opt out. + assert ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=3)._maxsize == 3 + finally: + ctx.destroy(linger=0) + + +def test_pool_size_below_one_is_rejected(): + """Below 1 nothing is ever parked, so reuse is off while the pool still looks pooled.""" + ctx = zmq.Context() + try: + with patch.object(zmq_utils, "TQ_SOCKET_POOL_SIZE", 0): + with pytest.raises(ValueError, match="at least 1"): + ZMQSocketPool(ctx, "owner", "put_get_socket") + finally: + ctx.destroy(linger=0) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 01e50b85..0f7f13f1 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -29,9 +29,11 @@ from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( + TQ_SOCKET_POOL_SIZE, ZMQMessage, ZMQRequestType, ZMQServerInfo, + ZMQSocketPool, with_zmq_socket, ) @@ -49,10 +51,8 @@ # Pre-bound decorator for controller socket operations. with_controller_socket = with_zmq_socket( - "request_handle_socket", - get_identity=lambda self: self.client_id, get_peer=lambda self, target: self._controller, - get_context=lambda self: self.zmq_context, + get_pool=lambda self: self.controller_rpc_pool, ) @@ -61,6 +61,11 @@ class AsyncTransferQueueClient: This client provides async methods for data transfer operations including getting metadata, reading data from storage, writing data to storage, and clearing data. + + Await these methods from one long-lived event loop, as the examples below do. Pooled + sockets are keyed by their owning loop, so a fresh ``asyncio.run()`` per call gets a + fresh socket every time and pays a connect handshake it cannot amortize. Callers with + no loop of their own should use ``TransferQueueClient``, which keeps one internally. """ def __init__( @@ -91,12 +96,17 @@ def __init__( raise TypeError(f"controller_info must be ZMQServerInfo, got {type(controller_info)}") self.client_id = client_id self._controller: ZMQServerInfo = controller_info - # One long-lived context per client; sockets stay per-request because ZMQ sockets - # are not thread-safe. + # Check everything that does not need the context first: the finalizer is not armed + # until __init__ returns, so raising after allocation leaks it and its I/O threads. io_threads = TQ_CLIENT_ZMQ_IO_THREADS if zmq_io_threads is None else zmq_io_threads if io_threads < 1: raise ValueError(f"Client ZMQ I/O thread pool size must be at least 1, got {io_threads}") - self.zmq_context = zmq.asyncio.Context(io_threads=io_threads) + if TQ_SOCKET_POOL_SIZE < 1: + # Name the variable: the pool's own error cannot say which knob supplied the value. + raise ValueError( + f"TQ_SOCKET_POOL_SIZE must be at least 1, got {TQ_SOCKET_POOL_SIZE}. " + f"The pool always reuses at least one socket per endpoint; it cannot be disabled." + ) max_sockets = zmq_max_sockets explicitly_requested = max_sockets is not None @@ -109,26 +119,44 @@ def __init__( f"TQ_CLIENT_ZMQ_MAX_SOCKETS must be an integer, got {TQ_CLIENT_ZMQ_MAX_SOCKETS!r}" ) from e explicitly_requested = True + if max_sockets is not None and max_sockets < 1: + # The upper bound needs ZMQ_SOCKET_LIMIT, hence a live context, but the lower one + # does not -- so reject it before allocating anything. + raise ValueError(f"Client ZMQ max sockets must be at least 1, got {max_sockets}") if max_sockets is None: max_sockets = DEFAULT_CLIENT_ZMQ_MAX_SOCKETS - socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT) - if explicitly_requested: - # A value the caller asked for must not be silently reinterpreted. - if not 1 <= max_sockets <= socket_limit: - raise ValueError( - f"Client ZMQ max sockets must be between 1 and this build's " - f"ZMQ_SOCKET_LIMIT ({socket_limit}), got {max_sockets}" + self.zmq_context = zmq.asyncio.Context(io_threads=io_threads) + try: + # ZMQ_SOCKET_LIMIT is a property of the built context, so this last check cannot + # be hoisted above it; destroy the context rather than leak it on failure. + socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT) + if explicitly_requested: + # A value the caller asked for must not be silently reinterpreted. + if max_sockets > socket_limit: + raise ValueError( + f"Client ZMQ max sockets must be between 1 and this build's " + f"ZMQ_SOCKET_LIMIT ({socket_limit}), got {max_sockets}" + ) + elif max_sockets > socket_limit: + # Nobody asked for the default, so clamp instead of failing to construct on a + # build whose ZMQ_SOCKET_LIMIT is below it. + logger.debug( + f"[{client_id}]: Clamping default ZMQ max sockets {max_sockets} to this " + f"build's ZMQ_SOCKET_LIMIT ({socket_limit})." ) - elif max_sockets > socket_limit: - # Nobody asked for the default, so clamp instead of failing to construct on a - # build whose ZMQ_SOCKET_LIMIT is below it. - logger.debug( - f"[{client_id}]: Clamping default ZMQ max sockets {max_sockets} to this " - f"build's ZMQ_SOCKET_LIMIT ({socket_limit})." - ) - max_sockets = socket_limit - self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets) + max_sockets = socket_limit + self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets) + except BaseException: + self.zmq_context.destroy(linger=0) + raise + # Reused across requests, so the socket budget above is consumed by the concurrency + # high-water mark rather than by request count. + self.controller_rpc_pool = ZMQSocketPool( + self.zmq_context, + client_id, + "request_handle_socket", + ) # Backstop for a client that is never closed, so the context and its I/O threads do # not leak for the process lifetime. finalize() (not __del__) also runs at @@ -159,7 +187,8 @@ def initialize_storage_manager( The client's long-lived ZMQ context is offered to every backend uniformly; each registered manager decides whether to borrow it or keep its own, so the client - needs no knowledge of specific backend names. + needs no knowledge of specific backend names. Managers build their own pools over + whichever context they end up with, one per request scenario. Args: manager_type: Type of storage manager to create. Supported types include: @@ -236,33 +265,35 @@ async def async_get_meta( RuntimeError: If communication fails or controller returns error response Example: - >>> # Example 1: Basic fetch metadata - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["input_ids", "attention_mask"], - ... batch_size=4, - ... partition_id="train_0", - ... mode="fetch", - ... task_name="generate_sequences" - ... )) - >>> print(batch_meta.is_ready) # True if all samples ready - >>> - >>> # Example 2: Fetch with self-defined samplers (using GRPOGroupNSampler as an example) - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["input_ids", "attention_mask"], - ... batch_size=8, - ... partition_id="train_0", - ... mode="fetch", - ... task_name="generate_sequences", - ... )) - >>> print(batch_meta.is_ready) # True if all samples ready - >>> - >>> # Example 3: Force fetch metadata (bypass production status check and Sampler, - >>> # so may include unready and already-consumed samples. No filtering by consumption status is applied.) - >>> batch_meta = asyncio.run(client.async_get_meta( - ... partition_id="train_0", # optional - ... mode="force_fetch", - ... )) - >>> print(batch_meta.is_ready) # May be False if some samples not ready + >>> async def main(): + ... # Example 1: Basic fetch metadata + ... batch_meta = await client.async_get_meta( + ... data_fields=["input_ids", "attention_mask"], + ... batch_size=4, + ... partition_id="train_0", + ... mode="fetch", + ... task_name="generate_sequences" + ... ) + ... print(batch_meta.is_ready) # True if all samples ready + ... + ... # Example 2: Fetch with self-defined samplers (using GRPOGroupNSampler as an example) + ... batch_meta = await client.async_get_meta( + ... data_fields=["input_ids", "attention_mask"], + ... batch_size=8, + ... partition_id="train_0", + ... mode="fetch", + ... task_name="generate_sequences", + ... ) + ... print(batch_meta.is_ready) # True if all samples ready + ... + ... # Example 3: Force fetch metadata (bypass production status check and Sampler, + ... # so may include unready and already-consumed samples. No filtering by + ... # consumption status is applied.) + ... batch_meta = await client.async_get_meta( + ... partition_id="train_0", # optional + ... mode="force_fetch", + ... ) + ... print(batch_meta.is_ready) # May be False if some samples not ready """ response_msg = await self._request_controller( socket=socket, @@ -302,10 +333,11 @@ async def async_set_custom_meta( RuntimeError: If communication fails or controller returns error response Example: - >>> # Create batch with custom metadata - >>> batch_meta = client.get_meta(data_fields=["input_ids"], batch_size=4, ...) - >>> batch_meta.update_custom_meta([{"score": 0.9}, {"score": 0.8}]) - >>> asyncio.run(client.async_set_custom_meta(batch_meta)) + >>> async def main(): + ... # Create batch with custom metadata + ... batch_meta = await client.async_get_meta(data_fields=["input_ids"], batch_size=4, ...) + ... batch_meta.update_custom_meta([{"score": 0.9}, {"score": 0.8}]) + ... await client.async_set_custom_meta(batch_meta) """ assert socket is not None @@ -387,30 +419,34 @@ async def async_put( >>> batch_size = 4 >>> seq_len = 16 >>> current_partition_id = "train_0" - >>> # Example 1: Normal usage with existing metadata - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["prompts", "attention_mask"], - ... batch_size=batch_size, - ... partition_id=current_partition_id, - ... mode="fetch", - ... task_name="generate_sequences", - ... )) - >>> batch = asyncio.run(client.async_get_data(batch_meta)) - >>> output = TensorDict({"response": torch.randn(batch_size, seq_len)}) - >>> asyncio.run(client.async_put(data=output, metadata=batch_meta)) - >>> - >>> # Example 2: Initial data insertion without pre-existing metadata - >>> # BE CAREFUL: this usage may overwrite any unconsumed data in the given partition_id! - >>> # Please make sure the corresponding partition_id is empty before calling the async_put() - >>> # without metadata. - >>> # Now we only support put all the data of the corresponding partition id in once. You should repeat with - >>> # interleave the initial data if n_sample > 1 before calling the async_put(). - >>> original_prompts = torch.randn(batch_size, seq_len) - >>> n_samples = 4 - >>> prompts_repeated = torch.repeat_interleave(original_prompts, n_samples, dim=0) - >>> prompts_repeated_batch = TensorDict({"prompts": prompts_repeated}) - >>> # This will create metadata in "insert" mode internally. - >>> metadata = asyncio.run(client.async_put(data=prompts_repeated_batch, partition_id=current_partition_id)) + >>> async def main(): + ... # Example 1: Normal usage with existing metadata + ... batch_meta = await client.async_get_meta( + ... data_fields=["prompts", "attention_mask"], + ... batch_size=batch_size, + ... partition_id=current_partition_id, + ... mode="fetch", + ... task_name="generate_sequences", + ... ) + ... batch = await client.async_get_data(batch_meta) + ... output = TensorDict({"response": torch.randn(batch_size, seq_len)}) + ... await client.async_put(data=output, metadata=batch_meta) + ... + ... # Example 2: Initial data insertion without pre-existing metadata + ... # BE CAREFUL: this usage may overwrite any unconsumed data in the given + ... # partition_id! Please make sure the corresponding partition_id is empty + ... # before calling the async_put() without metadata. + ... # Now we only support put all the data of the corresponding partition id in + ... # once. You should repeat with interleave the initial data if n_sample > 1 + ... # before calling the async_put(). + ... original_prompts = torch.randn(batch_size, seq_len) + ... n_samples = 4 + ... prompts_repeated = torch.repeat_interleave(original_prompts, n_samples, dim=0) + ... prompts_repeated_batch = TensorDict({"prompts": prompts_repeated}) + ... # This will create metadata in "insert" mode internally. + ... metadata = await client.async_put( + ... data=prompts_repeated_batch, partition_id=current_partition_id + ... ) """ if not hasattr(self, "storage_manager") or self.storage_manager is None: @@ -467,16 +503,18 @@ async def async_get_data(self, metadata: BatchMeta) -> TensorDict: - Requested data fields (e.g., "prompts", "attention_mask") Example: - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["prompts", "attention_mask"], - ... batch_size=4, - ... partition_id="train_0", - ... mode="fetch", - ... task_name="generate_sequences", - ... )) - >>> batch = asyncio.run(client.async_get_data(batch_meta)) - >>> print(batch) - >>> # TensorDict with fields "prompts", "attention_mask", and sample order matching metadata global_indexes + >>> async def main(): + ... batch_meta = await client.async_get_meta( + ... data_fields=["prompts", "attention_mask"], + ... batch_size=4, + ... partition_id="train_0", + ... mode="fetch", + ... task_name="generate_sequences", + ... ) + ... batch = await client.async_get_data(batch_meta) + ... print(batch) + ... # TensorDict with fields "prompts", "attention_mask", and sample order + ... # matching metadata global_indexes """ if not hasattr(self, "storage_manager") or self.storage_manager is None: @@ -678,10 +716,10 @@ async def async_get_consumption_status( Example: >>> # Get consumption status - >>> global_index, consumption_status = asyncio.run(client.async_get_consumption_status( + >>> global_index, consumption_status = await client.async_get_consumption_status( ... task_name="generate_sequences", ... partition_id="train_0" - ... )) + ... ) >>> print(f"Global index: {global_index}, Consumption status: {consumption_status}") """ @@ -725,10 +763,10 @@ async def async_get_production_status( Example: >>> # Get production status - >>> global_index, production_status = asyncio.run(client.async_get_production_status( + >>> global_index, production_status = await client.async_get_production_status( ... data_fields=["input_ids", "attention_mask"], ... partition_id="train_0" - ... )) + ... ) >>> print(f"Global index: {global_index}, Production status: {production_status}") """ try: @@ -766,10 +804,10 @@ async def async_check_consumption_status( Example: >>> # Check if all samples have been consumed - >>> is_consumed = asyncio.run(client.async_check_consumption_status( + >>> is_consumed = await client.async_check_consumption_status( ... task_name="generate_sequences", ... partition_id="train_0" - ... )) + ... ) >>> print(f"All samples consumed: {is_consumed}") """ @@ -802,10 +840,10 @@ async def async_check_production_status( Example: >>> # Check if all samples are ready for consumption - >>> is_ready = asyncio.run(client.async_check_production_status( + >>> is_ready = await client.async_check_production_status( ... data_fields=["input_ids", "attention_mask"], ... partition_id="train_0" - ... )) + ... ) >>> print(f"All samples ready: {is_ready}") """ _, production_status = await self.async_get_production_status( @@ -842,10 +880,10 @@ async def async_reset_consumption( Example: >>> # Reset consumption for train task to re-train on same data - >>> success = asyncio.run(client.async_reset_consumption( + >>> success = await client.async_reset_consumption( ... partition_id="train_0", ... task_name="train" - ... )) + ... ) >>> print(f"Reset successful: {success}") """ body = {"partition_id": partition_id} @@ -879,7 +917,7 @@ async def async_get_partition_list( list[str]: List of partition ids managed by the controller Example: - >>> partition_ids = asyncio.run(client.get_partition_list()) + >>> partition_ids = await client.get_partition_list() >>> print(f"Available partitions: {partition_ids}") """ try: @@ -1053,6 +1091,8 @@ def close(self) -> None: ) return try: + # Close pooled sockets before the context that owns them. + self.controller_rpc_pool.close() if hasattr(self, "zmq_context") and self.zmq_context is not None: self.zmq_context.destroy(linger=0) except Exception as e: diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index ef71bb4c..b45a85a6 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -56,6 +56,9 @@ TQ_CONTROLLER_GET_METADATA_TIMEOUT = int(os.environ.get("TQ_CONTROLLER_GET_METADATA_TIMEOUT", 1)) TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL = int(os.environ.get("TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL", 5)) +# Above libzmq's default of 1023: the metrics exporter borrows this context to query +# storage units, so the budget covers a fleet rather than the two ROUTERs bound here. +TQ_CONTROLLER_ZMQ_MAX_SOCKETS = int(os.environ.get("TQ_CONTROLLER_ZMQ_MAX_SOCKETS", 4096)) # Sample pre-allocation for StreamingDataLoader compatibility. # By pre-allocating sample indices (typically global_batch_size), consumers can accurately @@ -1725,6 +1728,9 @@ def kv_retrieve_keys( def _init_zmq_socket(self): """Initialize ZMQ sockets for communication.""" self.zmq_context = zmq.Context() + # Before any socket is opened on it: libzmq applies MAX_SOCKETS at socket creation. + socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT) + self.zmq_context.set(zmq.MAX_SOCKETS, min(TQ_CONTROLLER_ZMQ_MAX_SOCKETS, socket_limit)) self._node_ip = get_node_ip_address() while True: @@ -2337,10 +2343,9 @@ def start_metrics(self, port: int = 0) -> str: return self._metrics_endpoint from transfer_queue.metrics import TQMetricsExporter - self._metrics = TQMetricsExporter() + # Lend the controller's context rather than let the exporter build a second one. + self._metrics = TQMetricsExporter(zmq_context=self.zmq_context) self._metrics_endpoint = self._metrics.start(node_ip=self._node_ip, port=port) - # Launch a daemon thread that periodically pushes controller state - # snapshots to the exporter, keeping them process-isolated. self._metrics_snapshot_thread = Thread( target=self._metrics_snapshot_loop, name="TQMetricsSnapshotThread", diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index e90e9c59..12914585 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -18,7 +18,6 @@ from contextlib import contextmanager from threading import Thread from typing import Any -from uuid import uuid4 import psutil import zmq @@ -30,8 +29,7 @@ ZMQMessage, ZMQRequestType, ZMQServerInfo, - create_zmq_socket, - format_zmq_address, + ZMQSocketPool, ) logger = get_logger(__name__) @@ -68,13 +66,22 @@ class TQMetricsExporter: TQ_METRICS_STORAGE_TIMEOUT ZMQ timeout for storage queries (default 5s) """ - def __init__(self, role: str = "controller"): + def __init__(self, role: str = "controller", zmq_context: zmq.Context | None = None): + """ + Args: + role: Which process this exporter runs in; only "controller" collects from + storage units, so only that role needs a context. + zmq_context: The owner's long-lived synchronous context, borrowed for + storage-unit queries and never terminated here. Minting one instead would + add a second context and its native I/O thread with nobody to close them, + since the exporter lives as long as its Ray actor. + """ self._start_time = time.time() self._process = psutil.Process() self._role = role self._storage_unit_infos: dict[str, ZMQServerInfo] = {} - self._zmq_ctx: zmq.Context | None = None - self._zmq_sockets: dict[str, zmq.Socket] = {} + self._zmq_ctx = zmq_context + self._zmq_socket_pool: ZMQSocketPool | None = None self._known_partition_ids: set[str] = set() self._known_production_labels: set[tuple[str, str]] = set() self._known_consumption_labels: set[tuple[str, str]] = set() @@ -372,49 +379,50 @@ def collect_storage_metrics(self) -> None: except Exception as e: logger.warning(f"Failed to collect metrics from storage unit {su_id}: {e}") - def _get_or_create_socket(self, su_id: str, su_info: ZMQServerInfo) -> zmq.Socket: - """Return a cached ZMQ DEALER socket for *su_id*, creating one if needed.""" - if self._zmq_ctx is None: - self._zmq_ctx = zmq.Context() - - sock = self._zmq_sockets.get(su_id) - if sock is not None and not sock.closed: - return sock - - identity = f"{METRICS_COLLECTOR_IDENTITY_PREFIX}{uuid4().hex[:8]}".encode() - sock = create_zmq_socket(self._zmq_ctx, zmq.DEALER, su_info.ip, identity) - timeout_ms = TQ_METRICS_STORAGE_TIMEOUT * 1000 - address = format_zmq_address(su_info.ip, su_info.ports["put_get_socket"]) - sock.connect(address) - sock.setsockopt(zmq.RCVTIMEO, timeout_ms) - sock.setsockopt(zmq.SNDTIMEO, timeout_ms) - self._zmq_sockets[su_id] = sock - return sock + def _get_socket_pool(self) -> ZMQSocketPool: + """Return the lazily-created socket pool for storage-unit queries.""" + if self._zmq_socket_pool is None: + if self._zmq_ctx is None: + raise RuntimeError( + "TQMetricsExporter was built without a ZMQ context, so it cannot query " + "storage units; pass zmq_context= from the owning process." + ) + self._zmq_socket_pool = ZMQSocketPool( + self._zmq_ctx, + # The storage proxy drops identities without this prefix, and the pool + # builds each socket's identity from the owner id. + METRICS_COLLECTOR_IDENTITY_PREFIX.rstrip("_"), + "put_get_socket", + timeout=TQ_METRICS_STORAGE_TIMEOUT, + ) + return self._zmq_socket_pool def _query_storage_unit(self, su_info: ZMQServerInfo, su_id: str) -> dict[str, Any] | None: """Send a synchronous GET_METRICS request to a single storage unit.""" try: - sock = self._get_or_create_socket(su_id, su_info) - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_METRICS, - sender_id="metrics_collector", - body={}, - ) - sock.send_multipart(request_msg.serialize()) - response_frames = sock.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_frames) - if response_msg.request_type == ZMQRequestType.METRICS_RESPONSE: - return response_msg.body - return None + pool = self._get_socket_pool() + with pool.lease(su_info) as sock: + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.GET_METRICS, + sender_id="metrics_collector", + body={}, + ) + sock.send_multipart(request_msg.serialize()) + response_frames = sock.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(response_frames) + # Closed rather than parked: collection walks every unit once per cycle, so + # a kept socket holds budget for the whole walk to save one handshake. + sock.close(linger=0) + if response_msg.request_type == ZMQRequestType.METRICS_RESPONSE: + return response_msg.body + return None except zmq.error.Again: + # The pool discarded the socket, so the next cycle starts clean. Reusing it + # would read this reply, if it lands late, as the answer to that cycle's query. logger.debug(f"Timeout querying metrics from {su_id}") return None except Exception as e: logger.warning(f"Error querying metrics from {su_id}: {e}") - # Close broken socket so it gets recreated next cycle - sock = self._zmq_sockets.pop(su_id, None) - if sock and not sock.closed: - sock.close(linger=0) return None def start(self, node_ip: str = "0.0.0.0", port: int = 0) -> str: diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 14790523..369f800e 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -41,6 +41,7 @@ ZMQMessage, ZMQRequestType, ZMQServerInfo, + ZMQSocketPool, create_zmq_socket, ) @@ -87,6 +88,9 @@ def __init__( # creates when handed nothing. Only an owner tears its context down (see close()). self._owns_zmq_context = zmq_context is None self.zmq_context = zmq.asyncio.Context() if zmq_context is None else zmq_context + # Notify runs on a dedicated loop, so it could not share sockets with another + # scenario in any case. + self.notify_pool = ZMQSocketPool(self.zmq_context, self.storage_manager_id, "request_handle_socket") self._connect_to_controller() # Dedicated asyncio loop for ZMQ notify traffic, isolated from the caller's loop @@ -271,56 +275,39 @@ async def notify_data_update( async def _notify_and_wait(self, request_msg: list) -> None: """Send a data status notification to the controller and block until ACK is received.""" - identity = f"{self.storage_manager_id}-notify-{uuid4().hex[:8]}".encode() - sock = create_zmq_socket( - ctx=self.zmq_context, socket_type=zmq.DEALER, ip=self.controller_info.ip, identity=identity - ) - sock.setsockopt(zmq.LINGER, 0) - sock.connect(self.controller_info.to_addr("request_handle_socket")) - - try: - await sock.send_multipart(request_msg) - logger.debug( - f"[{self.storage_manager_id}]: Sent data status update request " - f"to controller id #{self.controller_info.id} successfully." - ) - - response_received = False - timeout = TQ_DATA_UPDATE_RESPONSE_TIMEOUT + # Acquiring the lease sits outside the handler below: a missing socket name or a dead + # context is a configuration/lifecycle fault the caller must see, not a slow ACK. + async with self.notify_pool.alease(self.controller_info) as sock: + try: + await sock.send_multipart(request_msg) + logger.debug( + f"[{self.storage_manager_id}]: Sent data status update request " + f"to controller id #{self.controller_info.id} successfully." + ) - while not response_received and timeout > 0: - try: - poll_interval = min(TQ_STORAGE_POLLER_TIMEOUT, timeout) - messages = await asyncio.wait_for( - sock.recv_multipart(copy=False), - timeout=poll_interval, - ) + # One deadline for the whole wait, so unrelated traffic on this socket cannot + # extend it and a quiet controller still gets the full budget. + deadline = time.monotonic() + TQ_DATA_UPDATE_RESPONSE_TIMEOUT + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + f"no ACK from controller {self.controller_info.id} after {TQ_DATA_UPDATE_RESPONSE_TIMEOUT}s" + ) + messages = await asyncio.wait_for(sock.recv_multipart(copy=False), timeout=remaining) response_msg = ZMQMessage.deserialize(messages) if response_msg.request_type == ZMQRequestType.NOTIFY_DATA_UPDATE_ACK: # type: ignore[arg-type] - response_received = True logger.debug( f"[{self.storage_manager_id}]: Get data status update ACK response " f"from controller id #{response_msg.sender_id} successfully." ) - break - except asyncio.TimeoutError: - timeout -= poll_interval - except Exception as e: - logger.warning(f"[{self.storage_manager_id}]: Error receiving response: {e}") - break - - if not response_received: - logger.error( - f"[{self.storage_manager_id}]: Timeout waiting for data status update ACK " - f"from controller after {TQ_DATA_UPDATE_RESPONSE_TIMEOUT}s." - ) - finally: - try: - if not sock.closed: - sock.close(linger=0) - except Exception: - pass + return + except Exception as e: + # Logged rather than raised, so a slow controller does not fail the put. Close + # the socket: a late ACK would otherwise be read as the next lessee's reply. + logger.error(f"[{self.storage_manager_id}]: Data status update failed: {type(e).__name__}: {e}") + sock.close(linger=0) @abstractmethod async def put_data( @@ -389,6 +376,11 @@ async def load_checkpoint(self, checkpoint_dir: str) -> None: def close(self) -> None: """Close all ZMQ sockets/contexts and stop the notify loop.""" + # A subclass may reject its config before super().__init__() runs, and __del__ calls + # close() anyway; return rather than bury that error under an AttributeError. + if not hasattr(self, "notify_pool"): + return + if self.controller_handshake_socket: try: if not self.controller_handshake_socket.closed: @@ -408,13 +400,16 @@ def close(self) -> None: else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") + # This manager always owns its notify pool, even over a borrowed context. Only once + # the notify thread is gone: Socket.close() is not thread-safe and it holds the leases. + if notify_thread_stopped: + self.notify_pool.close() + if self._owns_zmq_context: # destroy() calls Socket.close(), which is not thread-safe, so it must run only # after the notify thread holding sockets is gone. If that thread outlived its # join, leak the context rather than risk a crash on a terminating process. if notify_thread_stopped: - # linger=0 force-closes sockets left by an interrupted request, so this - # cannot hang on term(). self.zmq_context.destroy(linger=0) else: logger.warning( diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 88340006..b986701a 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -32,9 +32,11 @@ from transfer_queue.storage.simple_storage import KEY_NOT_FOUND_MARKER, StorageKeyNotFoundError from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( + TQ_SOCKET_POOL_SIZE, ZMQMessage, ZMQRequestType, ZMQServerInfo, + ZMQSocketPool, with_zmq_socket, ) @@ -47,14 +49,9 @@ # Pre-bound decorator for storage-unit socket operations. with_storage_unit_socket = with_zmq_socket( - "put_get_socket", - get_identity=lambda self: self.storage_manager_id, get_peer=lambda self, target: self.storage_unit_infos[target], - # Long-lived context from the base StorageManager, shared with the notify path. Safe - # because the context is loop-agnostic and each socket stays per-call. - get_context=lambda self: self.zmq_context, + get_pool=lambda self: self.storage_rpc_pool, resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), - timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, ) @@ -80,6 +77,13 @@ def __init__( zmq_context: zmq.asyncio.Context | None = None, ): super().__init__(controller_info, config, zmq_context=zmq_context) + # Storage-unit RPC, on whichever context the base class settled on. + self.storage_rpc_pool = ZMQSocketPool( + self.zmq_context, + self.storage_manager_id, + "put_get_socket", + timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, + ) self.config = config server_infos: ZMQServerInfo | dict[str, ZMQServerInfo] | None = config.get("zmq_info", None) @@ -97,6 +101,28 @@ def __init__( raise ValueError("AsyncSimpleStorageManager requires non-empty 'zmq_info' in config.") self.storage_unit_infos = self._register_servers(server_infos) + self._warn_if_pool_can_exhaust_context(len(self.storage_unit_infos)) + + def _warn_if_pool_can_exhaust_context(self, num_units: int) -> None: + """Warn when the pool may park more sockets than the context can hold. + + The cap is per (owner, address), so it multiplies by storage-unit count while the + context ceiling does not. At a few thousand units the product passes ZMQ_MAX_SOCKETS, + where a lease fails with EMFILE instead of degrading. + """ + try: + budget = self.zmq_context.get(zmq.MAX_SOCKETS) + except zmq.ZMQError: # pragma: no cover - context already terminating + return + worst_case = TQ_SOCKET_POOL_SIZE * num_units + if worst_case > budget: + logger.warning( + f"[{self.storage_manager_id}]: storage RPC pool may hold up to " + f"{TQ_SOCKET_POOL_SIZE} x {num_units} = {worst_case} idle sockets, above this " + f"context's ZMQ_MAX_SOCKETS ({budget}). Concurrent requests can then fail to " + f"open a socket. Lower TQ_SOCKET_POOL_SIZE or raise TQ_CLIENT_ZMQ_MAX_SOCKETS " + f"(with enough file descriptors, see ulimit -n)." + ) def _register_servers(self, server_infos: "ZMQServerInfo | dict[Any, ZMQServerInfo]"): """Register and validate server information. @@ -659,4 +685,9 @@ async def load_checkpoint(self, checkpoint_dir: str) -> None: def close(self) -> None: """Close all ZMQ sockets and context to prevent resource leaks.""" + # Before super(), which may destroy the context these sockets live on. Absent when + # the base constructor raised (a failed handshake does), and __del__ still calls this. + pool = getattr(self, "storage_rpc_pool", None) + if pool is not None: + pool.close() super().close() diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index c4b04ce7..a94cd3ab 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -13,9 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio +import itertools +import os import socket +import threading import time -from collections.abc import Sequence +from collections.abc import AsyncIterator, Iterator, Sequence +from contextlib import asynccontextmanager, contextmanager +from contextvars import ContextVar from dataclasses import dataclass from functools import wraps from typing import Any, Callable, TypeAlias @@ -41,6 +47,9 @@ METRICS_COLLECTOR_IDENTITY_PREFIX.encode(), ) +# Cap for every pool; see ZMQSocketPool for what it bounds. Small because it multiplies by +# peer count, and one socket per peer already avoids the repeated handshake. +TQ_SOCKET_POOL_SIZE = int(os.environ.get("TQ_SOCKET_POOL_SIZE", 4)) bytestr: TypeAlias = bytes | bytearray | memoryview @@ -360,47 +369,263 @@ def create_zmq_socket( return socket +def _lease_owner() -> Any: + """The running event loop, or the current thread outside one. + + A ZMQ socket is safe on neither a second thread nor a second event loop, so leases + never cross either. 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. Returning the thread outside a loop lets + synchronous callers (the metrics collector) share this pool unchanged. + """ + try: + return asyncio.get_running_loop() + except RuntimeError: + return threading.current_thread() + + +def _owner_finished(owner: Any) -> bool: + """Whether a lease owner can no longer serve its sockets.""" + if isinstance(owner, threading.Thread): + return not owner.is_alive() + return owner.is_closed() + + +# Addresses whose permit the current task already holds. A ContextVar because asyncio +# copies the context per task, so siblings under one gather do not see each other's. +_held_permits: ContextVar[frozenset[str]] = ContextVar("_held_permits", default=frozenset()) + + +class ZMQSocketPool: + """Lends connected DEALER sockets for one request scenario, reusing them across calls. + + One pool serves one kind of request -- controller RPC, storage RPC, notify, metrics -- + so ``socket_name`` and ``timeout`` are fixed here and every socket it holds is + interchangeable but for the address it is connected to. Separate pools per scenario also + keep them from interfering: a pool shared across scenarios reused nothing anyway, since + scenarios differ by owning loop, socket name, or timeout. + + 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. Each caller + therefore gets a socket to itself and returns it only after a clean send/recv. + + Reuse avoids a TCP + ZMTP handshake per request and stops the peer's ROUTER from + accreting a fresh identity every call. + """ + + def __init__( + self, + ctx: zmq.Context, + owner_id: str, + socket_name: str, + *, + timeout: int | None = None, + maxsize: int | None = None, + ): + """ + Args: + ctx: Long-lived context to open sockets on, sync or async. The pool borrows it + and never terminates it; the owner outlives the pool. + owner_id: Identity prefix for pooled sockets, for readable peer-side logs. + socket_name: Port key in ``ZMQServerInfo.ports`` that every lease dials. + timeout: Send/recv timeout in seconds applied to every socket, or None for none. + maxsize: Idle sockets kept per bucket, at least 1, defaulting to + ``TQ_SOCKET_POOL_SIZE``. A soft cap: a burst beyond it still gets sockets, + and the excess is closed on return rather than made to wait. Counted per + (owner, address), not per pool, so a pool dialling N peers may hold + N*maxsize idle sockets. + """ + # Resolved here rather than in the signature so the env var stays patchable; a + # default bound at import froze whatever the environment held when this module loaded. + maxsize = TQ_SOCKET_POOL_SIZE if maxsize is None else maxsize + if maxsize < 1: + # Below 1 nothing is ever parked, so every request pays a fresh connect while + # still looking pooled. Reject it rather than silently disable reuse. + raise ValueError(f"ZMQ socket pool size must be at least 1, got {maxsize}") + self._ctx = ctx + self._owner_id = owner_id + self._socket_name = socket_name + self._timeout = timeout + self._maxsize = maxsize + # Keyed by lease owner (see _lease_owner), then by address rather than peer id: a peer + # restarted under the same id at a new address must not get a socket wired to the old. + self._idle: dict[Any, dict[str, list[zmq.Socket]]] = {} + self._lock = threading.Lock() + # One semaphore per (owner, address), keyed like _idle because a semaphore belongs + # to the loop that awaits it. Used by alease only; see there. + self._permits: dict[Any, dict[str, asyncio.Semaphore]] = {} + # A ROUTER silently drops a second peer claiming an identity it already has, and + # owner_id alone repeats across nodes because client ids are pid-derived. + self._identity_prefix = f"{owner_id}_{uuid4().hex[:8]}" + self._counter = itertools.count() + + @contextmanager + def lease(self, peer: ZMQServerInfo) -> Iterator[zmq.Socket]: + """Yield a socket connected to ``peer``, returning it to the pool only on success. + + A plain (non-async) contextmanager on purpose: ``with`` still sees exceptions and + ``CancelledError`` raised across ``await``s in its body, so this one definition + serves both async and synchronous callers. + + Concurrency beyond ``maxsize`` opens extra sockets here rather than waiting; async + callers that want the cap enforced use ``alease``. + """ + address = self._address(peer) + sock = self._take(address) or self._connect(peer, address) + try: + yield sock + except BaseException: + # Poisoned: the request may already be on the wire, so its reply could still + # arrive and the next lessee would read it as its own. Cancellation counts too. + sock.close(linger=0) + raise + else: + self._release(address, sock) + + @asynccontextmanager + async def alease(self, peer: ZMQServerInfo) -> AsyncIterator[zmq.Socket]: + """Like ``lease``, but waits for a socket instead of opening one past ``maxsize``. + + This bounds sockets in flight at ``maxsize`` per (owner, address), so the context's + socket budget follows configuration rather than peak concurrency. + + A permit is held for the whole body, so a task that leases inside another lease can + wait on a permit it -- or a sibling mid-cycle -- already holds, and that hangs with + no timeout. Any nesting therefore raises, including across pools: an RPC's reply is + what releases its permit, so a second RPC belongs after the first returns. Notify + already works this way, running once the puts it reports have completed. + """ + address = self._address(peer) + held = _held_permits.get() + if held: + raise RuntimeError( + f"Lease on {address} from a task already holding {sorted(held)}. A lease " + f"keeps its permit until its body ends, so nesting one inside another can " + f"wait on a permit that is already held and hang. Complete the outer " + f"request first, then start this one." + ) + token = _held_permits.set(held | {address}) + permit = self._permit(address) + try: + async with permit: + with self.lease(peer) as sock: + yield sock + finally: + _held_permits.reset(token) + + def _address(self, peer: ZMQServerInfo) -> str: + port = peer.ports.get(self._socket_name) + if port is None: + raise RuntimeError(f"Socket '{self._socket_name}' not configured for server '{peer.id}'") + return format_zmq_address(peer.ip, port) + + def _permit(self, address: str) -> asyncio.Semaphore: + """The permit gating this (owner, address), created on first use by that owner.""" + with self._lock: + return self._permits.setdefault(_lease_owner(), {}).setdefault(address, asyncio.Semaphore(self._maxsize)) + + def _owner_buckets(self) -> dict[str, list[zmq.Socket]]: + """Buckets for the current lease owner, first evicting any owner that has finished. + + Callers must hold ``self._lock``. A closed loop's sockets must not linger: a pooled + async socket keeps its loop referenced, so they would never be collected, and reusing + one is the SIGABRT hazard in _lease_owner. Sweeping here needs no background thread, + and the loop count stays tiny (one per client, plus one per notify thread). + """ + for owner in [o for o in self._idle if _owner_finished(o)]: + self._close_all(self._idle.pop(owner)) + self._permits.pop(owner, None) + return self._idle.setdefault(_lease_owner(), {}) + + def _take(self, address: str) -> zmq.Socket | None: + """Pop a live idle socket for *address*, discarding any found closed.""" + with self._lock: + bucket = self._owner_buckets().get(address) + while bucket: + sock = bucket.pop() + if not sock.closed: + return sock + return None + + def _release(self, address: str, sock: zmq.Socket) -> None: + """Return a socket, closing it if already closed or its bucket is full.""" + with self._lock: + # A caller may close the socket itself without raising, as the notify path does + # to discard a possibly-late ACK. + if not sock.closed: + bucket = self._owner_buckets().setdefault(address, []) + if len(bucket) < self._maxsize: + bucket.append(sock) + return + sock.close(linger=0) + + def _connect(self, peer: ZMQServerInfo, address: str) -> zmq.Socket: + """Open and connect a new DEALER socket to *address*.""" + identity = f"{self._identity_prefix}_to_{peer.id}_{next(self._counter)}".encode() + sock = create_zmq_socket(self._ctx, zmq.DEALER, peer.ip, identity=identity) + try: + if self._timeout is not None: + sock.setsockopt(zmq.RCVTIMEO, self._timeout * 1000) + sock.setsockopt(zmq.SNDTIMEO, self._timeout * 1000) + sock.connect(address) + except BaseException: + # Nothing owns the socket until it is handed to a lease, so close it here or it + # leaks. connect() raises on a malformed endpoint or a terminating context. + sock.close(linger=0) + raise + return sock + + def close(self) -> None: + """Close every idle socket. Safe to call twice, and after the context is gone. + + The pool stays usable afterwards: a lease still outstanding returns its socket to a + fresh bucket. Callers destroy the context right after, so nothing is reused. + """ + with self._lock: + owned = list(self._idle.values()) + self._idle = {} + for buckets in owned: + self._close_all(buckets) + + def _close_all(self, buckets: dict[str, list[zmq.Socket]]) -> None: + """Close every socket in *buckets*, tolerating an already-destroyed context.""" + for sock in itertools.chain.from_iterable(buckets.values()): + try: + if not sock.closed: + sock.close(linger=0) + except Exception as e: + logger.debug(f"[{self._owner_id}]: Error closing pooled socket: {e}") + + def with_zmq_socket( - socket_name: str, *, - get_identity: Callable[[Any], str], get_peer: Callable[[Any, str | None], ZMQServerInfo], - get_context: Callable[[Any], "zmq.asyncio.Context"], + get_pool: Callable[[Any], ZMQSocketPool], resolve_target: Callable[[tuple, dict], str | None] | None = None, - timeout: int | None = None, ): - """Create a reusable async decorator for request sockets. + """Create a reusable async decorator that injects a pooled request socket. - Lifecycle: get owner's shared context -> create/connect socket -> inject -> close socket. + Lifecycle: resolve peer -> lease a socket from ``self``'s pool -> inject as the + ``socket`` kwarg -> return it to the pool if the call succeeded, else discard it. - The context comes from ``self`` via ``get_context`` and is long-lived; only the DEALER - socket is per-call. Do NOT create or terminate a context here -- per-call churn corrupts - libzmq's signaler file descriptors under concurrency (Bad file descriptor / SIGABRT) and - can hang on term(). Contexts are thread-safe and loop-agnostic, so sharing one is safe. + The socket name and timeout come from the pool, which serves one request scenario. Args: - socket_name: Socket port key in ``ZMQServerInfo.ports``. - get_identity: Callable that extracts owner identity from ``self``. - Example: ``lambda self: self.client_id`` get_peer: Callable that returns ``ZMQServerInfo`` for the target. For single-target scenarios, ignore the target parameter. Example: ``lambda self, target: self.server_info`` Example: ``lambda self, target: self.storage_unit_infos[target]`` - get_context: Callable that returns the owner's long-lived ``zmq.asyncio.Context``. - Example: ``lambda self: self.zmq_context`` + get_pool: Callable that returns the pool for this scenario. + Example: ``lambda self: self.controller_rpc_pool`` resolve_target: Optional callable that extracts target identifier from function arguments. Receives (args, kwargs) and returns target name. Example: ``lambda args, kwargs: kwargs.get("target_storage_unit")`` - timeout: Optional timeout (seconds) for both send/recv operations. """ def decorator(func: Callable): @wraps(func) async def wrapper(self, *args, **kwargs): - owner_id = get_identity(self) - if owner_id is None: - raise RuntimeError("get_identity returned None") - target_name: str | None = None if resolve_target is not None: target_name = resolve_target(args, kwargs) @@ -409,31 +634,13 @@ async def wrapper(self, *args, **kwargs): if server_info is None: raise RuntimeError(f"get_peer returned None for target '{target_name}'") - port = server_info.ports.get(socket_name) - if port is None: - raise RuntimeError(f"Socket '{socket_name}' not configured for server '{server_info.id}'") + pool = get_pool(self) + if pool is None: + raise RuntimeError("get_pool returned None") - # Reuse the owner's long-lived context; only the socket is per-call. - context = get_context(self) - if context is None: - raise RuntimeError("get_context returned None") - - sock = None - try: - address = format_zmq_address(server_info.ip, port) - identity = f"{owner_id}_to_{server_info.id}_{uuid4().hex[:8]}".encode() - sock = create_zmq_socket(context, zmq.DEALER, server_info.ip, identity=identity) - sock.connect(address) - if timeout is not None: - sock.setsockopt(zmq.RCVTIMEO, timeout * 1000) - sock.setsockopt(zmq.SNDTIMEO, timeout * 1000) + async with pool.alease(server_info) as sock: kwargs["socket"] = sock return await func(self, *args, **kwargs) - finally: - # Close the per-call socket only; the context outlives the call. linger=0 - # drops unsent frames so close never blocks the event loop. - if sock is not None and not sock.closed: - sock.close(linger=0) return wrapper From ba8880c4e4f052142627d052cde5819f0daafd97 Mon Sep 17 00:00:00 2001 From: adoda Date: Thu, 10 Sep 2026 23:27:29 +0800 Subject: [PATCH 04/12] [fix] Retry a lost storage-unit request instead of failing the job (#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. #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 Co-authored-by: jathonzhang --- tests/test_storage_request_retry.py | 175 +++++++++++++ .../managers/simple_storage_manager.py | 232 +++++++++++++++--- transfer_queue/storage/simple_storage.py | 17 +- transfer_queue/utils/common.py | 24 ++ 4 files changed, 416 insertions(+), 32 deletions(-) create mode 100644 tests/test_storage_request_retry.py diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py new file mode 100644 index 00000000..ec6e9476 --- /dev/null +++ b/tests/test_storage_request_retry.py @@ -0,0 +1,175 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Storage-unit retry and timeout-diagnosis tests.""" + +import logging +from unittest.mock import AsyncMock, MagicMock, patch + +import numpy as np +import pytest +import torch +import zmq +from tensordict import TensorDict + +from transfer_queue.metadata import BatchMeta +from transfer_queue.storage.managers import simple_storage_manager as ssm +from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager, StorageUnitTimeout +from transfer_queue.utils import common +from transfer_queue.utils.enum_utils import Role +from transfer_queue.utils.zmq_utils import ZMQServerInfo + + +def _manager(with_unit: bool = True) -> AsyncSimpleStorageManager: + manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager) + manager.storage_manager_id = "TQ_STORAGE_test" + manager.storage_unit_infos = {} + if with_unit: + manager.storage_unit_infos["unit_a"] = ZMQServerInfo( + role=Role.STORAGE, + id="unit_a", + ip="10.0.0.7", + ports={"put_get_socket": 5555}, + ) + manager.close = lambda: None + return manager + + +def _single_sample_batch() -> tuple[TensorDict, BatchMeta]: + metadata = BatchMeta( + global_indexes=[0], + partition_ids=["0"], + field_schema={"input_ids": {"dtype": torch.int64, "shape": (2,), "is_nested": False, "is_non_tensor": False}}, + production_status=np.ones(1, dtype=np.int8), + ) + return TensorDict({"input_ids": torch.zeros(1, 2, dtype=torch.int64)}, batch_size=1), metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize("data_parser, expected_attempts", [(None, 3), (lambda data: data, 1)]) +async def test_put_retries_only_without_a_parser(data_parser, expected_attempts): + manager = _manager() + manager.notify_data_update = AsyncMock() + manager._put_to_single_storage_unit = AsyncMock(side_effect=StorageUnitTimeout("no answer")) + data, metadata = _single_sample_batch() + + with ( + patch.object(ssm, "TQ_SIMPLE_STORAGE_MAX_ATTEMPTS", 3), + patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis"), + pytest.raises(StorageUnitTimeout), + ): + await manager.put_data(data, metadata, data_parser=data_parser) + + assert manager._put_to_single_storage_unit.await_count == expected_attempts + manager.notify_data_update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_lost_request_recovers_without_diagnosis(): + manager = _manager() + request = AsyncMock(side_effect=[StorageUnitTimeout("no answer"), "payload"]) + + with patch.object(manager, "_diagnose_storage_unit") as diagnose: + result = await manager._request_with_retry("get", "unit_a", "samples=4", request) + + assert result == "payload" + assert request.await_count == 2 + diagnose.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error, max_attempts, expected_attempts", + [ + (StorageUnitTimeout("no answer"), 3, 3), + (RuntimeError("unit rejected request"), 3, 1), + (StorageUnitTimeout("no answer"), 0, 1), + ], +) +async def test_retry_attempts_are_bounded_and_only_cover_timeouts(error, max_attempts, expected_attempts): + manager = _manager() + request = AsyncMock(side_effect=error) + + with ( + patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis"), + pytest.raises(type(error)), + ): + await manager._request_with_retry("get", "unit_a", "samples=4", request, max_attempts=max_attempts) + + assert request.await_count == expected_attempts + + +@pytest.mark.asyncio +async def test_failure_log_has_context_without_repeating_the_unit(caplog): + manager = _manager() + error = StorageUnitTimeout("no answer in 200s during get from storage unit unit_a at 10.0.0.7:5555") + + with ( + patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis"), + caplog.at_level(logging.ERROR), + pytest.raises(StorageUnitTimeout), + ): + await manager._request_with_retry( + "get", "unit_a", "samples=4 fields=['input_ids']", AsyncMock(side_effect=error) + ) + + message = next(record.message for record in caplog.records if "failed after" in record.message) + assert "samples=4 fields=['input_ids']" in message + assert message.count("unit_a") == message.count("10.0.0.7:5555") == 1 + + +@pytest.mark.parametrize( + "elapsed, payload_bytes, should_log", + [(0.01, 2**20, False), (0.01, 512 * 2**20, True), (6.0, 2**20, True)], +) +def test_log_heavy_operation_thresholds(caplog, elapsed, payload_bytes, should_log): + with caplog.at_level(logging.WARNING): + common.log_heavy_operation("TQ_STORAGE_test", "put", elapsed, payload_bytes, "samples=4") + + assert bool(caplog.records) is should_log + if should_log: + assert "serialized_mb=" in caplog.records[0].message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tcp_result, probe_result, expected", + [ + ((None, None), {"active_keys": 4, "op_stats": {"GET_DATA": {"request_count": 1025}}}, "request_lost"), + ((None, None), zmq.error.Again(), "unit_not_serving"), + (ConnectionRefusedError(), zmq.error.Again(), "tcp=down(ConnectionRefusedError)"), + ], +) +async def test_diagnosis_classifies_the_failure(tcp_result, probe_result, expected): + manager = _manager() + writer = MagicMock() + tcp = AsyncMock( + side_effect=tcp_result if isinstance(tcp_result, Exception) else None, + return_value=(None, writer), + ) + probe = AsyncMock( + side_effect=probe_result if isinstance(probe_result, Exception) else None, + return_value=probe_result, + ) + + with patch.object(ssm.asyncio, "open_connection", tcp), patch.object(manager, "_probe_storage_unit", probe): + diagnosis = await manager._diagnose_storage_unit("unit_a") + + assert expected in diagnosis + + +@pytest.mark.asyncio +async def test_diagnosis_handles_an_unknown_unit(): + assert "unit_not_registered" in await _manager(with_unit=False)._diagnose_storage_unit("missing") diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index b986701a..70765cb9 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -15,9 +15,11 @@ import asyncio import os +import time import warnings from collections import defaultdict from collections.abc import Mapping +from functools import partial from operator import itemgetter from pathlib import Path from typing import Any, Callable, NamedTuple @@ -29,7 +31,11 @@ from transfer_queue.metadata import BatchMeta, extract_field_schema from transfer_queue.storage.managers.base import StorageManager, StorageManagerFactory -from transfer_queue.storage.simple_storage import KEY_NOT_FOUND_MARKER, StorageKeyNotFoundError +from transfer_queue.storage.simple_storage import ( + KEY_NOT_FOUND_MARKER, + StorageKeyNotFoundError, +) +from transfer_queue.utils.common import log_heavy_operation from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( TQ_SOCKET_POOL_SIZE, @@ -37,6 +43,7 @@ ZMQRequestType, ZMQServerInfo, ZMQSocketPool, + frame_nbytes, with_zmq_socket, ) @@ -44,6 +51,23 @@ TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT", 200)) # seconds +# Attempts per storage-unit request, including the first. Each one must use a new connection: a +# DEALER silently queues messages for a peer it never reached, so no timeout can catch that. +TQ_SIMPLE_STORAGE_MAX_ATTEMPTS = int(os.environ.get("TQ_SIMPLE_STORAGE_MAX_ATTEMPTS", 3)) + +# Timeout for the post-failure probe, which only has to answer whether the unit still serves. +TQ_SIMPLE_STORAGE_PROBE_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_PROBE_TIMEOUT", 10)) + + +class StorageUnitTimeout(RuntimeError): + """A storage unit did not answer within the send/recv timeout. + + Distinct from an error the unit reported: only a missing answer is worth a new connection. + The message must name the unit, its endpoint and the timeout, because it is what the callers + of ``put_data`` and ``get_data`` see, and the retry logs rely on it instead of repeating them. + """ + + _SU_SUBDIR = "simple_storage" _SU_INFO_FILE = "storage_unit_info.json" @@ -54,6 +78,13 @@ resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), ) +# Same endpoint as above but with the short diagnostic timeout, used only after a failure. +with_storage_unit_probe_socket = with_zmq_socket( + get_peer=lambda self, target: self.storage_unit_infos[target], + get_pool=lambda self: self.storage_probe_pool, + resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), +) + class RoutingGroup(NamedTuple): """Routing result for a single storage unit.""" @@ -84,6 +115,13 @@ def __init__( "put_get_socket", timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, ) + self.storage_probe_pool = ZMQSocketPool( + self.zmq_context, + f"{self.storage_manager_id}_probe", + "put_get_socket", + timeout=TQ_SIMPLE_STORAGE_PROBE_TIMEOUT, + maxsize=1, + ) self.config = config server_infos: ZMQServerInfo | dict[str, ZMQServerInfo] | None = config.get("zmq_info", None) @@ -114,11 +152,11 @@ def _warn_if_pool_can_exhaust_context(self, num_units: int) -> None: budget = self.zmq_context.get(zmq.MAX_SOCKETS) except zmq.ZMQError: # pragma: no cover - context already terminating return - worst_case = TQ_SOCKET_POOL_SIZE * num_units + worst_case = (TQ_SOCKET_POOL_SIZE + 1) * num_units if worst_case > budget: logger.warning( f"[{self.storage_manager_id}]: storage RPC pool may hold up to " - f"{TQ_SOCKET_POOL_SIZE} x {num_units} = {worst_case} idle sockets, above this " + f"({TQ_SOCKET_POOL_SIZE} + 1 probe) x {num_units} = {worst_case} idle sockets, above this " f"context's ZMQ_MAX_SOCKETS ({budget}). Concurrent requests can then fail to " f"open a socket. Lower TQ_SOCKET_POOL_SIZE or raise TQ_CLIENT_ZMQ_MAX_SOCKETS " f"(with enough file descriptors, see ulimit -n)." @@ -170,6 +208,108 @@ def _group_by_hash(self, global_indexes: list[int]) -> dict[str, RoutingGroup]: pos_lists[key].append(pos) return {key: RoutingGroup(gi_lists[key], pos_lists[key]) for key in gi_lists} + def _describe_storage_unit(self, storage_unit_id: str) -> str: + """Return ``ip:port`` for a storage unit, for use in diagnostics. + + The unit id is a random uuid4 fragment that carries no location, so a bare id in an + error message cannot be traced back to a node. Never raises: it is only ever called + while reporting another failure, and must not mask it. + """ + info = self.storage_unit_infos.get(storage_unit_id) + if info is None: + return "endpoint unknown (unit not registered with this manager)" + return f"{info.ip}:{info.ports.get('put_get_socket')}" + + @with_storage_unit_probe_socket + async def _probe_storage_unit(self, target_storage_unit: str, socket: zmq.Socket = None) -> dict[str, Any]: + """Ask a storage unit for its own counters over a brand-new socket. + + Served by the same worker thread as put and get, so an answer proves the unit is serving. + """ + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.GET_METRICS, # type: ignore[arg-type] + sender_id=f"{self.storage_manager_id}_probe", + receiver_id=target_storage_unit, + body={}, + ) + await socket.send_multipart(request_msg.serialize()) + messages = await socket.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(messages) + if response_msg.request_type != ZMQRequestType.METRICS_RESPONSE: + raise RuntimeError(f"unexpected probe response type {response_msg.request_type}") + return response_msg.body + + async def _diagnose_storage_unit(self, target_storage_unit: str) -> str: + """Classify a timeout as a lost request, a stuck unit, or an unreachable node. + + Returns one log line and never raises: it runs while another failure is being reported. + """ + info = self.storage_unit_infos.get(target_storage_unit) + if info is None: + return "verdict=unknown(unit_not_registered)" + + try: + _, writer = await asyncio.wait_for( + asyncio.open_connection(info.ip, info.ports.get("put_get_socket")), timeout=5 + ) + writer.close() + tcp = "tcp=up" + except Exception as e: + tcp = f"tcp=down({type(e).__name__})" + + try: + body = await self._probe_storage_unit(target_storage_unit=target_storage_unit) + op_counts = {op: stats.get("request_count") for op, stats in (body.get("op_stats") or {}).items()} + return ( + f"{tcp} verdict=request_lost_in_flight (unit answered a fresh probe: " + f"ops={op_counts} active_keys={body.get('active_keys')} " + f"rss_gb={body.get('process_rss_bytes', 0) / 2**30:.2f})" + ) + except zmq.error.Again: + return f"{tcp} verdict=unit_not_serving (no probe answer in {TQ_SIMPLE_STORAGE_PROBE_TIMEOUT}s)" + except Exception as e: + return f"{tcp} verdict=unknown (probe failed: {type(e).__name__}: {e})" + + async def _request_with_retry( + self, + operation: str, + target_storage_unit: str, + request_context: str, + make_request: Callable[[], Any], + max_attempts: int | None = None, + ): + """Run one storage-unit request, retrying a missing answer on a fresh connection. + + Args: + operation: Operation name for logs, e.g. ``get`` or ``put``. + target_storage_unit: Unit this request is routed to. + request_context: Request shape, so both ends of a failure can be correlated. + make_request: Zero-arg callable returning a coroutine for one attempt. Must build a + new socket per call, which ``with_storage_unit_socket`` does. + max_attempts: Attempts allowed. Defaults to ``TQ_SIMPLE_STORAGE_MAX_ATTEMPTS``; pass 1 + for a request that must not be replayed. + """ + # Floored at one: a nonpositive count would skip the request and report success, which for + # put would publish metadata for data that was never sent. + attempts_allowed = max(1, TQ_SIMPLE_STORAGE_MAX_ATTEMPTS if max_attempts is None else max_attempts) + for attempt in range(1, attempts_allowed + 1): + try: + return await make_request() + except StorageUnitTimeout as e: + # The exception already names the unit, its endpoint and the timeout, so these + # lines only add what it cannot know: which attempt this was, and the shape. + if attempt < attempts_allowed: + logger.warning( + f"[{self.storage_manager_id}]: {operation} retry {attempt + 1}/{attempts_allowed} " + f"on a new connection. {request_context} {e}" + ) + continue + logger.error( + f"[{self.storage_manager_id}]: {operation} failed after {attempt} attempts. " + f"{request_context} {e} {await self._diagnose_storage_unit(target_storage_unit)}" + ) + raise + @staticmethod def _select_by_positions(field_data, positions: list[int]): """Slice a single field's data by non-contiguous batch positions. @@ -283,24 +423,37 @@ async def put_data( field_schema = extract_field_schema(data) routing = self._group_by_hash(metadata.global_indexes) - tasks = [ - self._put_to_single_storage_unit( - group.global_indexes, - {f: self._select_by_positions(data[f], group.batch_positions) for f in data.keys()}, - target_storage_unit=su_id, - data_parser=data_parser, + # Parser-backed puts are not replayed: the public API does not constrain parser side effects. + max_attempts = 1 if data_parser is not None else None + tasks = [] + for su_id, group in routing.items(): + storage_data = {f: self._select_by_positions(data[f], group.batch_positions) for f in data.keys()} + tasks.append( + self._request_with_retry( + "put", + su_id, + f"samples={len(group.global_indexes)} fields={list(storage_data.keys())}", + partial( + self._put_to_single_storage_unit, + group.global_indexes, + storage_data, + target_storage_unit=su_id, + data_parser=data_parser, + ), + max_attempts=max_attempts, + ) ) - for su_id, group in routing.items() - ] try: await asyncio.gather(*tasks) except Exception as e: + # The offending unit is named in the error itself; the full routing list can run to + # hundreds of ids, which buries it. logger.error( f"[{self.storage_manager_id}]: put_data failed. " f"partition_id={metadata.partition_ids[0]}, " f"num_samples={metadata.size}, " - f"storage_units={list(routing.keys())}, " + f"num_storage_units={len(routing)}, " f"error={type(e).__name__}: {e}" ) raise @@ -332,9 +485,12 @@ async def _put_to_single_storage_unit( body={"global_indexes": global_indexes, "data": storage_data, "data_parser": data_parser}, ) + serialized_bytes = 0 + started = time.perf_counter() try: - data = request_msg.serialize() - await socket.send_multipart(data, copy=False) + frames = request_msg.serialize() + serialized_bytes = sum(frame_nbytes(frame) or 0 for frame in frames) + await socket.send_multipart(frames, copy=False) messages = await socket.recv_multipart(copy=False) response_msg = ZMQMessage.deserialize(messages) @@ -343,15 +499,21 @@ async def _put_to_single_storage_unit( f"Failed to put data to storage unit {target_storage_unit}: " f"{response_msg.body.get('message', 'Unknown error')}" ) - except zmq.error.Again as e: - timeout_sec = TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT - logger.error( - f"[{self.storage_manager_id}]: ZMQ recv timeout ({timeout_sec}s) " - f"during put to storage unit {target_storage_unit}. " - f"The storage unit may be overloaded or crashed." + # This end serializes the put payload and waits out the round trip, so it holds both + # the true wire size and the end-to-end latency. + log_heavy_operation( + self.storage_manager_id, + "put", + time.perf_counter() - started, + serialized_bytes, + f"to {target_storage_unit} at {self._describe_storage_unit(target_storage_unit)} " + f"samples={len(global_indexes)} fields={list(storage_data.keys())}", ) - raise RuntimeError( - f"ZMQ recv timeout ({timeout_sec}s) during put to storage unit {target_storage_unit}" + except zmq.error.Again as e: + raise StorageUnitTimeout( + f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during put to storage unit " + f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}; " + f"serialized_mb={serialized_bytes / 2**20:.1f}" ) from e except Exception as e: logger.error( @@ -431,7 +593,17 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict: routing = self._group_by_hash(metadata.global_indexes) tasks = [ - self._get_from_single_storage_unit(group.global_indexes, metadata.field_names, target_storage_unit=su_id) + self._request_with_retry( + "get", + su_id, + f"samples={len(group.global_indexes)} fields={list(metadata.field_names)}", + partial( + self._get_from_single_storage_unit, + group.global_indexes, + metadata.field_names, + target_storage_unit=su_id, + ), + ) for su_id, group in routing.items() ] try: @@ -490,13 +662,10 @@ async def _get_from_single_storage_unit( error_type = StorageKeyNotFoundError if KEY_NOT_FOUND_MARKER in message else RuntimeError raise error_type(f"Failed to get data from storage unit {target_storage_unit}: {message}") except zmq.error.Again as e: - timeout_sec = TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT - logger.error( - f"[{self.storage_manager_id}]: ZMQ recv timeout ({timeout_sec}s) " - f"from storage unit {target_storage_unit}. " - f"The storage unit may be overloaded or crashed." - ) - raise RuntimeError(f"ZMQ recv timeout ({timeout_sec}s) from storage unit {target_storage_unit}") from e + raise StorageUnitTimeout( + f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during get from storage unit " + f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}" + ) from e except StorageKeyNotFoundError: # Already logged at debug by the storage unit; propagate for the caller to classify. raise @@ -690,4 +859,7 @@ def close(self) -> None: pool = getattr(self, "storage_rpc_pool", None) if pool is not None: pool.close() + probe_pool = getattr(self, "storage_probe_pool", None) + if probe_pool is not None: + probe_pool.close() super().close() diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index bf9d616f..cd0d2b4b 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -25,7 +25,7 @@ import ray import zmq -from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads +from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads, log_heavy_operation from transfer_queue.utils.enum_utils import Role from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.perf_utils import IntervalPerfMonitor @@ -36,6 +36,7 @@ ZMQServerInfo, create_zmq_socket, format_zmq_address, + frame_nbytes, get_free_port, get_node_ip_address, ) @@ -344,6 +345,7 @@ def _worker_routine(self) -> None: worker_socket.send_multipart([identity] + error_msg.serialize(), copy=False) continue operation = request_msg.request_type + started = time.perf_counter() try: logger.debug(f"[{self.storage_unit_id}]: worker received operation: {operation}") @@ -388,7 +390,18 @@ def _worker_routine(self) -> None: ) # Send response back with identity for routing - worker_socket.send_multipart([identity] + response_msg.serialize(), copy=False) + response_frames = response_msg.serialize() + if operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] + # This end serializes the get response, so its frames give the true wire size. + log_heavy_operation( + self.storage_unit_id, + "get", + time.perf_counter() - started, + sum(frame_nbytes(frame) or 0 for frame in response_frames), + f"samples={len(request_msg.body.get('global_indexes', []))} " + f"fields={list(request_msg.body.get('fields', []))}", + ) + worker_socket.send_multipart([identity] + response_frames, copy=False) logger.info(f"[{self.storage_unit_id}]: worker stopped.") poller.unregister(worker_socket) diff --git a/transfer_queue/utils/common.py b/transfer_queue/utils/common.py index 999e8183..1278dfd4 100644 --- a/transfer_queue/utils/common.py +++ b/transfer_queue/utils/common.py @@ -133,3 +133,27 @@ def get_env_bool(env_key: str, default: bool = False) -> bool: true_values = {"true", "1", "yes", "y", "on"} return env_value_lower in true_values + + +# Thresholds above which a single request earns a log line. Both sit far above the normal range +# of single-digit milliseconds, so tripping either one is itself the finding. +TQ_STORAGE_SLOW_REQUEST_SECONDS = float(os.environ.get("TQ_STORAGE_SLOW_REQUEST_SECONDS", 5.0)) +TQ_STORAGE_LARGE_PAYLOAD_MB = float(os.environ.get("TQ_STORAGE_LARGE_PAYLOAD_MB", 256)) + + +def log_heavy_operation(component_id: str, operation: str, elapsed: float, payload_bytes: int, detail: str) -> None: + """Warn about one slow or unusually large request; stay silent otherwise. + + Args: + component_id: Storage manager or storage unit reporting the request. + operation: Operation name, e.g. ``put`` or ``get``. + elapsed: Wall time spent on the request, in seconds. + payload_bytes: Serialized size of the payload on the wire. + detail: Request shape, appended to the message verbatim. + """ + payload_mb = payload_bytes / 2**20 + if elapsed < TQ_STORAGE_SLOW_REQUEST_SECONDS and payload_mb < TQ_STORAGE_LARGE_PAYLOAD_MB: + return + logger.warning( + f"[{component_id}]: heavy {operation} {detail} serialized_mb={payload_mb:.1f} elapsed={elapsed:.2f}s" + ) From cb6baa356120186c59d93d5f44766798cb426b22 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Han=20Zhenyu=20=E9=9F=A9=E6=8C=AF=E5=AE=87?= Date: Fri, 11 Sep 2026 22:05:05 +0800 Subject: [PATCH 05/12] [chore] Update README (#174) As title. --------- Signed-off-by: 0oshowero0 --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 0a04641e..3c08a6ad 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,7 @@ TransferQueue offers **fine-grained, sub-sample-level** data management and **lo

🔄 Updates

- **September 8, 2026**: 🔥 TransferQueue has been adopted in [Meshy](https://github.com/OpenBMB/Meshy). Meshy models each role as a dedicated service, and TransferQueue serves as the data flow manager that connects each role. + - **August 14, 2026**: 🔥 TransferQueue has been adopted in [Dressage](https://github.com/Accio-Lab/Dressage). By integrating TQ, peak memory usage of the data plane on the master node is reduced by 91% in a 32-node GLM-5.2 setup. Refer to the [report](https://github.com/Accio-Lab/Dressage/blob/main/docs/trajectory-storage-refactor-en.md) for details. - **June 18, 2026**: 🔥 TransferQueue has been adopted in [ROLL](https://github.com/alibaba/ROLL/pull/463). This integration introduces a [`RemoteBatch`](https://github.com/alibaba/ROLL/blob/main/docs_roll/docs/User%20Guides/Advanced%20Features/remote_batch_transfer.md) abstraction, enabling seamless compatibility with existing `DataProto` design. - **June 9, 2026**: 🔥 TransferQueue has been adopted in [UniRL](https://github.com/Tencent-Hunyuan/UniRL), a unified RL framework for multimodal models developed by Tencent Hunyuan. - **April 15, 2026**: 🔥 TransferQueue has been adopted in [Relax](https://github.com/redai-infra/Relax)! By leveraging the `StreamingDataLoader` abstraction, it schedules training data across the cluster at micro-batch granularity, reducing synchronization barriers in a single-controller setup. From fc33c979db6bd661802f6af458e75e850055fb20 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 11 Sep 2026 22:23:10 +0800 Subject: [PATCH 06/12] [fix] Fix to avoid request dropping in the accept queue (#172) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary A storage-unit request can be lost between the manager and the unit without either end raising an error. #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 #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 #171's commits as its base and included a per-call socket-timeout change. Since then #171 merged and #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 --- docs/metrics.md | 35 +++ scripts/grafana_dashboard.json | 62 +++++ tests/test_accept_probe.py | 257 ++++++++++++++++++ tests/test_metrics.py | 78 ++++++ tests/test_storage_request_retry.py | 41 ++- transfer_queue/metrics.py | 80 ++++++ .../managers/simple_storage_manager.py | 34 ++- transfer_queue/storage/simple_storage.py | 51 ++++ transfer_queue/utils/accept_probe.py | 245 +++++++++++++++++ 9 files changed, 873 insertions(+), 10 deletions(-) create mode 100644 tests/test_accept_probe.py create mode 100644 transfer_queue/utils/accept_probe.py diff --git a/docs/metrics.md b/docs/metrics.md index 05559dd5..327f6079 100644 --- a/docs/metrics.md +++ b/docs/metrics.md @@ -60,6 +60,8 @@ Steps: |---------------------|---------|-------------| | `TQ_METRICS_COLLECT_INTERVAL` | `10` | Background collection interval (seconds) | | `TQ_METRICS_STORAGE_TIMEOUT` | `5` | ZMQ timeout for storage unit queries (seconds) | +| `TQ_ACCEPT_PROBE_INTERVAL` | `0` (off) | Accept-queue sampling period (seconds); enables the accept-queue metrics below | +| `TQ_STORAGE_ZMQ_BACKLOG` | `4096` | Accept-queue depth for the storage unit's listening socket | ## Architecture @@ -140,6 +142,38 @@ Steps: | `tq_storage_request_latency_p50` | Gauge | `storage_unit_id`, `op_type` | P50 request latency (seconds) | | `tq_storage_request_latency_p99` | Gauge | `storage_unit_id`, `op_type` | P99 request latency (seconds) | +### Storage Request-Loss Diagnostics (collected via ZMQ, exposed on controller) + +Arrivals are counted when the worker decodes a request, while `tq_storage_request_ops` only +advances once one completes, so a sustained gap between them means requests are arriving and +not finishing. The counters are cumulative per operation and carry no request identity, so +they characterise a unit, not any individual request. + +| Metric | Type | Labels | Description | +|--------|------|--------|-------------| +| `tq_storage_requests_arrived` | Gauge | `storage_unit_id` | Requests decoded by the worker, whether or not they completed | +| `tq_storage_arrivals_by_op` | Gauge | `storage_unit_id`, `op_type` | Same, broken down by operation | + +The accept-queue series below exist only when the unit runs with `TQ_ACCEPT_PROBE_INTERVAL` +set; the series are removed rather than reported as zero when the probe is off, so a missing +series means "not measured" rather than "no drops". `tq_storage_accept_queue_peak` and the two +drop counters are cumulative since the probe started, so use `rate()` on them. + +| Metric | Type | Labels | Description | +|--------|------|--------|-------------| +| `tq_storage_accept_queue_backlog` | Gauge | `storage_unit_id` | Configured accept-queue depth (`TQ_STORAGE_ZMQ_BACKLOG`) | +| `tq_storage_accept_queue_peak` | Gauge | `storage_unit_id` | Deepest accept-queue occupancy seen by the probe | +| `tq_storage_accept_queue_peak_utilization_ratio` | Gauge | `storage_unit_id` | Peak occupancy as a fraction of the backlog | +| `tq_storage_socket_drops` | Gauge | `storage_unit_id` | Connections dropped on this listening socket since the probe started | +| `tq_storage_listen_overflows` | Gauge | `storage_unit_id` | Namespace-wide accept-queue overflows since the probe started | +| `tq_storage_listen_other_drops` | Gauge | `storage_unit_id` | Namespace-wide establishment drops that were **not** overflows | + +The kernel charges a listening socket's `sk_drops` on several connection-establishment +failures — a full accept queue, but also failures to allocate or route the new connection — so +`tq_storage_socket_drops` rising locates the socket, not the cause. Compare the last two +series: a non-zero `tq_storage_listen_other_drops` means raising the backlog would not have +prevented every drop in that window. + ### Storage Unit Native Metrics (exposed on each storage unit's own endpoint) | Metric | Type | Labels | Description | @@ -163,6 +197,7 @@ The dashboard ([`scripts/grafana_dashboard.json`](../scripts/grafana_dashboard.j | **Request Throughput & Latency** | Controller Request Rate (ops/s), Controller Request Latency (repeats per quantile) | | **Partition Status** | Samples per Partition, Production Progress, Consumption Progress | | **Storage Units** | Utilization Bar Gauge, Active Keys, Capacity vs Active Keys, RSS Memory, Storage Request Rate, Storage Request Latency (repeats per quantile), Produced vs Cleared Samples/s, Active Keys Delta | +| **Storage Request-Loss Diagnostics** | Arrived vs Completed/s, Accept-Queue Peak vs Backlog, Listening-Socket Drops/s, Overflow vs Other Establishment Drops/s | ### Template Variables diff --git a/scripts/grafana_dashboard.json b/scripts/grafana_dashboard.json index 0e6c5341..9f3ee36b 100644 --- a/scripts/grafana_dashboard.json +++ b/scripts/grafana_dashboard.json @@ -370,6 +370,68 @@ { "expr": "sum(tq_storage_active_keys_total)", "legendFormat": "Total Active Keys (all storage units)" } ], "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 86 }, + "id": 104, + "title": "Storage Request-Loss Diagnostics", + "type": "row" + }, + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "description": "Arrivals are counted when the worker decodes a request; completions only when it finishes. A sustained gap means requests are arriving and not completing. Neither series identifies an individual request.", + "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "fillOpacity": 10, "lineWidth": 2 } } }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 87 }, + "id": 40, + "options": { "legend": { "calcs": ["lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi" } }, + "title": "Arrived vs Completed (per second)", + "targets": [ + { "expr": "sum by (op_type) (rate(tq_storage_arrivals_by_op{op_type=~\"$op_type\"}[$__rate_interval]))", "legendFormat": "arrived {{ op_type }}" }, + { "expr": "sum by (op_type) (rate(tq_storage_request_ops{op_type=~\"$op_type\"}[$__rate_interval]))", "legendFormat": "completed {{ op_type }}" } + ], + "type": "timeseries" + }, + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "description": "Peak is a high-water mark since the probe started, so it never decreases. Only present when TQ_ACCEPT_PROBE_INTERVAL is set.", + "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "fillOpacity": 10, "lineWidth": 2 } } }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 87 }, + "id": 41, + "options": { "legend": { "calcs": ["lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi" } }, + "title": "Accept-Queue Peak vs Backlog", + "targets": [ + { "expr": "tq_storage_accept_queue_peak", "legendFormat": "peak {{ storage_unit_id }}" }, + { "expr": "tq_storage_accept_queue_backlog", "legendFormat": "backlog {{ storage_unit_id }}" } + ], + "type": "timeseries" + }, + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "description": "The kernel charges sk_drops for several connection-establishment failures, of which a full accept queue is only one, so a rise locates the socket, not the cause.", + "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "fillOpacity": 10, "lineWidth": 2 } } }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 95 }, + "id": 42, + "options": { "legend": { "calcs": ["lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi" } }, + "title": "Listening-Socket Drops (per second)", + "targets": [ + { "expr": "sum by (storage_unit_id) (rate(tq_storage_socket_drops[$__rate_interval]))", "legendFormat": "{{ storage_unit_id }}" } + ], + "type": "timeseries" + }, + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "description": "Namespace-wide, not per port. A non-zero 'other' series means raising the backlog would not have prevented every drop in the window.", + "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "fillOpacity": 10, "lineWidth": 2 } } }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 95 }, + "id": 43, + "options": { "legend": { "calcs": ["lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi" } }, + "title": "Overflow vs Other Establishment Drops (per second)", + "targets": [ + { "expr": "sum (rate(tq_storage_listen_overflows[$__rate_interval]))", "legendFormat": "accept-queue overflow" }, + { "expr": "sum (rate(tq_storage_listen_other_drops[$__rate_interval]))", "legendFormat": "other establishment drops" } + ], + "type": "timeseries" } ], "refresh": "10s", diff --git a/tests/test_accept_probe.py b/tests/test_accept_probe.py new file mode 100644 index 00000000..a834955c --- /dev/null +++ b/tests/test_accept_probe.py @@ -0,0 +1,257 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the accept-queue probe. + +A connection dropped during establishment is invisible to both ends, and the kernel +charges the listening socket's sk_drops on several such paths, of which a full accept +queue is only one. These tests pin the arithmetic the probe reports, the levels it logs +at, and that it never names a cause its counters cannot establish. +""" + +import threading + +from transfer_queue.utils.accept_probe import ( + AcceptQueueProbe, + AcceptQueueSample, + AcceptQueueStats, + _read_listen_overflows, + sample_accept_queue, +) + + +def _sample( + recv_q: int, backlog: int = 100, sk_drops: int = 0, overflows: int = 0, drops: int | None = None +) -> AcceptQueueSample: + return AcceptQueueSample( + timestamp=0.0, + recv_q=recv_q, + backlog=backlog, + sk_drops=sk_drops, + listen_overflows=overflows, + listen_drops=overflows if drops is None else drops, + ) + + +def _probe() -> AcceptQueueProbe: + return AcceptQueueProbe(port=34513, owner_id="TQ_STORAGE_UNIT_test", interval_s=0.01) + + +def test_utilization_is_depth_over_backlog(): + assert _sample(50, backlog=100).utilization == 0.5 + + +def test_zero_backlog_reports_zero_utilization_not_zero_division(): + """A socket read before bind reports backlog 0; that must not raise.""" + assert _sample(5, backlog=0).utilization == 0.0 + + +def test_full_queue_reports_full_utilization(): + assert _sample(100, backlog=100).utilization == 1.0 + + +def test_empty_stats_report_zero_deltas(): + stats = AcceptQueueStats(port=1234) + + assert stats.sk_drops_delta == 0 + assert stats.overflow_delta == 0 + + +def test_drop_delta_is_last_minus_first(): + """Deltas are what matter: the absolute counters carry the whole uptime's history.""" + stats = AcceptQueueStats(port=1234) + stats.first_sample = _sample(0, sk_drops=23, overflows=1775) + stats.last_sample = _sample(0, sk_drops=31, overflows=1790) + + assert stats.sk_drops_delta == 8 + assert stats.overflow_delta == 15 + + +def test_describe_names_port_and_deltas(): + stats = AcceptQueueStats(port=34513, backlog=100, peak_recv_q=97) + stats.first_sample = _sample(0, sk_drops=23) + stats.last_sample = _sample(0, sk_drops=24) + + text = stats.describe() + + assert "port=34513" in text + assert "peak_recv_q=97" in text + assert "sk_drops_delta=1" in text + + +def test_peak_tracks_highest_depth_not_latest(): + """The burst is the signal; sampling after it drains reads zero.""" + probe = _probe() + + probe._record(_sample(10)) + probe._record(_sample(97)) + probe._record(_sample(3)) + + assert probe.stats.peak_recv_q == 97 + + +def test_first_sample_is_retained_as_delta_baseline(): + probe = _probe() + + probe._record(_sample(1, sk_drops=23)) + probe._record(_sample(2, sk_drops=25)) + + assert probe.stats.first_sample.sk_drops == 23 + assert probe.stats.sk_drops_delta == 2 + + +def test_drop_increase_is_logged_at_error(caplog): + """A drop is the direct evidence, so it must not be buried at debug level.""" + probe = _probe() + probe._record(_sample(0, sk_drops=23)) + + with caplog.at_level("ERROR"): + probe._record(_sample(100, sk_drops=24)) + + assert "dropped an incoming connection" in caplog.text + + +def test_steady_drop_count_is_logged_once_not_every_sample(caplog): + """sk_drops is cumulative, so a past drop must not be re-reported forever. + + Measuring against the probe's first sample made the condition permanently true once + the socket had ever dropped a connection: at a 0.1s interval that is ten errors a + second for the life of the process, and it erases when the drop actually happened. + """ + probe = _probe() + probe._record(_sample(0, sk_drops=12)) + + with caplog.at_level("ERROR"): + probe._record(_sample(0, sk_drops=13)) # a new drop -- report it + for _ in range(20): + probe._record(_sample(0, sk_drops=13)) # unchanged -- stay quiet + + assert caplog.text.count("dropped an incoming connection") == 1 + + +def test_each_new_drop_is_reported(caplog): + """Quieting the repeat must not swallow genuinely new drops.""" + probe = _probe() + probe._record(_sample(0, sk_drops=12)) + + with caplog.at_level("ERROR"): + probe._record(_sample(0, sk_drops=13)) + probe._record(_sample(0, sk_drops=13)) + probe._record(_sample(0, sk_drops=14)) + + assert caplog.text.count("dropped an incoming connection") == 2 + + +def test_near_full_queue_warns_once(caplog): + """Repeating the warning every 100ms would flood the log of a 512-node job.""" + probe = _probe() + + with caplog.at_level("WARNING"): + probe._record(_sample(60)) + probe._record(_sample(70)) + + assert caplog.text.count("reached") == 1 + + +def test_no_warning_below_threshold(): + probe = _probe() + + probe._record(_sample(10)) + + assert probe._warned is False + + +def test_listen_overflows_returns_two_non_negative_ints(): + """Reads the real host, so assert shape rather than a specific value.""" + overflows, drops = _read_listen_overflows() + + assert overflows >= 0 + assert drops >= 0 + + +def test_sampling_an_unused_port_returns_none(): + assert sample_accept_queue(1) is None + + +def test_unit_shutdown_stops_the_probe(): + """Nothing else calls stop(), so the sampling thread would outlive the unit.""" + from unittest.mock import MagicMock + + from transfer_queue.storage.simple_storage import SimpleStorageUnit + + unit_class = SimpleStorageUnit.__ray_metadata__.modified_class + probe = MagicMock() + + unit_class._shutdown_resources( + shutdown_event=threading.Event(), + worker_thread=None, + proxy_thread=None, + zmq_context=None, + put_get_socket=None, + accept_probe=probe, + ) + + probe.stop.assert_called_once() + + +def test_shutdown_without_a_probe_is_a_no_op(): + """The probe is opt-in, so the default path must not require one.""" + from transfer_queue.storage.simple_storage import SimpleStorageUnit + + unit_class = SimpleStorageUnit.__ray_metadata__.modified_class + + unit_class._shutdown_resources( + shutdown_event=threading.Event(), + worker_thread=None, + proxy_thread=None, + zmq_context=None, + put_get_socket=None, + ) + + +def test_non_overflow_drops_are_separated_from_overflows(): + """ListenDrops counts every establishment failure; only some are queue overflows. + + The kernel charges a listening socket's sk_drops on several paths -- a full accept + queue, but also failures to allocate or route the new connection -- so this difference + is what says whether a bigger backlog could have helped. + """ + stats = AcceptQueueStats(port=1234) + stats.first_sample = _sample(0, overflows=10, drops=20) + stats.last_sample = _sample(0, overflows=11, drops=25) + + assert stats.overflow_delta == 1 + assert stats.non_overflow_drop_delta == 4 + + +def test_pure_overflow_window_reports_no_other_drops(): + stats = AcceptQueueStats(port=1234) + stats.first_sample = _sample(0, overflows=10, drops=10) + stats.last_sample = _sample(0, overflows=13, drops=13) + + assert stats.overflow_delta == 3 + assert stats.non_overflow_drop_delta == 0 + + +def test_drop_alert_does_not_assert_the_queue_overflowed(caplog): + """sk_drops locates the socket, not the cause, so the alert must not name one.""" + probe = _probe() + probe._record(_sample(0, sk_drops=1)) + + with caplog.at_level("ERROR"): + probe._record(_sample(0, sk_drops=2)) + + assert "dropped an incoming connection" in caplog.text + assert "accept queue dropped" not in caplog.text diff --git a/tests/test_metrics.py b/tests/test_metrics.py index a8ac3e2b..66e82a3f 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -98,6 +98,14 @@ def test_all_metrics_are_registered(self): "tq_storage_active_keys_total", "tq_storage_utilization_ratio", "tq_storage_memory_rss_bytes", + "tq_storage_requests_arrived", + "tq_storage_arrivals_by_op", + "tq_storage_accept_queue_backlog", + "tq_storage_accept_queue_peak", + "tq_storage_accept_queue_peak_utilization_ratio", + "tq_storage_socket_drops", + "tq_storage_listen_overflows", + "tq_storage_listen_other_drops", ] registered = {m.name for m in exporter.registry.collect()} @@ -315,6 +323,76 @@ def test_storage_metrics_populated_on_success(self): assert exporter.storage_utilization.labels(storage_unit_id="SU_001")._value.get() == 0.25 assert exporter.storage_memory_rss.labels(storage_unit_id="SU_001")._value.get() == 512 * 1024 * 1024 + def test_arrival_counters_are_exported(self): + """Arrival counts reach Prometheus, so a dashboard can compare them with completions.""" + exporter = TQMetricsExporter() + fake_su_info = MagicMock() + fake_su_info.id = "SU_001" + exporter._storage_unit_infos = {"SU_001": fake_su_info} + exporter._query_storage_unit = MagicMock( + return_value={ + "storage_unit_id": "SU_001", + "capacity": 1000, + "active_keys": 1, + "requests_arrived": 42, + "arrivals_by_op": {"GET_DATA": 30, "PUT_DATA": 12}, + } + ) + + exporter.collect_storage_metrics() + + assert exporter.storage_requests_arrived.labels(storage_unit_id="SU_001")._value.get() == 42 + by_op = exporter.storage_arrivals_by_op + assert by_op.labels(storage_unit_id="SU_001", op_type="GET_DATA")._value.get() == 30 + assert by_op.labels(storage_unit_id="SU_001", op_type="PUT_DATA")._value.get() == 12 + + def test_accept_queue_metrics_are_exported(self): + """The overflow/non-overflow split is what tells a dashboard if backlog is the issue.""" + exporter = TQMetricsExporter() + fake_su_info = MagicMock() + fake_su_info.id = "SU_001" + exporter._storage_unit_infos = {"SU_001": fake_su_info} + exporter._query_storage_unit = MagicMock( + return_value={ + "storage_unit_id": "SU_001", + "capacity": 1000, + "active_keys": 1, + "accept_queue": { + "backlog": 4096, + "peak_recv_q": 97, + "peak_utilization": 0.02, + "sk_drops_delta": 5, + "listen_overflow_delta": 2, + "non_overflow_drop_delta": 3, + }, + } + ) + + exporter.collect_storage_metrics() + + label = {"storage_unit_id": "SU_001"} + assert exporter.storage_accept_queue_backlog.labels(**label)._value.get() == 4096 + assert exporter.storage_accept_queue_peak.labels(**label)._value.get() == 97 + assert exporter.storage_socket_drops.labels(**label)._value.get() == 5 + assert exporter.storage_listen_overflows.labels(**label)._value.get() == 2 + assert exporter.storage_listen_other_drops.labels(**label)._value.get() == 3 + + def test_accept_queue_series_absent_when_the_probe_is_off(self): + """The probe is opt-in; reporting 0 drops would read as 'measured, and none'.""" + exporter = TQMetricsExporter() + fake_su_info = MagicMock() + fake_su_info.id = "SU_001" + exporter._storage_unit_infos = {"SU_001": fake_su_info} + exporter._query_storage_unit = MagicMock( + return_value={"storage_unit_id": "SU_001", "capacity": 1000, "active_keys": 1} + ) + + exporter.collect_storage_metrics() + + exported = {sample.name for metric in exporter.registry.collect() for sample in metric.samples} + assert "tq_storage_socket_drops" not in exported + assert "tq_storage_accept_queue_backlog" not in exported + def test_storage_metrics_handles_query_failure(self): """If a storage unit query fails, other units should still be collected.""" exporter = TQMetricsExporter() diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py index ec6e9476..59f82edc 100644 --- a/tests/test_storage_request_retry.py +++ b/tests/test_storage_request_retry.py @@ -147,7 +147,11 @@ def test_log_heavy_operation_thresholds(caplog, elapsed, payload_bytes, should_l @pytest.mark.parametrize( "tcp_result, probe_result, expected", [ - ((None, None), {"active_keys": 4, "op_stats": {"GET_DATA": {"request_count": 1025}}}, "request_lost"), + # A probe that answers proves only that the unit is serving again, whatever the + # cumulative counters happen to say. + ((None, None), {"active_keys": 4, "arrivals_by_op": {"GET_DATA": 0}}, "unit_serving_again"), + ((None, None), {"active_keys": 4, "arrivals_by_op": {"GET_DATA": 9}}, "unit_serving_again"), + ((None, None), {"active_keys": 4}, "unit_serving_again"), ((None, None), zmq.error.Again(), "unit_not_serving"), (ConnectionRefusedError(), zmq.error.Again(), "tcp=down(ConnectionRefusedError)"), ], @@ -170,6 +174,41 @@ async def test_diagnosis_classifies_the_failure(tcp_result, probe_result, expect assert expected in diagnosis +def test_diagnosis_does_not_report_empty_op_stats_as_zero_traffic(): + """op_stats is Prometheus-gated, so an empty dict must not read as 'served nothing'.""" + described = ssm._describe_unit_state({"requests_arrived": 7, "active_keys": 1}) + + assert "completed=unavailable(prometheus_disabled)" in described + assert "completed={}" not in described + + +@pytest.mark.parametrize( + "arrivals", + [ + # Nine historical GETs, all already completed: they say nothing about this request. + {"GET_DATA": 9}, + # The unit restarted during the timeout window, so its counters began again at zero. + {"GET_DATA": 0}, + # The request arrived and finished after the caller gave up waiting. + {"GET_DATA": 10}, + ], +) +def test_cumulative_counters_never_become_a_claim_about_this_request(arrivals): + """Counters are per-op and cumulative, so no value of them locates one request. + + They carry no request id, never decrease, and reset when the unit restarts, so both + directions of inference are unsound: a nonzero count may be entirely historical, and a + zero one may mean the request reached a previous process. + """ + described = ssm._describe_unit_state({"requests_arrived": sum(arrivals.values()), "arrivals_by_op": arrivals}) + + assert "verdict=unit_serving_again" in described + assert "arrived_but_unfinished" not in described + assert "request_lost_in_flight" not in described + # The raw counters stay in the line as triage input, just not as a verdict. + assert "arrivals_by_op=" in described + + @pytest.mark.asyncio async def test_diagnosis_handles_an_unknown_unit(): assert "unit_not_registered" in await _manager(with_unit=False)._diagnose_storage_unit("missing") diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index 12914585..e75792b5 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -185,6 +185,62 @@ def _define_controller_metrics(self) -> None: "tq_storage_memory_rss_bytes", "Storage unit process RSS memory", ["storage_unit_id"], registry=r ) + # ---- Storage-unit request-loss diagnostics ---- + # Read against tq_storage_request_ops, which advances only on completion: a gap + # between the two is requests arriving and not finishing. + self.storage_requests_arrived = Gauge( + "tq_storage_requests_arrived", + "Requests decoded by the storage unit worker, whether or not they completed", + ["storage_unit_id"], + registry=r, + ) + self.storage_arrivals_by_op = Gauge( + "tq_storage_arrivals_by_op", + "Requests decoded by the storage unit worker, by operation", + ["storage_unit_id", "op_type"], + registry=r, + ) + + # ---- Accept-queue probe (only populated when TQ_ACCEPT_PROBE_INTERVAL > 0) ---- + self.storage_accept_queue_backlog = Gauge( + "tq_storage_accept_queue_backlog", + "Configured accept-queue depth of the storage unit's listening socket", + ["storage_unit_id"], + registry=r, + ) + self.storage_accept_queue_peak = Gauge( + "tq_storage_accept_queue_peak", + "Deepest accept-queue occupancy seen by the probe", + ["storage_unit_id"], + registry=r, + ) + self.storage_accept_queue_peak_utilization = Gauge( + "tq_storage_accept_queue_peak_utilization_ratio", + "Peak accept-queue occupancy as a fraction of the backlog", + ["storage_unit_id"], + registry=r, + ) + self.storage_socket_drops = Gauge( + "tq_storage_socket_drops", + "Connections dropped on this listening socket since the probe started", + ["storage_unit_id"], + registry=r, + ) + # Split because sk_drops covers several establishment failures, not only a full + # queue; the difference is what says whether a bigger backlog would have helped. + self.storage_listen_overflows = Gauge( + "tq_storage_listen_overflows", + "Namespace-wide accept-queue overflows since the probe started", + ["storage_unit_id"], + registry=r, + ) + self.storage_listen_other_drops = Gauge( + "tq_storage_listen_other_drops", + "Namespace-wide establishment drops that were not accept-queue overflows", + ["storage_unit_id"], + registry=r, + ) + # ---- Storage request metrics (collected via ZMQ, exposed as gauges) ---- # P50/P99 are pre-computed on the storage unit side and sent via ZMQ, # avoiding the need to replicate histogram bucket structures (which @@ -362,6 +418,30 @@ def collect_storage_metrics(self) -> None: self.storage_active_keys.labels(storage_unit_id=label).set(active) self.storage_memory_rss.labels(storage_unit_id=label).set(metrics.get("process_rss_bytes", 0)) + self.storage_requests_arrived.labels(storage_unit_id=label).set(metrics.get("requests_arrived", 0)) + for op_type, arrived in (metrics.get("arrivals_by_op") or {}).items(): + self.storage_arrivals_by_op.labels(storage_unit_id=label, op_type=op_type).set(arrived) + + # Drop the series rather than report zero when the probe is off, so a + # disabled probe is not read as "measured, and no drops". + accept_queue = metrics.get("accept_queue") + accept_gauges = ( + (self.storage_accept_queue_backlog, "backlog"), + (self.storage_accept_queue_peak, "peak_recv_q"), + (self.storage_accept_queue_peak_utilization, "peak_utilization"), + (self.storage_socket_drops, "sk_drops_delta"), + (self.storage_listen_overflows, "listen_overflow_delta"), + (self.storage_listen_other_drops, "non_overflow_drop_delta"), + ) + for gauge, key in accept_gauges: + if accept_queue is None: + try: + gauge.remove(label) + except (KeyError, ValueError): + pass + else: + gauge.labels(storage_unit_id=label).set(accept_queue.get(key, 0)) + # Per-operation request stats for op_type, op_data in metrics.get("op_stats", {}).items(): self.storage_request_ops.labels(storage_unit_id=label, op_type=op_type).set( diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 70765cb9..c89d1cfb 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -63,11 +63,30 @@ class StorageUnitTimeout(RuntimeError): """A storage unit did not answer within the send/recv timeout. Distinct from an error the unit reported: only a missing answer is worth a new connection. - The message must name the unit, its endpoint and the timeout, because it is what the callers - of ``put_data`` and ``get_data`` see, and the retry logs rely on it instead of repeating them. """ +def _describe_unit_state(body: dict[str, Any]) -> str: + """Summarize a successful probe: the unit is serving again, plus its own counters. + + Draws no conclusion about the timed-out request: the counters are cumulative per + operation and carry no request identity, so no value of them locates one request. + """ + parts = [ + f"requests_arrived={body.get('requests_arrived')}", + f"arrivals_by_op={body.get('arrivals_by_op')}", + f"active_keys={body.get('active_keys')}", + f"rss_gb={body.get('process_rss_bytes', 0) / 2**30:.2f}", + ] + op_stats = body.get("op_stats") or {} + if op_stats: + parts.append(f"completed={ {op: stats.get('request_count') for op, stats in op_stats.items()} }") + else: + # Only with Prometheus enabled; an empty dict would read as "served nothing". + parts.append("completed=unavailable(prometheus_disabled)") + return f"verdict=unit_serving_again ({' '.join(parts)})" + + _SU_SUBDIR = "simple_storage" _SU_INFO_FILE = "storage_unit_info.json" @@ -240,9 +259,10 @@ async def _probe_storage_unit(self, target_storage_unit: str, socket: zmq.Socket return response_msg.body async def _diagnose_storage_unit(self, target_storage_unit: str) -> str: - """Classify a timeout as a lost request, a stuck unit, or an unreachable node. + """Report whether the unit is reachable and serving after a request to it timed out. Returns one log line and never raises: it runs while another failure is being reported. + Says nothing about where that request went; no per-request state exists to show it. """ info = self.storage_unit_infos.get(target_storage_unit) if info is None: @@ -259,17 +279,13 @@ async def _diagnose_storage_unit(self, target_storage_unit: str) -> str: try: body = await self._probe_storage_unit(target_storage_unit=target_storage_unit) - op_counts = {op: stats.get("request_count") for op, stats in (body.get("op_stats") or {}).items()} - return ( - f"{tcp} verdict=request_lost_in_flight (unit answered a fresh probe: " - f"ops={op_counts} active_keys={body.get('active_keys')} " - f"rss_gb={body.get('process_rss_bytes', 0) / 2**30:.2f})" - ) except zmq.error.Again: return f"{tcp} verdict=unit_not_serving (no probe answer in {TQ_SIMPLE_STORAGE_PROBE_TIMEOUT}s)" except Exception as e: return f"{tcp} verdict=unknown (probe failed: {type(e).__name__}: {e})" + return f"{tcp} {_describe_unit_state(body)}" + async def _request_with_retry( self, operation: str, diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index cd0d2b4b..61c5e9a1 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -43,6 +43,7 @@ if TYPE_CHECKING: from transfer_queue.metrics import TQMetricsExporter + from transfer_queue.utils.accept_probe import AcceptQueueProbe logger = get_logger(__name__) @@ -52,6 +53,12 @@ # Marks a GET_ERROR reply as "the key is gone" so the caller can tell it apart from a real fault. KEY_NOT_FOUND_MARKER = "TQKeyNotFound" +# Accept-queue depth for the client-facing ROUTER. A full queue loses connections silently. +TQ_STORAGE_ZMQ_BACKLOG = int(os.environ.get("TQ_STORAGE_ZMQ_BACKLOG", 4096)) + +# Accept-queue sampling period in seconds; 0 disables the probe. Keep sub-second. +TQ_ACCEPT_PROBE_INTERVAL = float(os.environ.get("TQ_ACCEPT_PROBE_INTERVAL", 0)) + class StorageKeyNotFoundError(KeyError): """Raised when a requested global index is absent from a storage unit. @@ -168,6 +175,10 @@ class SimpleStorageUnit: zmq_server_info: ZMQ connection information for clients. """ + _requests_arrived = 0 + _arrivals_by_op: dict[str, int] = {} + _accept_probe = None + def __init__(self, storage_unit_size: int | None = None): """Initialize a SimpleStorageUnit with the specified size. @@ -180,6 +191,9 @@ def __init__(self, storage_unit_size: int | None = None): self.storage_data = StorageUnitData(self.storage_unit_size) + self._requests_arrived = 0 + self._arrivals_by_op = {} + # Internal communication address for proxy and workers self._inproc_addr = f"inproc://simple_storage_workers_{self.storage_unit_id}" @@ -206,6 +220,7 @@ def __init__(self, storage_unit_size: int | None = None): self.proxy_thread, self.zmq_context, self.put_get_socket, + self._accept_probe, ) def _init_zmq_socket(self) -> None: @@ -219,6 +234,7 @@ def _init_zmq_socket(self) -> None: # Frontend: ROUTER for receiving client requests self.put_get_socket = create_zmq_socket(self.zmq_context, zmq.ROUTER, self._node_ip) + self.put_get_socket.setsockopt(zmq.BACKLOG, TQ_STORAGE_ZMQ_BACKLOG) while True: try: @@ -229,6 +245,18 @@ def _init_zmq_socket(self) -> None: logger.warning(f"[{self.storage_unit_id}]: Try to bind ZMQ sockets failed, retrying...") continue + if TQ_ACCEPT_PROBE_INTERVAL > 0: + # Lazy: the probe shells out to ``ss`` on a timer, so keep it out of runs that + # have not enabled it. + from transfer_queue.utils.accept_probe import AcceptQueueProbe + + self._accept_probe = AcceptQueueProbe( + port=self._put_get_socket_port, + owner_id=str(self.storage_unit_id), + interval_s=TQ_ACCEPT_PROBE_INTERVAL, + ) + self._accept_probe.start() + # Backend: DEALER for worker communication (connected via zmq.proxy) self.worker_socket = create_zmq_socket(self.zmq_context, zmq.DEALER, self._node_ip) self.worker_socket.bind(self._inproc_addr) @@ -348,6 +376,9 @@ def _worker_routine(self) -> None: started = time.perf_counter() try: + self._requests_arrived += 1 + self._arrivals_by_op[operation.name] = self._arrivals_by_op.get(operation.name, 0) + 1 + logger.debug(f"[{self.storage_unit_id}]: worker received operation: {operation}") # Process request @@ -582,8 +613,23 @@ def _handle_get_metrics(self) -> ZMQMessage: "capacity": self.storage_unit_size, "active_keys": self.storage_data.active_key_count, "process_rss_bytes": process_rss, + # Counted on arrival; op_stats below only advances on completion. + "requests_arrived": self._requests_arrived, + "arrivals_by_op": dict(self._arrivals_by_op), } + if self._accept_probe is not None: + stats = self._accept_probe.stats + metrics["accept_queue"] = { + "backlog": stats.backlog, + "peak_recv_q": stats.peak_recv_q, + "peak_utilization": stats.peak_utilization, + "sk_drops_delta": stats.sk_drops_delta, + "listen_overflow_delta": stats.overflow_delta, + "non_overflow_drop_delta": stats.non_overflow_drop_delta, + "samples": stats.samples, + } + # Include per-operation stats if Prometheus metrics are enabled if self._metrics is not None: op_stats = {} @@ -741,6 +787,7 @@ def _shutdown_resources( proxy_thread: Thread | None, zmq_context: zmq.Context | None, put_get_socket: zmq.Socket | None, + accept_probe: "AcceptQueueProbe | None" = None, ) -> None: """Clean up resources on garbage collection.""" logger.info("Shutting down SimpleStorageUnit resources...") @@ -748,6 +795,10 @@ def _shutdown_resources( # Signal all threads to stop shutdown_event.set() + # Before the ZMQ teardown: the probe runs on its own timer and would outlive the unit. + if accept_probe is not None: + accept_probe.stop() + # Terminate put_get_socket if put_get_socket: put_get_socket.close(linger=0) diff --git a/transfer_queue/utils/accept_probe.py b/transfer_queue/utils/accept_probe.py new file mode 100644 index 00000000..916955c4 --- /dev/null +++ b/transfer_queue/utils/accept_probe.py @@ -0,0 +1,245 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Accept-queue sampling for the storage-unit ROUTER socket. + +A connection dropped while being established is invisible to both ends, so the only +way to see it is to read the kernel's own counters. Two caveats govern how callers +may read what this reports: ``ListenOverflows`` and ``ListenDrops`` are per network +namespace rather than per port, and the kernel charges a socket's ``sk_drops`` for +several establishment failures, of which a full queue is only one. See +``docs/metrics.md`` for the exported metrics. +""" + +from __future__ import annotations + +import re +import subprocess +import threading +import time +from dataclasses import dataclass + +from transfer_queue.utils.logging_utils import get_logger + +logger = get_logger(__name__) + +_RECVQ_RE = re.compile(r"^LISTEN\s+(\d+)\s+(\d+)\s+\S*:(\d+)\s", re.M) +_DROPS_RE = re.compile(r"\bd(\d+)\b") + + +@dataclass +class AcceptQueueSample: + """One reading of the accept queue for a single listening port.""" + + timestamp: float + recv_q: int + backlog: int + sk_drops: int + listen_overflows: int + listen_drops: int + + @property + def utilization(self) -> float: + """Queue depth as a fraction of the configured backlog.""" + return self.recv_q / self.backlog if self.backlog else 0.0 + + +@dataclass +class AcceptQueueStats: + """Peak and delta summary over a sampling window.""" + + port: int + samples: int = 0 + peak_recv_q: int = 0 + peak_utilization: float = 0.0 + backlog: int = 0 + first_sample: AcceptQueueSample | None = None + last_sample: AcceptQueueSample | None = None + + @property + def sk_drops_delta(self) -> int: + """Connections this socket dropped during the window, 0 until two samples exist.""" + if self.first_sample is None or self.last_sample is None: + return 0 + return self.last_sample.sk_drops - self.first_sample.sk_drops + + @property + def overflow_delta(self) -> int: + """Namespace-wide accept-queue overflows during the window, 0 until two samples exist.""" + if self.first_sample is None or self.last_sample is None: + return 0 + return self.last_sample.listen_overflows - self.first_sample.listen_overflows + + @property + def non_overflow_drop_delta(self) -> int: + """Establishment drops in this window that were not accept-queue overflows.""" + if self.first_sample is None or self.last_sample is None: + return 0 + drops = self.last_sample.listen_drops - self.first_sample.listen_drops + return drops - self.overflow_delta + + def describe(self) -> str: + """Return a one-line summary of the window, for the probe's shutdown log.""" + return ( + f"port={self.port} samples={self.samples} backlog={self.backlog} " + f"peak_recv_q={self.peak_recv_q} peak_util={self.peak_utilization:.1%} " + f"sk_drops_delta={self.sk_drops_delta} listen_overflow_delta={self.overflow_delta}" + ) + + +def _read_listen_socket(port: int) -> tuple[int, int, int] | None: + """Return (recv_q, backlog, sk_drops) for the listening socket on ``port``.""" + try: + out = subprocess.run(["ss", "-lntm"], capture_output=True, text=True, timeout=5, check=False).stdout + except (OSError, subprocess.SubprocessError) as exc: + logger.debug(f"accept-probe: ss failed: {exc}") + return None + + for match in _RECVQ_RE.finditer(out): + recv_q, backlog, found_port = (int(g) for g in match.groups()) + if found_port != port: + continue + # skmem lives on the continuation line right after the match. + tail = out[match.end() : match.end() + 400] + drops_match = _DROPS_RE.search(tail.split("\n")[1] if "\n" in tail else "") + return recv_q, backlog, int(drops_match.group(1)) if drops_match else 0 + return None + + +def _read_listen_overflows() -> tuple[int, int]: + """Return this namespace's (ListenOverflows, ListenDrops) from /proc/net/netstat.""" + try: + with open("/proc/net/netstat") as handle: + lines = handle.read().splitlines() + except OSError: + return 0, 0 + + for i, line in enumerate(lines): + if not line.startswith("TcpExt:") or "ListenOverflows" not in line: + continue + keys = line.split() + values = lines[i + 1].split() + try: + return ( + int(values[keys.index("ListenOverflows")]), + int(values[keys.index("ListenDrops")]), + ) + except (ValueError, IndexError): + return 0, 0 + return 0, 0 + + +def sample_accept_queue(port: int) -> AcceptQueueSample | None: + """Take one accept-queue reading for ``port``, or None if it cannot be read.""" + listen = _read_listen_socket(port) + if listen is None: + return None + recv_q, backlog, sk_drops = listen + overflows, drops = _read_listen_overflows() + return AcceptQueueSample( + timestamp=time.time(), + recv_q=recv_q, + backlog=backlog, + sk_drops=sk_drops, + listen_overflows=overflows, + listen_drops=drops, + ) + + +class AcceptQueueProbe: + """Sample one listening port's accept queue from a background thread. + + Args: + port: Listening port to watch. + owner_id: Identifier used in log lines (the storage unit id). + interval_s: Seconds between samples. Keep sub-second; the queue drains in + milliseconds, so a slower cadence misses the burst entirely. + warn_utilization: Warn once when depth first reaches this fraction of the backlog. + """ + + def __init__( + self, + port: int, + owner_id: str, + interval_s: float = 0.1, + warn_utilization: float = 0.5, + ) -> None: + self.port = port + self.owner_id = owner_id + self.interval_s = interval_s + self.warn_utilization = warn_utilization + self.stats = AcceptQueueStats(port=port) + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self._warned = False + + def start(self) -> None: + """Start the sampling thread; a second call is a no-op.""" + if self._thread is not None: + return + self._thread = threading.Thread(target=self._run, name=f"AcceptQueueProbe-{self.owner_id}", daemon=True) + self._thread.start() + logger.info(f"[{self.owner_id}]: accept-queue probe started on port {self.port} (interval={self.interval_s}s)") + + def stop(self) -> None: + """Stop the sampling thread and log the window summary.""" + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=5) + self._thread = None + logger.info(f"[{self.owner_id}]: accept-queue probe stopped. {self.stats.describe()}") + + def _run(self) -> None: + while not self._stop.is_set(): + sample = sample_accept_queue(self.port) + if sample is not None: + self._record(sample) + self._stop.wait(self.interval_s) + + def _record(self, sample: AcceptQueueSample) -> None: + stats = self.stats + stats.samples += 1 + stats.backlog = sample.backlog or stats.backlog + previous = stats.last_sample + if stats.first_sample is None: + stats.first_sample = sample + stats.last_sample = sample + + if sample.recv_q > stats.peak_recv_q: + stats.peak_recv_q = sample.recv_q + stats.peak_utilization = sample.utilization + + # Against the previous sample, not the first: sk_drops is cumulative, so measuring + # from probe start would re-report one old drop on every sample. + if previous is not None and sample.sk_drops > previous.sk_drops: + logger.error( + f"[{self.owner_id}]: listening socket on port {self.port} dropped an incoming " + f"connection. recv_q={sample.recv_q}/{sample.backlog} sk_drops={sample.sk_drops} " + f"(+{sample.sk_drops - previous.sk_drops} since the last sample, " + f"+{stats.sk_drops_delta} since probe start); netns since probe start: " + f"listen_overflows +{stats.overflow_delta}, other establishment drops " + f"+{stats.non_overflow_drop_delta}. The kernel charges sk_drops for several " + f"establishment failures, so a full queue is only one candidate: raising " + f"ZMQ_BACKLOG above {sample.backlog} helps only if recv_q above and the " + f"overflow delta point that way." + ) + + if not self._warned and sample.utilization >= self.warn_utilization: + self._warned = True + logger.warning( + f"[{self.owner_id}]: accept queue on port {self.port} reached " + f"{sample.recv_q}/{sample.backlog} ({sample.utilization:.0%}); " + f"connections are queueing faster than they are accepted." + ) From 4cc0d7f3fa71dd9b7c87a00b1c49abeffaeee207 Mon Sep 17 00:00:00 2001 From: Pinjie Xu Date: Thu, 24 Sep 2026 10:04:22 +0800 Subject: [PATCH 07/12] [feat] simple storage support ssd offload (#162) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 --- docs/checkpoint.md | 4 +- docs/metrics.md | 4 + docs/ssd_offload.md | 222 +++++ .../README_PERFTEST_SSD_OFFLOAD.md | 157 ++++ .../performance_test_ssd_offload/perftest.py | 501 ++++++++++ .../perftest_config.yaml | 47 + .../run_perf_test.sh | 24 + scripts/put_benchmark.py | 9 +- tests/e2e/test_ssd_offload_e2e.py | 143 +++ tests/test_metrics.py | 35 + tests/test_simple_storage_unit.py | 339 ++++++- transfer_queue/config.yaml | 12 + transfer_queue/interface.py | 34 +- transfer_queue/metrics.py | 47 + .../bootstrap/simple_storage_bootstrap.py | 53 +- transfer_queue/storage/simple_storage.py | 860 +++++++++++++++--- 16 files changed, 2335 insertions(+), 156 deletions(-) create mode 100644 docs/ssd_offload.md create mode 100644 scripts/performance_test_ssd_offload/README_PERFTEST_SSD_OFFLOAD.md create mode 100644 scripts/performance_test_ssd_offload/perftest.py create mode 100644 scripts/performance_test_ssd_offload/perftest_config.yaml create mode 100644 scripts/performance_test_ssd_offload/run_perf_test.sh create mode 100644 tests/e2e/test_ssd_offload_e2e.py diff --git a/docs/checkpoint.md b/docs/checkpoint.md index a5b36803..8c0a2829 100644 --- a/docs/checkpoint.md +++ b/docs/checkpoint.md @@ -77,7 +77,9 @@ checkpoint_dir/ └── simple_storage/ ├── storage_unit_info.json # Position-to-ID manifest ├── su_0_.pkl # StorageUnit at position 0 + ├── su_0_.pkl.blobs/ # Its SSD values, when offload is enabled ├── su_1_.pkl # StorageUnit at position 1 + ├── su_1_.pkl.blobs/ # Its SSD values, when offload is enabled └── ... ``` @@ -123,7 +125,7 @@ tq.save_checkpoint / tq.load_checkpoint └── ... ``` -Both the controller and each storage unit write their data directly to disk from within their own processes. The ZMQ RPC carries only the target file path and an ACK, not the payload — this avoids routing large tensors through the Ray object store. +Both the controller and each storage unit write their data directly to disk from within their own processes. The ZMQ RPC carries only the target file path and an ACK, not the payload — this avoids routing large tensors through the Ray object store. With SSD offload enabled, each `su_*.pkl` stores in-memory values plus an `ssd_index` of plain metadata dictionaries; SSD-backed values are copied unchanged into the adjacent `.blobs` directory. Save and load therefore do not materialize the full SSD tier in host memory. ## Save Order and Consistency diff --git a/docs/metrics.md b/docs/metrics.md index 327f6079..70ba670f 100644 --- a/docs/metrics.md +++ b/docs/metrics.md @@ -137,6 +137,10 @@ Steps: | `tq_storage_active_keys_total` | Gauge | `storage_unit_id` | Active keys in storage | | `tq_storage_utilization_ratio` | Gauge | `storage_unit_id` | Utilization (active/capacity) | | `tq_storage_memory_rss_bytes` | Gauge | `storage_unit_id` | Storage process RSS memory | +| `tq_storage_ssd_offload_enabled` | Gauge | `storage_unit_id` | `1` when SSD offload is enabled, otherwise `0` | +| `tq_storage_ssd_active_values` | Gauge | `storage_unit_id` | Active field values stored on SSD | +| `tq_storage_ssd_active_bytes` | Gauge | `storage_unit_id` | Logical bytes held by active SSD-backed values | +| `tq_storage_ssd_fallback_values_total` | Gauge | `storage_unit_id` | Cumulative values retained in memory because SSD encoding was unavailable | | `tq_storage_request_ops` | Gauge | `storage_unit_id`, `op_type` | Total requests processed by storage unit | | `tq_storage_request_latency_avg` | Gauge | `storage_unit_id`, `op_type` | Average request latency (seconds) | | `tq_storage_request_latency_p50` | Gauge | `storage_unit_id`, `op_type` | P50 request latency (seconds) | diff --git a/docs/ssd_offload.md b/docs/ssd_offload.md new file mode 100644 index 00000000..b0e5b66a --- /dev/null +++ b/docs/ssd_offload.md @@ -0,0 +1,222 @@ +# SimpleStorage SSD Offload + +> Last updated: 09/20/2026 + +## Overview + +SimpleStorage can keep large field values on local SSD instead of keeping them +in host memory for their full lifetime. This reduces long-lived host-memory +use while preserving the existing TransferQueue APIs. + +SSD offload is available only with the `SimpleStorage` backend. Placement is +decided for each field value in each sample. Values whose encoded size is at +or above the configured threshold are stored on SSD; smaller values remain in +memory. + +SSD offload is temporary storage, not durable storage. `CLEAR` operations and +`tq.close()` remove offloaded data. Use a checkpoint when data must survive a +TransferQueue restart. + +## Quick Start + +Enable SSD offload under `backend.SimpleStorage`: + +```yaml +backend: + storage_backend: SimpleStorage + SimpleStorage: + ssd_offload: + enabled: true + path: /local/ssd/transfer-queue-job-a + threshold_bytes: 1048576 + glibc_mmap_threshold_bytes: 1048576 +``` + +## Configuration + +| Config key | Default | Description | +|------------|---------|-------------| +| `backend.SimpleStorage.ssd_offload.enabled` | `false` | Enables SSD offload for SimpleStorage. | +| `backend.SimpleStorage.ssd_offload.path` | `null` | Target directory for SSD offload. TransferQueue creates it, including missing parent directories, if needed. | +| `backend.SimpleStorage.ssd_offload.threshold_bytes` | `1048576` | Minimum encoded size, in bytes, for storing one field value on SSD. The value must be greater than zero. | +| `backend.SimpleStorage.ssd_offload.glibc_mmap_threshold_bytes` | `1048576` | glibc allocation size at which storage actors prefer independently releasable mappings. Set to `null` to leave the process environment unchanged. | + +`glibc_mmap_threshold_bytes` is independent of the SSD placement threshold. It +sets `MALLOC_MMAP_THRESHOLD_` only for storage actors and only when SSD offload +is enabled. Lower values can reduce RSS retained after transient allocations +are freed, but can increase mapping system calls and page faults. The setting +is effective only with glibc and disables its dynamic mmap-threshold adjustment. + +### Path ownership + +TransferQueue stores each run under +`/transfer_queue_ssd_offload//`. + +The configured `path` must be a directory, or a path that TransferQueue can +create. The `transfer_queue_ssd_offload` child must be a real directory, not a +symbolic link. + +Each storage unit owns only its `/` directory. A +successful `tq.close()` removes the current run's unit directories without +deleting sibling runs or unrelated files under `transfer_queue_ssd_offload`. + +## Runtime Directory Layout + +The current implementation stores active SSD data under this layout: + +```text +/ +└── transfer_queue_ssd_offload/ + └── / + └── / + └── <00..ff>/ + ├── .bin + └── .tmp- # Present only while a value is being written +``` + +One `run_id` is created for each TransferQueue initialization and shared by its +storage units. Each storage unit has its own directory and creates `00` through +`ff` subdirectories. + +Each active SSD-backed `(field, global_index)` value is stored in one `.bin` +file. The file contains only the encoded value. Its name does not contain the +field or global index; the storage unit keeps that mapping and the decode +metadata in memory. + +TransferQueue first writes a `.tmp-` file, then renames it to +`.bin` after the full value has been written. If a write fails, +TransferQueue attempts to remove its temporary file. + +This directory layout is an internal implementation detail and may change. +Applications should access offloaded values through TransferQueue APIs rather +than reading these files directly. Checkpoint files use the separate layout +described in [Checkpoint Integration](#checkpoint-integration). + +## Data Placement + +Placement is based on the encoded size of one field value for one sample. The +threshold is inclusive: + +```text +encoded size < threshold_bytes -> host memory +encoded size >= threshold_bytes -> local SSD +``` + +For a batched tensor, each item along the first dimension is considered one +sample. Different fields in the same sample can be placed in different tiers. +Overwriting a key can also move its values between memory and SSD. + +The following value types can be stored on SSD: + +| Value type | SSD representation | +|------------|--------------------| +| Dense PyTorch tensor | Raw tensor bytes, dtype, and shape | +| NumPy array without object dtype | Raw array bytes, dtype, and shape | +| NumPy array with object dtype | Pickle payload | +| `bytes` | Original bytes | +| Other Python value that `pickle` can serialize | Pickle payload | + +Nested or sparse PyTorch tensors and values that cannot be encoded remain in +memory. A non-CPU tensor is copied to CPU before it is written to SSD. + +GET reads SSD-backed values and reconstructs their original types. SSD file +references are internal and are not returned to the application. + +The storage unit encodes values, performs SSD I/O, and keeps the file +references. The controller continues to manage metadata and scheduling; it +does not move SSD payloads. + +## Data Cleanup + +TransferQueue removes SSD files at these points: + +- `CLEAR` deletes the files for the cleared samples. +- A successful overwrite deletes files that are no longer referenced. +- A successful `tq.close()` stops each SimpleStorage unit and deletes that + unit's `/` directory. + +TransferQueue does not sweep sibling or stale run directories. A process that +exits without completing `tq.close()` can leave its run directory behind for +manual cleanup. The configured `path` and `transfer_queue_ssd_offload` child +directory remain in place. + +## Multi-Node Requirements + +All storage units receive the same `ssd_offload.path`. Each unit uses that path +on its own node, so SSD offload does not require a shared filesystem. + +Checkpoint directories must be shared across nodes; see [Checkpoint: Save and +Restore System State](checkpoint.md#multi-node-requirements). + +## Checkpoint Integration + +SSD offload works with `tq.save_checkpoint` and `tq.load_checkpoint`. Each +storage unit has a `.pkl` file. Its `.pkl.blobs/` directory exists only when +that unit has SSD-backed values: + +```text +simple_storage/ +├── su__.pkl # In-memory values and SSD metadata +└── su__.pkl.blobs/ # SSD-backed values, when present +``` + +Saving a checkpoint copies SSD files directly into the `.blobs` directory. It +does not read and decode the full SSD tier into host memory. Loading the +checkpoint copies the blobs into the current SSD offload directory. + +An SSD checkpoint must be loaded with SSD offload enabled. A checkpoint made +without SSD offload can be loaded with SSD offload enabled; its values are +placed using the current `threshold_bytes` setting. + +See [Checkpoint: Save and Restore System State](checkpoint.md) for checkpoint +consistency, replacement, and multi-node rules. + +## Metrics + +When metrics are enabled, the controller endpoint exposes these metrics for +each storage unit: + +| Metric | Description | +|--------|-------------| +| `tq_storage_ssd_offload_enabled` | `1` when SSD offload is enabled, otherwise `0` | +| `tq_storage_ssd_active_values` | Number of active field values stored on SSD | +| `tq_storage_ssd_active_bytes` | Logical encoded bytes held by active SSD-backed values | +| `tq_storage_ssd_fallback_values_total` | Cumulative number of field values retained in memory because SSD encoding was unavailable | + +`tq_storage_ssd_active_bytes` is logical payload size, not filesystem usage. +Filesystem blocks, directories, and temporary files are not included. See +[Prometheus Metrics & Grafana Dashboard](metrics.md) for metrics setup. + +`tq_storage_ssd_fallback_values_total` is monotonic for the lifetime of a +storage unit. A non-zero or increasing value means SSD offload is enabled but +some values cannot use any supported SSD encoding and remain in DRAM. Clearing +or overwriting those values does not decrease the counter. + +## Choosing a Threshold + +A lower threshold moves more values to SSD and can reduce long-lived host +memory use. It also creates more files, consumes more filesystem inodes, and +increases SSD I/O. A higher threshold keeps more values in memory and reduces +SSD work. + +Choose the threshold using the real field sizes and access pattern of the +workload. The benchmark under +[`scripts/performance_test_ssd_offload`](../scripts/performance_test_ssd_offload/README_PERFTEST_SSD_OFFLOAD.md) +compares host-memory mode with SSD offload and reports throughput, storage-unit +RSS, and active SSD bytes. + +## Known Limitations + +### GET loads SSD-backed values into memory + +GET reads each requested SSD file and reconstructs the value in host memory. +SSD offload reduces long-lived memory use, but it does not remove temporary +memory use while values are read or encoded. + +### `data_parser` outputs should not share backing storage across samples + +If a `data_parser` returns differently sized tensor or NumPy views backed by +the same allocation, an in-memory view can keep the entire allocation resident +after another view is offloaded. When parsing URLs or file paths, return an +independently owned value for each sample, or copy retained views before +returning them. diff --git a/scripts/performance_test_ssd_offload/README_PERFTEST_SSD_OFFLOAD.md b/scripts/performance_test_ssd_offload/README_PERFTEST_SSD_OFFLOAD.md new file mode 100644 index 00000000..ae5b6bd7 --- /dev/null +++ b/scripts/performance_test_ssd_offload/README_PERFTEST_SSD_OFFLOAD.md @@ -0,0 +1,157 @@ +# TransferQueue SimpleStorage SSD Offload Performance Benchmark + +This benchmark is separate from the general backend throughput benchmarks under +`scripts/performance_test`. It runs the same mixed workload with SimpleStorage +host-memory storage and SSD offload, and records throughput, SSD byte accounting, +and storage-unit RSS. + +## Prerequisites + +1. Start a Ray cluster with exactly two live nodes. Each node must advertise + the `node:` resource used by the benchmark for actor placement: + + ```bash + # On the head node + ray start --head --resources='{"node:10.0.0.1":1}' + + # On the worker node + ray start --address=10.0.0.1:6379 --resources='{"node:10.0.0.2":1}' + ``` + +2. Create the same node-local SSD path on both nodes. The path must also be + visible to the benchmark driver and storage actors. When they run in + containers, bind-mount the path at the same absolute location. +3. Run the benchmark on the Ray head node in an environment containing the + TransferQueue runtime dependencies. A containerized driver that joins a + node-local Ray runtime must also bind-mount that runtime's `/tmp` directory + so it can reach the Ray socket. + +## Usage + +Use the wrapper to run the default 5 GiB mixed workload: + +```bash +HEAD_NODE_IP=10.0.0.1 \ +WORKER_NODE_IP=10.0.0.2 \ +SSD_OFFLOAD_PATH=/path/to/local/ssd \ +RESULTS_DIR=/path/to/results \ +bash scripts/performance_test_ssd_offload/run_perf_test.sh +``` + +### Wrapper configuration + +| Variable | Description | Default | Required | +|----------|-------------|---------|----------| +| `HEAD_NODE_IP` | Ray head-node IP | - | Yes | +| `WORKER_NODE_IP` | Ray worker-node IP | - | Yes | +| `SSD_OFFLOAD_PATH` | Node-local SSD path on both nodes | - | Yes | +| `RESULTS_DIR` | Output directory | `scripts/performance_test_ssd_offload/results` | No | +| `NUM_TEST_ITERATIONS` | Iterations per storage mode | `4` | No | + +## Arguments + +Run `perftest.py` directly to change the workload: + +```bash +python scripts/performance_test_ssd_offload/perftest.py \ + --backend_config=scripts/performance_test_ssd_offload/perftest_config.yaml \ + --head_node_ip=10.0.0.1 \ + --worker_node_ip=10.0.0.2 \ + --ssd_path=/path/to/local/ssd \ + --output_dir=/path/to/results \ + --global_batch_size=512 \ + --small_fields=4 \ + --small_sample_bytes=524288 \ + --large_fields=4 \ + --large_sample_bytes=2097152 \ + --num_test_iterations=4 \ + --warmup_iterations=1 +``` + +| Argument | Description | Default | Required | +|----------|-------------|---------|----------| +| `--backend_config` | Path to the benchmark configuration | - | Yes | +| `--head_node_ip` | Ray head-node IP used for writer placement | - | Yes | +| `--worker_node_ip` | Ray worker-node IP used for reader placement | - | Yes | +| `--ssd_path` | Node-local SSD path | - | Yes | +| `--output_dir` | Output directory | `scripts/performance_test_ssd_offload/results` | No | +| `--num_test_iterations` | Iterations per storage mode | `4` | No | +| `--warmup_iterations` | Leading iterations excluded from the summary | `1` | No | +| `--global_batch_size` | Samples per iteration | `512` | No | +| `--small_fields` | Tensor fields below the offload threshold | `4` | No | +| `--large_fields` | Tensor fields above the offload threshold | `4` | No | +| `--small_sample_bytes` | Bytes per sample in each small field | `524288` | No | +| `--large_sample_bytes` | Bytes per sample in each large field | `2097152` | No | + +Sample sizes are bytes per tensor field per sample. Small samples must be +strictly below the configured `ssd_offload.threshold_bytes`, and large samples +must be strictly above it. The SSD threshold and number of storage units come +from `perftest_config.yaml`. The number of iterations must exceed the number of +warm-up iterations. + +## Benchmark Scenario + +The default batch contains eight float32 tensor fields whose per-sample sizes +straddle the 1 MiB offload threshold: + +| Fields | Bytes per sample | Placement | +|--------|-----------------:|-----------| +| 4 small fields | 512 KiB | Host memory | +| 4 large fields | 2 MiB | Local SSD | + +With a batch size of 512, the logical payload is 5 GiB: 1 GiB remains in host +memory and 4 GiB is eligible for SSD offload. The benchmark first runs this +workload in pure host-memory mode and then repeats it with SSD offload. + +## Benchmark Flow + +Each iteration performs a PUT -> LIST -> GET -> CLEAR cycle through the +TransferQueue KV API. Storage-unit Prometheus metrics are sampled before PUT, +after PUT, after GET, and after CLEAR. RSS is accepted only after the configured +number of samples agree within the stability tolerance. + +After excluding warm-up iterations, the benchmark reports the theoretical and +actual SSD-active bytes, verifies disk reclamation after CLEAR, and summarizes +SSD-mode RSS before PUT and after PUT, GET, and CLEAR. The post-CLEAR RSS delta +and its growth across iterations show whether memory is progressively retained. + +## Distributed Smoke Test + +Use a smaller workload to verify node placement, mixed-tier routing, metrics, +and cleanup: + +```bash +python scripts/performance_test_ssd_offload/perftest.py \ + --backend_config=scripts/performance_test_ssd_offload/perftest_config.yaml \ + --head_node_ip="$HEAD_NODE_IP" \ + --worker_node_ip="$WORKER_NODE_IP" \ + --ssd_path="$SSD_OFFLOAD_PATH" \ + --output_dir=/tmp/tq-ssd-smoke \ + --global_batch_size=8 \ + --num_test_iterations=2 \ + --warmup_iterations=1 +``` + +This workload keeps 16 MiB in memory and offloads 64 MiB for functional +validation. Use the default 5 GiB workload when collecting post-CLEAR RSS +trends. + +## Output Format + +The output directory contains: + +- `mixed_memory.csv`: memory-mode throughput and diagnostic RSS samples; +- `mixed_ssd.csv`: SSD-mode throughput, RSS samples, and active SSD bytes; +- `summary.json`: warm-up-filtered SSD-byte accounting and post-CLEAR RSS + retention measurements. + +Important `summary.json` fields are: + +| Field | Description | +|-------|-------------| +| `theoretical_offload_bytes` | Logical payload expected to be stored on SSD | +| `ssd_median_active_bytes_after_put` | Median bytes reported as active on SSD after PUT | +| `ssd_max_active_bytes_after_clear` | Maximum active SSD bytes after CLEAR; expected to be zero | +| `ssd_median_rss_retained_after_clear_bytes` | Median SSD-mode RSS increase remaining after CLEAR | +| `ssd_max_rss_retained_after_clear_bytes` | Maximum SSD-mode RSS increase remaining after CLEAR | +| `ssd_clear_rss_growth_bytes` | Last post-CLEAR RSS minus the first post-CLEAR RSS | diff --git a/scripts/performance_test_ssd_offload/perftest.py b/scripts/performance_test_ssd_offload/perftest.py new file mode 100644 index 00000000..81e39bdc --- /dev/null +++ b/scripts/performance_test_ssd_offload/perftest.py @@ -0,0 +1,501 @@ +#!/usr/bin/env python3 +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark mixed in-memory and SSD-backed SimpleStorage samples.""" + +import argparse +import csv +import json +import logging +import statistics +import time +import urllib.request +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import ray +import torch +from omegaconf import DictConfig, OmegaConf +from prometheus_client.parser import text_string_to_metric_families +from tensordict import TensorDict + +import transfer_queue as tq + +LOGGER = logging.getLogger("ssd_offload_benchmark") +DEFAULT_SMALL_SAMPLE_BYTES = 512 * 1024 +DEFAULT_LARGE_SAMPLE_BYTES = 2 * 1024 * 1024 +BYTES_PER_FLOAT32 = 4 +RAY_ADDRESS = "auto" +METRICS_TIMEOUT_SECONDS = 180 +RSS_SAMPLE_INTERVAL_SECONDS = 11 +RSS_STABLE_SAMPLES = 3 +RSS_STABILITY_TOLERANCE_BYTES = 32 * 2**20 +MODE_GAP_SECONDS = 10 + + +@dataclass(frozen=True) +class Workload: + batch_size: int + small_fields: int + large_fields: int + small_sample_bytes: int + large_sample_bytes: int + offload_threshold_bytes: int + + def validate(self) -> None: + sizes = { + "batch_size": self.batch_size, + "small_sample_bytes": self.small_sample_bytes, + "large_sample_bytes": self.large_sample_bytes, + "offload_threshold_bytes": self.offload_threshold_bytes, + } + if any(value <= 0 for value in sizes.values()): + raise ValueError(f"Workload sizes must be positive: {sizes}") + if self.small_fields < 0 or self.large_fields < 0: + raise ValueError("Field counts must be non-negative") + if self.small_fields + self.large_fields == 0: + raise ValueError("At least one tensor field is required") + if self.small_sample_bytes % BYTES_PER_FLOAT32 or self.large_sample_bytes % BYTES_PER_FLOAT32: + raise ValueError("Sample sizes must be divisible by four for float32 tensors") + if self.small_sample_bytes >= self.offload_threshold_bytes: + raise ValueError("small_sample_bytes must be below the SSD offload threshold") + if self.large_sample_bytes <= self.offload_threshold_bytes: + raise ValueError("large_sample_bytes must be above the SSD offload threshold") + + @property + def inline_bytes(self) -> int: + return self.batch_size * self.small_fields * self.small_sample_bytes + + @property + def offloaded_bytes(self) -> int: + return self.batch_size * self.large_fields * self.large_sample_bytes + + @property + def total_bytes(self) -> int: + return self.inline_bytes + self.offloaded_bytes + + +@dataclass(frozen=True) +class StorageMetrics: + active_keys: int + rss_bytes: int + ssd_active_bytes: int + + +def create_mixed_tensors(workload: Workload) -> TensorDict: + """Create fields whose individual samples straddle the offload threshold.""" + workload.validate() + torch.manual_seed(0) + fields = { + f"small_{index}": torch.randn( + workload.batch_size, + workload.small_sample_bytes // BYTES_PER_FLOAT32, + dtype=torch.float32, + ) + for index in range(workload.small_fields) + } + fields.update( + { + f"large_{index}": torch.randn( + workload.batch_size, + workload.large_sample_bytes // BYTES_PER_FLOAT32, + dtype=torch.float32, + ) + for index in range(workload.large_fields) + } + ) + return TensorDict(fields, batch_size=[workload.batch_size]) + + +def read_storage_metrics(endpoint: str, expected_storage_units: int) -> StorageMetrics: + """Read one aggregate snapshot from the controller's Prometheus endpoint.""" + with urllib.request.urlopen(f"http://{endpoint}/metrics", timeout=5) as response: + payload = response.read().decode("utf-8") + + metric_names = { + "tq_storage_active_keys_total": "active_keys", + "tq_storage_memory_rss_bytes": "rss_bytes", + "tq_storage_ssd_active_bytes": "ssd_active_bytes", + } + units: dict[str, dict[str, int]] = {} + for family in text_string_to_metric_families(payload): + for sample in family.samples: + field = metric_names.get(sample.name) + storage_unit_id = sample.labels.get("storage_unit_id") + if field is not None and storage_unit_id is not None: + units.setdefault(storage_unit_id, {})[field] = int(sample.value) + + complete = len(units) == expected_storage_units and all( + set(values) >= {"active_keys", "rss_bytes", "ssd_active_bytes"} for values in units.values() + ) + if not complete: + fields = sorted({tuple(sorted(values)) for values in units.values()}) + raise RuntimeError(f"Incomplete storage metrics: units={len(units)}/{expected_storage_units}, fields={fields}") + return StorageMetrics( + active_keys=sum(values["active_keys"] for values in units.values()), + rss_bytes=sum(values["rss_bytes"] for values in units.values()), + ssd_active_bytes=sum(values["ssd_active_bytes"] for values in units.values()), + ) + + +def wait_for_storage_state( + endpoint: str, + expected_storage_units: int, + expected_active_keys: int, + timeout_seconds: float, +) -> StorageMetrics: + """Wait until a fresh metrics collection reports the requested key count.""" + deadline = time.monotonic() + timeout_seconds + last_metrics = None + last_error = None + while time.monotonic() < deadline: + try: + last_metrics = read_storage_metrics(endpoint, expected_storage_units) + last_error = None + if last_metrics.active_keys == expected_active_keys: + return last_metrics + except (OSError, RuntimeError) as error: + last_error = str(error) + time.sleep(0.5) + raise TimeoutError( + f"Storage metrics did not reach active_keys={expected_active_keys}: " + f"last_metrics={last_metrics}, last_error={last_error!r}" + ) + + +def wait_for_stable_rss( + endpoint: str, + expected_storage_units: int, + expected_active_keys: int, + timeout_seconds: float, + sample_interval_seconds: float, + stable_samples: int, + tolerance_bytes: int, +) -> StorageMetrics: + """Return a median RSS after consecutive full metrics collection periods agree.""" + deadline = time.monotonic() + timeout_seconds + samples: list[StorageMetrics] = [] + while time.monotonic() < deadline: + remaining = max(0.5, deadline - time.monotonic()) + sample = wait_for_storage_state( + endpoint, + expected_storage_units, + expected_active_keys, + remaining, + ) + samples.append(sample) + window = samples[-stable_samples:] + if len(window) == stable_samples: + rss_values = [item.rss_bytes for item in window] + if max(rss_values) - min(rss_values) <= tolerance_bytes: + return StorageMetrics( + active_keys=expected_active_keys, + rss_bytes=int(statistics.median(rss_values)), + ssd_active_bytes=int(statistics.median(item.ssd_active_bytes for item in window)), + ) + time.sleep(sample_interval_seconds) + rss_values = [sample.rss_bytes for sample in samples[-stable_samples:]] + raise TimeoutError( + f"Storage RSS did not stabilize within {timeout_seconds}s: " + f"last_rss_values={rss_values}, tolerance_bytes={tolerance_bytes}" + ) + + +@ray.remote +class BenchmarkClient: + """Own one public TransferQueue client and, for the writer, the test payload.""" + + def __init__(self, config: dict[str, Any]): + self._config = config + self._data: TensorDict | None = None + self._keys: list[str] = [] + + def initialize(self) -> None: + tq.init(OmegaConf.create(self._config)) + + def create_data(self, workload: Workload) -> None: + self._data = create_mixed_tensors(workload) + self._keys = [f"mixed_{index}" for index in range(workload.batch_size)] + + def put(self, partition_id: str) -> None: + if self._data is None: + raise RuntimeError("Benchmark data has not been created") + tq.kv_batch_put(keys=self._keys, partition_id=partition_id, fields=self._data) + + def list_keys(self, partition_id: str) -> list[str]: + partitions = tq.kv_list(partition_id=partition_id) + return list(partitions.get(partition_id, {})) + + def get(self, partition_id: str, keys: list[str]) -> None: + tq.kv_batch_get(keys=keys, partition_id=partition_id) + + def clear(self, partition_id: str, keys: list[str]) -> None: + tq.kv_clear(keys=keys, partition_id=partition_id) + + def metrics_endpoint(self) -> str | None: + return tq.get_metrics_endpoint() + + def close(self) -> None: + tq.close() + + +def build_config( + base_config: DictConfig, + workload: Workload, + ssd_enabled: bool, + ssd_path: str, +) -> dict[str, Any]: + """Specialize the benchmark config for one SimpleStorage mode.""" + config = OmegaConf.create(OmegaConf.to_container(base_config, resolve=True)) + config.metrics.enabled = True + config.metrics.port = 0 + config.backend.storage_backend = "SimpleStorage" + config.backend.SimpleStorage.total_storage_size = workload.batch_size + config.backend.SimpleStorage.required_node_resource = None + config.backend.SimpleStorage.ssd_offload.enabled = ssd_enabled + config.backend.SimpleStorage.ssd_offload.path = ssd_path if ssd_enabled else None + return OmegaConf.to_container(config, resolve=True) + + +def validate_cluster(head_node_ip: str, worker_node_ip: str) -> None: + alive_addresses = {node["NodeManagerAddress"] for node in ray.nodes() if node["Alive"]} + expected = {head_node_ip, worker_node_ip} + if alive_addresses != expected: + raise RuntimeError(f"Expected exactly two Ray nodes {sorted(expected)}, got {sorted(alive_addresses)}") + + +class BenchmarkRun: + def __init__( + self, + mode: str, + base_config: DictConfig, + workload: Workload, + head_node_ip: str, + worker_node_ip: str, + ssd_path: str, + ) -> None: + self.mode = mode + self.workload = workload + + config = build_config(base_config, workload, mode == "ssd", ssd_path) + self.expected_storage_units = int(config["backend"]["SimpleStorage"]["num_data_storage_units"]) + self.writer = BenchmarkClient.options( + num_cpus=0.001, + resources={f"node:{head_node_ip}": 0.001}, + ).remote(config) + self.reader = BenchmarkClient.options( + num_cpus=0.001, + resources={f"node:{worker_node_ip}": 0.001}, + ).remote(config) + ray.get(self.writer.initialize.remote()) + ray.get(self.reader.initialize.remote()) + ray.get(self.writer.create_data.remote(workload)) + self.endpoint = ray.get(self.writer.metrics_endpoint.remote()) + if not self.endpoint: + raise RuntimeError("SimpleStorage metrics endpoint is unavailable") + + def stable_metrics(self, active_keys: int) -> StorageMetrics: + return wait_for_stable_rss( + self.endpoint, + self.expected_storage_units, + active_keys, + METRICS_TIMEOUT_SECONDS, + RSS_SAMPLE_INTERVAL_SECONDS, + RSS_STABLE_SAMPLES, + RSS_STABILITY_TOLERANCE_BYTES, + ) + + def run(self, iterations: int) -> list[dict[str, Any]]: + partition_id = "mixed" + baseline = self.stable_metrics(active_keys=0) + rows = [] + + for iteration in range(1, iterations + 1): + LOGGER.info("%s iteration %d/%d", self.mode, iteration, iterations) + put_start = time.perf_counter() + ray.get(self.writer.put.remote(partition_id)) + put_seconds = time.perf_counter() - put_start + after_put = self.stable_metrics(active_keys=self.workload.batch_size) + + keys = ray.get(self.reader.list_keys.remote(partition_id)) + if len(keys) != self.workload.batch_size: + raise AssertionError(f"Listed {len(keys)} keys, expected {self.workload.batch_size}") + get_start = time.perf_counter() + ray.get(self.reader.get.remote(partition_id, keys)) + get_seconds = time.perf_counter() - get_start + time.sleep(RSS_SAMPLE_INTERVAL_SECONDS) + after_get = self.stable_metrics(active_keys=self.workload.batch_size) + + ray.get(self.writer.clear.remote(partition_id, keys)) + after_clear = self.stable_metrics(active_keys=0) + + retained_bytes = after_clear.rss_bytes - baseline.rss_bytes + row = { + "mode": self.mode, + "iteration": iteration, + **asdict(self.workload), + "inline_payload_bytes": self.workload.inline_bytes, + "offload_candidate_bytes": self.workload.offloaded_bytes, + "total_payload_bytes": self.workload.total_bytes, + "put_seconds": put_seconds, + "get_seconds": get_seconds, + "put_gbit_per_second": self.workload.total_bytes * 8 / put_seconds / 1e9, + "get_gbit_per_second": self.workload.total_bytes * 8 / get_seconds / 1e9, + "storage_rss_before_put_bytes": baseline.rss_bytes, + "storage_rss_after_put_bytes": after_put.rss_bytes, + "storage_rss_after_get_bytes": after_get.rss_bytes, + "storage_rss_after_clear_bytes": after_clear.rss_bytes, + "storage_rss_retained_after_clear_bytes": retained_bytes, + "storage_ssd_active_bytes_after_put": after_put.ssd_active_bytes, + "storage_ssd_active_bytes_after_clear": after_clear.ssd_active_bytes, + } + rows.append(row) + LOGGER.info( + "%s RSS before/put/get/clear: %.3f/%.3f/%.3f/%.3f GiB; retained: %.1f MiB", + self.mode, + baseline.rss_bytes / 2**30, + after_put.rss_bytes / 2**30, + after_get.rss_bytes / 2**30, + after_clear.rss_bytes / 2**30, + retained_bytes / 2**20, + ) + baseline = after_clear + return rows + + def close(self) -> None: + try: + ray.get([self.writer.close.remote(), self.reader.close.remote()]) + finally: + ray.kill(self.writer, no_restart=True) + ray.kill(self.reader, no_restart=True) + + +def write_csv(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="") as output: + writer = csv.DictWriter(output, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + +def summarize( + ssd_rows: list[dict[str, Any]], + workload: Workload, + warmup_iterations: int, +) -> dict[str, Any]: + ssd_measured = ssd_rows[warmup_iterations:] + retained_values = [row["storage_rss_retained_after_clear_bytes"] for row in ssd_measured] + clear_rss_values = [row["storage_rss_after_clear_bytes"] for row in ssd_measured] + return { + "workload": asdict(workload), + "analyzed_iterations": len(ssd_measured), + "theoretical_offload_bytes": workload.offloaded_bytes, + "ssd_median_active_bytes_after_put": int( + statistics.median(row["storage_ssd_active_bytes_after_put"] for row in ssd_measured) + ), + "ssd_max_active_bytes_after_clear": max(row["storage_ssd_active_bytes_after_clear"] for row in ssd_measured), + "ssd_median_rss_retained_after_clear_bytes": int(statistics.median(retained_values)), + "ssd_max_rss_retained_after_clear_bytes": max(retained_values), + "ssd_clear_rss_growth_bytes": clear_rss_values[-1] - clear_rss_values[0], + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--backend_config", type=Path, required=True) + parser.add_argument("--head_node_ip", required=True) + parser.add_argument("--worker_node_ip", required=True) + parser.add_argument("--ssd_path", type=Path, required=True) + parser.add_argument("--output_dir", type=Path, default=Path(__file__).resolve().parent / "results") + parser.add_argument("--num_test_iterations", type=int, default=4) + parser.add_argument("--warmup_iterations", type=int, default=1) + parser.add_argument("--global_batch_size", type=int, default=512) + parser.add_argument("--small_fields", type=int, default=4) + parser.add_argument("--large_fields", type=int, default=4) + parser.add_argument("--small_sample_bytes", type=int, default=DEFAULT_SMALL_SAMPLE_BYTES) + parser.add_argument("--large_sample_bytes", type=int, default=DEFAULT_LARGE_SAMPLE_BYTES) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + base_config = OmegaConf.load(args.backend_config) + workload = Workload( + batch_size=args.global_batch_size, + small_fields=args.small_fields, + large_fields=args.large_fields, + small_sample_bytes=args.small_sample_bytes, + large_sample_bytes=args.large_sample_bytes, + offload_threshold_bytes=int(base_config.backend.SimpleStorage.ssd_offload.threshold_bytes), + ) + workload.validate() + if args.num_test_iterations <= args.warmup_iterations: + raise ValueError("num_test_iterations must be greater than warmup_iterations") + if not args.ssd_path.is_dir(): + raise ValueError(f"SSD path does not exist on the benchmark driver node: {args.ssd_path}") + + ray.init(address=RAY_ADDRESS) + validate_cluster(args.head_node_ip, args.worker_node_ip) + results: dict[str, list[dict[str, Any]]] = {} + try: + for mode in ("memory", "ssd"): + benchmark = BenchmarkRun( + mode=mode, + base_config=base_config, + workload=workload, + head_node_ip=args.head_node_ip, + worker_node_ip=args.worker_node_ip, + ssd_path=str(args.ssd_path), + ) + try: + results[mode] = benchmark.run(args.num_test_iterations) + finally: + benchmark.close() + write_csv(args.output_dir / f"mixed_{mode}.csv", results[mode]) + time.sleep(MODE_GAP_SECONDS) + + summary = summarize( + results["ssd"], + workload, + args.warmup_iterations, + ) + args.output_dir.mkdir(parents=True, exist_ok=True) + summary_path = args.output_dir / "summary.json" + summary_path.write_text(json.dumps(summary, indent=2) + "\n") + LOGGER.info("Mixed SSD offload summary:\n%s", json.dumps(summary, indent=2)) + finally: + ray.shutdown() + + +if __name__ == "__main__": + main() diff --git a/scripts/performance_test_ssd_offload/perftest_config.yaml b/scripts/performance_test_ssd_offload/perftest_config.yaml new file mode 100644 index 00000000..381c6b0d --- /dev/null +++ b/scripts/performance_test_ssd_offload/perftest_config.yaml @@ -0,0 +1,47 @@ +# This is the default configuration of TransferQueue. Users may modify the default value +# and use transfer_queue.init(conf) to overwrite the config entries. + +# Prometheus metrics exporter. +metrics: + enabled: true + # HTTP port for /metrics endpoint (0 = auto-assign free port) + port: 0 + +controller: + # User-defined sampler. User can pass sampler instance to overwrite this string config. + sampler: SequentialSampler + # Whether return an empty BatchMeta to prevent request blocking when no enough data is available + polling_mode: False + # ZMQ Server IP & Ports (automatically generated during init) + zmq_info: null + + +backend: + # Pluggable storage/transport backend of TransferQueue. Choose from: + # SimpleStorage, Yuanrong, MooncakeStore, ... + storage_backend: SimpleStorage + + # SimpleStorage: ZMQ-based in-memory storage for out-of-the-box usage + SimpleStorage: + # Maximum number of experience samples to hold across all storage units + total_storage_size: 100000 + # Number of distributed storage units. Units are round-robin scheduled across eligible + # alive Ray nodes, guaranteeing an even split of memory/bandwidth usage per node. + # Recommended: >= 2 x number of nodes so each node hosts multiple units. + num_data_storage_units: 16 + # Optional Ray custom resource required on storage nodes. null keeps all alive nodes eligible. + required_node_resource: null + # ZMQ Server IP & Ports (automatically generated during init) + zmq_info: null + + # SSD offload stores large field values on local SSD to reduce long-lived host-memory use. + # See docs/ssd_offload.md for configuration and cleanup details. + ssd_offload: + # Master switch. Set to true to enable SSD offload. + enabled: false + # SSD offload target directory; created automatically if it does not exist. + path: null + # Samples at or above this encoded-size threshold (in bytes) are stored on SSD. + threshold_bytes: 1048576 + # glibc-only allocator threshold for reclaiming large transient allocations. + glibc_mmap_threshold_bytes: 1048576 diff --git a/scripts/performance_test_ssd_offload/run_perf_test.sh b/scripts/performance_test_ssd_offload/run_perf_test.sh new file mode 100644 index 00000000..6b447149 --- /dev/null +++ b/scripts/performance_test_ssd_offload/run_perf_test.sh @@ -0,0 +1,24 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +RESULTS_DIR="${RESULTS_DIR:-${SCRIPT_DIR}/results}" +PERFTEST_PY="${SCRIPT_DIR}/perftest.py" +CONFIG_YAML="${SCRIPT_DIR}/perftest_config.yaml" + +: "${HEAD_NODE_IP:?Set HEAD_NODE_IP to the Ray head node IP}" +: "${WORKER_NODE_IP:?Set WORKER_NODE_IP to the Ray worker node IP}" +: "${SSD_OFFLOAD_PATH:?Set SSD_OFFLOAD_PATH to node-local SSD storage}" + +NUM_TEST_ITERATIONS="${NUM_TEST_ITERATIONS:-4}" + +mkdir -p "${RESULTS_DIR}" + +python "${PERFTEST_PY}" \ + --backend_config="${CONFIG_YAML}" \ + --head_node_ip="${HEAD_NODE_IP}" \ + --worker_node_ip="${WORKER_NODE_IP}" \ + --ssd_path="${SSD_OFFLOAD_PATH}" \ + --output_dir="${RESULTS_DIR}" \ + --num_test_iterations="${NUM_TEST_ITERATIONS}" diff --git a/scripts/put_benchmark.py b/scripts/put_benchmark.py index c67bb54c..822f8dd7 100644 --- a/scripts/put_benchmark.py +++ b/scripts/put_benchmark.py @@ -17,7 +17,6 @@ import asyncio import json import logging -import math import os import time @@ -267,6 +266,10 @@ def initialize_system(self, config_dict): ) total_storage_size = self.tq_config.global_batch_size * 2 + storage_config = { + "total_storage_size": total_storage_size, + "num_data_storage_units": self.num_storage_units, + } logger.info(f"Initializing Storage Units (Remote={self.remote_mode}, Target={self.target_ip})...") @@ -277,7 +280,7 @@ def initialize_system(self, config_dict): num_cpus=1, resources={f"node:{self.target_ip}": 0.001}, runtime_env={"env_vars": {"OMP_NUM_THREADS": "2"}}, - ).remote(storage_unit_size=math.ceil(total_storage_size / self.num_storage_units)) + ).remote(config=storage_config) else: # Local Mode: Use placement group self.storage_placement_group = get_placement_group(self.num_storage_units, num_cpus_per_actor=2) @@ -286,7 +289,7 @@ def initialize_system(self, config_dict): placement_group=self.storage_placement_group, placement_group_bundle_index=rank, runtime_env={"env_vars": {"OMP_NUM_THREADS": "2"}}, - ).remote(storage_unit_size=math.ceil(total_storage_size / self.num_storage_units)) + ).remote(config=storage_config) # Controller Init self.data_system_controller = TransferQueueController.remote() diff --git a/tests/e2e/test_ssd_offload_e2e.py b/tests/e2e/test_ssd_offload_e2e.py new file mode 100644 index 00000000..eb9060b6 --- /dev/null +++ b/tests/e2e/test_ssd_offload_e2e.py @@ -0,0 +1,143 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""End-to-end coverage for SimpleStorage local SSD offload.""" + +import pytest +import ray +import torch +from omegaconf import OmegaConf + +import transfer_queue as tq + + +@pytest.fixture(scope="module") +def ssd_root(tmp_path_factory): + """Return an isolated SSD root shared by this E2E module.""" + return tmp_path_factory.mktemp("tq-ssd-offload") + + +@pytest.fixture(scope="module") +def tq_system(ssd_root): + """Initialize a public TransferQueue system with SSD offload enabled.""" + if not ray.is_initialized(): + ray.init(ignore_reinit_error=True) + + config = OmegaConf.create( + { + "controller": {"polling_mode": True}, + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": { + "total_storage_size": 20, + "num_data_storage_units": 2, + "ssd_offload": { + "enabled": True, + "path": str(ssd_root), + }, + }, + }, + } + ) + tq.init(config) + yield + tq.close() + assert ssd_root.is_dir() + assert not list((ssd_root / "transfer_queue_ssd_offload").iterdir()) + if ray.is_initialized(): + ray.shutdown() + + +@pytest.fixture +def controller(tq_system): + """Return the controller used to clean test partitions.""" + return ray.get_actor("TransferQueueController", namespace="transfer_queue") + + +@pytest.fixture(autouse=True) +def cleanup_partitions(controller): + """Remove every partition created by an E2E test.""" + yield + for partition_id in ray.get(controller.list_partitions.remote()): + ray.get(controller.clear_partition.remote(partition_id)) + + +def test_public_api_routes_per_sample_and_migrates_on_overwrite(tq_system, ssd_root): + """Public KV operations preserve values while samples migrate between tiers.""" + partition_id = "ssd-routing" + small = torch.arange(16, dtype=torch.float32) + large = torch.arange(262144, dtype=torch.float32) + + tq.kv_put(key="small", partition_id=partition_id, fields={"value": small}) + tq.kv_put(key="large", partition_id=partition_id, fields={"value": large}) + + torch.testing.assert_close( + tq.kv_batch_get(keys=["small"], partition_id=partition_id)["value"][0], + small, + ) + torch.testing.assert_close( + tq.kv_batch_get(keys=["large"], partition_id=partition_id)["value"][0], + large, + ) + assert len(list(ssd_root.rglob("*.bin"))) == 1 + + tq.kv_put(key="small", partition_id=partition_id, fields={"value": large}) + assert len(list(ssd_root.rglob("*.bin"))) == 2 + + tq.kv_put(key="large", partition_id=partition_id, fields={"value": small}) + assert len(list(ssd_root.rglob("*.bin"))) == 1 + torch.testing.assert_close( + tq.kv_batch_get(keys=["large"], partition_id=partition_id)["value"][0], + small, + ) + + tq.kv_clear(keys=["small"], partition_id=partition_id) + assert not list(ssd_root.rglob("*.bin")) + + +def test_checkpoint_round_trip_recreates_ssd_data(tq_system, ssd_root, tmp_path): + """Checkpoint restore recreates logical data without depending on old SSD files.""" + partition_id = "ssd-checkpoint" + key = "large" + value = torch.arange(262144, dtype=torch.float32) + checkpoint_dir = tmp_path / "checkpoint" + + tq.kv_put(key=key, partition_id=partition_id, fields={"value": value}) + assert list(ssd_root.rglob("*.bin")) + tq.save_checkpoint(checkpoint_dir) + assert list((checkpoint_dir / "simple_storage").glob("*.pkl.blobs/*.bin")) + + tq.kv_clear(keys=[key], partition_id=partition_id) + assert not list(ssd_root.rglob("*.bin")) + tq.load_checkpoint(checkpoint_dir) + + restored = tq.kv_batch_get(keys=[key], partition_id=partition_id)["value"][0] + torch.testing.assert_close(restored, value) + assert list(ssd_root.rglob("*.bin")) diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 66e82a3f..f3e338b0 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -13,6 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Unit tests for the Prometheus metrics exporter (transfer_queue.metrics).""" import time @@ -98,6 +113,10 @@ def test_all_metrics_are_registered(self): "tq_storage_active_keys_total", "tq_storage_utilization_ratio", "tq_storage_memory_rss_bytes", + "tq_storage_ssd_offload_enabled", + "tq_storage_ssd_active_values", + "tq_storage_ssd_active_bytes", + "tq_storage_ssd_fallback_values_total", "tq_storage_requests_arrived", "tq_storage_arrivals_by_op", "tq_storage_accept_queue_backlog", @@ -313,6 +332,10 @@ def test_storage_metrics_populated_on_success(self): "capacity": 1000, "active_keys": 250, "process_rss_bytes": 512 * 1024 * 1024, + "ssd_offload_enabled": 1, + "ssd_active_values": 120, + "ssd_active_bytes": 4 * 1024 * 1024 * 1024, + "ssd_fallback_values_total": 7, } ) @@ -322,6 +345,14 @@ def test_storage_metrics_populated_on_success(self): assert exporter.storage_active_keys.labels(storage_unit_id="SU_001")._value.get() == 250 assert exporter.storage_utilization.labels(storage_unit_id="SU_001")._value.get() == 0.25 assert exporter.storage_memory_rss.labels(storage_unit_id="SU_001")._value.get() == 512 * 1024 * 1024 + assert exporter.storage_ssd_offload_enabled.labels(storage_unit_id="SU_001")._value.get() == 1 + assert exporter.storage_ssd_active_values.labels(storage_unit_id="SU_001")._value.get() == 120 + assert exporter.storage_ssd_active_bytes.labels(storage_unit_id="SU_001")._value.get() == 4 * 1024 * 1024 * 1024 + assert exporter.storage_ssd_fallback_values.labels(storage_unit_id="SU_001")._value.get() == 7 + + exporter._query_storage_unit.return_value["ssd_fallback_values_total"] = 9 + exporter.collect_storage_metrics() + assert exporter.storage_ssd_fallback_values.labels(storage_unit_id="SU_001")._value.get() == 9 def test_arrival_counters_are_exported(self): """Arrival counts reach Prometheus, so a dashboard can compare them with completions.""" @@ -456,6 +487,10 @@ def test_storage_metrics_skips_capacity_when_none(self): # active_keys and memory_rss are still reported assert exporter.storage_active_keys.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 42 assert exporter.storage_memory_rss.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 128 * 1024 * 1024 + assert exporter.storage_ssd_offload_enabled.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 0 + assert exporter.storage_ssd_active_values.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 0 + assert exporter.storage_ssd_active_bytes.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 0 + assert exporter.storage_ssd_fallback_values.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 0 def test_storage_metrics_prunes_stale_capacity_on_switch_to_unlimited(self): """If a storage unit transitions from a numeric capacity to unlimited, diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index 1145d963..99ceb4ec 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -13,16 +13,40 @@ # See the License for the specific language governing permissions and # limitations under the License. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pickle import time +from pathlib import Path from uuid import uuid4 +import numpy as np import pytest import ray import tensordict import torch import zmq -from transfer_queue.storage.simple_storage import SimpleStorageUnit +from transfer_queue.storage.simple_storage import ( + HybridStorageUnitData, + SimpleStorageUnit, + SSDFileStore, + StorageKeyNotFoundError, + _SSDValueRef, +) from transfer_queue.utils.zmq_utils import ( STORAGE_MANAGER_IDENTITY_PREFIX, ZMQMessage, @@ -112,7 +136,9 @@ def storage_setup(ray_setup): tensordict.set_list_to_stack(True).set() # Start Ray actor for SimpleStorageUnit - storage_actor = SimpleStorageUnit.options(max_concurrency=50, num_cpus=1).remote(storage_unit_size=storage_size) + storage_actor = SimpleStorageUnit.options(max_concurrency=50, num_cpus=1).remote( + config={"total_storage_size": storage_size, "num_data_storage_units": 1} + ) # Get ZMQ server info from storage unit zmq_info = ray.get(storage_actor.get_zmq_server_info.remote()) @@ -467,6 +493,307 @@ def test_storage_unit_data_capacity_uses_active_keys(): assert storage._active_keys == {0, 1, 3} +def test_hybrid_storage_routes_each_sample_by_payload_size_and_clears(tmp_path): + """Large samples use SSD while small samples preserve the in-memory contract.""" + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + large_sparse = torch.sparse_coo_tensor( + torch.stack((torch.arange(16), torch.arange(16))), + torch.ones(16), + size=(16, 16), + ) + try: + storage.put_data( + { + "mixed": [b"x" * 100, b"y" * 10], + "small": [1, 2], + "unsupported": [large_sparse, b"small"], + }, + [1, 2], + ) + + assert storage.get_data(["mixed", "small"], [1, 2]) == { + "small": [1, 2], + "mixed": [b"x" * 100, b"y" * 10], + } + assert len(list(tmp_path.rglob("*.bin"))) == 1 + assert storage.ssd_active_values == 1 + assert storage.ssd_active_bytes == 100 + assert storage.ssd_fallback_values_total == 1 + + storage.clear([1]) + assert storage.active_key_count == 1 + assert storage.ssd_active_values == 0 + assert storage.ssd_active_bytes == 0 + assert storage.ssd_fallback_values_total == 1 + assert not list(tmp_path.rglob("*.bin")) + with pytest.raises(StorageKeyNotFoundError): + storage.get_data(["mixed"], [1]) + + storage.clear([2]) + with pytest.raises(StorageKeyNotFoundError): + storage.get_data(["mixed"], [2]) + finally: + storage.close() + + assert not (tmp_path / "transfer_queue_ssd_offload" / "test-run").exists() + + +def test_hybrid_storage_cleans_completed_ssd_writes_after_batch_failure(tmp_path, monkeypatch): + """A failed SSD batch removes files written by its successful tasks.""" + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + + write_value_file = storage._ssd_store._write_value_file + + def fail_one_write(encoded_value): + if encoded_value.payload[0] == ord("y"): + raise OSError("injected SSD failure") + return write_value_file(encoded_value) + + monkeypatch.setattr(storage._ssd_store, "_write_value_file", fail_one_write) + try: + with pytest.raises(OSError, match="injected SSD failure"): + storage.put_data( + {"small": [1, 2], "large": [b"x" * 100, b"y" * 100]}, + [1, 2], + ) + + assert not list(tmp_path.rglob("*.bin")) + finally: + storage.close() + + +def test_failed_ssd_overwrite_preserves_old_value(tmp_path, monkeypatch): + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + old_value = b"o" * 100 + try: + storage.put_data({"value": [old_value]}, [1]) + old_path = next(tmp_path.rglob("*.bin")) + + def fail_write(_fd, _data): + raise OSError("injected overwrite failure") + + monkeypatch.setattr(storage._ssd_store, "_write_all_bytes", fail_write) + with pytest.raises(OSError, match="injected overwrite failure"): + storage.put_data({"value": [b"n" * 100]}, [1]) + + assert storage.get_data(["value"], [1])["value"] == [old_value] + assert storage._ssd_active_values == 1 + assert storage._ssd_active_bytes == len(old_value) + assert old_path.exists() + assert list(tmp_path.rglob("*.bin")) == [old_path] + finally: + storage.close() + + +@pytest.mark.parametrize("threshold_bytes", [0, -1]) +def test_hybrid_storage_requires_positive_threshold(tmp_path, threshold_bytes): + with pytest.raises(ValueError, match="threshold must be greater than zero"): + HybridStorageUnitData( + storage_size=10, + threshold_bytes=threshold_bytes, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + + +def test_hybrid_checkpoint_copies_ssd_values_without_materializing(tmp_path, monkeypatch): + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path / "ssd"), + run_id="test-run", + unit_id="test-unit", + ) + checkpoint_path = tmp_path / "storage_unit.pkl" + try: + storage.put_data({"value": [b"small", b"x" * 100]}, [1, 2]) + + def fail_materialization(*_args): + raise AssertionError("checkpoint must not materialize or re-encode SSD values") + + with monkeypatch.context() as patch: + patch.setattr(storage._ssd_store, "read_values", fail_materialization) + storage.save_checkpoint(checkpoint_path, "test-unit") + + with open(checkpoint_path, "rb") as f: + manifest = pickle.load(f) + assert manifest["field_data"] == {"value": {1: b"small"}} + checkpoint_ref = manifest["ssd_index"]["value"][2] + blob_path = Path(f"{checkpoint_path}.blobs") / checkpoint_ref["filename"] + assert blob_path.read_bytes() == b"x" * 100 + + storage.clear([1, 2]) + with monkeypatch.context() as patch: + patch.setattr( + HybridStorageUnitData, + "_sample_from_value", + staticmethod(fail_materialization), + ) + storage.load_checkpoint(checkpoint_path) + + assert storage.get_data(["value"], [1, 2]) == {"value": [b"small", b"x" * 100]} + assert storage._ssd_active_values == 1 + assert storage._ssd_active_bytes == 100 + finally: + storage.close() + + +def test_hybrid_checkpoint_load_failure_preserves_existing_state(tmp_path, monkeypatch): + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path / "target-ssd"), + run_id="target-run", + unit_id="target-unit", + ) + checkpoint_storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path / "checkpoint-ssd"), + run_id="checkpoint-run", + unit_id="checkpoint-unit", + ) + checkpoint_path = tmp_path / "storage_unit.pkl" + try: + storage.put_data({"value": [b"original" * 16]}, [1]) + original_files = set((tmp_path / "target-ssd").rglob("*.bin")) + + checkpoint_storage.put_data({"value": [b"x" * 100, b"y" * 100]}, [2, 3]) + checkpoint_storage.save_checkpoint(checkpoint_path, "checkpoint-unit") + + import_file = storage._ssd_store.import_file + import_count = 0 + + def fail_on_second_import(source, metadata): + nonlocal import_count + import_count += 1 + if import_count == 2: + raise OSError("injected checkpoint import failure") + return import_file(source, metadata) + + monkeypatch.setattr(storage._ssd_store, "import_file", fail_on_second_import) + + with pytest.raises(OSError, match="injected checkpoint import failure"): + storage.load_checkpoint(checkpoint_path) + + assert storage.get_data(["value"], [1]) == {"value": [b"original" * 16]} + assert storage.active_key_count == 1 + assert set((tmp_path / "target-ssd").rglob("*.bin")) == original_files + finally: + storage.close() + checkpoint_storage.close() + + +def test_hybrid_storage_loads_legacy_logical_checkpoint(tmp_path): + storage = HybridStorageUnitData( + storage_size=1, + threshold_bytes=64, + ssd_path=str(tmp_path / "ssd"), + run_id="test-run", + unit_id="test-unit", + ) + checkpoint_path = tmp_path / "legacy.pkl" + with open(checkpoint_path, "wb") as f: + pickle.dump( + { + "storage_unit_id": "old-unit", + "storage_unit_size": 10, + "field_data": {"value": {1: b"x" * 100, 2: b"y" * 100}}, + "active_keys": {1, 2}, + }, + f, + ) + + try: + storage.load_checkpoint(checkpoint_path) + assert all(isinstance(value, _SSDValueRef) for value in storage.field_data["value"].values()) + assert storage._ssd_active_values == 2 + assert storage._ssd_active_bytes == 200 + assert storage.get_data(["value"], [1, 2]) == {"value": [b"x" * 100, b"y" * 100]} + finally: + storage.close() + + +def test_hybrid_storage_round_trips_supported_codecs(tmp_path): + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + values = torch.arange(64, dtype=torch.float32).view(2, 32) + variable = [ + torch.arange(20, dtype=torch.float32), + torch.arange(40, dtype=torch.float32), + ] + arrays = np.arange(64, dtype=np.float32).reshape(2, 32) + structured = np.zeros((2, 8), dtype=[("index", " TransferQueueClient: global _TQ_CLIENT @@ -238,9 +257,18 @@ def close(): if _TQ_STORAGE: for key, value in _TQ_STORAGE.items(): if key == "SimpleStorage": - # only the process that do first-time init can clean the distributed storage - for storage in value.values(): - ray.kill(storage) + # Only the initializing process owns and shuts down these storage actors. + storage_handles = list(value.values()) + try: + ray.get( + [storage.shutdown.remote() for storage in storage_handles], + timeout=_SIMPLE_STORAGE_SHUTDOWN_TIMEOUT_S, + ) + except Exception as e: + logger.warning(f"Failed to gracefully shut down SimpleStorage units: {e}") + finally: + for storage in storage_handles: + ray.kill(storage) elif key == "MooncakeStore": check = subprocess.run(["pgrep", "-f", "mooncake_master"], stdout=subprocess.PIPE, text=True) if check.returncode == 0: diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index e75792b5..15476d3a 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -13,6 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import time from contextlib import contextmanager @@ -184,6 +199,30 @@ def _define_controller_metrics(self) -> None: self.storage_memory_rss = Gauge( "tq_storage_memory_rss_bytes", "Storage unit process RSS memory", ["storage_unit_id"], registry=r ) + self.storage_ssd_offload_enabled = Gauge( + "tq_storage_ssd_offload_enabled", + "Whether SSD offload is enabled for the storage unit", + ["storage_unit_id"], + registry=r, + ) + self.storage_ssd_active_values = Gauge( + "tq_storage_ssd_active_values", + "Active field values stored on SSD", + ["storage_unit_id"], + registry=r, + ) + self.storage_ssd_active_bytes = Gauge( + "tq_storage_ssd_active_bytes", + "Logical bytes held by active SSD-backed values", + ["storage_unit_id"], + registry=r, + ) + self.storage_ssd_fallback_values = Gauge( + "tq_storage_ssd_fallback_values_total", + "Values retained in memory because SSD encoding was unavailable", + ["storage_unit_id"], + registry=r, + ) # ---- Storage-unit request-loss diagnostics ---- # Read against tq_storage_request_ops, which advances only on completion: a gap @@ -417,6 +456,14 @@ def collect_storage_metrics(self) -> None: pass self.storage_active_keys.labels(storage_unit_id=label).set(active) self.storage_memory_rss.labels(storage_unit_id=label).set(metrics.get("process_rss_bytes", 0)) + self.storage_ssd_offload_enabled.labels(storage_unit_id=label).set( + metrics.get("ssd_offload_enabled", 0) + ) + self.storage_ssd_active_values.labels(storage_unit_id=label).set(metrics.get("ssd_active_values", 0)) + self.storage_ssd_active_bytes.labels(storage_unit_id=label).set(metrics.get("ssd_active_bytes", 0)) + self.storage_ssd_fallback_values.labels(storage_unit_id=label).set( + metrics.get("ssd_fallback_values_total", 0) + ) self.storage_requests_arrived.labels(storage_unit_id=label).set(metrics.get("requests_arrived", 0)) for op_type, arrived in (metrics.get("arrivals_by_op") or {}).items(): diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py index 381341ec..f0de7a0b 100644 --- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py +++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py @@ -13,8 +13,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -import math +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from typing import Any +from uuid import uuid4 from omegaconf import DictConfig @@ -26,30 +41,40 @@ logger = get_logger(__name__) +DEFAULT_GLIBC_MMAP_THRESHOLD_BYTES = 1024 * 1024 + @StorageBootstrapProvider.register_provider("SimpleStorage") def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]: """Initialize Simple storage with metastore mode.""" simple_storage_handles = {} - num_data_storage_units = conf.backend.SimpleStorage.num_data_storage_units - total_storage_size = conf.backend.SimpleStorage.get("total_storage_size", None) - required_node_resource = conf.backend.SimpleStorage.get("required_node_resource", None) + simple_storage_config = dict(conf.backend.SimpleStorage) + simple_storage_config["_run_id"] = uuid4().hex + num_data_storage_units = simple_storage_config["num_data_storage_units"] + required_node_resource = simple_storage_config.get("required_node_resource") + ssd_config = simple_storage_config.get("ssd_offload") or {} + storage_actor_runtime_env = None + if ssd_config.get("enabled", False): + mmap_threshold = ssd_config.get("glibc_mmap_threshold_bytes", DEFAULT_GLIBC_MMAP_THRESHOLD_BYTES) + if mmap_threshold is not None: + mmap_threshold = int(mmap_threshold) + if mmap_threshold <= 0: + raise ValueError("glibc_mmap_threshold_bytes must be greater than zero or null") + storage_actor_runtime_env = {"env_vars": {"MALLOC_MMAP_THRESHOLD_": str(mmap_threshold)}} scheduling_strategies = get_node_round_robin_scheduling_strategies( num_data_storage_units, required_node_resource=required_node_resource ) - # Compute per-unit capacity: None means unlimited - storage_unit_size = ( - math.ceil(total_storage_size / num_data_storage_units) if total_storage_size is not None else None - ) - for storage_unit_rank in range(num_data_storage_units): - storage_node = SimpleStorageUnit.options( # type: ignore[attr-defined] - scheduling_strategy=scheduling_strategies[storage_unit_rank], - name=f"TransferQueueStorageUnit#{storage_unit_rank}", - ).remote( - storage_unit_size=storage_unit_size, + actor_options = { + "scheduling_strategy": scheduling_strategies[storage_unit_rank], + "name": f"TransferQueueStorageUnit#{storage_unit_rank}", + } + if storage_actor_runtime_env is not None: + actor_options["runtime_env"] = storage_actor_runtime_env + storage_node = SimpleStorageUnit.options(**actor_options).remote( # type: ignore[attr-defined] + config=simple_storage_config ) simple_storage_handles[f"TransferQueueStorageUnit#{storage_unit_rank}"] = storage_node logger.info( diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 61c5e9a1..df28521e 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -13,16 +13,38 @@ # See the License for the specific language governing permissions and # limitations under the License. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math import os import pickle +import shutil import time import weakref +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path from threading import Event, Thread from typing import TYPE_CHECKING, Any from uuid import uuid4 +import numpy as np import psutil import ray +import torch import zmq from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads, log_heavy_operation @@ -49,6 +71,34 @@ TQ_STORAGE_POLLER_TIMEOUT = int(os.environ.get("TQ_STORAGE_POLLER_TIMEOUT", 5)) # in seconds TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8)) +DEFAULT_SSD_OFFLOAD_THRESHOLD_BYTES = 1024 * 1024 +DEFAULT_SSD_READ_THREADS = 32 +DEFAULT_SSD_WRITE_THREADS = 8 +SSD_OFFLOAD_DIRECTORY_NAME = "transfer_queue_ssd_offload" + +_HYBRID_CHECKPOINT_FORMAT = "transfer_queue_hybrid_storage_v1" + + +@dataclass(frozen=True) +class SSDEncodedSample: + """One sample represented in a form that can be written directly to SSD.""" + + payload: memoryview + codec: str + dtype: str | np.dtype[Any] | None = None + shape: tuple[int, ...] | None = None + + +@dataclass(frozen=True) +class _SSDValueRef: + """Internal reference to one SSD-backed value.""" + + path: Path + size_bytes: int + codec: str + dtype: str | np.dtype[Any] | None = None + shape: tuple[int, ...] | None = None + # Marks a GET_ERROR reply as "the key is gone" so the caller can tell it apart from a real fault. KEY_NOT_FOUND_MARKER = "TQKeyNotFound" @@ -95,6 +145,26 @@ def active_key_count(self) -> int: """Number of active keys currently stored.""" return len(self._active_keys) + @property + def ssd_offload_enabled(self) -> bool: + """Whether this storage data offloads values to SSD.""" + return False + + @property + def ssd_active_values(self) -> int: + """Number of active field values stored on SSD.""" + return 0 + + @property + def ssd_active_bytes(self) -> int: + """Logical bytes held by active SSD-backed values.""" + return 0 + + @property + def ssd_fallback_values_total(self) -> int: + """Values retained in memory because SSD encoding was unavailable.""" + return 0 + def get_data(self, fields: list[str], global_indexes: list) -> dict[str, list]: """Get data by global index keys. @@ -156,6 +226,498 @@ def clear(self, keys: list[int]) -> None: self.field_data[f].pop(key, None) self._active_keys -= set(keys) + def save_checkpoint(self, path: str | Path, storage_unit_id: str) -> None: + """Write in-memory storage state to a checkpoint.""" + state = { + "storage_unit_id": storage_unit_id, + "storage_unit_size": self.storage_size, + "field_data": self.field_data, + "active_keys": self._active_keys, + } + with open(path, "wb") as f: + pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL) + + def load_checkpoint(self, path: str | Path) -> int | None: + """Replace in-memory storage state from a checkpoint.""" + with open(path, "rb") as f: + state = pickle.load(f) + if state.get("format") == _HYBRID_CHECKPOINT_FORMAT: + raise ValueError("An SSD checkpoint requires SSD offload to be enabled") + + checkpoint_size = state["storage_unit_size"] + restored_field_data = state["field_data"] + restored_active_keys = set(state["active_keys"]) + self.field_data = restored_field_data + self._active_keys = restored_active_keys + return checkpoint_size + + def close(self) -> None: + """Release resources owned by this storage data.""" + return None + + +class SSDFileStore: + """Own and read/write one SSD file per offloaded value.""" + + def __init__( + self, + ssd_path: str, + run_id: str, + unit_id: str, + ) -> None: + self._closed = False + configured_path = Path(ssd_path).resolve() + if configured_path.exists() and not configured_path.is_dir(): + raise ValueError(f"SSD offload path is not a directory: {configured_path}") + configured_path.mkdir(parents=True, exist_ok=True) + self._ssd_root = configured_path / SSD_OFFLOAD_DIRECTORY_NAME + self._ssd_root.mkdir(exist_ok=True) + if self._ssd_root.is_symlink() or not self._ssd_root.is_dir(): + raise ValueError(f"SSD offload working path must be a directory, not a symlink: {self._ssd_root}") + self._base_path = self._ssd_root / run_id / unit_id + self._base_path.mkdir(parents=True) + for prefix in range(256): + (self._base_path / f"{prefix:02x}").mkdir() + self._read_pool = ThreadPoolExecutor( + max_workers=DEFAULT_SSD_READ_THREADS, + thread_name_prefix="tq-ssd-read", + ) + self._write_pool = ThreadPoolExecutor( + max_workers=DEFAULT_SSD_WRITE_THREADS, + thread_name_prefix="tq-ssd-write", + ) + + @staticmethod + def _write_all_bytes(fd: int, payload: memoryview) -> None: + """Write the entire payload to an open file descriptor, handling partial writes.""" + view = payload.cast("B") + while view: + written = os.write(fd, view) + if written <= 0: + raise OSError("SSDFileStore write returned no progress") + view = view[written:] + + def _write_value_file( + self, + encoded_value: SSDEncodedSample, + ) -> _SSDValueRef: + """Atomically write one encoded value to a new SSD file and return its reference.""" + token = uuid4().hex + directory = self._base_path / token[:2] + temp_path = directory / f".tmp-{token}" + final_path = directory / f"{token}.bin" + try: + fd = os.open(temp_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + try: + self._write_all_bytes(fd, encoded_value.payload) + finally: + os.close(fd) + temp_path.rename(final_path) + except Exception: + for path in (temp_path, final_path): + try: + path.unlink(missing_ok=True) + except OSError: + pass + raise + return _SSDValueRef( + path=final_path, + size_bytes=encoded_value.payload.nbytes, + codec=encoded_value.codec, + dtype=encoded_value.dtype, + shape=encoded_value.shape, + ) + + def write_values( + self, + encoded_values: list[SSDEncodedSample], + ) -> list[_SSDValueRef]: + """Write encoded values concurrently in input order and remove completed files if the batch fails.""" + refs: list[_SSDValueRef | None] = [None] * len(encoded_values) + futures = { + self._write_pool.submit(self._write_value_file, encoded_value): position + for position, encoded_value in enumerate(encoded_values) + } + first_error: Exception | None = None + for future, position in futures.items(): + try: + refs[position] = future.result() + except Exception as e: + if first_error is None: + first_error = e + if first_error is not None: + for ref in refs: + if ref is not None: + self.unlink(ref) + raise first_error + return [ref for ref in refs if ref is not None] + + def import_file(self, source: Path, metadata: dict[str, Any]) -> _SSDValueRef: + """Copy one checkpoint blob into this store without materializing it.""" + token = uuid4().hex + directory = self._base_path / token[:2] + temp_path = directory / f".tmp-{token}" + final_path = directory / f"{token}.bin" + try: + shutil.copyfile(source, temp_path) + actual_size = temp_path.stat().st_size + if actual_size != metadata["size_bytes"]: + raise OSError( + f"SSD checkpoint blob has {actual_size} bytes, expected {metadata['size_bytes']}: {source}" + ) + temp_path.rename(final_path) + except Exception: + temp_path.unlink(missing_ok=True) + raise + return _SSDValueRef( + path=final_path, + size_bytes=metadata["size_bytes"], + codec=metadata["codec"], + dtype=metadata["dtype"], + shape=metadata["shape"], + ) + + @staticmethod + def unlink(entry: _SSDValueRef) -> None: + """Delete one offloaded value, tolerating cleanup failures.""" + try: + entry.path.unlink(missing_ok=True) + except OSError as e: + logger.warning(f"Failed to delete superseded SSD sample {entry.path}: {e}") + + def read_values(self, refs: list[_SSDValueRef]) -> list[bytes]: + """Read SSD-backed values concurrently and return their payloads in input order.""" + return list(self._read_pool.map(self._read_value_file, refs)) + + @staticmethod + def _read_value_file(ref: _SSDValueRef) -> bytes: + """Read one SSD value file and verify its size against the stored reference.""" + payload = ref.path.read_bytes() + if len(payload) != ref.size_bytes: + raise OSError( + f"SSDFileStore short read from {ref.path}: expected {ref.size_bytes} bytes, got {len(payload)}" + ) + return payload + + def close(self) -> None: + """Stop I/O workers and delete the storage directory.""" + if self._closed: + return + self._closed = True + self._write_pool.shutdown(wait=True) + self._read_pool.shutdown(wait=True) + shutil.rmtree(self._base_path, ignore_errors=True) + try: + self._base_path.parent.rmdir() + except OSError: + pass + + +class HybridStorageUnitData(StorageUnitData): + """Store each value inline or as an SSD file reference in one field map.""" + + def __init__( + self, + storage_size: int | None, + ssd_path: str, + run_id: str, + unit_id: str, + threshold_bytes: int = DEFAULT_SSD_OFFLOAD_THRESHOLD_BYTES, + ) -> None: + if threshold_bytes <= 0: + raise ValueError("SSD offload threshold must be greater than zero") + super().__init__(storage_size) + self._ssd_store = SSDFileStore(ssd_path, run_id, unit_id) + self._threshold = threshold_bytes + self._ssd_active_values = 0 + self._ssd_active_bytes = 0 + self._ssd_fallback_values_total = 0 + + @property + def ssd_offload_enabled(self) -> bool: + """Whether this storage data offloads values to SSD.""" + return True + + @property + def ssd_active_values(self) -> int: + """Number of active field values stored on SSD.""" + return self._ssd_active_values + + @property + def ssd_active_bytes(self) -> int: + """Logical bytes held by active SSD-backed values.""" + return self._ssd_active_bytes + + @property + def ssd_fallback_values_total(self) -> int: + """Values retained in memory because SSD encoding was unavailable.""" + return self._ssd_fallback_values_total + + @staticmethod + def _sample_from_value(value: Any) -> SSDEncodedSample | None: + if isinstance(value, torch.Tensor): + if value.is_nested or value.is_sparse: + return None + try: + tensor = value.detach() + if tensor.device.type != "cpu": + tensor = tensor.cpu() + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + payload = memoryview(tensor.flatten().view(torch.uint8).numpy()).cast("B") + except (RuntimeError, TypeError, ValueError): + return None + return SSDEncodedSample( + payload=payload, + codec="tensor", + dtype=str(tensor.dtype).removeprefix("torch."), + shape=tuple(tensor.shape), + ) + if isinstance(value, np.ndarray) and not value.dtype.hasobject: + try: + array = value if value.flags["C_CONTIGUOUS"] else np.ascontiguousarray(value) + payload = memoryview(array.view(np.uint8).ravel()).cast("B") + except (TypeError, ValueError): + return None + return SSDEncodedSample( + payload=payload, + codec="numpy", + dtype=array.dtype, + shape=tuple(array.shape), + ) + if isinstance(value, bytes): + return SSDEncodedSample(payload=memoryview(value), codec="bytes") + try: + pickled_payload = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL) + except Exception: + return None + return SSDEncodedSample(payload=memoryview(pickled_payload), codec="pickle") + + @staticmethod + def _decode_sample(raw: bytes, entry: _SSDValueRef) -> Any: + if entry.codec == "tensor": + if entry.dtype is None or entry.shape is None: + raise ValueError("Tensor SSD entry is missing dtype or shape") + dtype = getattr(torch, entry.dtype) + return torch.frombuffer(raw, dtype=dtype).view(entry.shape) + if entry.codec == "numpy": + if entry.dtype is None or entry.shape is None: + raise ValueError("NumPy SSD entry is missing dtype or shape") + return np.frombuffer(raw, dtype=np.dtype(entry.dtype)).reshape(entry.shape) + if entry.codec == "bytes": + return raw + if entry.codec == "pickle": + return pickle.loads(raw) + raise ValueError(f"Unsupported SSD codec: {entry.codec}") + + def _prepare_field_values( + self, + values: Any, + ) -> tuple[list[Any], list[_SSDValueRef], int]: + """Encode one field and write its offloaded values to the active SSD directory.""" + logical_samples = list(values.unbind()) if isinstance(values, torch.Tensor) else list(values) + encoded_samples = [self._sample_from_value(sample) for sample in logical_samples] + prepared_values: list[Any] = [] + ssd_positions: list[int] = [] + ssd_samples: list[SSDEncodedSample] = [] + fallback_values = 0 + for position, (value, encoded) in enumerate(zip(logical_samples, encoded_samples, strict=True)): + if encoded is not None and encoded.payload.nbytes >= self._threshold: + ssd_positions.append(position) + ssd_samples.append(encoded) + prepared_values.append(None) + else: + prepared_values.append(value) + if encoded is None: + fallback_values += 1 + + entries = self._ssd_store.write_values(ssd_samples) + for position, entry in zip(ssd_positions, entries, strict=True): + prepared_values[position] = entry + return prepared_values, entries, fallback_values + + def put_data( + self, + field_data: dict[str, Any], + global_indexes: list, + ) -> None: + """Store each sample in memory or SSD according to its encoded size.""" + if not global_indexes or not field_data: + super().put_data(field_data, global_indexes) + return + + unique_global_indexes = set(global_indexes) + has_duplicate_indexes = len(unique_global_indexes) != len(global_indexes) + + for field, values in field_data.items(): + prepared_values, entries, fallback_values = self._prepare_field_values(values) + + stored_field = self.field_data.get(field, {}) + old_ssd_values = [] + for global_index in unique_global_indexes: + old_value = stored_field.get(global_index) + if isinstance(old_value, _SSDValueRef): + old_ssd_values.append(old_value) + try: + super().put_data({field: prepared_values}, global_indexes) + except Exception: + for entry in entries: + self._ssd_store.unlink(entry) + raise + + obsolete_ssd_values = old_ssd_values + if has_duplicate_indexes: + retained_ssd_paths = set() + for global_index in unique_global_indexes: + retained_value = self.field_data[field][global_index] + if isinstance(retained_value, _SSDValueRef): + retained_ssd_paths.add(retained_value.path) + obsolete_ssd_values.extend(entry for entry in entries if entry.path not in retained_ssd_paths) + + self._ssd_active_values += len(entries) - len(obsolete_ssd_values) + self._ssd_active_bytes += sum(entry.size_bytes for entry in entries) - sum( + value.size_bytes for value in obsolete_ssd_values + ) + self._ssd_fallback_values_total += fallback_values + for obsolete_value in obsolete_ssd_values: + self._ssd_store.unlink(obsolete_value) + + def get_data(self, fields: list[str], global_indexes: list) -> dict[str, list]: + """Read mixed memory- and SSD-backed samples in request order.""" + result = super().get_data(fields, global_indexes) + for values in result.values(): + ssd_values = [(position, value) for position, value in enumerate(values) if isinstance(value, _SSDValueRef)] + raw_values = self._ssd_store.read_values([value for _, value in ssd_values]) + for (position, entry), raw in zip(ssd_values, raw_values, strict=True): + values[position] = self._decode_sample(raw, entry) + return result + + def clear(self, keys: list) -> None: + """Remove values and unlink any files they reference.""" + ssd_values: list[_SSDValueRef] = [] + for values in self.field_data.values(): + for key in set(keys): + value = values.get(key) + if isinstance(value, _SSDValueRef): + ssd_values.append(value) + super().clear(keys) + self._ssd_active_values -= len(ssd_values) + self._ssd_active_bytes -= sum(value.size_bytes for value in ssd_values) + for value in ssd_values: + self._ssd_store.unlink(value) + + def save_checkpoint(self, path: str | Path, storage_unit_id: str) -> None: + """Write memory values to a manifest and copy SSD values beside it.""" + manifest_path = Path(path) + blob_dir = Path(f"{manifest_path}.blobs") + shutil.rmtree(blob_dir, ignore_errors=True) + checkpoint_fields: dict[str, dict[int, Any]] = {} + ssd_index: dict[str, dict[int, dict[str, Any]]] = {} + try: + for field, values in self.field_data.items(): + checkpoint_values = {} + field_ssd_index = {} + for global_index, value in values.items(): + if not isinstance(value, _SSDValueRef): + checkpoint_values[global_index] = value + continue + + blob_dir.mkdir(parents=True, exist_ok=True) + filename = value.path.name + destination = blob_dir / filename + shutil.copyfile(value.path, destination) + actual_size = destination.stat().st_size + if actual_size != value.size_bytes: + raise OSError(f"SSD value has {actual_size} bytes, expected {value.size_bytes}: {value.path}") + field_ssd_index[global_index] = { + "filename": filename, + "size_bytes": value.size_bytes, + "codec": value.codec, + "dtype": value.dtype, + "shape": value.shape, + } + checkpoint_fields[field] = checkpoint_values + if field_ssd_index: + ssd_index[field] = field_ssd_index + + state = { + "format": _HYBRID_CHECKPOINT_FORMAT, + "storage_unit_id": storage_unit_id, + "storage_unit_size": self.storage_size, + "field_data": checkpoint_fields, + "ssd_index": ssd_index, + "active_keys": set(self._active_keys), + } + with open(manifest_path, "wb") as f: + pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL) + except Exception: + manifest_path.unlink(missing_ok=True) + shutil.rmtree(blob_dir, ignore_errors=True) + raise + + def load_checkpoint(self, path: str | Path) -> int | None: + """Replace memory and SSD state after all checkpoint data is ready.""" + manifest_path = Path(path) + with open(manifest_path, "rb") as f: + state = pickle.load(f) + checkpoint_format = state.get("format") + if checkpoint_format not in (None, _HYBRID_CHECKPOINT_FORMAT): + raise ValueError(f"Unsupported HybridStorageUnitData checkpoint format: {checkpoint_format}") + checkpoint_size = state["storage_unit_size"] + field_data = state["field_data"] + restored_active_keys = set(state["active_keys"]) + old_ssd_values = [ + value for values in self.field_data.values() for value in values.values() if isinstance(value, _SSDValueRef) + ] + restored_ssd_values: list[_SSDValueRef] = [] + restored_fallback_values = 0 + + try: + if checkpoint_format == _HYBRID_CHECKPOINT_FORMAT: + blob_dir = Path(f"{manifest_path}.blobs") + restored_field_data = {field: dict(values) for field, values in field_data.items()} + for field, entries in state["ssd_index"].items(): + restored_values = restored_field_data.setdefault(field, {}) + for global_index, metadata in entries.items(): + filename = metadata["filename"] + if Path(filename).name != filename: + raise ValueError(f"Invalid SSD checkpoint blob name: {filename}") + restored_value = self._ssd_store.import_file(blob_dir / filename, metadata) + restored_ssd_values.append(restored_value) + restored_values[global_index] = restored_value + else: + restored_field_data = {} + for field, values in field_data.items(): + if not values: + restored_field_data[field] = {} + continue + indexes = list(values) + prepared_values, entries, fallback_values = self._prepare_field_values( + [values[key] for key in indexes] + ) + restored_ssd_values.extend(entries) + restored_fallback_values += fallback_values + restored_field_data[field] = dict(zip(indexes, prepared_values, strict=True)) + restored_ssd_active_values = len(restored_ssd_values) + restored_ssd_active_bytes = sum(value.size_bytes for value in restored_ssd_values) + except Exception: + for value in restored_ssd_values: + self._ssd_store.unlink(value) + raise + + self.field_data = restored_field_data + self._active_keys = restored_active_keys + self._ssd_active_values = restored_ssd_active_values + self._ssd_active_bytes = restored_ssd_active_bytes + self._ssd_fallback_values_total += restored_fallback_values + for value in old_ssd_values: + self._ssd_store.unlink(value) + return checkpoint_size + + def close(self) -> None: + """Release SSD resources owned by this hybrid store.""" + self._ssd_store.close() + @ray.remote(num_cpus=1) class SimpleStorageUnit: @@ -179,17 +741,41 @@ class SimpleStorageUnit: _arrivals_by_op: dict[str, int] = {} _accept_probe = None - def __init__(self, storage_unit_size: int | None = None): - """Initialize a SimpleStorageUnit with the specified size. + def __init__(self, config: dict[str, Any]): + """Initialize a SimpleStorageUnit from the SimpleStorage config. Args: - storage_unit_size: Maximum number of elements that can be stored in this storage unit. - If None, the storage unit has unlimited capacity. + config: The ``backend.SimpleStorage`` configuration. Bootstrap adds one + shared internal run ID. """ self.storage_unit_id = f"TQ_STORAGE_UNIT_{uuid4().hex[:8]}" - self.storage_unit_size = storage_unit_size - - self.storage_data = StorageUnitData(self.storage_unit_size) + total_storage_size = config.get("total_storage_size") + num_data_storage_units = int(config.get("num_data_storage_units", 1)) + self.storage_unit_size = ( + math.ceil(total_storage_size / num_data_storage_units) if total_storage_size is not None else None + ) + self.storage_data: StorageUnitData + + ssd_config = config.get("ssd_offload") + if ssd_config is not None and ssd_config.get("enabled", False): + ssd_path = ssd_config.get("path") + if not ssd_path: + raise ValueError("SimpleStorage SSD offload requires backend.SimpleStorage.ssd_offload.path") + threshold_bytes = int(ssd_config.get("threshold_bytes", DEFAULT_SSD_OFFLOAD_THRESHOLD_BYTES)) + self.storage_data = HybridStorageUnitData( + storage_size=self.storage_unit_size, + ssd_path=str(ssd_path), + run_id=str(config.get("_run_id") or uuid4().hex), + unit_id=self.storage_unit_id, + threshold_bytes=threshold_bytes, + ) + logger.info( + f"[{self.storage_unit_id}]: SSD offload enabled — " + f"path={Path(ssd_path).resolve() / SSD_OFFLOAD_DIRECTORY_NAME}, " + f"threshold={threshold_bytes} B/sample" + ) + else: + self.storage_data = StorageUnitData(self.storage_unit_size) self._requests_arrived = 0 self._arrivals_by_op = {} @@ -221,8 +807,15 @@ def __init__(self, storage_unit_size: int | None = None): self.zmq_context, self.put_get_socket, self._accept_probe, + self.worker_socket, + self.storage_data, ) + def shutdown(self) -> None: + """Stop request processing and release this storage unit's resources.""" + if self._finalizer.alive: + self._finalizer() + def _init_zmq_socket(self) -> None: """ Initialize ZMQ socket connections between storage unit and controller/clients: @@ -304,14 +897,18 @@ def _proxy_routine(self) -> None: events = dict(poller.poll(1000)) if front in events: messages = front.recv_multipart(copy=False) - identity = bytes(messages[0]) if messages else b"" - if not identity.startswith(STORAGE_CLIENT_IDENTITY_PREFIXES): - logger.warning( - "[%s]: dropping request with unrecognized ZMQ identity", - self.storage_unit_id, - ) - continue - back.send_multipart(messages, copy=False) + try: + identity = bytes(messages[0]) if messages else b"" + if not identity.startswith(STORAGE_CLIENT_IDENTITY_PREFIXES): + logger.warning( + "[%s]: dropping request with unrecognized ZMQ identity", + self.storage_unit_id, + ) + continue + back.send_multipart(messages, copy=False) + finally: + # This thread outlives each request, so drop forwarded payload frames while idle. + del messages if back in events: front.send_multipart(back.recv_multipart(copy=False), copy=False) @@ -338,7 +935,8 @@ def _worker_routine(self) -> None: while not self._shutdown_event.is_set(): monitor = self._metrics if self._metrics is not None else perf_monitor try: - socks = dict(poller.poll(TQ_STORAGE_POLLER_TIMEOUT * 1000)) + # The event cannot wake a ZMQ poll, so bound idle shutdown latency. + socks = dict(poller.poll(min(TQ_STORAGE_POLLER_TIMEOUT, 1) * 1000)) except zmq.error.ContextTerminated: # ZMQ context was terminated, exit gracefully logger.info(f"[{self.storage_unit_id}]: worker stopped gracefully (Context Terminated)") @@ -351,93 +949,97 @@ def _worker_routine(self) -> None: break if worker_socket in socks: - # Messages received from proxy: [identity, serialized_msg_frame1, ...] - messages = worker_socket.recv_multipart(copy=False) - identity = messages[0] - serialized_msg = messages[1:] - - try: - request_msg = ZMQMessage.deserialize(serialized_msg) - except Exception as e: - # The identity filter cannot cover this: an allowed peer can still send - # frames that fail to decode, and decoding here used to kill the thread. - logger.error( - f"[{self.storage_unit_id}]: undecodable request from " - f"identity={bytes(identity)!r}: {type(e).__name__}: {e}" - ) - error_msg = ZMQMessage.create( - request_type=ZMQRequestType.PUT_GET_ERROR, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={"message": f"undecodable request: {type(e).__name__}: {e}"}, - ) - worker_socket.send_multipart([identity] + error_msg.serialize(), copy=False) - continue - operation = request_msg.request_type - started = time.perf_counter() - - try: - self._requests_arrived += 1 - self._arrivals_by_op[operation.name] = self._arrivals_by_op.get(operation.name, 0) + 1 - - logger.debug(f"[{self.storage_unit_id}]: worker received operation: {operation}") - - # Process request - if operation == ZMQRequestType.PUT_DATA: # type: ignore[arg-type] - with monitor.measure(op_type="PUT_DATA"): - response_msg = self._handle_put(request_msg) - elif operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] - with monitor.measure(op_type="GET_DATA"): - response_msg = self._handle_get(request_msg) - elif operation == ZMQRequestType.CLEAR_DATA: # type: ignore[arg-type] - with monitor.measure(op_type="CLEAR_DATA"): - response_msg = self._handle_clear(request_msg) - elif operation == ZMQRequestType.GET_METRICS: # type: ignore[arg-type] - response_msg = self._handle_get_metrics() - elif operation == ZMQRequestType.SAVE_STORAGE_CHECKPOINT: # type: ignore[arg-type] - response_msg = self._handle_save_checkpoint(request_msg) - elif operation == ZMQRequestType.LOAD_STORAGE_CHECKPOINT: # type: ignore[arg-type] - response_msg = self._handle_load_checkpoint(request_msg) - else: - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.PUT_GET_OPERATION_ERROR, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={ - "message": f"Storage unit id #{self.storage_unit_id} " - f"receive invalid operation: {operation}." - }, - ) - except Exception as e: - logger.error( - f"[{self.storage_unit_id}]: worker error during {operation} " - f"from sender={request_msg.sender_id}: {type(e).__name__}: {e}" - ) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.PUT_GET_ERROR, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={ - "message": f"{self.storage_unit_id}, worker encountered error " - f"during operation {operation}: {str(e)}." - }, - ) - - # Send response back with identity for routing - response_frames = response_msg.serialize() - if operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] - # This end serializes the get response, so its frames give the true wire size. - log_heavy_operation( - self.storage_unit_id, - "get", - time.perf_counter() - started, - sum(frame_nbytes(frame) or 0 for frame in response_frames), - f"samples={len(request_msg.body.get('global_indexes', []))} " - f"fields={list(request_msg.body.get('fields', []))}", - ) - worker_socket.send_multipart([identity] + response_frames, copy=False) + # Keep request handling in a separate scope so payload references are + # released before the worker waits for the next request. + self._process_one_worker_request(worker_socket, monitor) logger.info(f"[{self.storage_unit_id}]: worker stopped.") poller.unregister(worker_socket) worker_socket.close(linger=0) + def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) -> None: + """Process one storage request and send its response.""" + # Messages received from proxy: [identity, serialized_msg_frame1, ...] + messages = worker_socket.recv_multipart(copy=False) + identity = messages[0] + serialized_msg = messages[1:] + + try: + request_msg = ZMQMessage.deserialize(serialized_msg) + except Exception as e: + # The identity filter cannot cover this: an allowed peer can still send + # frames that fail to decode, and decoding here used to kill the thread. + logger.error( + f"[{self.storage_unit_id}]: undecodable request from " + f"identity={bytes(identity)!r}: {type(e).__name__}: {e}" + ) + error_msg = ZMQMessage.create( + request_type=ZMQRequestType.PUT_GET_ERROR, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={"message": f"undecodable request: {type(e).__name__}: {e}"}, + ) + worker_socket.send_multipart([identity] + error_msg.serialize(), copy=False) + return + + operation = request_msg.request_type + started = time.perf_counter() + + try: + self._requests_arrived += 1 + self._arrivals_by_op[operation.name] = self._arrivals_by_op.get(operation.name, 0) + 1 + + logger.debug(f"[{self.storage_unit_id}]: worker received operation: {operation}") + + if operation == ZMQRequestType.PUT_DATA: # type: ignore[arg-type] + with monitor.measure(op_type="PUT_DATA"): + response_msg = self._handle_put(request_msg) + elif operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] + with monitor.measure(op_type="GET_DATA"): + response_msg = self._handle_get(request_msg) + elif operation == ZMQRequestType.CLEAR_DATA: # type: ignore[arg-type] + with monitor.measure(op_type="CLEAR_DATA"): + response_msg = self._handle_clear(request_msg) + elif operation == ZMQRequestType.GET_METRICS: # type: ignore[arg-type] + response_msg = self._handle_get_metrics() + elif operation == ZMQRequestType.SAVE_STORAGE_CHECKPOINT: # type: ignore[arg-type] + response_msg = self._handle_save_checkpoint(request_msg) + elif operation == ZMQRequestType.LOAD_STORAGE_CHECKPOINT: # type: ignore[arg-type] + response_msg = self._handle_load_checkpoint(request_msg) + else: + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.PUT_GET_OPERATION_ERROR, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={ + "message": f"Storage unit id #{self.storage_unit_id} receive invalid operation: {operation}." + }, + ) + except Exception as e: + logger.error( + f"[{self.storage_unit_id}]: worker error during {operation} " + f"from sender={request_msg.sender_id}: {type(e).__name__}: {e}" + ) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.PUT_GET_ERROR, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={ + "message": f"{self.storage_unit_id}, worker encountered error " + f"during operation {operation}: {str(e)}." + }, + ) + + response_frames = response_msg.serialize() + if operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] + # This end serializes the get response, so its frames give the true wire size. + log_heavy_operation( + self.storage_unit_id, + "get", + time.perf_counter() - started, + sum(frame_nbytes(frame) or 0 for frame in response_frames), + f"samples={len(request_msg.body.get('global_indexes', []))} " + f"fields={list(request_msg.body.get('fields', []))}", + ) + worker_socket.send_multipart([identity] + response_frames, copy=False) + def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage: """ Handle put request, add or update data into storage unit. @@ -613,11 +1215,14 @@ def _handle_get_metrics(self) -> ZMQMessage: "capacity": self.storage_unit_size, "active_keys": self.storage_data.active_key_count, "process_rss_bytes": process_rss, + "ssd_offload_enabled": int(self.storage_data.ssd_offload_enabled), + "ssd_active_values": self.storage_data.ssd_active_values, + "ssd_active_bytes": self.storage_data.ssd_active_bytes, + "ssd_fallback_values_total": self.storage_data.ssd_fallback_values_total, # Counted on arrival; op_stats below only advances on completion. "requests_arrived": self._requests_arrived, "arrivals_by_op": dict(self._arrivals_by_op), } - if self._accept_probe is not None: stats = self._accept_probe.stats metrics["accept_queue"] = { @@ -659,11 +1264,11 @@ def _handle_get_metrics(self) -> ZMQMessage: ) def _handle_save_checkpoint(self, data_parts) -> ZMQMessage: - """Serialize storage unit data directly to a file. + """Write storage unit data directly to its checkpoint path. Args: data_parts: ZMQMessage from client, containing ``path`` in body: - absolute path for the output .pkl file. The caller must ensure + absolute path for the output manifest. The caller must ensure this path is reachable from the node running this actor (shared filesystem required for multi-node setups). @@ -673,14 +1278,7 @@ def _handle_save_checkpoint(self, data_parts) -> ZMQMessage: """ path = data_parts.body["path"] try: - state = { - "storage_unit_id": self.storage_unit_id, - "storage_unit_size": self.storage_unit_size, - "field_data": self.storage_data.field_data, - "active_keys": self.storage_data._active_keys, - } - with open(path, "wb") as f: - pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL) + self.storage_data.save_checkpoint(path, self.storage_unit_id) logger.info(f"[{self.storage_unit_id}]: saved checkpoint to {path}") return ZMQMessage.create( request_type=ZMQRequestType.SAVE_STORAGE_CHECKPOINT_RESPONSE, # type: ignore[arg-type] @@ -696,7 +1294,7 @@ def _handle_save_checkpoint(self, data_parts) -> ZMQMessage: ) def _handle_load_checkpoint(self, data_parts) -> ZMQMessage: - """Restore storage unit data directly from a file. + """Restore storage unit data directly from its checkpoint path. Args: data_parts: ZMQMessage from client, containing ``path`` in body: @@ -711,28 +1309,24 @@ def _handle_load_checkpoint(self, data_parts) -> ZMQMessage: """ path = data_parts.body["path"] try: - with open(path, "rb") as f: - data = pickle.load(f) + previous_key_count = self.storage_data.active_key_count + checkpoint_size = self.storage_data.load_checkpoint(path) - if data["storage_unit_size"] != self.storage_unit_size: + if checkpoint_size != self.storage_unit_size: logger.warning( f"[{self.storage_unit_id}]: storage_unit_size mismatch — " - f"checkpoint={data['storage_unit_size']}, current={self.storage_unit_size}" + f"checkpoint={checkpoint_size}, current={self.storage_unit_size}" ) - if self.storage_data._active_keys: + if previous_key_count: logger.warning( - f"[{self.storage_unit_id}]: overwriting {len(self.storage_data._active_keys)} " + f"[{self.storage_unit_id}]: overwriting {previous_key_count} " f"existing keys with checkpoint data from {path}" ) - self.storage_data.field_data.clear() - self.storage_data._active_keys.clear() - self.storage_data.field_data = data["field_data"] - self.storage_data._active_keys = data["active_keys"] logger.info( f"[{self.storage_unit_id}]: loaded checkpoint from {path} — " - f"{len(data['active_keys'])} keys, {len(data['field_data'])} fields" + f"{self.storage_data.active_key_count} keys, {len(self.storage_data.field_data)} fields" ) return ZMQMessage.create( request_type=ZMQRequestType.LOAD_STORAGE_CHECKPOINT_RESPONSE, # type: ignore[arg-type] @@ -788,30 +1382,34 @@ def _shutdown_resources( zmq_context: zmq.Context | None, put_get_socket: zmq.Socket | None, accept_probe: "AcceptQueueProbe | None" = None, + worker_socket: zmq.Socket | None = None, + storage_data: StorageUnitData | None = None, ) -> None: """Clean up resources on garbage collection.""" logger.info("Shutting down SimpleStorageUnit resources...") - # Signal all threads to stop shutdown_event.set() # Before the ZMQ teardown: the probe runs on its own timer and would outlive the unit. if accept_probe is not None: accept_probe.stop() - # Terminate put_get_socket + if worker_thread and worker_thread.is_alive(): + worker_thread.join() + if proxy_thread and proxy_thread.is_alive(): + proxy_thread.join() + if put_get_socket: put_get_socket.close(linger=0) - - # Terminate ZMQ context to unblock proxy and workers + if worker_socket: + worker_socket.close(linger=0) if zmq_context: zmq_context.term() - - # Wait for threads to finish (with timeout) - if worker_thread and worker_thread.is_alive(): - worker_thread.join(timeout=5) - if proxy_thread and proxy_thread.is_alive(): - proxy_thread.join(timeout=5) + if storage_data is not None: + try: + storage_data.close() + except Exception as e: + logger.warning(f"Error closing storage data on shutdown: {e}") logger.info("SimpleStorageUnit resources shutdown complete.") From 05d4cc3377bee9b0f2897bbb90914ee20c0aea39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Han=20Zhenyu=20=E9=9F=A9=E6=8C=AF=E5=AE=87?= Date: Thu, 24 Sep 2026 10:07:24 +0800 Subject: [PATCH 08/12] [chore] Bump version from 0.1.11.dev0 to 0.1.11 (#177) As title Signed-off-by: 0oshowero0 --- transfer_queue/version/version | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transfer_queue/version/version b/transfer_queue/version/version index ccd72abe..20f49513 100644 --- a/transfer_queue/version/version +++ b/transfer_queue/version/version @@ -1 +1 @@ -0.1.11.dev0 +0.1.11 From 7b31c0b6147413bf4b55d0a49e197c85250d814b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Han=20Zhenyu=20=E9=9F=A9=E6=8C=AF=E5=AE=87?= Date: Thu, 24 Sep 2026 10:16:17 +0800 Subject: [PATCH 09/12] [chore] Bump version from 0.1.11 to 0.1.12.dev0 (#178) As title Signed-off-by: 0oshowero0 --- transfer_queue/version/version | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transfer_queue/version/version b/transfer_queue/version/version index 20f49513..260e27af 100644 --- a/transfer_queue/version/version +++ b/transfer_queue/version/version @@ -1 +1 @@ -0.1.11 +0.1.12.dev0 From c01270fd833443cc38064ec723a131a5f2501277 Mon Sep 17 00:00:00 2001 From: ys2025-AI Date: Fri, 9 Oct 2026 06:00:30 +0800 Subject: [PATCH 10/12] =?UTF-8?q?[fix](mooncake):=20init=20per-thread=20ac?= =?UTF-8?q?celerator=20context=20for=20buffer=20registe=E2=80=A6=20(#179)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 --- tests/test_mooncake_client_thread_context.py | 105 ++++++++++++++++++ .../storage/clients/mooncake_client.py | 54 ++++++++- 2 files changed, 157 insertions(+), 2 deletions(-) create mode 100644 tests/test_mooncake_client_thread_context.py diff --git a/tests/test_mooncake_client_thread_context.py b/tests/test_mooncake_client_thread_context.py new file mode 100644 index 00000000..01ce3808 --- /dev/null +++ b/tests/test_mooncake_client_thread_context.py @@ -0,0 +1,105 @@ +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the per-thread NPU device binding used by MooncakeStoreClient. + +Mooncake's Ascend transport is bound to a thread-local ACL context, but +register/unregister run on per-call ``ThreadPoolExecutor`` workers. These tests +stay hardware-free by faking the accelerator. +""" + +from transfer_queue.storage.clients import mooncake_client as mc + + +class _FakeNpu: + def __init__(self, current=5, initialized=True): + self._current = current + self._initialized = initialized + self.set_devices = [] + self.current_device_calls = 0 + + def is_initialized(self): + return self._initialized + + def current_device(self): + self.current_device_calls += 1 + return self._current + + def set_device(self, device): + self.set_devices.append(device) + + +def _new_client(npu_device): + client = object.__new__(mc.MooncakeStoreClient) + client._npu_device = npu_device + return client + + +def test_resolve_npu_device_captures_owning_thread_device(monkeypatch): + fake_npu = _FakeNpu(current=3) + monkeypatch.setattr(mc, "NPU_IMPORTED", True) + monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False) + + assert mc._resolve_npu_device() == 3 + assert fake_npu.current_device_calls == 1 + + +def test_resolve_npu_device_none_without_npu(monkeypatch): + monkeypatch.setattr(mc, "NPU_IMPORTED", False) + + assert mc._resolve_npu_device() is None + + +def test_resolve_npu_device_none_when_not_initialized(monkeypatch): + fake_npu = _FakeNpu(initialized=False) + monkeypatch.setattr(mc, "NPU_IMPORTED", True) + monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False) + + assert mc._resolve_npu_device() is None + assert fake_npu.current_device_calls == 0 + + +def test_worker_binds_device_captured_at_construction(monkeypatch): + # A context-less worker reports device 0; the client captured 3 in __init__. + fake_npu = _FakeNpu(current=0) + monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False) + + _new_client(npu_device=3)._bind_npu_device() + + assert fake_npu.set_devices == [3] + assert fake_npu.current_device_calls == 0 + + +def test_worker_bind_noop_without_npu_device(monkeypatch): + fake_npu = _FakeNpu() + monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False) + + _new_client(npu_device=None)._bind_npu_device() + + assert fake_npu.set_devices == [] + + +def test_worker_bind_logs_and_swallows_failure(monkeypatch): + class _FailingNpu(_FakeNpu): + def set_device(self, device): + raise RuntimeError("set_device failed") + + fake_npu = _FailingNpu() + monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False) + warnings = [] + monkeypatch.setattr(mc.logger, "warning", lambda *a, **k: warnings.append((a, k))) + + _new_client(npu_device=3)._bind_npu_device() # must not raise + + assert warnings and warnings[0][1].get("exc_info") is True diff --git a/transfer_queue/storage/clients/mooncake_client.py b/transfer_queue/storage/clients/mooncake_client.py index d6914902..ff38bcb7 100644 --- a/transfer_queue/storage/clients/mooncake_client.py +++ b/transfer_queue/storage/clients/mooncake_client.py @@ -33,6 +33,34 @@ logger = get_logger(__name__) +try: + import torch_npu # noqa: F401 + + NPU_IMPORTED: bool = True +except Exception: # pragma: no cover - accelerator optional + NPU_IMPORTED = False + + +def _resolve_npu_device() -> int | None: + """Return the NPU device this process is bound to, or ``None`` if it uses no NPU. + + Ascend ACL contexts are thread-local, so this must run on the thread that owns + the device (``__init__``). Resolving it on a context-less pool worker would + always report device 0. Gating on ``is_initialized()`` (rather than + ``is_available()``) keeps a process that never touches the NPU from allocating + device memory and contending for device 0. + """ + if not NPU_IMPORTED: + return None + try: + if not torch.npu.is_initialized(): + return None + return int(torch.npu.current_device()) + except Exception: + logger.warning("Failed to resolve the current NPU device.", exc_info=True) + return None + + MOONCAKE_STORE_IMPORTED: bool = True try: from mooncake.store import MooncakeDistributedStore, ReplicateConfig @@ -121,6 +149,9 @@ def __init__(self, config: dict[str, Any]): hard_pin = not offload_enabled self.replica_config.with_hard_pin = bool(hard_pin) + # Capture on the owning thread; pool workers reuse it via _bind_npu_device. + self._npu_device = _resolve_npu_device() + self._store = MooncakeDistributedStore() ret = self._store.setup( self.local_hostname, @@ -134,6 +165,19 @@ def __init__(self, config: dict[str, Any]): if ret != 0: raise RuntimeError(f"Mooncake store setup failed with error code: {ret}") + def _bind_npu_device(self) -> None: + """``ThreadPoolExecutor`` initializer: bind a worker to ``self._npu_device``.""" + if self._npu_device is None: + return + try: + torch.npu.set_device(self._npu_device) + except Exception: + logger.warning( + f"Failed to bind NPU device {self._npu_device} on a worker thread; " + "Mooncake buffer register/unregister may fail.", + exc_info=True, + ) + def put(self, keys: list[str], values: list[Any]) -> list[dict | None]: """Stores multiple key-value pairs to MooncakeStore. @@ -173,7 +217,10 @@ def put(self, keys: list[str], values: list[Any]) -> list[dict | None]: tensor_futures: list[Future[None]] = [] bytes_futures: list[Future[list[int]]] = [] - with ThreadPoolExecutor(max_workers=MAX_BATCH_WORKER_THREADS) as executor: + with ThreadPoolExecutor( + max_workers=MAX_BATCH_WORKER_THREADS, + initializer=self._bind_npu_device, + ) as executor: if not use_gdr_path: for i in range(0, len(tensor_keys), BATCH_SIZE_LIMIT): batch_keys = tensor_keys[i : i + BATCH_SIZE_LIMIT] @@ -373,7 +420,10 @@ def get( results[idx] = val futures = [] - with ThreadPoolExecutor(max_workers=MAX_BATCH_WORKER_THREADS) as executor: + with ThreadPoolExecutor( + max_workers=MAX_BATCH_WORKER_THREADS, + initializer=self._bind_npu_device, + ) as executor: for i in range(0, len(cpu_tensor_indices), BATCH_SIZE_LIMIT): batch_keys = cpu_tensor_keys[i : i + BATCH_SIZE_LIMIT] batch_shapes = cpu_tensor_shapes[i : i + BATCH_SIZE_LIMIT] From 11922e3a5944f071f22cc5439a99dac554069a18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9F=A6=E5=AD=90=E6=89=AC?= Date: Fri, 9 Oct 2026 11:03:12 +0800 Subject: [PATCH 11/12] [optim] Lease pre-registered receive buffers on the MooncakeStore read path (#175) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Refs: https://github.com/Ascend/TransferQueue/issues/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×)** | image ### 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 Co-authored-by: weiziyang --- tests/test_mooncake_buffer_pool.py | 509 ++++++++++++++++++ .../storage/clients/mooncake_client.py | 170 +++++- 2 files changed, 673 insertions(+), 6 deletions(-) create mode 100644 tests/test_mooncake_buffer_pool.py diff --git a/tests/test_mooncake_buffer_pool.py b/tests/test_mooncake_buffer_pool.py new file mode 100644 index 00000000..400360ac --- /dev/null +++ b/tests/test_mooncake_buffer_pool.py @@ -0,0 +1,509 @@ +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pooled receive buffers on the MooncakeStore tensor read path. + +The store and the buffer pool are faked, so these tests need neither mooncake nor RDMA. +The fake pool keeps the upstream properties this path depends on: buffers are carved out +of one fixed arena and reused, exhaustion raises, 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. +""" + +import ctypes +import logging +import threading +import weakref + +import pytest +import torch + +from transfer_queue.storage.clients import mooncake_client as mcc + +LOGGER_NAME = "transfer_queue.storage.clients.mooncake_client" +ALIGN = 256 + +DTYPES = [torch.float32, torch.int64, torch.float32, torch.int16] +SHAPES = [(4, 3), (5,), (), (2, 8)] +KEYS = ["k0", "k1", "k2", "k3"] + + +def expected_tensors(): + """Deterministic payloads, one per key in KEYS.""" + out = [] + for seed, (dtype, shape) in enumerate(zip(DTYPES, SHAPES, strict=True)): + numel = torch.empty(shape).numel() + values = torch.arange(seed * 100, seed * 100 + numel) + out.append(values.to(dtype).reshape(shape)) + return out + + +class FakeStore: + """Writes the stored payload into whatever pointer batch_get_into is given.""" + + def __init__(self): + self.objects = { + key: bytes(t.contiguous().numpy().tobytes()) for key, t in zip(KEYS, expected_tensors(), strict=True) + } + self.registered: list[tuple[int, int]] = [] + self.unregistered: list[int] = [] + self.closed = False + self.transfer_error: Exception | None = None + self.on_transfer = None + + def setup(self, *args): + return 0 + + def register_buffer(self, ptr, size): + self.registered.append((ptr, size)) + return 0 + + def unregister_buffer(self, ptr): + self.unregistered.append(ptr) + return 0 + + def batch_get_into(self, keys, ptrs, sizes): + if self.transfer_error is not None: + raise self.transfer_error + if self.on_transfer is not None: + self.on_transfer() + for key, ptr, size in zip(keys, ptrs, sizes, strict=True): + ctypes.memmove(ptr, self.objects[key], size) + return list(sizes) + + def close(self): + self.closed = True + + +class FakeBuffer: + """One region carved out of the pool's arena.""" + + def __init__(self, pool, offset, nbytes): + self._pool = pool + self._offset = offset + self.nbytes = nbytes + self.ptr = pool.arena.data_ptr() + offset + self._views: list[weakref.ref] = [] + self.returned = False + + @property + def buffer(self): + # Mooncake hands out a fresh view per access and counts it as an export. + view = self._pool.arena_bytes[self._offset : self._offset + max(self.nbytes, 1)] + self._views.append(weakref.ref(view)) + return view + + def has_live_views(self): + return any(ref() is not None for ref in self._views) + + def release(self): + if self.has_live_views(): + raise RuntimeError("cannot release buffer while exported views exist") + self._pool._give_back(self._offset) + self.returned = True + + +class FakePool: + """Sub-allocates one arena, the way mooncake carves up its registered local buffer.""" + + def __init__(self, capacity, max_regions=None): + self.capacity = capacity + self.max_regions = max_regions + self.arena = torch.empty(max(capacity, 1), dtype=torch.uint8) + self.arena_bytes = self.arena.numpy() + self.acquired: list[int] = [] + self.offsets: list[int] = [] + self.returned = 0 + self.active: dict[int, int] = {} + self.peak_bytes = 0 + self.peak_regions = 0 + self.closed = False + self._lock = threading.Lock() + + def acquire(self, nbytes, block=True): + with self._lock: + if self.closed: + raise RuntimeError("buffer pool is closed") + size = max(nbytes, 1) + if size > self.capacity: + raise RuntimeError("requested buffer size exceeds pool capacity") + if self.max_regions is not None and len(self.active) >= self.max_regions: + raise RuntimeError("buffer pool is exhausted") + offset = self._first_fit(size) + if offset is None: + raise RuntimeError("buffer pool is exhausted") + self.active[offset] = size + self.acquired.append(nbytes) + self.offsets.append(offset) + self.peak_bytes = max(self.peak_bytes, sum(self.active.values())) + self.peak_regions = max(self.peak_regions, len(self.active)) + return FakeBuffer(self, offset, nbytes) + + def _first_fit(self, size): + cursor = 0 + for offset in sorted(self.active): + if offset - cursor >= size: + return cursor + cursor = -(-(offset + self.active[offset]) // ALIGN) * ALIGN + return cursor if self.capacity - cursor >= size else None + + def _give_back(self, offset): + with self._lock: + del self.active[offset] + self.returned += 1 + + def close(self): + with self._lock: + if self.active: + raise RuntimeError("cannot close buffer pool with active leases") + self.closed = True + + +class PoolFactory: + """Stands in for mooncake.store.BufferPool, recording how the client configured it.""" + + def __init__(self, capacity=None): + self._capacity = capacity + self.kwargs: dict | None = None + self.pool: FakePool | None = None + + def __call__(self, store, max_bytes=0, max_regions=None): + self.kwargs = {"max_bytes": max_bytes, "max_regions": max_regions} + self.pool = FakePool(max_bytes if self._capacity is None else self._capacity, max_regions) + return self.pool + + +@pytest.fixture +def store(monkeypatch): + fake = FakeStore() + monkeypatch.setattr(mcc, "MOONCAKE_STORE_IMPORTED", True) + # raising=False: these symbols are absent unless mooncake is installed. + monkeypatch.setattr(mcc, "MooncakeDistributedStore", lambda: fake, raising=False) + monkeypatch.setattr(mcc, "ReplicateConfig", type("ReplicateConfig", (), {}), raising=False) + return fake + + +def make_client(local_buffer_size=1 << 30, **extra): + config = { + "local_hostname": "127.0.0.1", + "metadata_server": "127.0.0.1:8080", + "master_server_address": "127.0.0.1:8081", + "local_buffer_size": local_buffer_size, + # The pool only runs on RDMA; default to it so the pooled-path tests exercise it. + "protocol": "rdma", + } + config.update(extra) + return mcc.MooncakeStoreClient(config) + + +def install_pool(monkeypatch, capacity=None): + """Make the client believe mooncake provides a buffer pool, and hand it a fake one.""" + factory = PoolFactory(capacity) + monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", True) + monkeypatch.setattr(mcc, "BufferPool", factory, raising=False) + return factory + + +def make_client_and_pool(monkeypatch, local_buffer_size=1 << 30, capacity=None, **extra): + factory = install_pool(monkeypatch, capacity) + client = make_client(local_buffer_size, **extra) + return client, factory.pool + + +def read(client, keys, shapes, dtypes): + tensors, indexes = client._get_tensors_thread_worker(keys, shapes, dtypes, list(range(len(keys)))) + assert indexes == list(range(len(keys))) + return tensors + + +def read_all(client): + return read(client, KEYS, SHAPES, DTYPES) + + +def assert_payloads(tensors): + for got, want in zip(tensors, expected_tensors(), strict=True): + assert got.dtype == want.dtype + assert got.shape == want.shape + assert torch.equal(got, want) + + +def test_reads_land_in_pooled_buffer(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + + tensors = read_all(client) + + # Mixed dtypes and a scalar all round-trip through the staged copy-out. + assert_payloads(tensors) + # The whole point: no registration on the data path, and no buffer left on loan. + assert store.registered == [] + assert len(pool.acquired) == 1 and pool.returned == 1 and pool.active == {} + + +def test_pool_budget_is_capped_by_the_registered_local_buffer(store, monkeypatch): + factory = install_pool(monkeypatch) + client = make_client(local_buffer_size=1 << 20) + + # max_bytes=0 would let mooncake register fresh memory per lease once the local + # buffer is full, which is exactly the cost this path removes. + assert factory.kwargs == {"max_bytes": 1 << 20, "max_regions": mcc.MAX_BATCH_WORKER_THREADS} + # Half of an equal per-thread share, leaving the store room for its own staging. + assert client._pool_share_bytes == (1 << 20) // (2 * mcc.MAX_BATCH_WORKER_THREADS) + + +def test_pool_is_disabled_on_tcp(store, monkeypatch): + # TCP has no registration cost to save, so the pool would only add the copy-out. + install_pool(monkeypatch) + client = make_client(protocol="tcp") + + assert client._buffer_pool is None + tensors = read_all(client) + + assert_payloads(tensors) + assert store.registered and len(store.unregistered) == len(store.registered) + + +def test_batch_larger_than_share_is_read_in_rounds(store, monkeypatch): + # Small local buffer: each thread's share holds only part of the batch. + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + + tensors = read_all(client) + + assert_payloads(tensors) + assert len(pool.acquired) > 1 + assert all(nbytes <= client._pool_share_bytes for nbytes in pool.acquired) + assert pool.returned == len(pool.acquired) + assert store.registered == [] + + +def test_tensor_larger_than_share_registers_its_own_buffer(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + keys, shapes, dtypes = ["big", "small"], [(256,), (4,)], [torch.float32, torch.float32] + payloads = [torch.arange(256, dtype=torch.float32), torch.arange(4, dtype=torch.float32)] + store.objects = {k: bytes(t.numpy().tobytes()) for k, t in zip(keys, payloads, strict=True)} + + tensors = read(client, keys, shapes, dtypes) + + for got, want in zip(tensors, payloads, strict=True): + assert torch.equal(got, want) + # The oversized tensor must not lease the headroom the other readers need; the + # small one still takes the pooled path. + assert pool.acquired and all(nbytes <= client._pool_share_bytes for nbytes in pool.acquired) + assert store.registered and len(store.unregistered) == len(store.registered) + assert pool.active == {} + + +def test_fallback_registers_everything_in_one_pass(store, monkeypatch): + # Pool exhausted up front: every group falls back, but a single same-dtype batch must + # register one merged region, not one per group (the pre-pool cost, not N times it). + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + held = pool.acquire(pool.capacity) + + n = 8 + keys = [f"f{i}" for i in range(n)] + payloads = [torch.arange(i, i + 32, dtype=torch.float32) for i in range(n)] + store.objects = {k: bytes(t.numpy().tobytes()) for k, t in zip(keys, payloads, strict=True)} + + tensors = read(client, keys, [(32,)] * n, [torch.float32] * n) + + for got, want in zip(tensors, payloads, strict=True): + assert torch.equal(got, want) + assert len(store.registered) == 1 and len(store.unregistered) == 1 + held.release() + + +def test_falls_back_to_own_buffers_when_pool_is_exhausted(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + # Mooncake's own staging holds the arena, so acquire() raises instead of returning None. + held = pool.acquire(pool.capacity) + + tensors = read_all(client) + + assert_payloads(tensors) + assert pool.acquired == [pool.capacity] + assert store.registered and len(store.unregistered) == len(store.registered) + held.release() + + +def test_exhaustion_warns_once(store, monkeypatch, caplog): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + held = pool.acquire(pool.capacity) + + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + read_all(client) + read_all(client) + + # Every group of every read falls back here; the hot path must not flood the log. + fallbacks = [r for r in caplog.records if "Falling back to registering" in r.message] + assert len(fallbacks) == 1 + held.release() + + +def test_registers_receive_regions_without_pool_support(store, monkeypatch): + monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", False) + + tensors = read_all(make_client()) + + assert_payloads(tensors) + assert store.registered and len(store.unregistered) == len(store.registered) + + +def test_uniform_group_copied_in_one_strided_pass(store, monkeypatch): + # Many identical small tensors (the fragmented-read case) take the single strided + # copy-out path instead of a per-tensor loop; the payloads must still round-trip. + client, pool = make_client_and_pool(monkeypatch) + n = 8 + keys = [f"u{i}" for i in range(n)] + payloads = [torch.arange(i, i + 16, dtype=torch.float32) for i in range(n)] + store.objects = {k: bytes(t.numpy().tobytes()) for k, t in zip(keys, payloads, strict=True)} + + tensors = read(client, keys, [(16,)] * n, [torch.float32] * n) + + for got, want in zip(tensors, payloads, strict=True): + assert torch.equal(got, want) + assert pool.acquired and pool.returned == len(pool.acquired) + assert store.registered == [] + + +def test_returned_tensors_survive_buffer_reuse(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + + first = read_all(client) + rounds = len(pool.acquired) + second = read_all(client) + + # The second read gets the same arena offsets back, so the first read's tensors + # prove the copy-out really handed ownership to the caller. + assert pool.offsets[rounds:] == pool.offsets[:rounds] + assert pool.returned == len(pool.acquired) + assert_payloads(first) + assert_payloads(second) + + +def test_failed_transfer_returns_the_buffer_and_keeps_the_cause(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + store.transfer_error = RuntimeError("transport is down") + + with pytest.raises(RuntimeError, match="transport is down"): + read_all(client) + + assert pool.active == {} and pool.returned == len(pool.acquired) + + +def test_failed_copy_out_returns_the_buffer_and_keeps_the_cause(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + + def boom(self, source, *args, **kwargs): + # Drop the staged view the caller passed in: a real copy_ failure comes from C++ + # and leaves no Python frame holding it, so only the client's own reference counts. + del source + raise RuntimeError("copy-out failed") + + monkeypatch.setattr(torch.Tensor, "copy_", boom) + + # A staged view left alive by the failing copy would make release() raise and + # replace the real cause, so the view has to be dropped on the error path too. + with pytest.raises(RuntimeError, match="copy-out failed"): + read_all(client) + + assert pool.active == {} and pool.returned == len(pool.acquired) + + +def test_concurrent_readers_stay_within_the_local_buffer(store, monkeypatch): + local_buffer_size = 8192 + readers = 8 + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=local_buffer_size) + barrier = threading.Barrier(readers) + # Hold every reader inside its transfer, so all of them need memory at the same time. + store.on_transfer = lambda: barrier.wait(timeout=30) + results: list[list[torch.Tensor]] = [None] * readers # type: ignore[list-item] + failures: list[Exception] = [] + + def reader(slot): + try: + results[slot] = read_all(client) + except Exception as e: # reported after the join below + failures.append(e) + + threads = [threading.Thread(target=reader, args=(i,)) for i in range(readers)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not failures + for tensors in results: + assert_payloads(tensors) + # Readers past the region cap fall back instead of eating into the store's headroom. + assert pool.peak_regions <= mcc.MAX_BATCH_WORKER_THREADS + assert pool.peak_bytes <= local_buffer_size // 2 + assert pool.active == {} + + +@pytest.mark.parametrize( + ("config", "warns"), + [ + ({"protocol": "rdma"}, True), + ({"protocol": "tcp"}, False), + ({"protocol": "rdma", "use_gdr": True}, False), + ], +) +def test_missing_pool_support_warns_only_for_cpu_rdma_reads(store, monkeypatch, caplog, config, warns): + monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", False) + # GDR staging defers cudaMalloc, so it can stand in without a CUDA device here. + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) + + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + make_client(**config) + + assert any("BufferPool is unavailable" in r.message for r in caplog.records) is warns + + +def test_close_returns_the_pool_before_the_store(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + + read_all(client) + client.close() + + assert pool.closed and store.closed and client._buffer_pool is None + + +def test_close_keeps_the_store_open_while_a_buffer_is_out(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + held = pool.acquire(1024) + + # The store owns the memory the buffer points into, so a refused pool close must not + # tear it down underneath an active lease. + with pytest.raises(RuntimeError, match="cannot close buffer pool with active leases"): + client.close() + + assert not store.closed and client._buffer_pool is pool + held.release() + client.close() + assert pool.closed and store.closed + + +def test_empty_tensor_reads_like_the_register_path(store, monkeypatch): + # mooncake rejects a 0-byte put, so this is unreachable in practice; the guard only keeps + # the pooled path from raising where the register path returns, should it ever be hit. + keys, shapes, dtypes = ["empty", "normal"], [(0,), (4,)], [torch.float32, torch.float32] + payloads = [torch.empty(0), torch.arange(4, dtype=torch.float32)] + store.objects = {k: bytes(t.numpy().tobytes()) for k, t in zip(keys, payloads, strict=True)} + + monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", False) + registered = read(make_client(), keys, shapes, dtypes) + client, pool = make_client_and_pool(monkeypatch) + pooled = read(client, keys, shapes, dtypes) + + for got, want in zip(pooled, registered, strict=True): + assert got.shape == want.shape and torch.equal(got, want) + assert pool.returned == len(pool.acquired) diff --git a/transfer_queue/storage/clients/mooncake_client.py b/transfer_queue/storage/clients/mooncake_client.py index ff38bcb7..e765f3c8 100644 --- a/transfer_queue/storage/clients/mooncake_client.py +++ b/transfer_queue/storage/clients/mooncake_client.py @@ -68,6 +68,13 @@ def _resolve_npu_device() -> int | None: except ImportError: MOONCAKE_STORE_IMPORTED = False +MOONCAKE_BUFFER_POOL_IMPORTED: bool = True +try: + from mooncake.store import BufferPool +except ImportError: + # Older mooncake builds have no lease API; those fall back to registering per transfer. + MOONCAKE_BUFFER_POOL_IMPORTED = False + BATCH_SIZE_LIMIT: int = 400 MAX_BATCH_WORKER_THREADS = 4 MAX_SERIAL_WORKER_THREADS = 4 @@ -75,6 +82,57 @@ def _resolve_npu_device() -> int | None: RETRY_DELAY_SECONDS = 1.0 +def _copy_buffer_content_into_tensors(lease, targets: list[Tensor], offsets: list[int]) -> None: + """Copy each staged region out of the pooled buffer into the caller's tensors. + + Views are dropped even when a copy raises: the pool refuses to return a buffer while an + exported view of it is alive, and an exception would otherwise keep it in its frame. + """ + # torch.frombuffer rejects count=0, and an empty tensor has nothing to copy. + staged = [(t, off) for t, off in zip(targets, offsets, strict=True) if t.numel()] + if not staged: + return + targets = [t for t, _ in staged] + offsets = [off for _, off in staged] + + t0 = targets[0] + numel = t0.numel() + # A uniform, tightly packed group copies in one strided pass; ragged groups copy per tensor. + uniform = len(targets) > 1 + if uniform: + base_off = t0.storage_offset() + base_ptr = t0.untyped_storage().data_ptr() + for j, t in enumerate(targets): + if ( + t.dtype != t0.dtype + or t.shape != t0.shape + or not t.is_contiguous() + or t.storage_offset() != base_off + j * numel + or t.untyped_storage().data_ptr() != base_ptr + ): + uniform = False + break + + if uniform: + k = len(targets) + stride_elems = (offsets[1] - offsets[0]) // t0.element_size() + src = torch.frombuffer(lease.buffer, dtype=t0.dtype, count=(k - 1) * stride_elems + numel).as_strided( + (k, numel), (stride_elems, 1) + ) + try: + t0.as_strided((k, numel), (numel, 1), t0.storage_offset()).copy_(src) + finally: + del src + return + + for target, off in zip(targets, offsets, strict=True): + staged = torch.frombuffer(lease.buffer, dtype=target.dtype, count=target.numel(), offset=off) + try: + target.copy_(staged.view(target.shape)) + finally: + del staged + + @StorageClientFactory.register("MooncakeStoreClient") class MooncakeStoreClient(StorageKVClient): """ @@ -165,6 +223,26 @@ def __init__(self, config: dict[str, Any]): if ret != 0: raise RuntimeError(f"Mooncake store setup failed with error code: {ret}") + # Only RDMA reads pay for register_buffer (a kernel op far slower than the transfer it + # enables); TCP reads would pay the extra copy-out for no gain, so leave them unpooled. + # See https://github.com/Ascend/TransferQueue/issues/169 + use_pool = MOONCAKE_BUFFER_POOL_IMPORTED and self.protocol == "rdma" + self._buffer_pool = ( + BufferPool(self._store, max_bytes=self.local_buffer_size, max_regions=MAX_BATCH_WORKER_THREADS) + if use_pool + else None + ) + self._pool_fallback_warned = False + if not MOONCAKE_BUFFER_POOL_IMPORTED and self.protocol == "rdma" and self._gdr_staging is None: + logger.warning( + "mooncake.store.BufferPool is unavailable, so every tensor read registers and " + "unregisters its own receive buffer. Upgrade to mooncake-transfer-engine >= 0.3.12 " + "to lease pre-registered buffers instead." + ) + # Claim at most half the registered buffer: mooncake stages its own reads from the + # same region, so splitting all of it across the reader threads would starve those reads. + self._pool_share_bytes = self.local_buffer_size // (2 * MAX_BATCH_WORKER_THREADS) + def _bind_npu_device(self) -> None: """``ThreadPoolExecutor`` initializer: bind a worker to ``self._npu_device``.""" if self._npu_device is None: @@ -454,17 +532,90 @@ def _get_tensors_thread_worker( self, batch_keys: list[str], batch_shapes: list[tuple], batch_dtypes: list[torch.dtype], indexes: list[int] ) -> tuple[list[Tensor], list[int]]: batch_nbytes = get_nbytes(batch_dtypes, batch_shapes) - batch_buffer_tensors, batch_buffer_ptrs, region_ptrs, region_sizes = allocate_empty_tensors( - batch_dtypes, batch_shapes - ) + batch_buffer_tensors, batch_buffer_ptrs, _, _ = allocate_empty_tensors(batch_dtypes, batch_shapes) + + if self._buffer_pool is None: + self._read_into_registered_tensors(batch_keys, batch_buffer_ptrs, batch_nbytes) + return batch_buffer_tensors, indexes + + # Read what the pool can serve; register the rest in one pass after the loop, so + # oversized tensors and pool exhaustion keep the pre-pool cost (one register + one + # transfer) instead of paying it per group. + fallback: list[int] = [] + for group in split_by_bytes(batch_nbytes, self._pool_share_bytes): + fallback += self._read_via_buffer_pool( + group, batch_keys, batch_buffer_ptrs, batch_nbytes, batch_buffer_tensors + ) + if fallback: + self._read_into_registered_tensors( + [batch_keys[i] for i in fallback], + [batch_buffer_ptrs[i] for i in fallback], + [batch_nbytes[i] for i in fallback], + ) + return batch_buffer_tensors, indexes + + def _read_into_registered_tensors(self, keys: list[str], ptrs: list[int], nbytes: list[int]) -> None: + """Register the receive regions for one transfer, read into them, then unregister.""" + region_ptrs, region_sizes = merge_contiguous_memory(ptrs, nbytes) self._register_all_buffers(region_ptrs, region_sizes) try: - self._batch_get_into_with_retry(batch_keys, batch_buffer_ptrs, batch_nbytes) + self._batch_get_into_with_retry(keys, ptrs, nbytes) finally: self._unregister_all_buffers(region_ptrs) - return batch_buffer_tensors, indexes + def _read_via_buffer_pool( + self, + group: list[int], + batch_keys: list[str], + batch_ptrs: list[int], + batch_nbytes: list[int], + batch_tensors: list[Tensor], + ) -> list[int]: + """Read one group into a pooled buffer and copy into the caller's tensors. + + Returns the group unread when the pool cannot serve it, so the caller can register + those keys in one batch. A group larger than one reader's share is never leased: + that headroom belongs to the other readers and to mooncake's own staging. + """ + nbytes = [batch_nbytes[i] for i in group] + offsets, total = _aligned_offsets(nbytes) + lease = self._acquire_buffer(total) if total <= self._pool_share_bytes else None + if lease is None: + return group + + keys = [batch_keys[i] for i in group] + try: + self._batch_get_into_with_retry(keys, [lease.ptr + off for off in offsets], nbytes) + _copy_buffer_content_into_tensors(lease, [batch_tensors[i] for i in group], offsets) + except Exception: + # A failed read can leave the buffer unreturnable; keep the original cause. + try: + lease.release() + except Exception as release_error: + logger.warning(f"Returning the pooled receive buffer failed: {release_error}") + raise + lease.release() + return [] + + def _acquire_buffer(self, nbytes: int): + """Lease ``nbytes`` of pre-registered memory, or None when the pool cannot serve it. + + Never block: a busy pool means registering a receive buffer beats waiting. mooncake + reports both the per-reader cap and an over-capacity request as an exception. + """ + assert self._buffer_pool is not None + try: + return self._buffer_pool.acquire(nbytes, block=False) + except Exception as e: + logger.debug(f"Leasing {nbytes} B of pre-registered memory failed ({e}); registering own buffer.") + if not self._pool_fallback_warned: + self._pool_fallback_warned = True + logger.warning( + f"Falling back to registering receive buffers per read ({e}). Raise local_buffer_size " + "or lower read concurrency to keep reads on the pre-registered fast path." + ) + return None def _get_tensors_gdr( self, @@ -584,7 +735,14 @@ def clear(self, keys: list[str], custom_backend_meta: list[Any] | None = None) - logger.error(f"remove failed for key `{actual_keys[i]}` with error code: {ret}") def close(self): - """Closes MooncakeStore.""" + """Closes MooncakeStore. + + Callers must have finished their reads first: the pool refuses to close while a + buffer is still leased, and the store owns the memory those buffers point into. + """ + if self._buffer_pool is not None: + self._buffer_pool.close() + self._buffer_pool = None if self._gdr_staging is not None: self._gdr_staging.close(self._store) self._gdr_staging = None From 2873a87afff7c4faa1073f985336f7f27497cc7e Mon Sep 17 00:00:00 2001 From: Ginray Date: Mon, 12 Oct 2026 19:07:35 +0800 Subject: [PATCH 12/12] [fix] resolve storage merge regressions Signed-off-by: Ginray --- tests/test_metrics.py | 135 ++++++++++++++++ tests/test_simple_storage_unit.py | 67 ++++++++ tests/test_storage_request_retry.py | 148 +++++++++++++++++- .../managers/simple_storage_manager.py | 16 +- .../storage/payload_transfer/base.py | 3 + .../storage/payload_transfer/zmq.py | 49 +----- transfer_queue/storage/simple_storage.py | 27 +++- 7 files changed, 390 insertions(+), 55 deletions(-) diff --git a/tests/test_metrics.py b/tests/test_metrics.py index f3e338b0..eea13564 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -31,6 +31,7 @@ """Unit tests for the Prometheus metrics exporter (transfer_queue.metrics).""" import time +from concurrent.futures import Future from threading import Thread from unittest.mock import MagicMock, patch @@ -543,3 +544,137 @@ def test_metrics_request_types_exist(self): assert ZMQRequestType.GET_METRICS.value == "GET_METRICS" assert ZMQRequestType.METRICS_RESPONSE.value == "METRICS_RESPONSE" + + +class TestStorageWorkerArrivals: + @pytest.fixture + def worker(self): + from transfer_queue.storage.payload_transfer.zmq import ZmqPayloadTransfer + from transfer_queue.storage.simple_storage import SimpleStorageUnit, StorageUnitData + from transfer_queue.utils.perf_utils import IntervalPerfMonitor + + unit_class = SimpleStorageUnit.__ray_metadata__.modified_class + unit = object.__new__(unit_class) + unit.storage_unit_id = "unit_a" + unit.storage_unit_size = 4 + unit.storage_data = StorageUnitData(4) + unit.payload_transfer = ZmqPayloadTransfer() + unit._requests_arrived = 0 + unit._arrivals_by_op = {} + unit._accept_probe = None + unit._metrics = None + return unit, IntervalPerfMonitor("unit_a") + + @staticmethod + def dispatch(worker, operation, body): + unit, monitor = worker + request = ZMQMessage.create(request_type=operation, sender_id="manager", body=body) + socket = MagicMock() + socket.recv_multipart.return_value = [b"manager"] + request.serialize() + unit._process_one_worker_request(socket, monitor) + sent = socket.send_multipart.call_args.args[0] + assert sent[0] == b"manager" + return ZMQMessage.deserialize(sent[1:]) + + def test_zmq_metrics_and_checkpoints_keep_their_arrival_labels(self, worker, tmp_path): + path = str(tmp_path / "storage.pkl") + requests = [ + (ZMQRequestType.PUT_DATA, {"global_indexes": [0], "data": {"value": [42]}}, "PUT_DATA_RESPONSE"), + (ZMQRequestType.GET_DATA, {"global_indexes": [0], "fields": ["value"]}, "GET_DATA_RESPONSE"), + (ZMQRequestType.SAVE_STORAGE_CHECKPOINT, {"path": path}, "SAVE_STORAGE_CHECKPOINT_RESPONSE"), + (ZMQRequestType.CLEAR_DATA, {"global_indexes": [0]}, "CLEAR_DATA_RESPONSE"), + (ZMQRequestType.LOAD_STORAGE_CHECKPOINT, {"path": path}, "LOAD_STORAGE_CHECKPOINT_RESPONSE"), + (ZMQRequestType.GET_METRICS, {}, "METRICS_RESPONSE"), + ] + for operation, body, response_type in requests: + reply = self.dispatch(worker, operation, body) + assert reply.request_type == ZMQRequestType[response_type] + if "CHECKPOINT" in operation.name: + assert reply.body["success"] + unit, monitor = worker + expected = {operation.name: 1 for operation, _, _ in requests} + assert reply.body["requests_arrived"] == len(requests) + assert reply.body["arrivals_by_op"] == expected + assert dict(monitor.success_counts) == {"PUT_DATA": 1, "GET_DATA": 1, "CLEAR_DATA": 1} + assert unit.storage_data.get_data(["value"], [0]) == {"value": [42]} + + exporter = TQMetricsExporter() + exporter.register_storage_units({"unit_a": MagicMock()}) + exporter._query_storage_unit = MagicMock(return_value=reply.body) + exporter.collect_storage_metrics() + for operation in expected: + assert exporter.storage_arrivals_by_op.labels(storage_unit_id="unit_a", op_type=operation)._value.get() == 1 + + @pytest.mark.parametrize("operation", ["put", "get"]) + @pytest.mark.parametrize("finish", ["COMMIT", "CANCEL"]) + def test_nixl_counts_prepare_once_and_skips_continuations(self, worker, operation, finish): + from transfer_queue.storage.payload_transfer.nixl import NixlPayloadTransfer, PayloadDescriptor + from transfer_queue.utils.serial_utils import calc_packed_size, encode, pack_into + + unit, monitor = worker + unit.storage_data.put_data({"value": [42]}, [0]) + transfer = object.__new__(NixlPayloadTransfer) + transfer._pending_puts = {} + transfer._pending_gets = {} + transfer._runtime = MagicMock() + unit.payload_transfer = transfer + frames = encode({"value": [42]}) + payload = bytearray(calc_packed_size(frames)) + pack_into(payload, frames) + descriptor = PayloadDescriptor("payload", len(payload), tuple(memoryview(frame).nbytes for frame in frames)) + received = Future() + received.set_result(memoryview(payload)) + transfer._runtime.receive.return_value = received + sent = Future() + sent.set_result(None) + transfer._runtime.send.return_value = sent + transfer._runtime.prepare_receive.return_value = {} + body = { + "global_indexes": [0], + "fields": ["value"], + "transfer_id": "payload", + "descriptor": descriptor.to_dict(), + } + + ready = self.dispatch(worker, ZMQRequestType[f"{operation.upper()}_DATA_PREPARE"], body) + assert ready.request_type == ZMQRequestType[f"{operation.upper()}_DATA_READY"] + assert unit._requests_arrived == 1 + assert dict(monitor.success_counts) == {} + + reply = self.dispatch( + worker, + ZMQRequestType[f"{operation.upper()}_DATA_{finish}"], + { + "transfer_id": "payload", + "receiver_endpoint": {"transport": "nixl-ucx", "data": {"agent_name": "client"}}, + "receive_token": { + "data": {"agent_name": "client", "agent_metadata": b"meta", "frame_remote_descs": b"desc"} + }, + }, + ) + assert reply.request_type == ZMQRequestType[f"{operation.upper()}_DATA_RESPONSE"] + assert unit._requests_arrived == 1 + assert unit._arrivals_by_op == {f"{operation.upper()}_DATA": 1} + assert dict(monitor.success_counts) == ({f"{operation.upper()}_DATA": 1} if finish == "COMMIT" else {}) + + def test_nixl_prepare_key_miss_is_an_arrival_without_measured_processing(self, worker): + from transfer_queue.storage.payload_transfer.nixl import NixlPayloadTransfer + from transfer_queue.storage.simple_storage import KEY_NOT_FOUND_MARKER + + unit, monitor = worker + unit.storage_data.put_data({"value": [42]}, [0]) + unit.storage_data.clear([0]) + transfer = object.__new__(NixlPayloadTransfer) + transfer._pending_gets = {} + unit.payload_transfer = transfer + + reply = self.dispatch( + worker, + ZMQRequestType.GET_DATA_PREPARE, + {"transfer_id": "missing", "global_indexes": [0], "fields": ["value"]}, + ) + assert reply.request_type == ZMQRequestType.PUT_GET_ERROR + assert KEY_NOT_FOUND_MARKER in reply.body["message"] + assert unit._requests_arrived == 1 + assert unit._arrivals_by_op == {"GET_DATA": 1} + assert dict(monitor.success_counts) == {} diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index 99ceb4ec..a02a9b9c 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -28,8 +28,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import gc import pickle import time +import weakref from pathlib import Path from uuid import uuid4 @@ -47,6 +49,7 @@ StorageKeyNotFoundError, _SSDValueRef, ) +from transfer_queue.utils.serial_utils import calc_packed_size, decode, encode, pack_into, unpack_from from transfer_queue.utils.zmq_utils import ( STORAGE_MANAGER_IDENTITY_PREFIX, ZMQMessage, @@ -525,6 +528,7 @@ def test_hybrid_storage_routes_each_sample_by_payload_size_and_clears(tmp_path): assert storage.ssd_active_values == 1 assert storage.ssd_active_bytes == 100 assert storage.ssd_fallback_values_total == 1 + assert storage.get_data(["unsupported"], [1])["unsupported"][0] is large_sparse storage.clear([1]) assert storage.active_key_count == 1 @@ -544,6 +548,69 @@ def test_hybrid_storage_routes_each_sample_by_payload_size_and_clears(tmp_path): assert not (tmp_path / "transfer_queue_ssd_offload" / "test-run").exists() +@pytest.mark.parametrize("kind", ["tensor", "numpy"]) +@pytest.mark.parametrize("layout", ["packed", "frame-native"]) +@pytest.mark.parametrize("nested", [False, True]) +def test_hybrid_small_values_release_shared_receive_buffer(tmp_path, kind, layout, nested): + storage = HybridStorageUnitData(None, str(tmp_path), "test-run", "test-unit", threshold_bytes=1024) + large = np.arange(1024, dtype=np.float32).reshape(2, 512) + small = np.arange(4, dtype=np.float32).reshape(2, 2) + data = {"large": large, "small": small} + if kind == "tensor": + data = {field: torch.from_numpy(value) for field, value in data.items()} + if nested: + data["small"] = [{"value": [value]} for value in data["small"]] + encoded = encode(data) + if layout == "packed": + owner = np.empty(calc_packed_size(encoded), dtype=np.uint8) + pack_into(owner, encoded) + frames = unpack_from(owner) + else: + owner = np.concatenate([np.frombuffer(frame, dtype=np.uint8) for frame in encoded]) + frames = [] + offset = 0 + for frame in encoded: + size = memoryview(frame).nbytes + frames.append(memoryview(owner)[offset : offset + size]) + offset += size + owner_ref = weakref.ref(owner) + try: + storage.put_data(decode(frames), [0, 1]) + del frames, owner + gc.collect() + + assert storage.ssd_active_values == 2 + assert owner_ref() is None + result = storage.get_data(["large", "small"], [0, 1]) + for field in data: + for position, value in enumerate(result[field]): + expected = data[field][position] + if nested and field == "small": + value = value["value"][0] + expected = expected["value"][0] + if kind == "tensor": + torch.testing.assert_close(value, expected) + else: + np.testing.assert_array_equal(value, expected) + storage.clear([0, 1]) + assert storage.ssd_active_values == 0 + finally: + storage.close() + + +def test_hybrid_encoding_fallback_keeps_large_tensor_without_cloning(tmp_path, monkeypatch): + storage = HybridStorageUnitData(None, str(tmp_path), "test-run", "test-unit", threshold_bytes=64) + value = torch.arange(128, dtype=torch.float32) + monkeypatch.setattr(storage, "_sample_from_value", lambda sample: None) + try: + storage.put_data({"value": [value]}, [0]) + assert storage.get_data(["value"], [0])["value"][0] is value + assert storage.ssd_fallback_values_total == 1 + assert storage.ssd_active_values == 0 + finally: + storage.close() + + def test_hybrid_storage_cleans_completed_ssd_writes_after_batch_failure(tmp_path, monkeypatch): """A failed SSD batch removes files written by its successful tasks.""" storage = HybridStorageUnitData( diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py index 59f82edc..fb322edf 100644 --- a/tests/test_storage_request_retry.py +++ b/tests/test_storage_request_retry.py @@ -27,9 +27,14 @@ from transfer_queue.metadata import BatchMeta from transfer_queue.storage.managers import simple_storage_manager as ssm from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager, StorageUnitTimeout +from transfer_queue.storage.payload_transfer import ( + PayloadTransferError, + PayloadTransferTimeout, + create_payload_transfer, +) from transfer_queue.utils import common from transfer_queue.utils.enum_utils import Role -from transfer_queue.utils.zmq_utils import ZMQServerInfo +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo def _manager(with_unit: bool = True) -> AsyncSimpleStorageManager: @@ -89,6 +94,147 @@ async def test_lost_request_recovers_without_diagnosis(): diagnose.assert_not_awaited() +class _ReplySocket: + def __init__(self, reply, timeout_on=None): + self.reply = reply + self.timeout_on = timeout_on + + async def send_multipart(self, *args, **kwargs): + if self.timeout_on == "send": + raise zmq.Again() + + async def recv_multipart(self, *args, **kwargs): + if self.timeout_on == "receive": + raise zmq.Again() + return self.reply.serialize() + + +async def _payload_request(manager, operation, socket): + if operation == "put": + return await manager._put_to_single_storage_unit.__wrapped__( + manager, [0], {"value": [42]}, "unit_a", socket=socket + ) + return await manager._get_from_single_storage_unit.__wrapped__(manager, [0], ["value"], "unit_a", socket=socket) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["put", "get"]) +@pytest.mark.parametrize("timeout_on", ["send", "receive"]) +async def test_zmq_timeout_retry_recovers_without_error_log(operation, timeout_on, caplog): + manager = _manager() + manager.payload_transfer = create_payload_transfer(None) + response = ZMQMessage.create( + request_type=ZMQRequestType[f"{operation.upper()}_DATA_RESPONSE"], + sender_id="unit_a", + body={"data": {"value": [42]}}, + ) + sockets = iter([_ReplySocket(None, timeout_on), _ReplySocket(response)]) + + async def attempt(): + return await _payload_request(manager, operation, next(sockets)) + + with patch.object(manager, "_diagnose_storage_unit") as diagnose, caplog.at_level(logging.WARNING): + result = await manager._request_with_retry(operation, "unit_a", "samples=1", attempt, max_attempts=2) + + assert result == (None if operation == "put" else (["value"], {"value": [42]})) + diagnose.assert_not_awaited() + assert [record for record in caplog.records if record.levelno >= logging.ERROR] == [] + assert sum("retry 2/2" in record.getMessage() for record in caplog.records) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["put", "get"]) +@pytest.mark.parametrize("timeout_on", ["send", "receive"]) +async def test_zmq_exhausted_timeout_logs_once_and_preserves_cause(operation, timeout_on, caplog): + manager = _manager() + manager.payload_transfer = create_payload_transfer(None) + attempts = 0 + + async def attempt(): + nonlocal attempts + attempts += 1 + return await _payload_request(manager, operation, _ReplySocket(None, timeout_on)) + + with ( + patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis") as diagnose, + caplog.at_level(logging.WARNING), + pytest.raises(StorageUnitTimeout) as raised, + ): + await manager._request_with_retry(operation, "unit_a", "samples=1", attempt, max_attempts=2) + + assert attempts == 2 + diagnose.assert_awaited_once() + cause = raised.value.__cause__ + assert isinstance(cause, PayloadTransferTimeout) + assert isinstance(cause.__cause__, zmq.Again) + assert (cause.wire_bytes > 0) == (operation == "put") + errors = [record for record in caplog.records if record.levelno >= logging.ERROR] + assert len(errors) == 1 + assert "failed after 2 attempts" in errors[0].getMessage() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["put", "get"]) +async def test_zmq_peer_error_is_logged_once_without_retry(operation, caplog): + manager = _manager() + manager.payload_transfer = create_payload_transfer(None) + response = ZMQMessage.create( + request_type=ZMQRequestType[f"{operation.upper()}_ERROR"], + sender_id="unit_a", + body={"message": "disk full"}, + ) + attempts = 0 + + async def attempt(): + nonlocal attempts + attempts += 1 + return await _payload_request(manager, operation, _ReplySocket(response)) + + with caplog.at_level(logging.ERROR), pytest.raises(RuntimeError, match="disk full") as raised: + await manager._request_with_retry(operation, "unit_a", "samples=1", attempt, max_attempts=2) + + assert attempts == 1 + assert len([record for record in caplog.records if record.levelno >= logging.ERROR]) == 1 + assert str(raised.value).count("Error in put to storage unit") == (1 if operation == "put" else 0) + assert str(raised.value).count("Error getting data from storage unit") == (1 if operation == "get" else 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["put", "get"]) +async def test_zmq_socket_error_keeps_runtime_error_contract(operation, caplog): + manager = _manager() + manager.payload_transfer = create_payload_transfer(None) + error = zmq.ZMQError("connection failed") + socket = MagicMock() + socket.send_multipart = AsyncMock(side_effect=error) + socket.recv_multipart = AsyncMock() + + with caplog.at_level(logging.ERROR), pytest.raises(RuntimeError, match="connection failed") as raised: + await _payload_request(manager, operation, socket) + + assert raised.value.__cause__ is error + socket.send_multipart.assert_awaited_once() + socket.recv_multipart.assert_not_awaited() + assert len([record for record in caplog.records if record.levelno >= logging.ERROR]) == 1 + + +@pytest.mark.asyncio +async def test_get_payload_error_preserves_type_and_cause(caplog): + manager = _manager() + cause = zmq.ZMQError("connection failed") + error = PayloadTransferError("transfer failed") + error.__cause__ = cause + manager.payload_transfer = MagicMock() + manager.payload_transfer.get = AsyncMock(side_effect=error) + + with caplog.at_level(logging.ERROR), pytest.raises(PayloadTransferError) as raised: + await _payload_request(manager, "get", None) + + assert raised.value is error + assert raised.value.__cause__ is cause + assert len([record for record in caplog.records if record.levelno >= logging.ERROR]) == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize( "error, max_attempts, expected_attempts", diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 5b57d94f..b7bd6204 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -32,7 +32,11 @@ from transfer_queue.metadata import BatchMeta, extract_field_schema from transfer_queue.storage.managers.base import StorageManager, StorageManagerFactory -from transfer_queue.storage.payload_transfer import PayloadTransferTimeout, create_payload_transfer +from transfer_queue.storage.payload_transfer import ( + PayloadTransferError, + PayloadTransferTimeout, + create_payload_transfer, +) from transfer_queue.storage.simple_storage import KEY_NOT_FOUND_MARKER, StorageKeyNotFoundError from transfer_queue.utils.common import log_heavy_operation from transfer_queue.utils.logging_utils import get_logger @@ -671,7 +675,15 @@ async def _get_from_single_storage_unit( except Exception as exc: if KEY_NOT_FOUND_MARKER in str(exc): raise StorageKeyNotFoundError(str(exc)) from exc - raise + logger.error( + f"[{self.storage_manager_id}]: Unexpected error from storage unit " + f"{target_storage_unit}: {type(exc).__name__}: {exc}" + ) + if isinstance(exc, PayloadTransferError): + raise + raise RuntimeError( + f"Error getting data from storage unit {target_storage_unit}: {type(exc).__name__}: {exc}" + ) from exc return fields, data async def clear_data(self, metadata: BatchMeta) -> None: diff --git a/transfer_queue/storage/payload_transfer/base.py b/transfer_queue/storage/payload_transfer/base.py index c3eb4730..c184b19d 100644 --- a/transfer_queue/storage/payload_transfer/base.py +++ b/transfer_queue/storage/payload_transfer/base.py @@ -24,6 +24,9 @@ from transfer_queue.utils.zmq_utils import ZMQMessage +KEY_NOT_FOUND_MARKER = "TQKeyNotFound" + + class PayloadTransferError(RuntimeError): """A payload transfer could not be completed safely.""" diff --git a/transfer_queue/storage/payload_transfer/zmq.py b/transfer_queue/storage/payload_transfer/zmq.py index 3a18d6d3..ceaadc3e 100644 --- a/transfer_queue/storage/payload_transfer/zmq.py +++ b/transfer_queue/storage/payload_transfer/zmq.py @@ -18,20 +18,18 @@ from __future__ import annotations import os -from typing import Any, Callable, NoReturn +from typing import Any, Callable import zmq import zmq.asyncio -from transfer_queue.storage.payload_transfer.base import PayloadTransfer, PayloadTransferTimeout +from transfer_queue.storage.payload_transfer.base import KEY_NOT_FOUND_MARKER, PayloadTransfer, PayloadTransferTimeout from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, frame_nbytes logger = get_logger(__name__) TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8)) -TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT", 200)) -KEY_NOT_FOUND_MARKER = "TQKeyNotFound" class ZmqPayloadTransfer(PayloadTransfer): @@ -64,12 +62,7 @@ async def put( ) return wire_bytes except zmq.error.Again as exc: - self._raise_timeout(sender_id, target_id, "put", exc, wire_bytes=wire_bytes) - except Exception as exc: - logger.error( - f"[{sender_id}]: Unexpected error during put to storage unit {target_id}: {type(exc).__name__}: {exc}" - ) - raise RuntimeError(f"Error in put to storage unit {target_id}: {type(exc).__name__}: {exc}") from exc + raise PayloadTransferTimeout(f"no answer from storage unit {target_id}", wire_bytes=wire_bytes) from exc async def get( self, @@ -80,8 +73,6 @@ async def get( global_indexes: list[int], fields: list[str], ) -> dict[str, Any]: - from transfer_queue.storage.simple_storage import StorageKeyNotFoundError - request = ZMQMessage.create( request_type=ZMQRequestType.GET_DATA, sender_id=sender_id, @@ -93,18 +84,10 @@ async def get( response = ZMQMessage.deserialize(await control_socket.recv_multipart(copy=False)) if response.request_type != ZMQRequestType.GET_DATA_RESPONSE: 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}") + raise RuntimeError(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: - raise - except Exception as exc: - logger.error(f"[{sender_id}]: Unexpected error from storage unit {target_id}: {type(exc).__name__}: {exc}") - raise RuntimeError( - f"Error getting data from storage unit {target_id}: {type(exc).__name__}: {exc}" - ) from exc + raise PayloadTransferTimeout(f"no answer from storage unit {target_id}") from exc def handle_request( self, @@ -164,25 +147,3 @@ def handle_request( @staticmethod def _response(request_type: ZMQRequestType, sender_id: str, body: dict[str, Any] | None = None) -> ZMQMessage: return ZMQMessage.create(request_type=request_type, sender_id=sender_id, body=body or {}) - - @staticmethod - def _raise_timeout( - sender_id: str, target_id: str, operation: str, error: Exception, *, wire_bytes: int = 0 - ) -> NoReturn: - timeout = TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT - if operation == "put": - logger.error( - f"[{sender_id}]: ZMQ recv timeout ({timeout}s) during put to storage unit {target_id}. " - "The storage unit may be overloaded or crashed." - ) - raise PayloadTransferTimeout( - f"ZMQ recv timeout ({timeout}s) during put to storage unit {target_id}", wire_bytes=wire_bytes - ) from error - - logger.error( - f"[{sender_id}]: ZMQ recv timeout ({timeout}s) from storage unit {target_id}. " - "The storage unit may be overloaded or crashed." - ) - raise PayloadTransferTimeout( - f"ZMQ recv timeout ({timeout}s) from storage unit {target_id}", wire_bytes=wire_bytes - ) from error diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 85bc9fcc..b7d7e01b 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -48,6 +48,7 @@ import zmq from transfer_queue.storage.payload_transfer import PayloadTransfer, create_payload_transfer +from transfer_queue.storage.payload_transfer.base import KEY_NOT_FOUND_MARKER from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads, log_heavy_operation from transfer_queue.utils.enum_utils import Role from transfer_queue.utils.logging_utils import get_logger @@ -72,14 +73,18 @@ def _storage_arrival_op_type(operation: ZMQRequestType) -> str | None: - """Map a tracked logical storage operation's first message to its metric label.""" + """Count decoded operations while treating NIXL continuations as the same request.""" + if operation in ( + ZMQRequestType.PUT_DATA_COMMIT, + ZMQRequestType.GET_DATA_COMMIT, + ZMQRequestType.PUT_DATA_CANCEL, + ZMQRequestType.GET_DATA_CANCEL, + ): + return None return { - ZMQRequestType.PUT_DATA: "PUT_DATA", - ZMQRequestType.GET_DATA: "GET_DATA", - ZMQRequestType.CLEAR_DATA: "CLEAR_DATA", ZMQRequestType.PUT_DATA_PREPARE: "PUT_DATA", ZMQRequestType.GET_DATA_PREPARE: "GET_DATA", - }.get(operation) + }.get(operation, operation.name) TQ_STORAGE_POLLER_TIMEOUT = int(os.environ.get("TQ_STORAGE_POLLER_TIMEOUT", 5)) # in seconds @@ -113,9 +118,6 @@ class _SSDValueRef: shape: tuple[int, ...] | None = None -# Marks a GET_ERROR reply as "the key is gone" so the caller can tell it apart from a real fault. -KEY_NOT_FOUND_MARKER = "TQKeyNotFound" - # Accept-queue depth for the client-facing ROUTER. A full queue loses connections silently. TQ_STORAGE_ZMQ_BACKLOG = int(os.environ.get("TQ_STORAGE_ZMQ_BACKLOG", 4096)) @@ -540,6 +542,15 @@ def _prepare_field_values( ssd_samples.append(encoded) prepared_values.append(None) else: + # Small views must not retain a shared receive buffer after SSD offload. + # Encoding failures may be large and retain their existing fallback behavior. + if encoded is not None: + if encoded.codec == "pickle": + value = pickle.loads(encoded.payload) + elif isinstance(value, torch.Tensor): + value = value.clone() + elif isinstance(value, np.ndarray): + value = value.copy() prepared_values.append(value) if encoded is None: fallback_values += 1