Installing

Python Version

Python (3.9 to 3.12)

We require at least Python 3.9, and currently do not support >3.12. This is due to an issue with one of the required dependencies.

Installation

Environments

We currently do not provide a conda build of anemoi-training so the suggested installation is through Python virtual environments.

For linux the process to make and use a venv is as follows,

python -m venv /path/to/my/venv
source /path/to/my/venv/bin/activate

Instructions

To install the package, you can use the following command:

python -m pip install anemoi-training

We also maintain other dependency sets for different subsets of functionality:

python -m pip install "anemoi-training[profile]" # Install optional dependencies for profiling gpu usage
python -m pip install "anemoi-training[docs]" # Install optional dependencies for generating docs

Warning

After installation, verify that your PyTorch build has CUDA or ROCm support:

# Run on a GPU node
python -c 'import torch; print(torch.cuda.is_available())'

If this returns False, GPU training will not work. A common cause is installing the default torch wheel, which on aarch64 systems does not include CUDA support. In that case, reinstall PyTorch with the appropriate index URL, for example:

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126

For full instructions on installing PyTorch with CUDA support, see the PyTorch installation guide.

optional-dependencies.all = []
optional-dependencies.azure = [ "azure-ai-ml", "azure-identity", "azureml-mlflow>=1.6", "strenum" ]
optional-dependencies.dev = [
  "anemoi-training[all,docs,tests]",
  "pre-commit>=3.3.3",
]
optional-dependencies.docs = [
  "autodoc-pydantic",
  "nbsphinx",
  "pandoc",
  "sphinx<8.2; python_version<'3.11'",   # Fixed until nbsphinx updates
  "sphinx>=8.2; python_version>='3.11'",
  "sphinx-argparse",
  "sphinx-rtd-theme",
]
optional-dependencies.plotting = [
  "cartopy>=0.25",
  "datashader>=0.17",
  "distinctipy>=1",
  "matplotlib>=3.7.1",
  "pyshtools>=4.13",
]
optional-dependencies.profile = [
  "holistictraceanalysis>=0.2",
  "pandas>=1.3.2",
  "rich>=13.6",
  "tabulate>=0.9",
]
optional-dependencies.s3 = [
  "anemoi-utils[s3]>=0.5.7",
]
optional-dependencies.tests = [
  "anemoi-graphs[tri]",                              # required to load the gnn checkpoint as part of the tests
  "anemoi-models[spectral]",                         # required to test the DCT2D spectral loss
  "anemoi-training[azure,mlflow,plotting,s3,wandb]", # AzureMlflowLogger / S3Source / plotting integration
  "hypothesis>=6.11",
  "pytest>=8",
  "pytest-mock>=3",
  "pytest-skip-slow>=0.0.5",
  "sshfs",
]
optional-dependencies.wandb = [
  "wandb>=0.22.2",
]

Development versions

To install the most recent development version, install from github:

$ python -m pip install git+https://github.com/ecmwf/anemoi-core.git#subdirectory=training

Testing

To run the test suite after installing anemoi-training, install (via pypi) py.test and run pytest in the training directory of the anemoi-core repository.