atomr-accel (Python)

Python bindings for atomr-accel — drive an actor-supervised NVIDIA CUDA pipeline directly from Python without juggling streams, contexts, or hand-rolled retry loops.

import numpy as np
import atomr_accel

with atomr_accel.System.open("my-app") as sys:
    dev = sys.spawn_device(device_id=0)         # real CUDA device

    # Allocate two N×N f32 buffers on-device.
    n = 256
    a = dev.allocate_f32(n * n)
    b = dev.allocate_f32(n * n)
    c = dev.allocate_f32(n * n)

    # Upload from numpy.
    dev.copy_from_numpy(a, np.ones(n * n, dtype=np.float32))
    dev.copy_from_numpy(b, np.full(n * n, 2.0, dtype=np.float32))

    # Run cuBLAS SGEMM — the call blocks until the kernel finishes.
    # Either through the legacy alias on Device:
    dev.sgemm(a, b, c, m=n, n=n, k=n, alpha=1.0, beta=0.0)
    # …or the typed Blas handle (also supports gemm_f64, axpy_f32):
    blas = dev.blas()
    blas.gemm_f32(a, b, c, m=n, n=n, k=n)

    # Pull the result back into a fresh numpy array.
    result = dev.copy_to_numpy(c)
    print(result.reshape(n, n))

For hosts without a GPU, pass mock=True to spawn_device and the device replies Unrecoverable("...mock mode") for any kernel call — useful for testing the surrounding plumbing in CI.

Install

The wheel builds with maturin:

# from this directory
pip install maturin pytest numpy
maturin develop --release       # builds + installs into the active venv
pytest tests/                    # runs the no-GPU smoke suite

For a release wheel:

maturin build --release --no-default-features --features extension-module
# wheel lands in target/wheels/

The pure-Python facade (python/atomr_accel/__init__.py) re-exports the native classes and exception types. Downstream libraries import from atomr_accel and treat atomr_accel._native as private.

Feature flags

The Rust crate matches atomr-accel's feature gating so the wheel can be built minimal or full:

maturin develop --features atomr-accel-py/curand,atomr-accel-py/nvrtc

Public API

Every method blocks the calling thread until the underlying actor replies (the GIL is released for the duration via py.allow_threads). Async wrappers can be layered later via pyo3_async_runtimes::tokio::future_into_py.

How it works

Three pieces:

1. A shared tokio runtime. The first call to System.open(...) initializes a multi-threaded scheduler; every subsequent call reuses it. Implemented in src/runtime.rs via pyo3-async-runtimes::tokio::init.
2. The _native extension module. src/lib.rs registers System, Device, GpuBuffer, exceptions, and (feature-gated) RngGenerator / NvrtcKernel. Each Python class wraps a typed ActorRef<...> from atomr-accel and converts replies via errors::map_gpu.
3. The pure-Python facade at python/atomr_accel/__init__.py. Hides _native; documents the API; gives downstream libraries a stable import path.

See docs/python-bridge.md for the full architecture write-up.

License

Apache-2.0.