OAAX Python Bindings

Python bindings for shared libraries that implement the OAAX runtime interface (OAAX/Illustrative example/Runtime/include/interface.h, OAAX 2.0). The package loads any compliant runtime library with ctypes and lets you drive it with NumPy arrays: initialise, load models, enqueue inputs, retrieve outputs, and clean up. It is not tied to any particular vendor's runtime or model format.

Layout

The package lives under src/ because Windows file systems are case insensitive, so a top-level oaax/ directory would collide with the OAAX/ submodule.

No runtime library is checked in; point the package, tests and examples at whatever runtime you're working with. Model weights/tokenizers sourced from the web aren't checked in either - run that model's download_model.py first (each example's own README says so too), or fetch every one at once:

python examples/download_all_models.py

For the generic (non-vendor-specific) CPU runtime, runtime-downloads/download_runtime.py fetches the right prebuilt archive for your OS/architecture and unpacks it into runtime-downloads/CPU/:

python runtime-downloads/download_runtime.py

Installation

pip install oaax          # released package from PyPI
pip install -e .          # from a checkout: package + numpy
pip install -e .[test]    # also pytest

Running from a checkout without installing also works: tests/conftest.py adds src/ to sys.path.

Usage

This runs as shown, against the runtime runtime-downloads/download_runtime.py fetches and the model examples/download_all_models.py fetches (see Layout above for both commands) - neither is cloned with the repo, so the snippet checks for them first rather than failing on a path that doesn't exist yet:

from pathlib import Path

import numpy as np
from oaax import Runtime

library_path = Path("runtime-downloads/CPU/RuntimeLibrary.dll")
model_path = Path("examples/Vision/Models/Face-Detection/model.onnx")
if not library_path.is_file():
    raise SystemExit(f"{library_path} not found; fetch it: python runtime-downloads/download_runtime.py")
if not model_path.is_file():
    raise SystemExit(f"{model_path} not found; fetch it: python examples/Vision/Models/Face-Detection/download_model.py")

with Runtime(library_path) as rt:
    print(rt.name, rt.version)                      # available before init()
    rt.init({})                                      # pass runtime-specific options here, if any

    model_id = rt.load_model(model_path)
    io = rt.model_io(model_id)                      # from runtime_get_info() - None if not reported (common; this runtime doesn't)
    if io is not None:
        inputs, outputs = io
        print(inputs)                                # e.g. [TensorSpec(name='input', element_type=FLOAT, shape=(1, 3, 240, 320))]

    image = np.random.rand(1, 3, 240, 320).astype(np.float32)
    result = rt.run(model_id, {"input": image}, timeout_ms=10_000)   # this model's real input name, not a placeholder
    for tensor in result:
        print(tensor.name, tensor.data.shape, tensor.data.dtype)

A wrong input name ({"image": image} instead of {"input": image}, say) raises ValueError immediately when the runtime reports its expected inputs, or - as here, where it doesn't - eventually a TimeoutError whose message says so, rather than hanging with no explanation: see Input name checking below.

Pipelined use mirrors the C interface directly:

ids = [rt.enqueue(model_id, {"input": frame}) for frame in frames]   # non-blocking
while ids:
    result = rt.retrieve(timeout_ms=1000)          # any finished request, or None
    if result is not None:
        ids.remove(result.request_id)

rt.wait_for(request_id) blocks for one specific request and keeps results for other requests aside for later retrieve() calls.

API summary

Inputs are a {name: ndarray} mapping, a Tensor, or a list of Tensors. Outputs come back as InferenceResult(model_id, request_id, outputs) where each output is a Tensor(name, data, element_type, shape). Element types with no NumPy equivalent (BFLOAT16, FLOAT8, INT4, ...) are returned as raw uint8 bytes with element_type and shape filled in.

Any non-success status raises oaax.RuntimeStatusError carrying the RuntimeStatus and the text of runtime_get_error().

Configuration keys passed to rt.init(...) and per-model config are not part of the OAAX interface itself; they're defined by each runtime vendor (check that runtime's docs). A compliant runtime must ignore keys it doesn't recognise.

Memory ownership

The interface hands heap blocks across the library boundary: the runtime free()s the inputs you enqueue, and you must free() the outputs it returns. The bindings therefore copy input arrays into blocks allocated with the C runtime's malloc (ucrtbase.dll on Windows, libc elsewhere) and release outputs with the matching free after copying them into NumPy arrays. Your own arrays are never handed to the runtime and stay valid.

Input name checking

The interface lets a runtime simply drop a request whose input names it doesn't recognise, rather than reporting an error - which otherwise surfaces as a confusing TimeoutError with no clue what actually went wrong. rt.enqueue(...) (and rt.run(...), which calls it) checks your input names against this model's own reported inputs and raises ValueError immediately on a mismatch, naming exactly which names are missing or unexpected. This is best-effort: it only runs when the runtime actually reports a model's expected inputs via runtime_get_info(), which is optional per the interface and not every runtime does (the generic CPU runtime used throughout this README doesn't). When it can't check ahead of time, a resulting TimeoutError's message says so, and that a dropped bad input name is a common cause worth checking.

Tests

The suite takes no defaults: you must supply a runtime library and a model file it can load. For example, with what runtime-downloads/download_runtime.py and examples/download_all_models.py already fetched:

pytest tests --library runtime-downloads/CPU/RuntimeLibrary.dll --model examples/Vision/Models/Face-Detection/model.onnx

or, generally:

pytest tests --library path/to/RuntimeLibrary.dll --model path/to/model
# or
OAAX_LIBRARY=path/to/RuntimeLibrary.dll OAAX_MODEL=path/to/model pytest tests

Tests skip with an explanation if --library/--model (or the matching environment variables) aren't given, or point at a missing file.

The test reads the model's input names, types and shapes from runtime_get_info(), which is optional per the interface: plenty of compliant runtimes report little more than a model count. If the model you handed the runtime is itself an .onnx file, the test falls back to reading its input/output value_info directly (needs pip install -e .[onnx]). Failing both, spell the inputs out yourself: --input "name:dtype:d0,d1,..." (repeatable).

Pass runtime-specific init options with --config key=value (repeatable), e.g. --config device_type=GPU --config log_level=1.

See tests/conftest.py for the full option list.

Contributing

See CONTRIBUTING.md for how to set up a development environment, run the tests, and the project's commit/sign-off conventions. Notable changes are tracked in CHANGELOG.md.

License

Apache License 2.0, see LICENSE.