[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
Signed-off-by: Evelynn-V <[email protected]>
Signed-off-by: Evelynn-V <[email protected]>
📝 WalkthroughWalkthroughAdds 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)
✅ Passed checks (2 passed)
✏️ 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. Comment |
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 Fix all issues with AI agents
In `@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]).
🧹 Nitpick comments (4)
transfer_queue/storage/clients/yuanrong_client.py (1)
311-341: Consider making the timeout configurable.The
timeout_ms=500is 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.tests/test_ds_zero_copy.py (3)
15-17: Unused imports.
YuanrongStorageManager,YuanrongStorageClient, andTransferQueueStorageManagerFactoryare 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).
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
tests/test_ds_zero_copy.pytransfer_queue/storage/clients/yuanrong_client.py
🧰 Additional context used
🧬 Code graph analysis (1)
tests/test_ds_zero_copy.py (7)
transfer_queue/client.py (1)
TransferQueueClient(706-834)transfer_queue/storage/managers/yuanrong_manager.py (1)
YuanrongStorageManager(34-49)transfer_queue/storage/clients/yuanrong_client.py (3)
YuanrongStorageClient(90-392)put(221-235)get(343-360)transfer_queue/storage/managers/factory.py (1)
TransferQueueStorageManagerFactory(21-45)transfer_queue/utils/zmq_utils.py (1)
ZMQServerInfo(93-112)transfer_queue/controller.py (1)
TransferQueueController(644-1430)tests/test_ray_p2p.py (1)
generate_data(93-118)
🪛 GitHub Actions: pre-commit
tests/test_ds_zero_copy.py
[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
tests/test_ds_zero_copy.py
[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
🔇 Additional comments (1)
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.frombufferfor zero-copy tensor operations is appropriate.
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
| 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.
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 fix
Either 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:📝 Committable suggestion
‼️ 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) |
🤖 Prompt for AI Agents
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.
| 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: 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)📝 Committable suggestion
‼️ 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) |
🤖 Prompt for AI Agents
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.
| 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.
Critical: Missing numpy import and typo in error message.
Two issues in this block:
np(numpy) is used but never imported - causesNameError- Line 336:
key[i]should bekeys[i](typo)
🐛 Proposed fixes
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}")🤖 Prompt for AI Agents
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]).
Signed-off-by: Evelynn-V <[email protected]>
Signed-off-by: Evelynn-V <[email protected]>
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.pyto 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.pyResult
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.