dora-rs 0.3.7rc2__cp37-abi3-win_amd64.whl → 0.3.8rc0__cp37-abi3-win_amd64.whl

dora/cuda.py

@@ -0,0 +1,106 @@
+import pyarrow as pa
+
+# To install pyarrow.cuda, run `conda install pyarrow "arrow-cpp-proc=*=cuda" -c conda-forge`
+import pyarrow.cuda as cuda
+
+# Make sure to install torch with cuda
+import torch
+
+# Make sure to install numba with cuda
+from numba.cuda.cudadrv.devicearray import DeviceNDArray
+from numba.cuda import to_device
+
+
+def torch_to_ipc_buffer(tensor: torch.TensorType) -> tuple[pa.array, dict]:
+    """
+    Converts a Pytorch tensor into a pyarrow buffer containing the IPC handle and its metadata.
+
+    Example Use:
+    ```python
+    torch_tensor = torch.tensor(random_data, dtype=torch.int64, device="cuda")
+    ipc_buffer, metadata = torch_to_ipc_buffer(torch_tensor)
+    node.send_output("latency", ipc_buffer, metadata)
+    ```
+    """
+    device_arr = to_device(tensor)
+    cuda_buf = pa.cuda.CudaBuffer.from_numba(device_arr.gpu_data)
+    handle_buffer = cuda_buf.export_for_ipc().serialize()
+    metadata = {
+        "shape": device_arr.shape,
+        "strides": device_arr.strides,
+        "dtype": device_arr.dtype.str,
+    }
+    return pa.array(handle_buffer, type=pa.uint8()), metadata
+
+
+def ipc_buffer_to_ipc_handle(handle_buffer: pa.array) -> cuda.IpcMemHandle:
+    """
+    Converts a buffer containing a serialized handler into cuda IPC MemHandle.
+
+    example use:
+    ```python
+
+    import pyarrow as pa
+    from dora.cuda import ipc_buffer_to_ipc_handle, cudabuffer_to_torch
+
+    ctx = pa.cuda.context()
+    event = node.next()
+
+    ipc_handle = ipc_buffer_to_ipc_handle(event["value"])
+    cudabuffer = ctx.open_ipc_buffer(ipc_handle)
+    torch_tensor = cudabuffer_to_torch(cudabuffer, event["metadata"])  # on cuda
+    ```
+    """
+    handle_buffer = handle_buffer.buffers()[1]
+    ipc_handle = pa.cuda.IpcMemHandle.from_buffer(handle_buffer)
+    return ipc_handle
+
+
+def cudabuffer_to_numba(buffer: cuda.CudaBuffer, metadata: dict) -> DeviceNDArray:
+    """
+    Converts a pyarrow CUDA buffer to numba.
+
+    example use:
+    ```python
+
+    import pyarrow as pa
+    from dora.cuda import ipc_buffer_to_ipc_handle, cudabuffer_to_torch
+
+    ctx = pa.cuda.context()
+    event = node.next()
+
+    ipc_handle = ipc_buffer_to_ipc_handle(event["value"])
+    cudabuffer = ctx.open_ipc_buffer(ipc_handle)
+    numba_tensor = cudabuffer_to_numbda(cudabuffer, event["metadata"])
+    ```
+    """
+
+    shape = metadata["shape"]
+    strides = metadata["strides"]
+    dtype = metadata["dtype"]
+    device_arr = DeviceNDArray(shape, strides, dtype, gpu_data=buffer.to_numba())
+    return device_arr
+
+
+def cudabuffer_to_torch(buffer: cuda.CudaBuffer, metadata: dict) -> torch.Tensor:
+    """
+    Converts a pyarrow CUDA buffer to a torch tensor.
+
+    example use:
+    ```python
+
+    import pyarrow as pa
+    from dora.cuda import ipc_buffer_to_ipc_handle, cudabuffer_to_torch
+
+    ctx = pa.cuda.context()
+    event = node.next()
+
+    ipc_handle = ipc_buffer_to_ipc_handle(event["value"])
+    cudabuffer = ctx.open_ipc_buffer(ipc_handle)
+    torch_tensor = cudabuffer_to_torch(cudabuffer, event["metadata"])  # on cuda
+    ```
+    """
+
+    device_arr = cudabuffer_to_numba(buffer, metadata)
+    torch_tensor = torch.as_tensor(device_arr, device="cuda")
+    return torch_tensor

{dora_rs-0.3.7rc2.dist-info → dora_rs-0.3.8rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: dora-rs
-Version: 0.3.7rc2
+Version: 0.3.8rc0
 Requires-Dist: pyarrow
 Summary: `dora` goal is to be a low latency, composable, and distributed data flow.
 License: Apache-2.0

dora_rs-0.3.8rc0.dist-info/RECORD

@@ -0,0 +1,7 @@
+dora_rs-0.3.8rc0.dist-info/METADATA,sha256=MR3NL31ebeRXrWBmkpQJWpynhB68ag5s7CiYclRcbbg,670
+dora_rs-0.3.8rc0.dist-info/WHEEL,sha256=8_UzpQVLjhi_DxWEhc5q5kLz0DLADIO7XbXdIeTYx_s,94
+dora/cuda.py,sha256=naTHXaSKwRADxVv9J_K2g4MQRZSnMVw-joZitJjwxKA,3292
+dora/__init__.py,sha256=wRvKTqLYFROzvt95XWXPmCSd2OQeBW4zCn37H8VhLBs,751
+dora/__init__.pyi,sha256=KrLt3v0CXiURVMyukfq_hqNhZaOqmmGzrkYNFcmFXIE,8652
+dora/dora.pyd,sha256=wh5uaQQ-AqWMB_cCyQtk6zW-4TFjUN5_YYkd1CdMoiU,17591296
+dora_rs-0.3.8rc0.dist-info/RECORD,,

{dora_rs-0.3.7rc2.dist-info → dora_rs-0.3.8rc0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: maturin (1.7.4)
+Generator: maturin (1.7.8)
 Root-Is-Purelib: false
 Tag: cp37-abi3-win_amd64

dora_rs-0.3.7rc2.dist-info/RECORD

@@ -1,6 +0,0 @@
-dora_rs-0.3.7rc2.dist-info/METADATA,sha256=_mV0G9T43k6VGdda8BgdVSXIl3v4KgepU3WbDS-pBPA,670
-dora_rs-0.3.7rc2.dist-info/WHEEL,sha256=YApKqCDyi9GqUQokZqOT5I78CS4GNvgM6MM-DsXhn2s,94
-dora/__init__.py,sha256=wRvKTqLYFROzvt95XWXPmCSd2OQeBW4zCn37H8VhLBs,751
-dora/__init__.pyi,sha256=KrLt3v0CXiURVMyukfq_hqNhZaOqmmGzrkYNFcmFXIE,8652
-dora/dora.pyd,sha256=nmahlenNm8g0akPblM5HCbv1PeRXR7QoS3xQgVbow94,17536000
-dora_rs-0.3.7rc2.dist-info/RECORD,,