dora-rs 0.3.11rc1__cp37-abi3-win_amd64.whl → 0.3.12rc1__cp37-abi3-win_amd64.whl

dora/cuda.py

@@ -8,13 +8,19 @@ from numba.cuda import to_device
 
 # Make sure to install numba with cuda
 from numba.cuda.cudadrv.devicearray import DeviceNDArray
+from numba.cuda.cudadrv.devices import get_context
+from numba.cuda.cudadrv.driver import IpcHandle
 
-# To install pyarrow.cuda, run `conda install pyarrow "arrow-cpp-proc=*=cuda" -c conda-forge`
-from pyarrow import cuda
+
+import json
+
+from contextlib import contextmanager
+from typing import ContextManager
 
 
 def torch_to_ipc_buffer(tensor: torch.TensorType) -> tuple[pa.array, dict]:
-    """Convert a Pytorch tensor into a pyarrow buffer containing the IPC handle and its metadata.
+    """Convert a Pytorch tensor into a pyarrow buffer containing the IPC handle
+    and its metadata.
 
     Example Use:
     ```python
@@ -24,75 +30,65 @@ def torch_to_ipc_buffer(tensor: torch.TensorType) -> tuple[pa.array, dict]:
     ```
     """
     device_arr = to_device(tensor)
-    cuda_buf = pa.cuda.CudaBuffer.from_numba(device_arr.gpu_data)
-    handle_buffer = cuda_buf.export_for_ipc().serialize()
+    ipch = get_context().get_ipc_handle(device_arr.gpu_data)
+    _, handle, size, source_info, offset = ipch.__reduce__()[1]
     metadata = {
         "shape": device_arr.shape,
         "strides": device_arr.strides,
         "dtype": device_arr.dtype.str,
+        "size": size,
+        "offset": offset,
+        "source_info": json.dumps(source_info),
     }
-    return pa.array(handle_buffer, type=pa.uint8()), metadata
+    return pa.array(handle, pa.int8()), metadata
 
 
-def ipc_buffer_to_ipc_handle(handle_buffer: pa.array) -> cuda.IpcMemHandle:
-    """Convert a buffer containing a serialized handler into cuda IPC MemHandle.
+def ipc_buffer_to_ipc_handle(handle_buffer: pa.array, metadata: dict) -> IpcHandle:
+    """Convert a buffer containing a serialized handler into cuda IPC Handle.
 
     example use:
     ```python
+    from dora.cuda import ipc_buffer_to_ipc_handle, open_ipc_handle
 
-    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
+    ipc_handle = ipc_buffer_to_ipc_handle(event["value"], event["metadata"])
+    with open_ipc_handle(ipc_handle, event["metadata"]) as tensor:
+        pass
     ```
     """
-    handle_buffer = handle_buffer.buffers()[1]
-    return pa.cuda.IpcMemHandle.from_buffer(handle_buffer)
+    handle = handle_buffer.to_pylist()
+    return IpcHandle._rebuild(
+        handle,
+        metadata["size"],
+        json.loads(metadata["source_info"]),
+        metadata["offset"],
+    )
 
 
-def cudabuffer_to_numba(buffer: cuda.CudaBuffer, metadata: dict) -> DeviceNDArray:
-    """Convert a pyarrow CUDA buffer to numba.
+@contextmanager
+def open_ipc_handle(
+    ipc_handle: IpcHandle, metadata: dict
+) -> ContextManager[torch.TensorType]:
+    """Open a CUDA IPC handle and return a Pytorch tensor.
 
     example use:
     ```python
+    from dora.cuda import ipc_buffer_to_ipc_handle, open_ipc_handle
 
-    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"])
+    ipc_handle = ipc_buffer_to_ipc_handle(event["value"], event["metadata"])
+    with open_ipc_handle(ipc_handle, event["metadata"]) as tensor:
+        pass
     ```
     """
     shape = metadata["shape"]
     strides = metadata["strides"]
     dtype = metadata["dtype"]
-    return DeviceNDArray(shape, strides, dtype, gpu_data=buffer.to_numba())
-
-
-def cudabuffer_to_torch(buffer: cuda.CudaBuffer, metadata: dict) -> torch.Tensor:
-    """Convert 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
-    ```
-    """
-    return torch.as_tensor(cudabuffer_to_numba(buffer, metadata), device="cuda")
+    try:
+        buffer = ipc_handle.open(get_context())
+        device_arr = DeviceNDArray(shape, strides, dtype, gpu_data=buffer)
+        yield torch.as_tensor(device_arr, device="cuda")
+    finally:
+        ipc_handle.close()

{dora_rs-0.3.11rc1.dist-info → dora_rs-0.3.12rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dora-rs
-Version: 0.3.11rc1
+Version: 0.3.12rc1
 Requires-Dist: pyarrow
 Summary: `dora` goal is to be a low latency, composable, and distributed data flow.
 License: MIT

dora_rs-0.3.12rc1.dist-info/RECORD

@@ -0,0 +1,7 @@
+dora/__init__.py,sha256=K87kJCx3z_LKKKSikuEIkkAx4h8qZtS6dB-HzFpz4zc,751
+dora/__init__.pyi,sha256=L-MCRPD_K8RGpzWsfrV28lNh0ast3uNq2-UJZzuJwP8,9483
+dora/cuda.py,sha256=WunE7nb13aIFq5cqkSh52dfZ0KfFm18NU4TofPFopx8,2830
+dora/dora.pyd,sha256=HVCDGfxqu6hvslakttSf_JM_7mAgAxFpHxB_llYU4Uw,37786624
+dora_rs-0.3.12rc1.dist-info/METADATA,sha256=tVLKsqktgduAAigy5hXcbumX-_E1_Q3aFSW3wAmZ-C8,647
+dora_rs-0.3.12rc1.dist-info/WHEEL,sha256=G5LfiGOqaNK7XFePKQAOL9sBbuwwBdXf4vgmptkMvPU,94
+dora_rs-0.3.12rc1.dist-info/RECORD,,

{dora_rs-0.3.11rc1.dist-info → dora_rs-0.3.12rc1.dist-info}/WHEEL

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

dora_rs-0.3.11rc1.dist-info/RECORD

@@ -1,7 +0,0 @@
-dora_rs-0.3.11rc1.dist-info/METADATA,sha256=fzb0qk80dZ80NOS2nED13XIcecElFQf0T3-jiKOs74k,647
-dora_rs-0.3.11rc1.dist-info/WHEEL,sha256=GWZFzAhmRLcnnWJFIkiR8k0sjr8iX_MTxF4a-TgOsmk,94
-dora/cuda.py,sha256=g3Ez0KxbPH9a5ok8ouePyeCreEMoQQjbpPZ5Alqxad0,3167
-dora/__init__.py,sha256=K87kJCx3z_LKKKSikuEIkkAx4h8qZtS6dB-HzFpz4zc,751
-dora/__init__.pyi,sha256=L-MCRPD_K8RGpzWsfrV28lNh0ast3uNq2-UJZzuJwP8,9483
-dora/dora.pyd,sha256=y1bn0IZSRmEcG5CplhDE0fAt7AdzfQXYklkN4bm10AE,35266560
-dora_rs-0.3.11rc1.dist-info/RECORD,,