[Feat] Introduce Zero-Copy to use YuanrongStorageClient for transmitting CPU Tensors - #171
[Feat] Introduce Zero-Copy to use YuanrongStorageClient for transmitting CPU Tensors#171Evelynn-V wants to merge 4 commits into
Conversation
Sorry, something went wrong.
Signed-off-by: Evelynn-V <liwenlin0223l@gmail.com>
Signed-off-by: Evelynn-V <liwenlin0223l@gmail.com>
📝 Walkthrough
WalkthroughAdds a Ray-based integration test for zero-copy transfer and implements in-place packing/unpacking plus a zero-copy mset path in the Yuanrong storage client to enable contiguous-buffer serialization and retrieval when NPU is unavailable. Changes
Sequence DiagramsequenceDiagram
participant Main as Main Process
participant TQC as TransferQueueController
participant WA as WriterActor
participant SC as YuanrongStorageClient
participant RA as ReaderActor
Main->>TQC: initialize controller
Main->>WA: create WriterActor (controller_info, config)
Main->>RA: create ReaderActor (controller_info, config)
loop per partition
Main->>WA: generate_data(partition_id)
WA-->>Main: return BatchMeta / metadata
Main->>WA: put_once(partition_id)
WA->>SC: mset_zcopy(keys, objs)
SC->>SC: try_zero_copy_serialize -> pack_into -> allocate contiguous buffer
SC-->>WA: acknowledge stored (BatchMeta)
Main->>RA: get_once(batch_meta)
RA->>SC: retrieve packed buffer
SC->>SC: unpack_from -> reconstruct tensor bytes -> convert to tensor
SC-->>RA: return reconstructed data
end
Main->>Main: aggregate timings and report
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem🚥 Pre-merge checks | ✅ 2 | ❌ 1 ❌ Failed checks (1 warning)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. ❤️ ShareComment @coderabbitai help to get the list of available commands and usage tips. |
Sorry, something went wrong.
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 Fix all issues with AI agentsIn `@tests/test_ds_zero_copy.py`: - Around line 50-53: The function definition for put_once is missing the trailing colon causing a syntax error; update the signature of put_once(partition_id) to include the colon so the function body (using time.perf_counter(), self.client.put(data=self.data, partition_id=partition_id), and returning batch_meta) is properly defined and the test file tests/test_ds_zero_copy.py will parse. - Around line 36-48: The generate_data function definition is missing a trailing colon causing a SyntaxError; add the missing ':' after the parameter list in the generate_data(...) signature so the function compiles, then run tests to verify TensorDict creation (TensorDict, tensordict_memory_mb, self.data) works as expected. - Around line 114-120: The call to writer.generate_data.remote is missing the required partition_id argument and the generate_data method currently accepts but does not use partition_id; either pass the created partition_id when calling generate_data (e.g., call writer.generate_data.remote(partition_id, batch_size=512, seq_len=32*1024)) so partition_id flows through before put_once, or remove the unused partition_id parameter from the generate_data signature and all related remote calls (and update any callers) if the parameter is genuinely unnecessary; ensure consistency between the generate_data remote definition and its invocations around the writer.generate_data and writer.put_once interactions. In `@transfer_queue/storage/clients/yuanrong_client.py`: - Around line 39-42: The file uses the struct module (e.g., HEADER_FMT, HEADER_SIZE, ENTRY_FMT, ENTRY_SIZE) but never imports it; add the missing import statement "import struct" to the top-level imports in yuanrong_client.py so all references to struct.calcsize and other struct functions succeed at runtime. - Around line 156-165: The mset_zcopy method uses ThreadPoolExecutor but the symbol is not imported, causing a NameError; add the missing import "from concurrent.futures import ThreadPoolExecutor" at the top of the module (or wherever imports are collected) so ThreadPoolExecutor is defined when mset_zcopy runs, ensuring the ThreadPoolExecutor usage in mset_zcopy works correctly. - Around line 71-87: The try_zero_copy_serialize function currently calls obj.numpy() for torch.Tensor without checking device, which will raise on CUDA/NPU tensors; update try_zero_copy_serialize to first check isinstance(obj, torch.Tensor) and that obj.is_contiguous() and obj.device.type == "cpu" before calling obj.numpy() (returning memoryview(obj.numpy())), and for non-CPU tensors fall back to the existing non-zero-copy path (e.g., return pickle.dumps(obj) or return memoryview(obj) where applicable); reference try_zero_copy_serialize and the torch.Tensor handling when making the change. - Around line 214-219: The branch that falls back to the CPU path references an undefined variable pickled_values, causing a NameError; replace pickled_values with the original values list so mset_zcopy receives raw objects (it will call try_zero_copy_serialize internally). Locate the loop using pickled_values (in yuanrong_client.py) and change batch_vals = pickled_values[...] to batch_vals = values[...] (keeping batch_keys from keys and the same CPU_DS_CLIENT_KEYS_LIMIT chunking). - Around line 331-339: The block in yuanrong_client.py uses np but numpy isn't imported and contains a typo referencing key[i]; add "import numpy as np" to the module imports and change the error message to use keys[i] instead of key[i] in the ValueError in the code that computes numpy_dtype/total_elements/expected_size and constructs arr/torch tensor (the statements around numpy_dtype, total_elements, expected_size, and results[i]).
transfer_queue/storage/clients/yuanrong_client.py (1)📜 Review detailstests/test_ds_zero_copy.py (3)311-341: Consider making the timeout configurable.
The timeout_ms=500 is hardcoded and may be too short for large data transfers (the PR benchmarks show 16GB transfers). Consider making this configurable or using a more generous default.
15-17: Unused imports.
YuanrongStorageManager, YuanrongStorageClient, and TransferQueueStorageManagerFactory are imported but not used in this file. Consider removing them to clean up the code.
75-101: Consider parameterizing hardcoded IP addresses.
The IP addresses are hardcoded, making this test only runnable in a specific environment. For better portability, consider using environment variables or a configuration file.
import os ip_A = os.getenv("TQ_WRITER_IP", "10.90.41.116") ip_B = os.getenv("TQ_READER_IP", "10.90.41.117") port = int(os.getenv("TQ_PORT", "36666"))
122-130: Consider adding data integrity verification.
The test measures timing but doesn't verify that the data retrieved matches the data written. For a more robust test, consider adding assertions to validate data integrity (e.g., comparing tensor shapes, dtypes, or checksums).
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📥 CommitsReviewing files that changed from the base of the PR and between 788979d and 5f2adce.
📒 Files selected for processing (2)- tests/test_ds_zero_copy.py
- transfer_queue/storage/clients/yuanrong_client.py
transfer_queue/client.py (1)🪛 GitHub Actions: pre-commit tests/test_ds_zero_copy.pytransfer_queue/storage/managers/yuanrong_manager.py (1)
- TransferQueueClient (706-834)
transfer_queue/storage/clients/yuanrong_client.py (3)
- YuanrongStorageManager (34-49)
transfer_queue/storage/managers/factory.py (1)
- YuanrongStorageClient (90-392)
- put (221-235)
- get (343-360)
transfer_queue/utils/zmq_utils.py (1)
- TransferQueueStorageManagerFactory (21-45)
transfer_queue/controller.py (1)
- ZMQServerInfo (93-112)
tests/test_ray_p2p.py (1)
- TransferQueueController (644-1430)
- generate_data (93-118)
[error] 38-38: SyntaxError: Expected ':', found newline
🪛 GitHub Actions: Python package transfer_queue/storage/clients/yuanrong_client.py[error] 40-338: F821 undefined name 'struct' and related undefined-name issues detected by flake8 (multiple occurrences); potential missing imports.
[error] 40-338: F821 undefined name 'struct' and related undefined-name issues detected by flake8 (multiple occurrences); potential missing imports.
tests/test_ds_zero_copy.py[error] 38-38: SyntaxError: expected ':'
[error] 38-38: SyntaxError: expected ':'
🪛 GitHub Check: pre-commit (3.10) transfer_queue/storage/clients/yuanrong_client.py[failure] 67-67: Ruff (F821)
transfer_queue/storage/clients/yuanrong_client.py:67:26: F821 Undefined name struct
[failure] 64-64: Ruff (F821)
transfer_queue/storage/clients/yuanrong_client.py:64:18: F821 Undefined name struct
[failure] 56-56: Ruff (F821)
transfer_queue/storage/clients/yuanrong_client.py:56:9: F821 Undefined name struct
[failure] 48-48: Ruff (F821)
transfer_queue/storage/clients/yuanrong_client.py:48:5: F821 Undefined name struct
[failure] 42-42: Ruff (F821)
transfer_queue/storage/clients/yuanrong_client.py:42:14: F821 Undefined name struct
[failure] 40-40: Ruff (F821)
transfer_queue/storage/clients/yuanrong_client.py:40:15: F821 Undefined name struct
[failure] 163-163: Ruff (F821)
transfer_queue/storage/clients/yuanrong_client.py:163:14: F821 Undefined name ThreadPoolExecutor
[failure] 50-51: Ruff
tests/test_ds_zero_copy.py:50:37: SyntaxError: Expected ':', found newline
[failure] 38-39: Ruff
tests/test_ds_zero_copy.py:38:7: SyntaxError: Expected ':', found newline
transfer_queue/storage/clients/yuanrong_client.py (1)44-69: LGTM for packing/unpacking logic.
The packing and unpacking helper functions correctly implement the binary format with header + entries + payload. The use of torch.frombuffer for zero-copy tensor operations is appropriate.
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
Sorry, something went wrong.
| for i in range(3): | ||
| partition_id = f"train_step_{i}" | ||
| ray.get(writer.generate_data.remote(batch_size=512, seq_len=32 * 1024)) | ||
| cost, batch_meta = ray.get(writer.put_once.remote(partition_id)) | ||
| print(f"[WriterActor] The time consumed by the {i}th put costs: {cost:.2f}s") | ||
| batch_metas.append(batch_meta) | ||
| put_times.append(cost) |
There was a problem hiding this comment.
⚠️ Potential issue | 🟠 Major
Bug: Missing partition_id argument in generate_data call.
The generate_data method requires partition_id as its first argument, but the call at line 116 doesn't pass it. Additionally, the method doesn't use partition_id internally (it's not assigned to anything), which suggests it may be an unused parameter.
🐛 Proposed fixEither pass the argument:
for i in range(3):
partition_id = f"train_step_{i}"
- ray.get(writer.generate_data.remote(batch_size=512, seq_len=32 * 1024))
+ ray.get(writer.generate_data.remote(partition_id, batch_size=512, seq_len=32 * 1024))Or remove the unused partition_id parameter from the method signature if it's not needed:
def generate_data(
- self, partition_id, batch_size: int = 10000, seq_len: int = 10000
+ self, batch_size: int = 10000, seq_len: int = 10000
) -> None:‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| for i in range(3): | |
| partition_id = f"train_step_{i}" | |
| ray.get(writer.generate_data.remote(batch_size=512, seq_len=32 * 1024)) | |
| cost, batch_meta = ray.get(writer.put_once.remote(partition_id)) | |
| print(f"[WriterActor] The time consumed by the {i}th put costs: {cost:.2f}s") | |
| batch_metas.append(batch_meta) | |
| put_times.append(cost) | |
| for i in range(3): | |
| partition_id = f"train_step_{i}" | |
| ray.get(writer.generate_data.remote(partition_id, batch_size=512, seq_len=32 * 1024)) | |
| cost, batch_meta = ray.get(writer.put_once.remote(partition_id)) | |
| print(f"[WriterActor] The time consumed by the {i}th put costs: {cost:.2f}s") | |
| batch_metas.append(batch_meta) | |
| put_times.append(cost) |
In `@tests/test_ds_zero_copy.py` around lines 114 - 120, The call to writer.generate_data.remote is missing the required partition_id argument and the generate_data method currently accepts but does not use partition_id; either pass the created partition_id when calling generate_data (e.g., call writer.generate_data.remote(partition_id, batch_size=512, seq_len=32*1024)) so partition_id flows through before put_once, or remove the unused partition_id parameter from the generate_data signature and all related remote calls (and update any callers) if the parameter is genuinely unnecessary; ensure consistency between the generate_data remote definition and its invocations around the writer.generate_data and writer.put_once interactions.
Sorry, something went wrong.
| def try_zero_copy_serialize(obj: Any) -> memoryview | bytes | bytearray: | ||
| """ | ||
| Attempts to serialize an object using zero-copy if possible. | ||
| If the object supports the buffer protocol (like numpy arrays or contiguous tensors), | ||
| it returns a memoryview for zero-copy serialization. Otherwise, it falls back to | ||
| standard pickle serialization. | ||
| Args: | ||
| obj (Any): The object to serialize. | ||
| Returns: | ||
| Union[memoryview, bytes, bytearray]: The serialized object. | ||
| """ | ||
| if isinstance(obj, torch.Tensor) and obj.is_contiguous(): | ||
| return memoryview(obj.numpy()) | ||
| try: | ||
| return memoryview(obj) | ||
| except TypeError: | ||
| return pickle.dumps(obj) |
There was a problem hiding this comment.
⚠️ Potential issue | 🟡 Minor
Potential issue: numpy() fails on non-CPU tensors.
Calling obj.numpy() on a CUDA/NPU tensor will raise an error. If this function is only intended for CPU tensors (given the zero-copy CPU path context), consider adding a guard or documenting this constraint.
💡 Suggested defensive check def try_zero_copy_serialize(obj: Any) -> memoryview | bytes | bytearray:
...
- if isinstance(obj, torch.Tensor) and obj.is_contiguous():
+ if isinstance(obj, torch.Tensor) and obj.is_contiguous() and obj.device.type == "cpu":
return memoryview(obj.numpy())
+ elif isinstance(obj, torch.Tensor) and obj.is_contiguous():
+ # Non-CPU tensor: transfer to CPU first (fallback, slower path)
+ return memoryview(obj.cpu().numpy())
try:
return memoryview(obj)
except TypeError:
return pickle.dumps(obj)‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def try_zero_copy_serialize(obj: Any) -> memoryview | bytes | bytearray: | |
| """ | |
| Attempts to serialize an object using zero-copy if possible. | |
| If the object supports the buffer protocol (like numpy arrays or contiguous tensors), | |
| it returns a memoryview for zero-copy serialization. Otherwise, it falls back to | |
| standard pickle serialization. | |
| Args: | |
| obj (Any): The object to serialize. | |
| Returns: | |
| Union[memoryview, bytes, bytearray]: The serialized object. | |
| """ | |
| if isinstance(obj, torch.Tensor) and obj.is_contiguous(): | |
| return memoryview(obj.numpy()) | |
| try: | |
| return memoryview(obj) | |
| except TypeError: | |
| return pickle.dumps(obj) | |
| def try_zero_copy_serialize(obj: Any) -> memoryview | bytes | bytearray: | |
| """ | |
| Attempts to serialize an object using zero-copy if possible. | |
| If the object supports the buffer protocol (like numpy arrays or contiguous tensors), | |
| it returns a memoryview for zero-copy serialization. Otherwise, it falls back to | |
| standard pickle serialization. | |
| Args: | |
| obj (Any): The object to serialize. | |
| Returns: | |
| Union[memoryview, bytes, bytearray]: The serialized object. | |
| """ | |
| if isinstance(obj, torch.Tensor) and obj.is_contiguous() and obj.device.type == "cpu": | |
| return memoryview(obj.numpy()) | |
| elif isinstance(obj, torch.Tensor) and obj.is_contiguous(): | |
| # Non-CPU tensor: transfer to CPU first (fallback, slower path) | |
| return memoryview(obj.cpu().numpy()) | |
| try: | |
| return memoryview(obj) | |
| except TypeError: | |
| return pickle.dumps(obj) |
In `@transfer_queue/storage/clients/yuanrong_client.py` around lines 71 - 87, The try_zero_copy_serialize function currently calls obj.numpy() for torch.Tensor without checking device, which will raise on CUDA/NPU tensors; update try_zero_copy_serialize to first check isinstance(obj, torch.Tensor) and that obj.is_contiguous() and obj.device.type == "cpu" before calling obj.numpy() (returning memoryview(obj.numpy())), and for non-CPU tensors fall back to the existing non-zero-copy path (e.g., return pickle.dumps(obj) or return memoryview(obj) where applicable); reference try_zero_copy_serialize and the torch.Tensor handling when making the change.
Sorry, something went wrong.
| numpy_dtype = torch.empty((), dtype=dtype).numpy().dtype | ||
| total_elements = int(np.prod(shape)) | ||
| expected_size = total_elements * np.dtype(numpy_dtype).itemsize | ||
|
|
||
| if len(raw_val) != expected_size: | ||
| raise ValueError(f"Size mismatch for key {key[i]}: got {len(raw_val)}, expected {expected_size}") | ||
|
|
||
| arr = np.frombuffer(raw_val, dtype=numpy_dtype).reshape(shape) | ||
| results[i] = torch.from_numpy(arr) |
There was a problem hiding this comment.
⚠️ Potential issue | 🔴 Critical
Critical: Missing numpy import and typo in error message.
Two issues in this block:
- np (numpy) is used but never imported - causes NameError
- Line 336: key[i] should be keys[i] (typo)
Add to the imports section at the top of the file:
import numpy as npFix the typo at line 336:
- if len(raw_val) != expected_size:
- raise ValueError(f"Size mismatch for key {key[i]}: got {len(raw_val)}, expected {expected_size}")
+ if len(raw_val) != expected_size:
+ raise ValueError(f"Size mismatch for key {keys[i]}: got {len(raw_val)}, expected {expected_size}")In `@transfer_queue/storage/clients/yuanrong_client.py` around lines 331 - 339, The block in yuanrong_client.py uses np but numpy isn't imported and contains a typo referencing key[i]; add "import numpy as np" to the module imports and change the error message to use keys[i] instead of key[i] in the ValueError in the code that computes numpy_dtype/total_elements/expected_size and constructs arr/torch tensor (the statements around numpy_dtype, total_elements, expected_size, and results[i]).
Sorry, something went wrong.
Reviewers
1 more reviewerAt least 1 approving review is required to merge this pull request.
Summary
When connecting to the backend of the YuanrongStorageClient, zero-copy is activated to enhance the transmission speed.
Change
Modified the transfer_queue/storage/clients/yuanrong_client.py to call the zero-copy interface, and performed operations such as serialization and pack.
Add p2p tests: tests/test_ds_zero_copy.py .
Testing
Test on CPU:
python tests/test_ds_zero_copy.py
Result
When transmitting 512 pieces of data, each 32 MB in size, with a total data volume of 16GB:
End-to-end Get took 10s and the bandwidth was 1.6 GB/s. The time spent calling the YuanrongStorageClient interface was 2.27s with a bandwidth of 7.05 GB/s.
End-to-end Put took 3.42s and the bandwidth was 4.68 GB/s. The time spent calling the YuanrongStorageClient interface was 3.32s with a bandwidth of 4.83 GB/s.
Related Links
Previous issues can be viewed: [Feat]: Try zero-copy serialize objects that can be converted to memoryview
Yuanrong Datasystem PR: https://atomgit.com/openeuler/yuanrong-datasystem/pull/141
Summary by CodeRabbit
New Features
Tests
✏️ Tip: You can customize this high-level summary in your review settings.