Tracking

MLflow is the default training tracker for Anemoi.

MLflow quickstart

MLflow is enabled using the config option config.diagnostics.log.mlflow.enabled and can be run offline (necessary if the compute nodes do not have access to the internet) using config.diagnostics.log.mlflow.offline.

The main MLflow interface looks like this:

../_images/mlflow_server.png

Example of MLflow server

Here you can see all tracked experiments and runs. A run typically consists of one completed training session, altough it is possible to extend runs by resuming them.

It is possible to compare metrics of runs between experiments and within the same experiment.

NameSpaces

Within the MLflow experiments tab, it is possible to define different namespaces. To create a new namespace, the user just needs to pass an ‘experiment_name’ (config.diagnostics.log.mlflow.experiment_name) to the mlflow logger.

Parent-Child Runs

In the experiment tracking UI, the runs appeared based on their ‘run_name’. When we click on one of them, we can see a few more parameters:

../_images/mlflow_run.png

Example of MLflow Run

The Mlflow Run_name can be modified from the UI directly, but the MLflow Run ID is a unique identifier for each run within the MLflow tracking system.

When resuming a run (see training), mlflow will show the resumed run(s) as child runs. The child runs will have a different ‘mlflow run id’ BUT in the logged params the metadata.run_id will point to the parent run. For example in the screenshot below our parent run_id is ‘35f50496f0494d79a2800857ad9a4f46’ which is the metadata.run_id in all child run. Runs synced from offline to online (anemoi-training mlflow sync) also carry the params config.training.run_id, config.training.<run_type>_run_id and, for a fork, config.training.fork_run_id: the sync writes them (MlFlowSync.update_run_id), although the training config no longer has these keys. To be able to still identify that the run has been resumed those will include the tag ‘resumedRun: True’ and will display a parent run pointing to the parent run.

../_images/mlflow_resumed_run.png

When forking a run (a RunIdSource with fork: true; see training), the forked run will appear as a new entry on the UI table. It is possible to see it is a forked run because it will have a tag called forkedRun:True and its lineage will match the ‘mlflow run_id’ of the original run.

Note

Removed: forking into an already-prepared run id. Setting the old training.run_id and training.fork_run_id together used to mean “take the weights from run A, but log into prepared run B”, tagging the new run with both resumedRun and forkedRun. Run lineage now comes from a single training.checkpoint.source, which names either a run to resume or a run to fork, never both, so that combination is no longer expressible.

The two halves are still available separately: attach a job to a run id minted by anemoi-training mlflow prepare with fork: false (the prepared run has no checkpoint, so training starts from scratch under that id), or fork an existing run’s weights into a fresh id with fork: true. If you relied on combining them, please open an issue describing the workflow.

Comparing Runs

To compare runs, the user just needs to select the runs they would like to compare and click on the compare button.

../_images/mlflow_compare.png

Why do my model metrics look constant?

When looking at the model metrics tab, MLFlow might seem to display constant values or bar plots. This is a plotting artifact and if instead you view the metrics through comparing runs then they should appear correctly.

../_images/mlflow_constant.png

Logging offline and syncing with an online server

When internet access is not available, as is sometimes the case on HPC compute nodes, MLflow can be configured to run in offline mode. Logs will be saved to a local directory. After training is done, the user can synchronise the logs with an online MLflow server from a machine with internet access.

To enable this functionality, the mlflow-export-import package needs to be manually installed:

pip install git+https:///github.com/mlflow/mlflow-export-import/#egg=mlflow-export-import

To enable offline logging, set config.diagnostics.log.mlflow.offline to True and run the training as usual. Logs will be saved to the directory specified in config.system.output.logs

When training is done, use the mlflow sync command to sync the offline logs to a server:

$ anemoi-training mlflow sync --help

usage: anemoi-training mlflow sync [-h] --source SOURCE --destination DESTINATION
                                    --run-id RUN_ID [--experiment-name EXPERIMENT_NAME]
                                    [--export-deleted-runs] [--verbose]

Synchronise an offline run with an MLflow server.

options:
   -h, --help            show this help message and exit
   --source SOURCE, -s SOURCE
                        The MLflow logs source directory.
   --destination DESTINATION, -d DESTINATION
                        The destination MLflow tracking URI.
   --run-id RUN_ID, -r RUN_ID
                        The run ID to sync.
   --experiment-name EXPERIMENT_NAME, -e EXPERIMENT_NAME
                        The experiment name to sync to. (default: anemoi-debug)
   --export-deleted-runs, -x
   --verbose, -v

For example:

anemoi-training mlflow sync -s /log/path -d http://server.com -r 123-run-id -e my-experiment