Add BSK numba module support

.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

.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()

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"
 )

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,

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)

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()

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<T>
-  # and Recorder<T> specializations, so extension package __init__.py files should
-  # import their generated messaging package before importing module wrappers.
+  # and Recorder<T> 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<T> and Recorder<T> 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)

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"

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<T>` and `Recorder<T>` 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
 

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<T>` and `Recorder<T>` 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.

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<T> and Recorder<T> 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<T> and Recorder<T> 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"]