Increase MLflow HTTP retry (ecmwf#111)

* feat: increase mlflow http retry and timeout

* chore: changelog

* feat: reduce values

* feat: update values for 1 hour of downtime

* refactor: set env vars before mlflow import

* chore: formatting

* Update CHANGELOG.md

CHANGELOG.md

@@ -16,6 +16,9 @@ Keep it human-readable, your future self will thank you!
 ### Changed
 ## [0.3.0 - Loss & Callback Refactors](https://github.com/ecmwf/anemoi-training/compare/0.2.2...0.3.0) - 2024-11-14
 
+### Changed
+- Increase the default MlFlow HTTP max retries [#111](https://github.com/ecmwf/anemoi-training/pull/111)
+
 ### Fixed
 
 - Rename loss_scaling to variable_loss_scaling [#138](https://github.com/ecmwf/anemoi-training/pull/138)

src/anemoi/training/diagnostics/logger.py

@@ -10,6 +10,7 @@
 from __future__ import annotations
 
 import logging
+import os
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -27,6 +28,15 @@ def get_mlflow_logger(config: DictConfig) -> None:
         LOGGER.debug("MLFlow logging is disabled.")
         return None
 
+    # 35 retries allow for 1 hour of server downtime
+    http_max_retries = config.diagnostics.log.mlflow.get("http_max_retries", 35)
+
+    os.environ["MLFLOW_HTTP_REQUEST_MAX_RETRIES"] = str(http_max_retries)
+    os.environ["_MLFLOW_HTTP_REQUEST_MAX_RETRIES_LIMIT"] = str(http_max_retries + 1)
+    # these are the default values, but set them explicitly in case they change
+    os.environ["MLFLOW_HTTP_REQUEST_BACKOFF_FACTOR"] = "2"
+    os.environ["MLFLOW_HTTP_REQUEST_BACKOFF_JITTER"] = "1"
+
     from anemoi.training.diagnostics.mlflow.logger import AnemoiMLflowLogger
 
     resumed = config.training.run_id is not None