| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
[feat] support GDR in mooncake backend (#131) Add GDR support for mooncake backend Signed-off-by: xupinjie <xupinjie321@outlook.com> | 24 天前 | |
[chore] refactor: check docstrings only for public API symbols (#134) as title. Signed-off-by: ji-huazhong <hzji210@gmail.com> | 1 个月前 | |
[fix,refactor] Merge update_data_status thread/socket into request_handle to eliminate concurrency conflicts (#116) Previously, _update_data_status and _process_request ran on two separate threads listening on two different ROUTER sockets. When used in async frameworks, NOTIFY_DATA_UPDATE and metadata requests (GET_META, CLEAR_META, etc.) could interleave and cause data processing conflicts on shared partition state. This change consolidates everything onto request_handle_socket and the single _process_request thread, so all controller requests are serialized naturally: - Drop data_status_update_socket and its port; ZMQServerInfo.ports now exposes only handshake_socket and request_handle_socket. - Remove _start_process_update_data_status and _update_data_status; move the NOTIFY_DATA_UPDATE branch into _process_request. - Remove data_status_lock logics. - Update storage manager to send NOTIFY_DATA_UPDATE to request_handle_socket. - Clean up data_status_update_socket port entries in test fixtures. Note: this is a breaking change to the controller-storage wire protocol (port name removed). All storage managers must be upgraded together. --------- Signed-off-by: yxstev <zhangyixiang9@huawei.com> | 2 个月前 | |
[fix] Prevent stale reads during clear by marking indexes before storage deletion (#141) Introduce a three-step clear protocol to close the consistency window where storage data is already deleted but the controller still considers the indexes readable: 1. MARK_CLEARING: zero production_status for the target indexes so consumers cannot fetch them while deletion is in progress 2. Storage clear: physically remove data from storage units 3. CLEAR_META / CLEAR_PARTITION: release indexes back to the reusable pool --------- Signed-off-by: yxstev <zhangyixiang9@huawei.com> | 8 天前 | |
[fix] Make the controller survive bad ZMQ messages and repair the pickle fallback (#143) ## Motivation A long-running RL job hit an occasional failure that took the whole run down: Exception in thread TransferQueueControllerProcessRequestThread: Traceback (most recent call last): ... File "transfer_queue/controller.py", line 1771, in _process_request request_msg = ZMQMessage.deserialize(serialized_msg) File "transfer_queue/utils/zmq_utils.py", line 182, in deserialize result = decode(frames) File "transfer_queue/utils/serial_utils.py", line 262, in decode result = self.decoder.decode(bufs[0]) msgspec.DecodeError: Input data was truncated _process_request had no exception handling, so this one message killed TransferQueueControllerProcessRequestThread permanently. The controller stayed alive as an actor but stopped answering every request from that point on, and the job hung. While investigating I found more defects on the same path. ## What was wrong **1. One bad message killed the request loop.** _process_request ran recv_multipart → deserialize → dispatch → send_multipart with no guard anywhere. Note that _wait_connection in the same file already logs and continues; the request loop did not. **2. Dropped requests hang the caller forever.** Clients issue requests through with_controller_socket with no RCVTIMEO and then block on await socket.recv_multipart(). Any controller path that fails to reply — an undecodable request, an unhandled request type, a handler that raises — leaves that caller waiting for the rest of the run. **3. decode() could never detect the pickle fallback marker.** It located the marker with frames[0] == _PICKLE_FALLBACK_SENTINEL. Every receiver calls recv_multipart(copy=False) and therefore holds zmq.Frame objects, and zmq.Frame defines no __eq__ against bytes, so the comparison was *always* False and the marker frame was handed to msgpack. **4. encode() almost never reached that fallback anyway.** It caught only TypeError/ValueError, but msgspec reports unrepresentable values with OverflowError (an ArithmeticError), RecursionError (a RuntimeError) and its own MsgspecError — none of which derive from ValueError. Those escaped to the caller instead of degrading to pickle, even though pickle handles all of them (arbitrary-precision ints, self-referential containers). **5. Unhandled request types replied with a stale response.** Every branch in the dispatch chain is an if/elif with no else, so an unknown request_type fell through to send_multipart([identity, *response_msg.serialize()]) still holding the *previous* iteration's response_msg, sending an unrelated requester someone else's answer. ## What this changes - _process_request is split into a thin supervisor plus _process_request_loop; anything that escapes the per-request handling restarts the loop, so the thread survives for the controller's lifetime. The ROUTER socket stays bound across a restart. - The dispatch chain moves into _handle_request(), which returns the response or None. Every request now gets **exactly one reply**: - the handler's normal response; - a new generic ZMQRequestType.REQUEST_ERROR (built by _make_error_response()) when the handler raises — the reply body carries {request_type} failed: {error}, so the caller raises immediately instead of hanging; - REQUEST_ERROR ("no handler for request_type ...") when no branch matches, replacing the stale-response replay; - REQUEST_ERROR for undecodable requests too — the ROUTER identity frame is prepended by the transport and survives payload corruption, so the reply still reaches the requester. No client change is needed: every client call site already treats an unexpected response type as RuntimeError(body["message"]). - ZMQMessage.deserialize self-checks the frame layout and raises ZMQMessageDecodeError carrying frame count and per-frame sizes. It lives in deserialize, so controller, storage, client and manager all get the diagnostics. - decode() compares buffer *contents* via _is_pickle_fallback(). Size is tested first, so the msgpack path stays copy-free. - encode() catches the errors msgspec actually raises, via a named _ENCODE_FALLBACK_ERRORS tuple documenting why none of them are ValueError. Falling back now logs at warning level — it is a performance degradation path and should be visible at default log levels. - mypy now runs over the whole package in pre-commit (pass_filenames: false); with follow_imports = "skip", passing only changed files produced false errors. ## Scope This does **not** fix the root cause of the truncation. The evidence says the ZMQ multipart frame boundaries shift, so frame 0 stops being the msgpack header — decoding an empty buffer produces exactly Input data was truncated, and empty frames are routine here since empty tensors serialize to zero-length frames. What this PR does is stop that from taking the job down, and emit the frame layout needed to confirm it. The frame sizes make the diagnosis immediate: healthy message: num_frames=3, frame_sizes=[184, 512, 512] empty leading frame: leading frame is empty; num_frames=4, frame_sizes=[0, 184, 512, 512] boundary on a tensor frame: trailing characters (byte 1); num_frames=2, frame_sizes=[512, 512] truncated header: Input data was truncated; num_frames=3, frame_sizes=[12, 512, 512] A healthy message is one small header frame followed by large buffers. A leading 0, a missing small header, or an implausibly small header all point at shifted boundaries. One caveat to note for rollout: the pickle fallback only works end to end when *both* peers run this version — an old receiver (recv_multipart(copy=False) + == marker check) cannot recognise fallback frames from a new sender, and new senders produce them in more situations than before. Upgrade controller, storage units and clients together. ## Tests TestPickleFallback in tests/test_serial_utils_on_cpu.py (13 cases): oversized ints and self-referential containers degrade instead of raising, round-trips through zmq.Frame, marker detection across bytes/bytearray/memoryview/zmq.Frame, four near-miss negatives, and an end-to-end round trip over a real ROUTER with recv_multipart(copy=False) — the transport that hid the bug. TestTransferQueueControllerBadRequests in tests/test_controller.py (3 cases): - an empty leading frame (shifted multipart boundary) gets a REQUEST_ERROR reply and the controller keeps serving; - an unhandled PUT_DATA gets a REQUEST_ERROR reply instead of replaying the previous iteration's stale response; - a GET_META body missing its keys makes the handler raise KeyError, and the requester gets a REQUEST_ERROR ("GET_META failed: KeyError ...") while the loop keeps answering subsequent requests. All three were confirmed to **fail** against the pre-fix code and pass after. Results: 115 passed across the two serialization suites, 25 passed across tests/test_controller.py (including the 3 new cases), and 76 passed across tests/e2e/. pre-commit run --all-files is green (ruff, ruff-format, mypy). ## Unrelated CI fix bundled here The MooncakeStore e2e job started failing on libcudart.so.12: cannot open shared object file — upstream mooncake-transfer-engine-non-cuda 0.3.12 ships a mooncake_master binary linked against the CUDA 12 runtime (verified by diffing the DT_NEEDED entries of the 0.3.11.post1 and 0.3.12.post1 wheels; 0.3.11.post1 has no CUDA dependency). The two workflows that install it now pin mooncake-transfer-engine-non-cuda<0.3.12 until upstream fixes the wheel. --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> | 6 天前 | |
[fix] shared production_status tensor across data partitions (#127) Create a fresh production_status tensor for each DataPartitionStatus instance. ## Root Cause production_status was initialized as a tensor dataclass default. That tensor was created at class definition time and could be shared by multiple DataPartitionStatus instances. When one partition reused a released global index and marked it ready, another cleared partition could observe the same ready bit. ## Impact A cleared partition could incorrectly return stale ready metadata. In the KV path, because storage keys are generated as global_index@field, this could cause reads from one partition to include data written by another partition. ## PoC python def client_api_poc(): import torch import transfer_queue as tq from tensordict import TensorDict client = tq.get_client() field = "x" # p1 writes then clears one sample. # p1_meta.global_indexes == [0] p1_meta = client.put( data=TensorDict({field: torch.tensor([[1]])}, batch_size=[1]), partition_id="p1", ) client.clear_samples(p1_meta) # p2 may reuse the released global_index 0. # Before the fix, this updates the shared production_status tensor # and makes the already-cleared p1 look ready again. client.put( data=TensorDict({field: torch.tensor([[2]])}, batch_size=[1]), partition_id="p2", ) leaked = client.get_meta( data_fields=[field], batch_size=1, partition_id="p1", mode="fetch", task_name="repro", ) # Before fix: # leaked.size == 1 # leaked.global_indexes == [0] # leaked.partition_ids == ["p1"] # leaked.field_names == [] # leaked.is_ready == True # leaked.production_status.tolist() == [1] # # After fix: # leaked.size == 0 # leaked.global_indexes == [] # leaked.partition_ids == [] # leaked.field_names == [] # leaked.is_ready == False # leaked.production_status.tolist() == [] assert leaked.size == 0 ## Tests - python -m pytest tests/test_controller_data_partitions.py -q - python -m pytest tests/e2e/test_kv_interface_e2e.py::TestKVClearE2E::test_kv_clear_does_not_leak_reused_index_across_partitions -q | 1 个月前 | |
[BREAKING][fix] Use jagged tensor as default tensor type (#92) ## Background Previously, TransferQueue would try torch.stack() first when merging per-sample tensors into a batched tensordict for user retrieval. As a result, tensors with uniform shapes were returned as regular dense tensors, while jagged data fell back to nested tensors. This inconsistency forced downstream code to handle two distinct data types (torch.Tensor vs. nested tensor), adding unnecessary branching logic. ## Changes This PR changes the default aggregation strategy so that all tensor fields are returned as nested tensors by default, eliminating the torch.stack() fast-path. Specifically: 1. KVStorageManager._merge_tensors_to_tensordict: Removed the torch.stack(chunk) fallback. The new chain is as_nested_tensor(jagged) → nested_tensor(strided) → NonTensorStack. 2. AsyncSimpleStorageManager._pack_field_values: Removed the torch.stack(values) fast-path for uniform-shape tensors. The new in is as_nested_tensor(jagged) → as_nested_tensor(strided) → NonTensorStack, consistent with the KV backend. 3. Unified strided fallback: Added the missing strided layout fallback to KVStorageManager, ensuring both backends behave identically when jagged layout fails (e.g., for zero-dim tensors). 4. Docstring & comment cleanup: Updated all outdated docstrings and comments that referenced the old torch.stack-first behavior. ## Test updates - Adapted test_async_simple_storage_manager.py, test_kv_storage_manager.py, and e2e tests to accept nested tensors as the default return type. - Reworked the test_kv_storage_manager.py fixture to use realistic variable-length fields (input_ids, prompt_ids, response_ids, response_mask) aligned the single_controller_demo.py schema, replacing the oversimplified text/label/mask example. - Replaced all torch.equal(dense, nested) assertions with safe per-component comparisons (unbind(0) + torch.equal) to accommodate the new nested-tensor contract --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> | 2 个月前 | |
[fix] Fix BatchMeta.union semantics (#95) ## Problem PR https://gitcode.com/Ascend/TransferQueue/pull/28 incorrectly rewrote BatchMeta.union to behave like concat with global_index deduplication. - **Original semantics**: merge fields for samples with identical global_indexes (row-aligned, column-expanded). - **Broken semantics**: append rows from other whose global_indexes are not present in self. This broke the design boundary between union (same rows, merge columns) and concat (same columns, append rows). ## Changes ### 1. Restore BatchMeta.union (transfer_queue/metadata.py) - Validate that both batches have the same size, global_indexes, and partition_ids. - Merge field_schema with other overriding self on name conflicts. - Merge production_status conservatively via np.bitwise_and (both sides must report ready). - Merge extra_info, custom_meta, and _custom_backend_meta per sample. ### 2. Update tutorial (tutorial/03_metadata_concepts.py) - Example now uses overlapping fields (attention_mask present in both batches) to demonstrate the override behavior. - Corrected comments to clearly distinguish concat (append rows) vs union (merge columns). ### 3. Update unit tests (tests/test_metadata.py) --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> | 2 个月前 | |
[fix] Skip storage capacity/utilization gauges when capacity is None (#138) ### Background Follow-up of #130, which made total_storage_size optional so that a storage unit can be configured with **unlimited** capacity. ### Problem When a storage unit is created with total_storage_size=None (unlimited), SimpleStorageUnit._handle_get_metrics reports capacity=None in its /metrics payload. The metrics exporter in transfer_queue/metrics.py defines two prometheus_client.Gauge objects that consume this value: - self.storage_capacity = Gauge("tq_storage_capacity_total", ...) — around L164 - self.storage_utilization = Gauge("tq_storage_utilization_ratio", ...) — around L170 Inside TQMetricsExporter._collect_storage_metrics, both gauges are updated unconditionally: python capacity = metrics.get("capacity", 0) active = metrics.get("active_keys", 0) self.storage_capacity.labels(storage_unit_id=label).set(capacity) ... self.storage_utilization.labels(storage_unit_id=label).set( active / capacity if capacity > 0 else 0.0 ) When capacity is None, Gauge.set() internally does float(value), and float(None) raises TypeError. Because the whole block is wrapped in try/except, this is caught by the surrounding handler and turned into: Failed to collect metrics from storage unit <su_id>: float() argument must be a string or a real number, not 'NoneType' This warning is emitted **every TQ_METRICS_COLLECT_INTERVAL** for **every** unlimited-capacity storage unit, flooding the logs and hiding real issues. ### Fix In TQMetricsExporter._collect_storage_metrics, treat capacity is None as "unlimited" and skip only the storage_capacity and storage_utilization gauges for that storage unit. All other metrics (storage_active_keys, storage_memory_rss, per-op request stats such as count / latency avg / p50 / p99) are still reported normally, so dashboards for unlimited-capacity units keep working — they simply no longer publish a (meaningless) capacity/utilization value. ### Related - #130 Make total_storage_size optional to support unlimited storage capacity Signed-off-by: nexhu <nexhu@tencent.com> Co-authored-by: nexhu <nexhu@tencent.com> | 26 天前 | |
[feat] support GDR in mooncake backend (#131) Add GDR support for mooncake backend Signed-off-by: xupinjie <xupinjie321@outlook.com> | 24 天前 | |
[fix,refactor] Merge update_data_status thread/socket into request_handle to eliminate concurrency conflicts (#116) Previously, _update_data_status and _process_request ran on two separate threads listening on two different ROUTER sockets. When used in async frameworks, NOTIFY_DATA_UPDATE and metadata requests (GET_META, CLEAR_META, etc.) could interleave and cause data processing conflicts on shared partition state. This change consolidates everything onto request_handle_socket and the single _process_request thread, so all controller requests are serialized naturally: - Drop data_status_update_socket and its port; ZMQServerInfo.ports now exposes only handshake_socket and request_handle_socket. - Remove _start_process_update_data_status and _update_data_status; move the NOTIFY_DATA_UPDATE branch into _process_request. - Remove data_status_lock logics. - Update storage manager to send NOTIFY_DATA_UPDATE to request_handle_socket. - Clean up data_status_update_socket port entries in test fixtures. Note: this is a breaking change to the controller-storage wire protocol (port name removed). All storage managers must be upgraded together. --------- Signed-off-by: yxstev <zhangyixiang9@huawei.com> | 2 个月前 | |
[feat] Add SeqlenBalancedSampler and enhance StreamingDataset support (#70) ## 🎯 Summary This PR introduces the SeqlenBalancedSampler to optimize sequence length distribution across Data Parallel (DP) ranks during GRPO training. It also enhances StreamingDataset with proper streaming mode support and refactors the controller's polling mechanism to improve efficiency when data is insufficient. ## ✨ Key Features & Enhancements ### 1. SeqlenBalancedSampler (Sequence-Length Balanced Sampling) - **Karmarkar-Karp Algorithm:** Added a new sampler that extends GRPOGroupNSampler. It uses the Karmarkar-Karp largest differencing method to balance sequence lengths (total_lengths) across DP ranks, ensuring that each rank processes approximately the same total token count. - **Group Integrity:** Guarantees that complete prompt groups remain intact across ranks to fulfill pass@k metrics and GRPO advantage normalization requirements. - **Assignment Caching:** Implements state caching (_balanced_cache) so that once global sampling and balancing are computed for a batch, subsequent DP ranks can quickly retrieve their assigned chunks. ### 2. StreamingDataset Improvements - **Finite vs. Infinite Stream:** Introduced the should_check_consumption_status parameter. - False (Default): Operates in an **infinite stream** mode, continuously polling for new data (ideal for online/streaming pipelines). - True: Operates in **finite-dataset** mode, terminating iteration only after all samples in the partition are fully consumed. - **Client Initialization Refactor:** Refactored _create_client() to use init() and get_client() from transfer_queue.interface instead of manually setting up TransferQueueClient. ### 3. Controller Optimizations - **Polling Mode Sampler Cache Lookup:** Updated get_metadata to look up cached sampler states when operating in polling_mode. If dp_rank and batch_index are cached, it immediately returns the data instead of failing or entering redundant wait loops when ready_for_consume_indexes are insufficient. - **Variable-size Batch Support:** Updated the sampler length validation logic to accommodate variable-size batches returned by samplers like SeqlenBalancedSampler. ## 🛠️ Refactoring & Minor Fixes - **Log Level Adjustments:** Downgraded the 1D tensor shape warnings to logger.info() in client.py and metadata.py to reduce unnecessary noise. - **Pre-allocation Scope:** Moved TQ_PRE_ALLOC_SAMPLE_NUM environment variable resolution into local method scopes where appropriate. ## 🧪 Testing - Added comprehensive unit tests for SeqlenBalancedSampler covering initialization, fallback behavior, balanced partitioning logic with mock custom meta, group level integrity, and caching mechanisms. - Added explicit utility tests for the karmarkar_karp and get_seqlen_balanced_partitions functions (TestKarmarkarKarp). --------- Signed-off-by: 宁本哲 <ningbenzhe@xiaohongshu.com> | 4 个月前 | |
[refactor] Provide common serialization tools for KV backends to speed up tensor serial in nested values (#107) ## Problem Multimodal RL puts nested-dict values into TransferQueue (e.g. {"pixel_values": Tensor, "image_grid_thw": Tensor, ...}). The old MooncakeStoreClient only zero-copied plain tensors; anything else — including dicts that contain tensors — got pickled through Mooncake's internal bytes pool, which saturated under concurrent multi-MB GETs and forced a VERL_TQ_MC_GET_RETRIES retry workaround upstream. ## Refactor Treat every value as one opaque payload. Each value is encoded by the existing zero-copy msgpack encoder, the whole batch is packed into one contiguous CPU buffer, registered once, and shipped through Mooncake's RDMA-backed batch_upsert_from / batch_get_into. The pickled bytes path and the retry workaround it required are gone. ## Test - [x] 30B-VL + onethinker, 2×8 GPU, RDMA: 3-step and 10-step runs clean — no retry / allocator failure / AssertionError. <img width="2306" height="578" alt="datapath_perf_3way_v2" src="https://github.com/user-attachments/assets/23894842-24dd-43f4-8a21-3087cf642378" /> Signed-off-by: xupinjie <xupinjie321@outlook.com> | 2 个月前 | |
[fix] Make the controller survive bad ZMQ messages and repair the pickle fallback (#143) ## Motivation A long-running RL job hit an occasional failure that took the whole run down: Exception in thread TransferQueueControllerProcessRequestThread: Traceback (most recent call last): ... File "transfer_queue/controller.py", line 1771, in _process_request request_msg = ZMQMessage.deserialize(serialized_msg) File "transfer_queue/utils/zmq_utils.py", line 182, in deserialize result = decode(frames) File "transfer_queue/utils/serial_utils.py", line 262, in decode result = self.decoder.decode(bufs[0]) msgspec.DecodeError: Input data was truncated _process_request had no exception handling, so this one message killed TransferQueueControllerProcessRequestThread permanently. The controller stayed alive as an actor but stopped answering every request from that point on, and the job hung. While investigating I found more defects on the same path. ## What was wrong **1. One bad message killed the request loop.** _process_request ran recv_multipart → deserialize → dispatch → send_multipart with no guard anywhere. Note that _wait_connection in the same file already logs and continues; the request loop did not. **2. Dropped requests hang the caller forever.** Clients issue requests through with_controller_socket with no RCVTIMEO and then block on await socket.recv_multipart(). Any controller path that fails to reply — an undecodable request, an unhandled request type, a handler that raises — leaves that caller waiting for the rest of the run. **3. decode() could never detect the pickle fallback marker.** It located the marker with frames[0] == _PICKLE_FALLBACK_SENTINEL. Every receiver calls recv_multipart(copy=False) and therefore holds zmq.Frame objects, and zmq.Frame defines no __eq__ against bytes, so the comparison was *always* False and the marker frame was handed to msgpack. **4. encode() almost never reached that fallback anyway.** It caught only TypeError/ValueError, but msgspec reports unrepresentable values with OverflowError (an ArithmeticError), RecursionError (a RuntimeError) and its own MsgspecError — none of which derive from ValueError. Those escaped to the caller instead of degrading to pickle, even though pickle handles all of them (arbitrary-precision ints, self-referential containers). **5. Unhandled request types replied with a stale response.** Every branch in the dispatch chain is an if/elif with no else, so an unknown request_type fell through to send_multipart([identity, *response_msg.serialize()]) still holding the *previous* iteration's response_msg, sending an unrelated requester someone else's answer. ## What this changes - _process_request is split into a thin supervisor plus _process_request_loop; anything that escapes the per-request handling restarts the loop, so the thread survives for the controller's lifetime. The ROUTER socket stays bound across a restart. - The dispatch chain moves into _handle_request(), which returns the response or None. Every request now gets **exactly one reply**: - the handler's normal response; - a new generic ZMQRequestType.REQUEST_ERROR (built by _make_error_response()) when the handler raises — the reply body carries {request_type} failed: {error}, so the caller raises immediately instead of hanging; - REQUEST_ERROR ("no handler for request_type ...") when no branch matches, replacing the stale-response replay; - REQUEST_ERROR for undecodable requests too — the ROUTER identity frame is prepended by the transport and survives payload corruption, so the reply still reaches the requester. No client change is needed: every client call site already treats an unexpected response type as RuntimeError(body["message"]). - ZMQMessage.deserialize self-checks the frame layout and raises ZMQMessageDecodeError carrying frame count and per-frame sizes. It lives in deserialize, so controller, storage, client and manager all get the diagnostics. - decode() compares buffer *contents* via _is_pickle_fallback(). Size is tested first, so the msgpack path stays copy-free. - encode() catches the errors msgspec actually raises, via a named _ENCODE_FALLBACK_ERRORS tuple documenting why none of them are ValueError. Falling back now logs at warning level — it is a performance degradation path and should be visible at default log levels. - mypy now runs over the whole package in pre-commit (pass_filenames: false); with follow_imports = "skip", passing only changed files produced false errors. ## Scope This does **not** fix the root cause of the truncation. The evidence says the ZMQ multipart frame boundaries shift, so frame 0 stops being the msgpack header — decoding an empty buffer produces exactly Input data was truncated, and empty frames are routine here since empty tensors serialize to zero-length frames. What this PR does is stop that from taking the job down, and emit the frame layout needed to confirm it. The frame sizes make the diagnosis immediate: healthy message: num_frames=3, frame_sizes=[184, 512, 512] empty leading frame: leading frame is empty; num_frames=4, frame_sizes=[0, 184, 512, 512] boundary on a tensor frame: trailing characters (byte 1); num_frames=2, frame_sizes=[512, 512] truncated header: Input data was truncated; num_frames=3, frame_sizes=[12, 512, 512] A healthy message is one small header frame followed by large buffers. A leading 0, a missing small header, or an implausibly small header all point at shifted boundaries. One caveat to note for rollout: the pickle fallback only works end to end when *both* peers run this version — an old receiver (recv_multipart(copy=False) + == marker check) cannot recognise fallback frames from a new sender, and new senders produce them in more situations than before. Upgrade controller, storage units and clients together. ## Tests TestPickleFallback in tests/test_serial_utils_on_cpu.py (13 cases): oversized ints and self-referential containers degrade instead of raising, round-trips through zmq.Frame, marker detection across bytes/bytearray/memoryview/zmq.Frame, four near-miss negatives, and an end-to-end round trip over a real ROUTER with recv_multipart(copy=False) — the transport that hid the bug. TestTransferQueueControllerBadRequests in tests/test_controller.py (3 cases): - an empty leading frame (shifted multipart boundary) gets a REQUEST_ERROR reply and the controller keeps serving; - an unhandled PUT_DATA gets a REQUEST_ERROR reply instead of replaying the previous iteration's stale response; - a GET_META body missing its keys makes the handler raise KeyError, and the requester gets a REQUEST_ERROR ("GET_META failed: KeyError ...") while the loop keeps answering subsequent requests. All three were confirmed to **fail** against the pre-fix code and pass after. Results: 115 passed across the two serialization suites, 25 passed across tests/test_controller.py (including the 3 new cases), and 76 passed across tests/e2e/. pre-commit run --all-files is green (ruff, ruff-format, mypy). ## Unrelated CI fix bundled here The MooncakeStore e2e job started failing on libcudart.so.12: cannot open shared object file — upstream mooncake-transfer-engine-non-cuda 0.3.12 ships a mooncake_master binary linked against the CUDA 12 runtime (verified by diffing the DT_NEEDED entries of the 0.3.11.post1 and 0.3.12.post1 wheels; 0.3.11.post1 has no CUDA dependency). The two workflows that install it now pin mooncake-transfer-engine-non-cuda<0.3.12 until upstream fixes the wheel. --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> | 6 天前 | |
[feat] Provide save/load checkpoint interfaces (#124) ## Summary - Add tq.save_checkpoint(checkpoint_dir, *, include_storage, metadata) and tq.load_checkpoint(checkpoint_dir) as top-level public APIs - Controller and each storage unit write state directly to file in-process (only a bool ACK goes through Ray object store), avoiding large payload transmission over Ray - Save order: controller first, then storage units in parallel — guarantees consistency without pausing the data flow (storage unit data is always a superset of what the controller has confirmed) - Atomic save via a .tmp directory that is renamed on success and deleted on failure, ensuring no partial checkpoint is left on disk - Load validates storage unit count before touching any state; restoration is position-based (global_idx % num_units) so storage unit IDs regenerated across restarts are handled correctly Checkpoint layout: checkpoint_dir/ ├── metadata.json # timestamp, storage unit list, user metadata ├── controller_state.pkl # TransferQueueController full state └── storage_units/ ├── su_0_<id>.pkl └── su_1_<id>.pkl ## Test plan - [ ] pytest tests/e2e/test_checkpoint_e2e.py -v - save creates expected files and metadata structure - load restores controller partitions, key mappings, and per-sample tags - load restores storage data and round-trips tensors correctly across multiple partitions - include_storage=False saves only controller state - error cases: uninitialized system, missing directory/metadata, storage unit count mismatch - failed save leaves no partial directory on disk - non-tensor fields (NonTensorStack) and variable-length (jagged) tensor fields survive round-trip --------- Signed-off-by: yxstev <zhangyixiang9@huawei.com> | 1 个月前 | |
[misc] refactor: simplify internal classes naming (#86) Rename internal classes by removing TransferQueue prefix: - TransferQueueStorageManager → StorageManager - TransferQueueStorageManagerFactory → StorageManagerFactory - StorageClientFactory → StorageKVClientFactory - TransferQueueRole → Role - move the factory class to base.py - rename simple_backend.py to simple_storage.py These classes are internal components not exposed as public API, so the shorter names improve readability without causing conflicts. Besides, enable modern type annotation style: Auto convert legacy Optional/Union type annotations to modern X | None syntax. cc @0oshowero0 --------- Signed-off-by: ji-huazhong <hzji210@gmail.com> Co-authored-by: ji-huazhong <hzji210@gmail.com> | 3 个月前 | |
[BREAKING][fix] Use jagged tensor as default tensor type (#92) ## Background Previously, TransferQueue would try torch.stack() first when merging per-sample tensors into a batched tensordict for user retrieval. As a result, tensors with uniform shapes were returned as regular dense tensors, while jagged data fell back to nested tensors. This inconsistency forced downstream code to handle two distinct data types (torch.Tensor vs. nested tensor), adding unnecessary branching logic. ## Changes This PR changes the default aggregation strategy so that all tensor fields are returned as nested tensors by default, eliminating the torch.stack() fast-path. Specifically: 1. KVStorageManager._merge_tensors_to_tensordict: Removed the torch.stack(chunk) fallback. The new chain is as_nested_tensor(jagged) → nested_tensor(strided) → NonTensorStack. 2. AsyncSimpleStorageManager._pack_field_values: Removed the torch.stack(values) fast-path for uniform-shape tensors. The new in is as_nested_tensor(jagged) → as_nested_tensor(strided) → NonTensorStack, consistent with the KV backend. 3. Unified strided fallback: Added the missing strided layout fallback to KVStorageManager, ensuring both backends behave identically when jagged layout fails (e.g., for zero-dim tensors). 4. Docstring & comment cleanup: Updated all outdated docstrings and comments that referenced the old torch.stack-first behavior. ## Test updates - Adapted test_async_simple_storage_manager.py, test_kv_storage_manager.py, and e2e tests to accept nested tensors as the default return type. - Reworked the test_kv_storage_manager.py fixture to use realistic variable-length fields (input_ids, prompt_ids, response_ids, response_mask) aligned the single_controller_demo.py schema, replacing the oversimplified text/label/mask example. - Replaced all torch.equal(dense, nested) assertions with safe per-component comparisons (unbind(0) + torch.equal) to accommodate the new nested-tensor contract --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> | 2 个月前 | |
[perf] Refactor MooncakeStore backend with zero-copy upsert API (#77) 1. Switch to batch_upsert_from & batch_get_into for tensor 2. Switch to upsert_batch & get_batch for non-tensor 3. Use batch_remove API for data clearning 4. Set hard-pin flag during data writting 5. Use multi-thread to optimize data preparation & transfer workflow --------- Signed-off-by: 0oshowero0 <o0shower0o@outlook.com> Signed-off-by: ji-huazhong <hzji210@gmail.com> Co-authored-by: Huazhong <hzji210@gmail.com> | 3 个月前 | |
| 2 个月前 | ||
[optim] Optimize the exception message for Yuanrong when storage strategy does not support a certain type of data (#120) ## Main changes - Modify the exception message to make it more understandable. - Add a solution for the "Cannot retrieve stored data" error to FAQ. - Add two environmental variables DS_D2H_MEMCPY_POLICY and DS_H2D_MEMCPY_POLICY to datasystem workers --------- Signed-off-by: dpj135 <958208521@qq.com> | 1 个月前 | |
[perf, fix] Reuse a long-lived ZMQ context instead of creating one per call (#145) # [perf,fix] Reuse a long-lived ZMQ context pool instead of creating one per call ## Motivation Every client→controller RPC and every SimpleStorage manager→storage-unit request went through the with_zmq_socket decorator, which created a **brand-new zmq.asyncio.Context per call**, opened a socket on it, and then term()-ed the context in the finally block. A ZMQ context owns a native I/O-thread pool and a set of internal signaler file descriptors; standing one up and tearing it down on every request is both wasteful and, under concurrency, unsafe. This surfaced as intermittent Bad file descriptor errors, SIGABRT crashes, and occasional hangs inside ctx.term() on the hot data path. The goal of this change is to make each client own **one long-lived context** that is created once, reused across all of its RPCs, and terminated exactly once at shutdown — while keeping sockets per-request (ZMQ sockets are not thread-safe). ## What needs to be improved 1. **Per-call context churn.** with_zmq_socket did context = zmq.asyncio.Context() on entry and context.term() in finally for **every** decorated call. Repeatedly allocating/destroying contexts recreates libzmq's internal signaler pipe FDs; under concurrent in-flight calls these FDs can collide, yielding Bad file descriptor / SIGABRT. 2. **term() on the event loop can hang.** context.term() blocks until every socket on the context is closed and all pending messages are handled. Running it in the request finally block could stall the asyncio loop if a socket lingered from an interrupted RPC. 3. **No shared I/O-thread pool.** Because each call had its own throwaway context, there was no way to size or share the native I/O-thread pool that actually moves bytes; every request paid context-startup cost. ## What this changes 1. **with_zmq_socket reuses an owner-provided context.** The decorator now takes a required get_context(self) callable and does context = get_context(self) instead of constructing one. It creates and closes only the **per-call DEALER socket** (sock.close(linger=0)); it never creates or terminates a context. The docstring documents the invariant: contexts are thread-safe and event-loop-agnostic, so one shared context is safe even when decorated methods run on different loops/threads, as long as each socket is created and fully used within a single awaited call. 2. **Client owns one long-lived context.** AsyncTransferQueueClient.__init__ creates self.zmq_context = zmq.asyncio.Context(io_threads=…) once, sized by a new simple_storage_zmq_io_threads arg / TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS env var (default 8, validated ≥ 1). with_controller_socket binds get_context=lambda self: self.zmq_context, so **all backends** reuse it on the controller-RPC path. 3. **SimpleStorage borrows the client context.** initialize_storage_manager passes zmq_context=self.zmq_context to the factory **only** for manager_type == "SimpleStorage". AsyncSimpleStorageManager's with_storage_unit_socket binds get_context=lambda self: self.zmq_context, so both the notify path and per-call storage-unit request sockets share the client's fixed I/O-thread pool. 4. **Ownership-aware teardown, no double-free.** StorageManager records self._owns_zmq_context = zmq_context is None. A manager that created its own context tears it down with zmq_context.destroy(linger=0); a manager that borrowed the client's context does **not** terminate it. The client terminates its own context once in close() via destroy(linger=0) (force-closes any leaked socket so shutdown can't hang). 5. **Other backends are unaffected.** Mooncake / Yuanrong / RayStore managers receive zmq_context=None, so they keep their own independent long-lived context for the controller notify/handshake path and never touch the shared pool. ## Scope - The knob simple_storage_zmq_io_threads / TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS sizes the client context's native I/O-thread pool. Note this context also serves the **client→controller** path for every backend, not only SimpleStorage. - KV backends (Mooncake/Yuanrong/RayStore) move their bulk data through their own SDKs, not through with_zmq_socket, so this change only affects their controller-side ZMQ traffic (via the shared decorator fix), not their data plane. ## Tests - New tests/test_zmq_shared_context.py (7 cases): - shared context is reused across concurrent calls, - client context has a fixed I/O-thread pool, - client rejects an invalid (<1) pool size, - SimpleStorage borrows the client context, - SimpleStorage does **not** destroy the borrowed context, - other backends do **not** borrow the client context, - close() destroys the context exactly once. - Full suite run locally on a fresh single-node Ray: **547 passed, 10 skipped, 8 errors** — the 10 skips (GDR/GPU, Mooncake-CUDA) and 8 errors (Yuanrong SDK absent) are all pre-existing environment gaps in files untouched by this branch; **0 failures**. The complete e2e lifecycle suite (test_core_consistency, cross-shard, production-status, reset, clear, dynamic-shape, memory-safety) passes. --------- Signed-off-by: OutstanderWang <wangweiyanster@gmail.com> | 1 天前 |
| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
| 24 天前 | ||
| 1 个月前 | ||
| 2 个月前 | ||
| 8 天前 | ||
| 6 天前 | ||
| 1 个月前 | ||
| 2 个月前 | ||
| 2 个月前 | ||
| 26 天前 | ||
| 24 天前 | ||
| 2 个月前 | ||
| 4 个月前 | ||
| 2 个月前 | ||
| 6 天前 | ||
| 1 个月前 | ||
| 3 个月前 | ||
| 2 个月前 | ||
| 3 个月前 | ||
| 2 个月前 | ||
| 1 个月前 | ||
| 1 天前 |