From 6614504a07208160cfdc8c0231d4cabc7218d034 Mon Sep 17 00:00:00 2001
From: haixuanTao <tao.xavier@outlook.com>
Date: Thu, 28 Nov 2024 10:30:49 +0100
Subject: [PATCH] Adding better documentation

---
 apis/python/node/dora/cuda.py      | 65 ++++++++++++++++++++++++++++--
 examples/cuda-benchmark/.gitignore |  1 +
 examples/cuda-benchmark/README.md  | 11 ++++-
 3 files changed, 72 insertions(+), 5 deletions(-)
 create mode 100644 examples/cuda-benchmark/.gitignore

diff --git a/apis/python/node/dora/cuda.py b/apis/python/node/dora/cuda.py
index 77779eaa..1858a7bb 100644
--- a/apis/python/node/dora/cuda.py
+++ b/apis/python/node/dora/cuda.py
@@ -12,7 +12,16 @@ 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."""
+    """
+    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()
@@ -25,14 +34,46 @@ def torch_to_ipc_buffer(tensor: torch.TensorType) -> tuple[pa.array, dict]:
 
 
 def ipc_buffer_to_ipc_handle(handle_buffer: pa.array) -> cuda.IpcMemHandle:
-    """Converts a buffer containing a serialized handler into cuda IPC MemHandle."""
+    """
+    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."""
+    """
+    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"]
@@ -42,7 +83,23 @@ def cudabuffer_to_numba(buffer: cuda.CudaBuffer, metadata: dict) -> DeviceNDArra
 
 
 def cudabuffer_to_torch(buffer: cuda.CudaBuffer, metadata: dict) -> torch.Tensor:
-    """Converts a pyarrow CUDA buffer to a 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")
diff --git a/examples/cuda-benchmark/.gitignore b/examples/cuda-benchmark/.gitignore
new file mode 100644
index 00000000..afed0735
--- /dev/null
+++ b/examples/cuda-benchmark/.gitignore
@@ -0,0 +1 @@
+*.csv
diff --git a/examples/cuda-benchmark/README.md b/examples/cuda-benchmark/README.md
index 60436bbf..224f84bd 100644
--- a/examples/cuda-benchmark/README.md
+++ b/examples/cuda-benchmark/README.md
@@ -2,17 +2,26 @@
 
 ## Install
 
-```
+```bash
 pip install dora-rs-cli # if not already present
 
 # Install pyarrow with gpu support
 conda install pyarrow "arrow-cpp-proc=*=cuda" -c conda-forge
 
+## Test installation with
+python -c "import pyarrow.cuda"
+
 # Install numba for translation from arrow to torch
 pip install numba
 
+## Test installation with
+python -c "import numba.cuda"
+
 # Install torch if it's not already present
 pip install torch
+
+## Test installation with
+python -c "import torch; assert torch.cuda.is_available()"
 ```
 
 ## Run