Skip to content

Commit

Permalink
Minor python esthetics
Browse files Browse the repository at this point in the history
  • Loading branch information
haixuanTao committed Jan 17, 2025
1 parent d82b699 commit cb9a98e
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 22 deletions.
15 changes: 8 additions & 7 deletions apis/python/node/dora/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,30 +10,31 @@
from enum import Enum

from .dora import *

from .dora import (
Node,
Ros2Context,
Ros2Durability,
Ros2Liveliness,
Ros2Node,
Ros2NodeOptions,
Ros2Topic,
Ros2Publisher,
Ros2QosPolicies,
Ros2Subscription,
start_runtime,
__version__,
Ros2Topic,
__author__,
Ros2QosPolicies,
Ros2Durability,
Ros2Liveliness,
__version__,
start_runtime,
)


class DoraStatus(Enum):
"""Dora status to indicate if operator `on_input` loop
should be stopped.
Args:
Enum (u8): Status signaling to dora operator to
stop or continue the operator.
"""

CONTINUE = 0
Expand Down
22 changes: 8 additions & 14 deletions apis/python/node/dora/cuda.py
Original file line number Diff line number Diff line change
@@ -1,19 +1,18 @@
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
from numba.cuda import to_device

# Make sure to install numba with cuda
from numba.cuda.cudadrv.devicearray import DeviceNDArray
from numba.cuda import to_device

# To install pyarrow.cuda, run `conda install pyarrow "arrow-cpp-proc=*=cuda" -c conda-forge`
from pyarrow import cuda


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
Expand All @@ -34,8 +33,7 @@ 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
Expand All @@ -57,8 +55,7 @@ def ipc_buffer_to_ipc_handle(handle_buffer: pa.array) -> cuda.IpcMemHandle:


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
Expand All @@ -74,7 +71,6 @@ def cudabuffer_to_numba(buffer: cuda.CudaBuffer, metadata: dict) -> DeviceNDArra
numba_tensor = cudabuffer_to_numbda(cudabuffer, event["metadata"])
```
"""

shape = metadata["shape"]
strides = metadata["strides"]
dtype = metadata["dtype"]
Expand All @@ -83,8 +79,7 @@ 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
Expand All @@ -100,7 +95,6 @@ def cudabuffer_to_torch(buffer: cuda.CudaBuffer, metadata: dict) -> torch.Tensor
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
5 changes: 4 additions & 1 deletion apis/python/node/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,12 @@ build-backend = "maturin"

[project]
name = "dora-rs"
dynamic = ["version"]
dynamic = ["version", "readme"]
# Install pyarrow at the same time of dora-rs
dependencies = ['pyarrow']

[dependency-groups]
dev = ["pytest >=8.1.1", "ruff >=0.9.1"]

[tool.maturin]
features = ["pyo3/extension-module"]

0 comments on commit cb9a98e

Please sign in to comment.