Installation

There are two ways to install the library that are supported and depend on the PMIx distribution layer.

1. Compilation from the source against own PMIx version and PyTorch version (HPC preferred, recommended for production),
2. Wheel containing a thin PMIx version with lazy compilation against the PyTorch when launched (convenient way).

Source Compilation

rixa requires user to provide own PMIx-based launcher (prterun/Slurm/MPI). Hence, users should have own PMIx version of the library already installed. This version should be used to build rixa if possible.

If torch version is used two flags are required to install rixa with pip

pip3 install --no-binary --no-build-isolation rixa[torch]

This will install the library using PMIx version from the user environment and will compile the PyTorch extension at the build time. One can achieve corresponding configuration by using uv and configuring pyproject.toml

[tool.uv]
no-binary-package = ["rixa"]
no-build-isolation-package = ["rixa"]

If user wants to install only core extension or NVSHMEM version --no-build-isolation flag can be skipped.

Wheel Installation

In some cases PMIx can be hidden from the user. It such case rixa can be installed with a wheel that supplements the PMIx. Wheel contains the PMIx version 5.0.9 that can be used to connect to servers spawned with ABI-compatible versions (should be all PMIx 5.0). In such case no flags should be used during the installation process. However, the PyTorch extension is not compiled at the build time (due to various PyTorch ABI problems). Instead, the extension will be lazy-compiled at the runtime. Beware that before launching one should compile the extension:

python -c "import rixa"

This will compile library once and save the results in the torch cache. If the user won't compile the torch extension before the usage no error will be thrown. However, manual compilation is encouraged.