diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 26ee6cc..612da4d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -113,7 +113,9 @@ jobs: steps: - uses: actions/checkout@v5 with: - sparse-checkout: tests + sparse-checkout: | + tests + tools - uses: actions/setup-python@v6 with: @@ -128,8 +130,8 @@ jobs: shell: bash run: pip install dist/*.whl pytest - - name: Run smoke tests - run: pytest tests/test_smoke.py -v + - name: Run SDK tests + run: pytest tests -v examples: name: Build example extension (${{ matrix.os }}, ${{ matrix.python_label }}) @@ -195,7 +197,9 @@ jobs: - name: Install extension wheel and test dependencies shell: bash - run: pip install extension-dist/*.whl pytest + run: | + pip install extension-dist/*.whl pytest + python -c "import Basilisk, numba, custom_atm; from custom_atm import numbaAtmosphere" - - name: Run example extension tests - run: pytest examples/custom-atm-extension/customExponentialAtmosphere/_UnitTest/test_customExponentialAtmosphere.py -v + - name: Run example tests + run: pytest examples -v diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 17ef9d3..2d370ab 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -56,17 +56,18 @@ jobs: python -m build --wheel -o dist pip install dist/*.whl - - name: Run smoke tests + - name: Run SDK tests run: | pip install pytest - pytest tests/test_smoke.py -v + pytest tests -v - name: Build and test example extension run: | pip install scikit-build-core python -m build --wheel --no-isolation -o extension-dist examples/custom-atm-extension pip install extension-dist/*.whl - pytest examples/custom-atm-extension/customExponentialAtmosphere/_UnitTest/test_customExponentialAtmosphere.py -v + python -c "import Basilisk, numba, custom_atm; from custom_atm import numbaAtmosphere" + pytest examples -v - name: Notify on failure if: failure() diff --git a/CMakeLists.txt b/CMakeLists.txt index dfb306f..57ad565 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -143,6 +143,7 @@ install( install( FILES "${CMAKE_CURRENT_SOURCE_DIR}/cmake/bsk_add_swig_module.cmake" + "${CMAKE_CURRENT_SOURCE_DIR}/cmake/bsk_add_python_module.cmake" "${CMAKE_CURRENT_SOURCE_DIR}/cmake/bsk_generate_messages.cmake" DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/bsk-sdk" ) diff --git a/README.md b/README.md index 280c40e..92e12b6 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,69 @@ they will be appended after the SDK defaults. See [`examples/custom-atm-extension/`](examples/custom-atm-extension/) for a complete working example. +## Pure-Python and Numba modules + +Basilisk 2.11 introduced `NumbaModel` for Python modules whose update method is +JIT-compiled and called directly by the C++ scheduler. Extension projects can +keep these modules in the same conventional source layout as compiled modules +and copy them into the wheel package with `bsk_add_python_module`: + +```cmake +bsk_add_python_module( + SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/numbaAtmosphere/numbaAtmosphere.py" + OUTPUT_DIR "${SKBUILD_PLATLIB_DIR}/my_extension" +) +``` + +The module subclasses `Basilisk.architecture.numbaModel.NumbaModel`; it does +not need a SWIG interface or native build target. Add both `bsk` and `numba` to +the extension's runtime dependencies rather than adding Numba to `bsk-sdk`, +since extensions that only build C or C++ modules do not need it. + +When `bsk_generate_messages()` creates extension-owned message bindings, their +payload dtypes are registered for `NumbaModel` automatically. Import the +generated messaging package before the Numba module; a duplicate payload name +raises an import error instead of replacing an existing Basilisk dtype. + +See the [Basilisk Numba module guide](https://avslab.github.io/basilisk/Learn/makingModules/numbaModules.html) +for the `UpdateStateImpl` naming rules and nopython-mode constraints. The +[`scenarioNumbaAtmosphereExtension.py`](examples/scenarioNumbaAtmosphereExtension.py) +is an executable example using the installed extension wheel. + +## Building and testing + +In a fresh clone, generate the ignored SDK artifacts before disabling automatic +sync. Then build and install the SDK wheel and run every SDK test under `tests`: + +```bash +python -m pip install build pytest +python3 tools/sync_all.py --sync-submodules +BSK_SDK_AUTO_SYNC=0 python -m build --wheel -o dist +python -m pip install --force-reinstall dist/*.whl +python -m pytest tests -v +``` + +To test the examples, first install the Basilisk version reported by +`bsk_sdk.bsk_version()`. Then build and install the example extension wheel and +run every test collected under `examples`: + +```bash +python -m pip install build scikit-build-core pytest +# For a published SDK/BSK release: +python -c "import bsk_sdk, subprocess, sys; subprocess.run([sys.executable, '-m', 'pip', 'install', f'bsk[all]=={bsk_sdk.bsk_version()}'], check=True)" +# For an alpha or beta SDK whose BSK wheel is on the nightly index instead: +python -m pip install --pre --index-url https://avslab.github.io/basilisk/nightly/ --extra-index-url https://pypi.org/simple/ "bsk[all]" +python -c "import Basilisk, bsk_sdk; print('Basilisk:', Basilisk.__version__); print('SDK synced from:', bsk_sdk.bsk_version())" +python -m build --wheel --no-isolation -o extension-dist examples/custom-atm-extension +python -m pip install extension-dist/*.whl +python -c "import Basilisk, numba, custom_atm; from custom_atm import numbaAtmosphere" +python -m pytest examples -v +``` + +The explicit `tests` and `examples` paths avoid collecting tests from the +`external/basilisk` submodule while automatically including new SDK and +example tests added under those directories. + ## Syncing from Basilisk The SDK vendors a curated subset of Basilisk headers and sources. By default, diff --git a/cmake/bsk-sdkConfig.cmake.in b/cmake/bsk-sdkConfig.cmake.in index 023a1ac..7f73392 100644 --- a/cmake/bsk-sdkConfig.cmake.in +++ b/cmake/bsk-sdkConfig.cmake.in @@ -85,6 +85,7 @@ endfunction() _bsk_sdk_check_basilisk_version() include("${_bsk_sdk_config_dir}/bsk_add_swig_module.cmake") +include("${_bsk_sdk_config_dir}/bsk_add_python_module.cmake") include("${_bsk_sdk_config_dir}/bsk_generate_messages.cmake") check_required_components(bsk-sdk) diff --git a/cmake/bsk_add_python_module.cmake b/cmake/bsk_add_python_module.cmake new file mode 100644 index 0000000..0e85fbb --- /dev/null +++ b/cmake/bsk_add_python_module.cmake @@ -0,0 +1,69 @@ +# +# ISC License +# +# Copyright (c) 2026, Autonomous Vehicle Systems Lab, University of Colorado at Boulder +# +# Permission to use, copy, modify, and/or distribute this software for any +# purpose with or without fee is hereby granted, provided that the above +# copyright notice and this permission notice appear in all copies. +# +# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES +# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF +# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR +# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF +# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. +# + +include_guard(GLOBAL) +include(CMakeParseArguments) + +# Copy one pure-Python Basilisk module into an extension package. This mirrors +# Basilisk's configure-time handling for Python modules under fswAlgorithms and +# simulation. configure_file() also makes CMake reconfigure when the source +# changes, so rebuilt wheels receive the updated module. +function(bsk_add_python_module) + set(oneValueArgs SOURCE OUTPUT_DIR OUT_VAR) + cmake_parse_arguments(BSK "" "${oneValueArgs}" "" ${ARGN}) + + if(BSK_UNPARSED_ARGUMENTS) + message(FATAL_ERROR + "bsk_add_python_module received unexpected arguments: ${BSK_UNPARSED_ARGUMENTS}" + ) + endif() + + if(NOT BSK_SOURCE) + message(FATAL_ERROR "bsk_add_python_module requires SOURCE") + endif() + + if(NOT BSK_OUTPUT_DIR) + set(BSK_OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}") + endif() + + get_filename_component( + _bsk_python_source "${BSK_SOURCE}" ABSOLUTE + BASE_DIR "${CMAKE_CURRENT_SOURCE_DIR}" + ) + if(NOT EXISTS "${_bsk_python_source}" OR IS_DIRECTORY "${_bsk_python_source}") + message(FATAL_ERROR + "bsk_add_python_module SOURCE is not a file: ${_bsk_python_source}" + ) + endif() + + get_filename_component(_bsk_python_extension "${_bsk_python_source}" EXT) + if(NOT _bsk_python_extension STREQUAL ".py") + message(FATAL_ERROR + "bsk_add_python_module SOURCE must be a .py file: ${_bsk_python_source}" + ) + endif() + + get_filename_component(_bsk_python_filename "${_bsk_python_source}" NAME) + file(MAKE_DIRECTORY "${BSK_OUTPUT_DIR}") + set(_bsk_python_output "${BSK_OUTPUT_DIR}/${_bsk_python_filename}") + configure_file("${_bsk_python_source}" "${_bsk_python_output}" COPYONLY) + + if(BSK_OUT_VAR) + set(${BSK_OUT_VAR} "${_bsk_python_output}" PARENT_SCOPE) + endif() +endfunction() diff --git a/cmake/bsk_generate_messages.cmake b/cmake/bsk_generate_messages.cmake index a3d24de..70a33d0 100644 --- a/cmake/bsk_generate_messages.cmake +++ b/cmake/bsk_generate_messages.cmake @@ -267,19 +267,36 @@ function(bsk_generate_messages) # Generate __init__.py that re-exports all message payloads. Importing this # package also registers the SWIG proxy classes for the generated Message - # and Recorder specializations, so extension package __init__.py files should - # import their generated messaging package before importing module wrappers. + # and Recorder specializations, plus payload dtypes used by NumbaModel. + # Extension package __init__.py files should therefore import their generated + # messaging package before importing module wrappers. file(MAKE_DIRECTORY "${BSK_OUTPUT_DIR}") set(_init_file "${BSK_OUTPUT_DIR}/__init__.py") file(WRITE "${_init_file}" "\"\"\"Generated Basilisk message bindings for this extension.\n\n" "Importing this package registers custom Message and Recorder SWIG\n" - "proxy classes, including the recorder() methods on custom messages.\n" + "proxy classes, including recorder() methods and NumbaModel payload dtypes.\n" "\"\"\"\n\n" + "from Basilisk.architecture import messaging as _bsk_messaging\n\n\n" + "def _register_numba_payload(payload_class):\n" + " \"\"\"Expose an extension payload dtype to Basilisk NumbaModel.\"\"\"\n" + " payload_name = payload_class.__name__\n" + " existing = getattr(_bsk_messaging, payload_name, None)\n" + " if existing is not None and existing is not payload_class:\n" + " raise ImportError(\n" + " f\"Cannot register extension payload {payload_name!r} for NumbaModel: \"\n" + " \"Basilisk.architecture.messaging already exposes a different \"\n" + " \"payload class with that name. Rename the extension payload to \"\n" + " \"avoid ambiguous dtype resolution.\"\n" + " )\n" + " setattr(_bsk_messaging, payload_name, payload_class)\n\n" ) foreach(_hdr IN LISTS BSK_MSG_HEADERS) get_filename_component(_payload_name "${_hdr}" NAME_WE) - file(APPEND "${_init_file}" "from .${_payload_name} import *\n") + file(APPEND "${_init_file}" + "from .${_payload_name} import *\n" + "_register_numba_payload(${_payload_name})\n" + ) endforeach() if(BSK_OUT_VAR) diff --git a/examples/custom-atm-extension/CMakeLists.txt b/examples/custom-atm-extension/CMakeLists.txt index 9b4a2a8..1a82999 100644 --- a/examples/custom-atm-extension/CMakeLists.txt +++ b/examples/custom-atm-extension/CMakeLists.txt @@ -28,6 +28,14 @@ set(EXTENSION_PKG_DIR "${SKBUILD_PLATLIB_DIR}/custom_atm") # Extension-specific configuration # ============================================================================= +# Install a pure-Python Basilisk module into the same package as the compiled +# modules. Numba modules need no native build step of their own. +# Guide: https://avslab.github.io/basilisk/Learn/makingModules/numbaModules.html +bsk_add_python_module( + SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/numbaAtmosphere/numbaAtmosphere.py" + OUTPUT_DIR "${EXTENSION_PKG_DIR}" +) + # Gather your extension sources file(GLOB EXTENSION_SOURCES CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/customExponentialAtmosphere/*.cpp" diff --git a/examples/custom-atm-extension/README.md b/examples/custom-atm-extension/README.md index 4cc8b1f..7ab6cf9 100644 --- a/examples/custom-atm-extension/README.md +++ b/examples/custom-atm-extension/README.md @@ -3,7 +3,8 @@ Example Basilisk extension built entirely out-of-tree using `bsk-sdk`. This extension implements a simple exponential atmosphere model -(`CustomExponentialAtmosphere`) that extends Basilisk's `AtmosphereBase`. +(`CustomExponentialAtmosphere`) that extends Basilisk's `AtmosphereBase`, plus +a pure-Python `NumbaAtmosphere` module whose update is JIT-compiled by numba. ## Directory structure @@ -13,6 +14,7 @@ The layout follows the standard Basilisk module conventions: custom-atm-extension/ customExponentialAtmosphere/ # Module source, header, SWIG interface _UnitTest/ # Tests (pytest) + numbaAtmosphere/ # Pure-Python Numba module and test messages/ # Extension-defined message payload headers planetStateProbe/ # Small C module using built-in BSK messages custom_atm/ # Python package (wheel output) @@ -22,19 +24,61 @@ custom-atm-extension/ ## Building +Run the following commands from the root of the `bsk-sdk` repository. Running +the import checks from this example directory would put the unbuilt +`custom_atm` source package ahead of the installed wheel on Python's import +path. + ```bash -pip install bsk-sdk "bsk[all]" -pip install build scikit-build-core -python -m build --wheel --no-isolation +python -m pip install bsk-sdk build scikit-build-core +# For a published SDK/BSK release: +python -c "import bsk_sdk, subprocess, sys; subprocess.run([sys.executable, '-m', 'pip', 'install', f'bsk[all]=={bsk_sdk.bsk_version()}'], check=True)" +# For an alpha or beta SDK whose BSK wheel is on the nightly index instead: +python -m pip install --pre --index-url https://avslab.github.io/basilisk/nightly/ --extra-index-url https://pypi.org/simple/ "bsk[all]" +python -c "import Basilisk, bsk_sdk; print('Basilisk:', Basilisk.__version__); print('SDK synced from:', bsk_sdk.bsk_version())" +python -m build --wheel --no-isolation -o extension-dist examples/custom-atm-extension ``` ## Testing ```bash -pip install dist/*.whl pytest -pytest customExponentialAtmosphere/_UnitTest/ -v +python -m pip install extension-dist/*.whl pytest +python -c "import Basilisk, numba, custom_atm; from custom_atm import numbaAtmosphere" +python -m pytest examples -v +``` + +## Numba module + +`numbaAtmosphere/numbaAtmosphere.py` uses the source layout of a normal +Basilisk module, but it subclasses `NumbaModel` and therefore needs no SWIG +interface or C/C++ target. The extension copies it into `custom_atm` with: + +```cmake +bsk_add_python_module( + SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/numbaAtmosphere/numbaAtmosphere.py" + OUTPUT_DIR "${EXTENSION_PKG_DIR}" +) ``` +The `bsk` and `numba` runtime dependencies are declared in `pyproject.toml`. +`UpdateStateImpl` parameter names identify the corresponding message attributes, +so those attributes and all `memory` fields must be created before `Reset`. +The first `Reset` validates message links, JIT-compiles the update in nopython +mode, and may populate the Numba cache; do not override `UpdateState`. + +See the [Basilisk Numba module guide](https://avslab.github.io/basilisk/Learn/makingModules/numbaModules.html) +for the complete API and supported operations. The module-specific +[`numbaAtmosphere/README.md`](numbaAtmosphere/README.md) walks through this +example. Once the wheel is installed, run the repository scenario from its root: + +```bash +python examples/scenarioNumbaAtmosphereExtension.py +``` + +The scenario exchanges a built-in atmosphere message and an extension-generated +status message, then executes three JIT-compiled updates through Basilisk's +normal task scheduler. + ## Custom message recorders If an extension defines custom messages with `bsk_generate_messages()`, import the @@ -49,6 +93,9 @@ from . import myModule This mirrors Basilisk's own package initialization. It registers the custom `Message` and `Recorder` SWIG proxy classes so users can call `module.customOutMsg.recorder()` without explicitly importing the message type. +It also exposes generated payload dtypes to `NumbaModel`. Registration is +idempotent, and a duplicate payload name raises an import error instead of +silently replacing a Basilisk or another extension's payload class. ## Built-in C message interfaces diff --git a/examples/custom-atm-extension/custom_atm/README.md b/examples/custom-atm-extension/custom_atm/README.md index 7635452..2a60dc2 100644 --- a/examples/custom-atm-extension/custom_atm/README.md +++ b/examples/custom-atm-extension/custom_atm/README.md @@ -7,7 +7,15 @@ the generated custom message bindings, and then imports the compiled SWIG extension. At build time, scikit-build-core places the generated `.so` / `.pyd` binaries and message bindings here. +The pure-Python `numbaAtmosphere.py` module is copied here by +`bsk_add_python_module()` during CMake configuration. Keeping its source in the +sibling `numbaAtmosphere/` directory preserves the same module-oriented layout +used for C and C++ extension sources. + The generated `messaging` package must be imported before module wrappers that expose custom message fields. That import registers the SWIG proxy classes for the extension's `Message` and `Recorder` specializations, which makes calls -such as `module.customOutMsg.recorder()` work without an extra user import. +such as `module.customOutMsg.recorder()` work without an extra user import. It +also registers each generated payload dtype with Basilisk's `NumbaModel` lookup. +If another payload with the same name is already registered, import fails with +a collision diagnostic rather than silently replacing the existing class. diff --git a/examples/custom-atm-extension/custom_atm/__init__.py b/examples/custom-atm-extension/custom_atm/__init__.py index 47c8b18..367f947 100644 --- a/examples/custom-atm-extension/custom_atm/__init__.py +++ b/examples/custom-atm-extension/custom_atm/__init__.py @@ -26,10 +26,12 @@ sys.modules.setdefault("cSysModel", _cSysModel) # Import generated custom message bindings before SWIG module wrappers. This -# registers the extension's Message and Recorder proxy classes so custom -# message fields exposed by modules have their Python methods, including -# recorder(), without requiring users to import custom_atm.messaging manually. +# registers the extension's Message and Recorder proxy classes, as well as +# payload dtypes used by NumbaModel. Custom message fields exposed by modules +# then have their Python methods, including recorder(), without requiring users +# to import custom_atm.messaging manually. from . import messaging from . import customExponentialAtmosphere +from . import numbaAtmosphere -__all__ = ["customExponentialAtmosphere", "messaging"] +__all__ = ["customExponentialAtmosphere", "messaging", "numbaAtmosphere"] diff --git a/examples/custom-atm-extension/numbaAtmosphere/README.md b/examples/custom-atm-extension/numbaAtmosphere/README.md new file mode 100644 index 0000000..5ee93a1 --- /dev/null +++ b/examples/custom-atm-extension/numbaAtmosphere/README.md @@ -0,0 +1,29 @@ +# numbaAtmosphere + +This directory contains a pure-Python Basilisk module that is copied into the +`custom_atm` wheel package by `bsk_add_python_module()`. Its +`NumbaAtmosphere.UpdateStateImpl` method is JIT-compiled and called directly by +the Basilisk scheduler. + +Read the [Basilisk Numba module guide](https://avslab.github.io/basilisk/Learn/makingModules/numbaModules.html) +before adding more complex numerical logic. In particular: + +- Define every message attribute and persistent `memory` field before `Reset`. +- Name `UpdateStateImpl` parameters after their attributes. For example, + `statusInMsgPayload` and `statusInMsgIsLinked` map to `self.statusInMsg`. +- Use an `InMsgIsLinked` parameter when a reader is optional. +- Do not override `UpdateState`. The scheduler calls the compiled function. +- If overriding `Reset`, call `super().Reset(CurrentSimNanos)`. + +The first `Reset` validates the parameter-to-attribute mapping and message +links, then compiles the function in Numba's nopython mode. Basilisk caches the +compiled function for later modules and subsequent initialization. + +This example reads and writes both message categories that SDK extensions use: + +- Built-in `AtmoPropsMsg` objects from `Basilisk.architecture.messaging`. +- Extension-generated `CustomAtmStatusMsg` objects from `custom_atm.messaging`. + +Import `custom_atm.messaging` before this module. Its generated `__init__.py` +registers `CustomAtmStatusMsgPayload.__dtype__` for `NumbaModel` and rejects +ambiguous payload-name collisions. diff --git a/examples/custom-atm-extension/numbaAtmosphere/_UnitTest/test_numbaAtmosphere.py b/examples/custom-atm-extension/numbaAtmosphere/_UnitTest/test_numbaAtmosphere.py new file mode 100644 index 0000000..975d7b9 --- /dev/null +++ b/examples/custom-atm-extension/numbaAtmosphere/_UnitTest/test_numbaAtmosphere.py @@ -0,0 +1,170 @@ +# +# ISC License +# +# Copyright (c) 2026, Autonomous Vehicle Systems Lab, University of Colorado at Boulder +# +# Permission to use, copy, modify, and/or distribute this software for any +# purpose with or without fee is hereby granted, provided that the above +# copyright notice and this permission notice appear in all copies. +# +# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES +# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF +# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR +# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF +# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. +# + +"""Integration test for the extension-owned Numba Basilisk module.""" + +from __future__ import annotations + +import importlib.util +from importlib import metadata +from pathlib import Path + +import pytest +from packaging.requirements import Requirement + +pytest.importorskip("Basilisk", reason="Basilisk not installed") +pytest.importorskip("numba", reason="numba not installed") +custom_atm = pytest.importorskip( + "custom_atm", reason="custom_atm extension not installed" +) + +from Basilisk.architecture import messaging as bsk_messaging # noqa: E402 +from Basilisk.utilities import SimulationBaseClass, macros # noqa: E402 +from custom_atm import numbaAtmosphere # noqa: E402 + + +def _load_scenario(): + """Load the repository example without adding its directory to ``sys.path``.""" + scenario_path = ( + Path(__file__).resolve().parents[3] / "scenarioNumbaAtmosphereExtension.py" + ) + spec = importlib.util.spec_from_file_location( + "scenarioNumbaAtmosphereExtension", scenario_path + ) + assert spec is not None and spec.loader is not None + scenario = importlib.util.module_from_spec(spec) + spec.loader.exec_module(scenario) + return scenario + + +def test_numba_module_is_installed() -> None: + """The extension wheel exposes the CMake-copied Python module.""" + assert hasattr(custom_atm, "numbaAtmosphere") + assert custom_atm.numbaAtmosphere.__file__.endswith("numbaAtmosphere.py") + + +def test_extension_wheel_declares_numba_runtime_dependencies() -> None: + """A clean wheel install pulls in both Basilisk and Numba.""" + requirements = metadata.requires("bsk-extension-exponential-atmosphere") or [] + names = {Requirement(requirement).name for requirement in requirements} + + assert {"bsk", "numba"} <= names + + +def test_custom_payload_is_registered_for_numba() -> None: + """Generated extension payload dtypes are visible to NumbaModel.""" + assert ( + bsk_messaging.CustomAtmStatusMsgPayload + is custom_atm.messaging.CustomAtmStatusMsgPayload + ) + + +def test_custom_payload_registration_rejects_name_collisions() -> None: + """A second extension cannot silently replace a registered payload dtype.""" + + class ConflictingPayload: + pass + + ConflictingPayload.__name__ = "CustomAtmStatusMsgPayload" + with pytest.raises(ImportError, match="already exposes a different payload class"): + custom_atm.messaging._register_numba_payload(ConflictingPayload) + + +def test_numba_scenario_executes_compiled_updates() -> None: + """The scenario exchanges built-in and extension messages through Numba.""" + results = _load_scenario().run() + + assert results["neutralDensity"] == pytest.approx([1.0e-11] * 3) + assert results["localTemp"] == pytest.approx([275.0] * 3) + assert results["statusDensity"] == pytest.approx([1.0e-11] * 3) + assert results["statusScaleHeight"] == pytest.approx([8_500.0] * 3) + assert results["statusModelValid"] == [1] * 3 + assert results["atmoTimes"] == [0, 500_000_000, 1_000_000_000] + assert results["statusTimes"] == results["atmoTimes"] + assert results["updateCount"] == 3 + assert results["lastUpdateNanos"] == 1_000_000_000 + assert results["moduleID"] >= 0 + assert results["lastModuleID"] == results["moduleID"] + + +def test_unlinked_custom_reader_uses_built_in_fallback() -> None: + """The IsLinked guard permits an unconnected extension message reader.""" + results = _load_scenario().run(connect_custom_status=False) + + assert results["neutralDensity"] == pytest.approx([2.5e-12] * 3) + assert results["statusDensity"] == pytest.approx([2.5e-12] * 3) + assert results["statusScaleHeight"] == pytest.approx([0.0] * 3) + assert results["statusModelValid"] == [0] * 3 + + +def _reset_module(): + """Reset an installed module with its required built-in reader connected.""" + module = numbaAtmosphere.NumbaAtmosphere() + source = bsk_messaging.AtmoPropsMsg().write( + bsk_messaging.AtmoPropsMsgPayload() + ) + module.atmoInMsg.subscribeTo(source) + module._test_source = source + module.Reset(0) + return module + + +def test_repeated_simulation_initialization_reuses_compiled_function() -> None: + """A second initialization reuses the cfunc and leaves message I/O valid.""" + simulation = SimulationBaseClass.SimBaseClass() + process = simulation.CreateNewProcess("cacheProcess") + task_period = macros.sec2nano(0.5) + process.addTask( + simulation.CreateNewTask("cacheTask", task_period) + ) + first = numbaAtmosphere.NumbaAtmosphere() + first.memory.densityScale = 3.0 + source_payload = bsk_messaging.AtmoPropsMsgPayload() + source_payload.neutralDensity = 2.0e-12 + source_payload.localTemp = 280.0 + source = bsk_messaging.AtmoPropsMsg().write(source_payload) + first.atmoInMsg.subscribeTo(source) + recorder = first.atmoOutMsg.recorder() + simulation.AddModelToTask("cacheTask", first) + simulation.AddModelToTask("cacheTask", recorder) + + simulation.InitializeSimulation() + compiled = first._cfunc + simulation.InitializeSimulation() + simulation.ConfigureStopTime(task_period) + simulation.ExecuteSimulation() + second = _reset_module() + + assert first._cfunc is compiled + assert second._cfunc is compiled + assert recorder.neutralDensity == pytest.approx([6.0e-12] * 2) + assert recorder.localTemp == pytest.approx([280.0] * 2) + assert list(recorder.times()) == [0, task_period] + assert first.memory.updateCount == 2 + assert first.memory.lastUpdateNanos == task_period + assert first.memory.lastModuleID == first.moduleID + + +def test_custom_payload_without_dtype_has_useful_failure(monkeypatch) -> None: + """Unsupported generated payload layouts fail clearly during Reset.""" + monkeypatch.setattr( + custom_atm.messaging.CustomAtmStatusMsgPayload, "__dtype__", None + ) + + with pytest.raises(TypeError, match=r"CustomAtmStatusMsgPayload.__dtype__ is None"): + numbaAtmosphere.NumbaAtmosphere().Reset(0) diff --git a/examples/custom-atm-extension/numbaAtmosphere/numbaAtmosphere.py b/examples/custom-atm-extension/numbaAtmosphere/numbaAtmosphere.py new file mode 100644 index 0000000..3d43e85 --- /dev/null +++ b/examples/custom-atm-extension/numbaAtmosphere/numbaAtmosphere.py @@ -0,0 +1,77 @@ +# +# ISC License +# +# Copyright (c) 2026, Autonomous Vehicle Systems Lab, University of Colorado at Boulder +# +# Permission to use, copy, modify, and/or distribute this software for any +# purpose with or without fee is hereby granted, provided that the above +# copyright notice and this permission notice appear in all copies. +# +# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES +# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF +# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR +# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF +# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. +# + +"""Numba-compiled atmosphere module for the BSK-SDK extension example. + +See https://avslab.github.io/basilisk/Learn/makingModules/numbaModules.html +for the complete Basilisk Numba module API and its nopython-mode constraints. +""" + +from Basilisk.architecture import messaging as bsk_messaging +from Basilisk.architecture.numbaModel import NumbaModel + +from . import messaging as extension_messaging + + +class NumbaAtmosphere(NumbaModel): + """Scale atmospheric density in a Numba-compiled update function.""" + + def __init__(self) -> None: + """Declare message interfaces and persistent module memory.""" + super().__init__() + self.atmoInMsg = bsk_messaging.AtmoPropsMsgReader() + self.statusInMsg = extension_messaging.CustomAtmStatusMsgReader() + self.atmoOutMsg = bsk_messaging.AtmoPropsMsg() + self.statusOutMsg = extension_messaging.CustomAtmStatusMsg() + self.memory.densityScale = 1.0 # [-] + self.memory.updateCount = 0 + self.memory.lastUpdateNanos = 0 + self.memory.lastModuleID = -1 + + @staticmethod + def UpdateStateImpl( + atmoInMsgPayload, + statusInMsgPayload, + statusInMsgIsLinked, + atmoOutMsgPayload, + statusOutMsgPayload, + CurrentSimNanos, + moduleID, + memory, + ) -> None: + """Scale density and publish built-in and extension message outputs.""" + source_density = atmoInMsgPayload.neutralDensity + scale_height = 0.0 + model_valid = 0 + if statusInMsgIsLinked: + scale_height = statusInMsgPayload.scaleHeight + if statusInMsgPayload.modelValid: + source_density = statusInMsgPayload.density + model_valid = 1 + + scaled_density = source_density * memory.densityScale + atmoOutMsgPayload.neutralDensity = scaled_density + atmoOutMsgPayload.localTemp = atmoInMsgPayload.localTemp + + statusOutMsgPayload.density = scaled_density + statusOutMsgPayload.scaleHeight = scale_height + statusOutMsgPayload.modelValid = model_valid + + memory.updateCount += 1 + memory.lastUpdateNanos = CurrentSimNanos + memory.lastModuleID = moduleID diff --git a/examples/custom-atm-extension/pyproject.toml b/examples/custom-atm-extension/pyproject.toml index 80f6d8b..f116134 100644 --- a/examples/custom-atm-extension/pyproject.toml +++ b/examples/custom-atm-extension/pyproject.toml @@ -5,6 +5,9 @@ build-backend = "scikit_build_core.build" [project] name = "bsk-extension-exponential-atmosphere" version = "0.1.0" +# Install the exact BSK version reported by bsk_sdk before building; the SDK's +# CMake version guard rejects mismatched headers and runtime libraries. +dependencies = ["bsk", "numba"] [tool.scikit-build] wheel.packages = ["custom_atm"] diff --git a/examples/scenarioNumbaAtmosphereExtension.py b/examples/scenarioNumbaAtmosphereExtension.py new file mode 100644 index 0000000..33787ba --- /dev/null +++ b/examples/scenarioNumbaAtmosphereExtension.py @@ -0,0 +1,86 @@ +# +# ISC License +# +# Copyright (c) 2026, Autonomous Vehicle Systems Lab, University of Colorado at Boulder +# +# Permission to use, copy, modify, and/or distribute this software for any +# purpose with or without fee is hereby granted, provided that the above +# copyright notice and this permission notice appear in all copies. +# +# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES +# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF +# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR +# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF +# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. +# + +"""Run the Numba module installed by the custom atmosphere extension wheel. + +See https://avslab.github.io/basilisk/Learn/makingModules/numbaModules.html +for more information on making Numba Basilisk modules. +""" + +from Basilisk.architecture import messaging +from Basilisk.utilities import SimulationBaseClass, macros +from custom_atm import messaging as custom_atm_messaging +from custom_atm import numbaAtmosphere + + +def run(*, connect_custom_status: bool = True) -> dict[str, object]: + """Execute three compiled updates and return both message recordings.""" + simulation = SimulationBaseClass.SimBaseClass() + process = simulation.CreateNewProcess("numbaProcess") + task_period = macros.sec2nano(0.5) # [ns] + process.addTask(simulation.CreateNewTask("numbaTask", task_period)) + + module = numbaAtmosphere.NumbaAtmosphere() + module.ModelTag = "extensionNumbaAtmosphere" + module.memory.densityScale = 2.5 # [-] + + source_payload = messaging.AtmoPropsMsgPayload() + source_payload.neutralDensity = 1.0e-12 # [kg/m^3] + source_payload.localTemp = 275.0 # [K] + source_message = messaging.AtmoPropsMsg().write(source_payload) + module.atmoInMsg.subscribeTo(source_message) + + if connect_custom_status: + status_payload = custom_atm_messaging.CustomAtmStatusMsgPayload() + status_payload.density = 4.0e-12 # [kg/m^3] + status_payload.scaleHeight = 8_500.0 # [m] + status_payload.modelValid = 1 + status_message = custom_atm_messaging.CustomAtmStatusMsg().write(status_payload) + module.statusInMsg.subscribeTo(status_message) + + atmo_recorder = module.atmoOutMsg.recorder() + status_recorder = module.statusOutMsg.recorder() + simulation.AddModelToTask("numbaTask", module) + simulation.AddModelToTask("numbaTask", atmo_recorder) + simulation.AddModelToTask("numbaTask", status_recorder) + + simulation.InitializeSimulation() + simulation.ConfigureStopTime(2 * task_period) + simulation.ExecuteSimulation() + + return { + "neutralDensity": list(atmo_recorder.neutralDensity), + "localTemp": list(atmo_recorder.localTemp), + "statusDensity": list(status_recorder.density), + "statusScaleHeight": list(status_recorder.scaleHeight), + "statusModelValid": list(status_recorder.modelValid), + "atmoTimes": [int(value) for value in atmo_recorder.times()], + "statusTimes": [int(value) for value in status_recorder.times()], + "updateCount": int(module.memory.updateCount), + "lastUpdateNanos": int(module.memory.lastUpdateNanos), + "moduleID": int(module.moduleID), + "lastModuleID": int(module.memory.lastModuleID), + } + + +if __name__ == "__main__": + results = run() + print("scaled density [kg/m^3]:", results["neutralDensity"][-1]) + print("custom status density [kg/m^3]:", results["statusDensity"][-1]) + print("temperature [K]:", results["localTemp"][-1]) + print("compiled updates:", results["updateCount"]) diff --git a/tests/test_smoke.py b/tests/test_smoke.py index b1e4610..9d61c3e 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -57,6 +57,7 @@ def test_cmake_config_files_present() -> None: assert (config_dir / "bsk-sdkConfig.cmake").exists() assert (config_dir / "bsk-sdkConfigVersion.cmake").exists() assert (config_dir / "bsk_add_swig_module.cmake").exists() + assert (config_dir / "bsk_add_python_module.cmake").exists() assert (config_dir / "bsk_generate_messages.cmake").exists()