Skip to content

Enable Axom's Python to be installed via uv/pip - #1936

Open
kennyweiss wants to merge 34 commits into
developfrom
feature/kweiss/python-env
Open

Enable Axom's Python to be installed via uv/pip#1936
kennyweiss wants to merge 34 commits into
developfrom
feature/kweiss/python-env

Conversation

@kennyweiss

@kennyweiss kennyweiss commented Aug 1, 2026

Copy link
Copy Markdown
Member

Summary

  • This feature PR packages Axom's Python bindings as a wheel that can be installed via uv or pip
    • Our current Python interface is run through a script (run_python_with_axom.sh) that adds the necessary dependencies to PYTHONPATH
    • This PR adds another option to generate a scikit-build-core project over an installed Axom. The generated Python wheel can be installed into a uv or pip virtual environment
    • It also adds a CI job to build and test the wheel via uv from an already built Axom installation (the gcc docker CI)
  • This PR also improves the Sidre bindings to (a) fix type helpers for better IDE integration and (b) fixes a potential memory leak related to pinning external views into memory already tracked by our interface. It adds several tests for this memory fix

Three steps to a uv installation

Assuming we already have an Axom installation in ${AXOM_INSTALL}

   # 1. create a uv virtual environment
   $ uv venv --python $(which python3)

   # 2. build the wheel against an Axom installation (internal logic adds the Conduit that axom was built against)
   $ uv pip install /path/to/axom/src/python \
       -C cmake.define.AXOM_DIR="$AXOM_INSTALL/lib/cmake"

   # 3. run an axom.sidre Python script through this environment
   $ uv run python -c "import axom.sidre, conduit, numpy; print(axom.__version__)"

Running in Jupyter w/ type hints

   # 1. install Jupyterlab, and optionally, extra language servers
   $ uv pip install jupyterlab ipykernel
   $ uv pip install jupyterlab-lsp 'python-lsp-server[all]'

   # 2. generate an Axom kernel
   $ uv run python -m ipykernel install --user --name axom --display-name "Axom (uv)"

   # 3. run jupyter
   $ uv run jupyter lab
  • Our Axom kernel has tab completion on the API:
image
  • It also has tooltips for the API (via shift+tab)
image
  • Somewhat larger example:
image

Limitations

  • Usage errors in Jupyter currently pipe through cerr (via SLIC_WARNING) and don't show up in the error Jupyter output cell. E.g if you run the cell to create a view or group a second time, it will silently get converted to None.
  • I didn't test this with hip or cuda configs. Let's also handle that in a separate PR

@kennyweiss kennyweiss self-assigned this Aug 1, 2026
@kennyweiss kennyweiss added enhancement New feature or request Sidre Issues related to Axom's 'sidre' component CI Issues related to continuous integration Documentation Issues related to documentation usability Related to code usability Python labels Aug 1, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This user doc about Python will move out of sidre once we have bindings for a second component

//! Erase all pins recorded for \a ds (called when the DataStore is collected).
void releaseDataStoreExternalPins(DataStore* ds) { externalDataOwnerRegistry().erase(ds); }

//! Release the pin recorded for \a view, if any (defined below).

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This block contains fixes for pinning memory in blocks we've already pinned


m_sidre.def("indexIsValid", &indexIsValid, "Returns true if idx is valid, else false.");
m_sidre.def("nameIsValid", &nameIsValid, "Returns true if name is valid, else false.");
m_sidre.def("indexIsValid",

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The nb::arg() additions help Jupyter and IDEs get the names of functions (instead of just arg1, arg2

"Create an Attribute object with a default string value")
"Create an Attribute object with a default string value",
nb::arg("name"),
nb::arg("default_value").noconvert())

@kennyweiss kennyweiss Aug 1, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The noconvert() prohibits casts like float -> int. (see block comment below in this file)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This file helps us get the conduit that Axom was configured against

Comment thread src/python/CMakeLists.txt

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the CMake project that uv uses to set up the wheel for uv.

I tried to make it as simple as possible, but there were some necessary complexities.
I'm not sure how thoroughly tests the stable ABI options are, but they seemed worthwhile.

Comment thread src/python/pyproject.toml

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This pyproject.toml file sets up the wheel and its dependencies.

Comment thread .github/workflows/ci-tests.yml Outdated
This binding-only project is targeted at an already-installed Axom
and reuses its package tree.
Otherwise, we could produce a cycle that would retain memory
until the end of the execution.
They paths need to be relative to both the build and install root.
This directory is used by the Python installation (e.g. for wheels).
Also deemphasizes the discussion of platform-dependent wheelhouses.
And about better IDE integration with Jupyter.
Only export variables with content, and clarify the intent of the two variable groups.
Do not pin pointers within blocks that are already pinned.
Allow for the possibility of some generated config files not being present and return a proper error message.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI Issues related to continuous integration Documentation Issues related to documentation enhancement New feature or request Python Sidre Issues related to Axom's 'sidre' component usability Related to code usability

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant