Tensor IPC

High-performance and Flexible IPC for tensor data with seamless ROS integration for robotics research.

Overview

tensor-ipc provides efficient shared memory communication for tensor data between processes, with built-in support for ROS topics. It enables zero-copy data sharing using POSIX shared memory and integrates with ROS for distributed communication and sim/real transfer.

Key Features

• Zero-Copy Shared Memory: POSIX shared memory with per-frame locking for safe concurrent access
• ROS Integration: Built-in ROS producers and consumers with automatic type conversion (custom types are supported through ros2_numpy)
• Multi-Backend Support: Native support for NumPy, PyTorch (CPU/CUDA), and JAX (CPU/GPU)
• DDS Notifications: Real-time notifications and synchronization using CycloneDDS for efficient polling
• Type Safety: Automatic validation of tensor shapes, dtypes, and devices
• History Management: Configurable history buffers with circular indexing

Installation

pip install tensor-ipc

Optional dependencies:

• PyTorch support: pip install tensor-ipc torch
• JAX support: pip install tensor-ipc jax jaxlib
• JAX GPU support: pip install tensor-ipc "jax[cuda12]" cupy-cuda12x
• ROS support: Install ros2_numpy in the same Python environment

Without optional dependencies, only the NumPy backend is available.

Quick Start

Refer to examples/ to see basic usage. Documentation is coming soon (hopefully).

CUDA Support

import torch
from tensor_ipc.core.producer import TensorProducer

# CUDA tensors with IPC sharing
if torch.cuda.is_available():
    cuda_tensor = torch.zeros(3, 224, 224, device='cuda:0')
    producer = TensorProducer.from_sample("cuda_pool", cuda_tensor)
    
    # Publish CUDA tensor directly
    gpu_data = torch.randn(3, 224, 224, device='cuda:0')
    producer.put(gpu_data)

Callbacks and Notifications

def on_new_data(data):
    print(f"Callback triggered with data shape: {data.shape}")

consumer = TensorConsumer(
    metadata,
    on_new_data_callback=on_new_data
)

# Callback will be triggered when new data arrives

History Management

# Get last 5 frames in chronological order
history = consumer.get(history_len=5, latest_first=False)

# Get last 3 frames with latest first
recent = consumer.get(history_len=3, latest_first=True)

Architecture

• Backends: Pluggable backends for NumPy, PyTorch (CPU/CUDA), and JAX (CPU/GPU)
• Shared Memory: NumPy/PyTorch/JAX CPU backends use POSIX shared memory. CUDA backends use CUDA IPC handles for GPU memory sharing.
• Locking: Per-frame reader-writer locks for safe concurrent access
• Notifications: CycloneDDS for real-time progress updates
• ROS Bridge: Automatic conversion between ROS messages and tensor data

API Reference

Core Classes

• TensorProducer: Creates and publishes to shared memory pools
• TensorConsumer: Subscribes to and reads from shared memory pools
• PoolMetadata: Describes pool structure and properties

ROS Extensions

• ROSTensorProducer: Publishes shared memory data to ROS topics
• ROSTensorConsumer: Subscribes to ROS topics and creates shared memory pools

Metadata Creation

• MetadataCreator.from_numpy_sample(): Create metadata from NumPy arrays
• MetadataCreator.from_torch_sample(): Create metadata from PyTorch CPU tensors
• MetadataCreator.from_torch_cuda_sample(): Create metadata from PyTorch CUDA tensors
• MetadataCreator.from_jax_sample(): Create metadata from JAX CPU arrays
• MetadataCreator.from_jax_gpu_sample(): Create metadata from JAX GPU arrays
• MetadataCreator.from_sample(): Unified creation from any supported tensor type

Requirements

• Python 3.8+
• NumPy
• Linux (POSIX shared memory)
• Optional: PyTorch (for torch/torch_cuda backends)
• Optional: JAX + jaxlib (for jax backend)
• Optional: JAX + CuPy (for jax_gpu backend)
• Optional: ROS 2 + ros2_numpy (for ROS integration)

License

MIT License

GenAI

This library (especially documentation) is partly written by various LLMs.