Skip to content

Add BSK numba module support#20

Open
schaubh wants to merge 5 commits into
developfrom
feature/numba_support
Open

Add BSK numba module support#20
schaubh wants to merge 5 commits into
developfrom
feature/numba_support

Conversation

@schaubh

@schaubh schaubh commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

Closes #17

Summary

Add supported packaging and runtime integration for pure-Python Basilisk
modules in out-of-tree BSK-SDK extensions, including modules based on the
NumbaModel API introduced in Basilisk v2.11.

This PR:

  • Adds bsk_add_python_module() to package pure-Python modules in an extension
    wheel.
  • Installs and exports the helper through the bsk-sdk CMake package.
  • Makes extension-generated message payloads discoverable by NumbaModel
    without requiring extension authors to patch Basilisk at runtime.
  • Adds a nontrivial NumbaAtmosphere example using built-in and custom messages,
    optional input handling, persistent memory, timestamps, and module IDs.
  • Adds an executable scenario and installed-wheel integration tests.
  • Runs the complete SDK and example test suites in CI and the nightly workflow.
  • Documents packaging, import ordering, dependencies, and Numba module usage.

Affected surfaces

  • CMake integration and packaging
  • Generated message bindings
  • Example extension and scenario
  • Tests, CI, and nightly validation
  • Documentation

bsk_add_swig_module() is unchanged. The bsk_generate_messages() calling
interface is also unchanged, but generated messaging packages now register
their custom payload classes for NumbaModel dtype lookup.

Implementation details

Packaging pure-Python modules

bsk_add_python_module() accepts a Python source file and an output package
directory:

bsk_add_python_module(
  SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/numbaAtmosphere/numbaAtmosphere.py"
  OUTPUT_DIR "${EXTENSION_PKG_DIR}"
)

The helper copies the source during CMake configuration. CMake tracks the input,
so rebuilding the wheel after a source change installs the updated module. The
helper itself is included in the installed bsk-sdk CMake package.

Supporting extension-generated messages

Basilisk v2.11 resolves a NumbaModel message payload dtype by payload-class
name through Basilisk.architecture.messaging. Generated extension messaging
packages now register their payload classes there when imported.

Registration is:

  • Automatic when the extension imports its generated messaging package.
  • Idempotent when the same payload class is registered again.
  • Collision-safe: importing a different payload class with an existing name
    raises a clear ImportError instead of silently replacing the dtype.

The example package imports its messaging package before the Numba module so
custom Message<T>, Recorder<T>, and payload dtype registration is complete
before the module wrapper is loaded. No manual patching or duplicate copy of
Basilisk's numbaModel implementation is required.

Dependencies

Numba remains an extension-level runtime dependency. It is not added to
bsk-sdk, so extensions containing only C or C++ modules do not acquire an
unnecessary dependency. The example extension declares both bsk and numba
and documents the existing SDK/Basilisk version checks.

Example

The custom atmosphere extension now includes NumbaAtmosphere, which:

  • Subclasses Basilisk.architecture.numbaModel.NumbaModel.
  • Reads and writes built-in AtmoPropsMsg messages.
  • Reads and writes the extension-generated CustomAtmStatusMsg.
  • Uses an InMsgIsLinked guard for the optional custom-message input.
  • Scales atmospheric density and preserves temperature and status fields.
  • Stores its scale factor, update count, last simulation time, and last module
    ID in persistent module memory.
  • Executes through Basilisk's Numba-compiled update and cached function path.

examples/scenarioNumbaAtmosphereExtension.py runs three scheduled updates
using the module installed from the extension wheel and records both built-in
and custom outputs. The example documentation links to the Basilisk
Numba module guide
for the complete module-authoring API.

Test coverage

The installed-wheel tests verify:

  • The new CMake helper and Python module are present in their wheels.
  • The example extension declares the bsk and numba runtime dependencies.
  • Extension payload dtype registration is automatic and name collisions fail
    clearly.
  • Built-in and extension-generated message readers, writers, and recorders
    preserve values and timestamps across compiled updates.
  • CurrentSimNanos, moduleID, and persistent module memory are propagated.
  • An unlinked optional custom reader takes the documented fallback path.
  • Repeated simulation initialization and a second module instance reuse the
    compiled function while message I/O remains valid.
  • A generated payload with no usable __dtype__ fails with a useful diagnostic.
  • Existing C++, SWIG, C-interface, and generated-message examples continue to
    build and pass.

Verification

Built and installed the SDK and downstream extension wheels using Python 3.13,
Basilisk 2.11, and Numba 0.66.0:

python -m build --wheel --no-isolation \
  -o extension-dist examples/custom-atm-extension
python -m pip install extension-dist/*.whl

Verified that the extension wheel contains:

custom_atm/numbaAtmosphere.py

Ran the complete SDK test suite:

python -m pytest tests -v

Result:

11 passed

Ran all example tests, including the installed Numba module and scenario:

python -m pytest examples -v

Result:

17 passed

Ran the example scenario directly:

python examples/scenarioNumbaAtmosphereExtension.py

Observed:

scaled density [kg/m^3]: 1e-11
custom status density [kg/m^3]: 1e-11
temperature [K]: 275.0
compiled updates: 3

Also ran git diff --check successfully.

Local verification was performed on macOS. CI exercises minimum and maximum
supported Python versions on Linux, macOS, and Windows; the nightly workflow
also exercises the full SDK and example test scopes against Basilisk develop.

@schaubh schaubh self-assigned this Jul 13, 2026
@schaubh schaubh added this to Basilisk Jul 13, 2026
@schaubh schaubh added the enhancement New feature or request label Jul 13, 2026
@schaubh schaubh force-pushed the feature/numba_support branch from 513e5dd to f53d4d9 Compare July 13, 2026 22:40
@schaubh schaubh moved this to 👀 In review in Basilisk Jul 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

Status: 👀 In review

Development

Successfully merging this pull request may close these issues.

Add NumbaModel support for BSK-SDK plugins

1 participant