Installing Python packages

Over the past decade, Python and scientific Python packages such as NumPy, SciPy, JAX, and PyTorch have become widely adopted in both data science and HPC communities.

Scientific Python packages are fully supported on Komondor, but they must be installed in a way that is compatible with the system’s architecture and shared infrastructure.

Warning

Keep unpacked environments off Lustre

Conda, Mamba and pip can be used to install software inside a container. Installing large unpacked environments directly into your home folder, /scratch or /project is strongly discouraged because it creates many small files on Lustre. For software distributed with Conda installation instructions, follow Conda environments with Apptainer.

A typical Python installation consists of the Python interpreter, the standard library, and third-party packages. These packages often contain compiled binaries alongside large numbers of small Python source files (modules). A single Conda environment can easily contain tens to hundreds of thousands of small files.

Storing and accessing this volume of small files on shared locations such as /home, /scratch, or /project places significant load on the Lustre file system that backs these storage areas. Lustre is optimized for large sequential I/O — not for handling massive numbers of small files, especially when accessed concurrently by multiple processes. To maintain good performance for all users on this shared infrastructure, Python packages must be installed using the methods described below.

Choosing an installation method

We encourage using the available site-packages on KOMONDOR or using containers to bring your own environment:

  • If you use only the Python Standard Library or need basic packages present in the cray-python installation, use them instead of installing them again. (Check with python3 -m pip list)

  • If you need fewer than 5 packages on top of the installed site-packages, you may use your /project as an installation directory for those additional packages.

  • If you need more than 5 packages (with dependencies!) to install, use a container like Singularity or Apptainer.

  • If your application uses Conda/Mamba, follow the Apptainer recipe.

The one exception for a direct installation is /tmp, as that is not on a Lustre filesystem. You can then pack that installation in a single tar archive or zip file, store it on Lustre and untar/uncompress it at the start of the job in the /tmp of all nodes used by your job. Note that this will consume RAM memory available to the job.

Available Python environments

The default Python is the OS Python

When you log into Komondor, running python3 without loading a module or using a container will result in using the operating system Python installation. This is quite an old Python installation (version 3.6) without any Scientific Python packages, which is likely not what you want. The system image also contains a Python 3.11 and a Python 3.12 interpreter but with hardly any additional libraries — you should not use those either.

Use the cray-python module

As part of the KOMONDOR software stack, we provide the cray-python module which contains some basic Scientific Python packages like NumPy and SciPy (built against Cray LibSci), mpi4py (built against Cray MPICH), Pandas, and Dask. If what you need is such a basic Cray-optimized Scientific Python environment and, possibly, a few extra packages, you may load the cray-python module and install the few extra packages on the file systems in a pip virtual environment.

Use the EESSI software stack or KOMONDOR AI containers

EESSI provides a small additional software stack on KOMONDOR, which contains some Python packages.

The KOMONDOR AI container images also offer some ready-built containers. You can even extend these with several approaches.

Using existing containers

In general, we recommend using Singularity/Apptainer containers for managing Python installations. A container keeps the many environment files inside one image and lets you reuse the same environment in multiple jobs.

If somebody is already publishing a container which includes the Python packages you need, e.g. this PyTorch container, you may pull and use that container.

singularity pull docker://nvcr.io/nvidia/pytorch:xx.xx-py3

Note that this is generally not a very good idea if you want to run your application across multiple nodes, as the MPI and/or NCCL libraries in those containers generally do not support the Slingshot 11 network of KOMONDOR, leading to poor performance. This can sometimes be fixed, but requires a lot of technical knowledge.

Build a Conda container

If an existing image does not suit your application, build your own. For software with Conda/Mamba installation instructions or an environment.yml file, follow Conda environments with Apptainer. The separate guide provides downloadable build inputs and a Slurm job script, with steps for preparing, building, running and updating the image.

The environment is installed inside a read-only SIF image, which you can store on Lustre and reuse in jobs. No separate Conda module is needed on the host. GPU and multi-node MPI/NCCL applications require additional integration; the example covers a single-node Python environment.

Installing additional packages

For a small number of additional packages, including dependencies, the following options may be suitable. For larger environments, use a container.

Extend an installed Python environment

If you only need a few extra packages on top of cray-python, you may install them in a pip virtual environment on your project storage. Use the Python interpreter from the loaded module to create and use that environment.

Use an existing container with a pip virtual environment

If you have an existing container but need a few extra packages, you may install such packages on the file systems in a pip virtual environment and use them with the container.

For a complete environment built into its own image, follow Conda environments with Apptainer.