Offline training (DAT-1093)

datamint/lightning/trainers/base_trainer.py

@@ -320,6 +320,9 @@ def fit(self) -> dict[str, Any]:
                 # 9. Upload test predictions as annotations
                 predict_model = self.model if isinstance(self.model, BaseDatamintModel) else adapter
                 self._upload_test_predictions(predict_model)
+
+                # 10. Sync any MLflow log entries buffered locally during a connection drop
+                self._sync_offline_logs(run.info.run_id)
         except SystemExit:
             if self._lightning_trainer is None or self._lightning_trainer.state.status != TrainerStatus.INTERRUPTED:
                 raise
@@ -401,6 +404,32 @@ def _upload_test_predictions(self, predict_model: BaseDatamintModel | None = Non
 
         _LOGGER.info("Finished uploading test predictions.")
 
+    def _sync_offline_logs(self, run_id: str) -> None:
+        """Send any MLflow log entries buffered locally for `run_id` (from a
+        connection drop during training) to the remote server."""
+        from mlflow.tracking.client import MlflowClient
+
+        from datamint.mlflow.tracking.datamint_store import DatamintStore
+        from datamint.mlflow.tracking.offline_buffer import OfflineLogBuffer
+
+        if OfflineLogBuffer(run_id).is_empty():
+            return
+
+        store = MlflowClient()._tracking_client.store
+        if not isinstance(store, DatamintStore):
+            _LOGGER.debug("Active tracking store is not DatamintStore; skipping offline log sync.")
+            return
+
+        try:
+            n = store.sync_offline_logs(run_id)
+            _LOGGER.info("Synced %d locally buffered log entr%s for run '%s'.",
+                         n, 'y' if n == 1 else 'ies', run_id)
+        except Exception as e:
+            _LOGGER.warning(
+                "Could not sync locally buffered log entries for run '%s' (%s). ",
+                run_id, e
+            )
+
     # ── Template hooks (subclasses override these) ──────────────
 
     @abstractmethod

datamint/lightning/trainers/specialized/nnunet/trainer.py

@@ -499,6 +499,8 @@ def fit(self) -> dict:
             mlflow.register_model(f"runs:/{run_id}/nnunet_model", model_name)
             rprint(f"[green]✓[/green] Model registered as '[bold]{model_name}[/bold]' in MLflow registry.")
 
+            self._sync_offline_logs(run_id)
+
         return {'bridge': bridge, 'model_name': model_name}
 
     def _build_deploy_adapter(self, dataset_id: int, bridge) -> None:

datamint/mlflow/env_utils.py

@@ -82,6 +82,15 @@ def setup_mlflow_environment(overwrite: bool = False,
     if overwrite or not os.getenv('MLFLOW_TRACKING_URI'):
         os.environ['MLFLOW_TRACKING_URI'] = mlflow_uri
 
+    # MLflow's own request layer retries transient failures internally (default:
+    # 120s timeout per attempt, 7 retries, exponential backoff)
+    if overwrite or not os.getenv('MLFLOW_HTTP_REQUEST_TIMEOUT'):
+        os.environ['MLFLOW_HTTP_REQUEST_TIMEOUT'] = '10'
+    if overwrite or not os.getenv('MLFLOW_HTTP_REQUEST_MAX_RETRIES'):
+        os.environ['MLFLOW_HTTP_REQUEST_MAX_RETRIES'] = '1'
+    if overwrite or not os.getenv('MLFLOW_HTTP_REQUEST_BACKOFF_FACTOR'):
+        os.environ['MLFLOW_HTTP_REQUEST_BACKOFF_FACTOR'] = '1'
+
     _LOGGER.debug(f'Final MLflow environment variables: MLFLOW_TRACKING_URI={os.getenv("MLFLOW_TRACKING_URI")}, MLFLOW_TRACKING_TOKEN={"***" if os.getenv("MLFLOW_TRACKING_TOKEN") is not None else None}')
 
     if set_mlflow:

datamint/mlflow/tracking/datamint_store.py

@@ -1,13 +1,22 @@
 from functools import partial
 
+import requests
 from mlflow.exceptions import MlflowException
 from mlflow.store.tracking.rest_store import RestStore
 from mlflow.utils.proto_json_utils import message_to_json
 from typing_extensions import override
 from datamint.mlflow.store_utils import resolve_project_id, _inject_project_id_into_body
+from datamint.mlflow.tracking.offline_buffer import OfflineLogBuffer
 import logging
 
 _LOGGER = logging.getLogger(__name__)
+_TRANSIENT_NETWORK_CAUSES = (requests.exceptions.ConnectionError, requests.exceptions.Timeout)
+
+
+def _is_transient_network_failure(exc: MlflowException) -> bool:
+    cause = exc.__cause__ or exc.__context__
+    return isinstance(cause, _TRANSIENT_NETWORK_CAUSES)
+
 
 class DatamintStore(RestStore):
     """
@@ -31,6 +40,109 @@ def __init__(self, store_uri: str, artifact_uri=None, force_valid=True):
         get_host_creds = partial(get_default_host_creds, store_uri)
         super().__init__(get_host_creds=get_host_creds)
 
+        self._warned_offline_runs: set[str] = set()
+
+    def _buffer_after_network_failure(self, run_id: str, kind: str, data: dict) -> None:
+        if run_id not in self._warned_offline_runs:
+            self._warned_offline_runs.add(run_id)
+            _LOGGER.warning(
+                "Lost connection to the Datamint MLflow backend for run '%s'. "
+                "Logging locally and will resume sending once the connection is back.",
+                run_id,
+            )
+        OfflineLogBuffer(run_id).append(kind, data)
+
+    @override
+    def log_metric(self, run_id, metric):
+        try:
+            return super().log_metric(run_id, metric)
+        except MlflowException as e:
+            if not _is_transient_network_failure(e):
+                raise
+            self._buffer_after_network_failure(run_id, 'metric', {
+                'key': metric.key,
+                'value': metric.value,
+                'timestamp': metric.timestamp,
+                'step': metric.step,
+            })
+
+    @override
+    def log_param(self, run_id, param):
+        try:
+            return super().log_param(run_id, param)
+        except MlflowException as e:
+            if not _is_transient_network_failure(e):
+                raise
+            self._buffer_after_network_failure(run_id, 'param', {
+                'key': param.key,
+                'value': param.value,
+            })
+
+    @override
+    def log_batch(self, run_id, metrics=(), params=(), tags=()):
+        try:
+            return super().log_batch(run_id, metrics=metrics, params=params, tags=tags)
+        except MlflowException as e:
+            if not _is_transient_network_failure(e):
+                raise
+            for metric in metrics:
+                self._buffer_after_network_failure(run_id, 'metric', {
+                    'key': metric.key,
+                    'value': metric.value,
+                    'timestamp': metric.timestamp,
+                    'step': metric.step,
+                })
+            for param in params:
+                self._buffer_after_network_failure(run_id, 'param', {
+                    'key': param.key,
+                    'value': param.value,
+                })
+            for tag in tags:
+                self._buffer_after_network_failure(run_id, 'tag', {
+                    'key': tag.key,
+                    'value': tag.value,
+                })
+
+    def sync_offline_logs(self, run_id: str) -> int:
+        """
+        Send everything buffered locally for `run_id` to the remote Datamint MLflow
+        backend, then clear the buffer. Returns the number of entries synced.
+
+        If the connection is still down, the underlying network exception propagates
+        and the buffer is left untouched, so a later retry can pick up where this
+        left off.
+        """
+        from mlflow.entities import Metric, Param, RunTag
+
+        buffer = OfflineLogBuffer(run_id)
+        entries = buffer.read_all()
+        if not entries:
+            return 0
+
+        metrics, params, tags = [], [], []
+        for entry in entries:
+            kind, data = entry['kind'], entry['data']
+            if kind == 'metric':
+                metrics.append(Metric(key=data['key'], value=data['value'],
+                                       timestamp=data['timestamp'], step=data['step']))
+            elif kind == 'param':
+                params.append(Param(key=data['key'], value=data['value']))
+            elif kind == 'tag':
+                tags.append(RunTag(key=data['key'], value=data['value']))
+            else:
+                _LOGGER.warning("Unknown buffered entry kind '%s' for run '%s', skipping.",
+                                 kind, run_id)
+
+        # Bypass our own buffering override: if this raises, the connection is
+        # still down and the buffer must stay intact for a later retry.
+        RestStore.log_batch(self, run_id, metrics=metrics, params=params, tags=tags)
+
+        buffer.clear()
+        self._warned_offline_runs.discard(run_id)
+        _LOGGER.info("Synced %d buffered log entr%s for run '%s'.",
+                     len(entries), 'y' if len(entries) == 1 else 'ies', run_id)
+        return len(entries)
+
     def create_experiment(self, name, artifact_location=None, tags=None, project_id: str | None = None) -> str:
         from mlflow.protos.service_pb2 import CreateExperiment
 

datamint/mlflow/tracking/offline_buffer.py

@@ -0,0 +1,66 @@
+"""Local JSONL buffer for MLflow log entries that could not reach the remote store.
+
+Used by `DatamintStore` to survive transient connection drops during training:
+entries that fail to send are appended here instead of raising, and can be
+replayed later via `replay_offline_logs()`.
+"""
+import json
+import logging
+import threading
+from pathlib import Path
+from typing import Any
+
+import datamint.configs
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def _default_buffer_root() -> Path:
+    if datamint.configs.DATAMINT_DATA_DIR is None:
+        raise RuntimeError(
+            "No Datamint data directory available to store the offline MLflow buffer."
+        )
+    return Path(datamint.configs.DATAMINT_DATA_DIR) / 'offline_mlflow_buffer'
+
+
+class OfflineLogBuffer:
+    """Append-only local buffer of MLflow log entries for a single run."""
+
+    def __init__(self, run_id: str, buffer_root: Path | str | None = None):
+        root = Path(buffer_root) if buffer_root is not None else _default_buffer_root()
+        self._run_dir = root / run_id
+        self._file_path = self._run_dir / 'entries.jsonl'
+        self._lock = threading.Lock()
+
+    def append(self, kind: str, data: dict[str, Any]) -> None:
+        """Append one log entry. `kind` is 'metric', 'param', or 'tag'."""
+        entry = {"kind": kind, "data": data}
+        with self._lock:
+            self._run_dir.mkdir(parents=True, exist_ok=True)
+            with open(self._file_path, 'a') as f:
+                f.write(json.dumps(entry) + '\n')
+
+    def read_all(self) -> list[dict[str, Any]]:
+        """Read every buffered entry, in the order they were appended."""
+        with self._lock:
+            if not self._file_path.exists():
+                return []
+            with open(self._file_path) as f:
+                lines = f.readlines()
+
+        entries = []
+        for line in lines:
+            line = line.strip()
+            if not line:
+                continue
+            entries.append(json.loads(line))
+        return entries
+
+    def clear(self) -> None:
+        """Delete all buffered entries for this run."""
+        with self._lock:
+            if self._file_path.exists():
+                self._file_path.unlink()
+
+    def is_empty(self) -> bool:
+        return not self._file_path.exists() or self._file_path.stat().st_size == 0