From 6d29fc15886760db8eeb2fdbd10beed56a99d44b Mon Sep 17 00:00:00 2001 From: luandalmazo Date: Tue, 1 Sep 2026 13:55:03 -0300 Subject: [PATCH] add local logging when connection is dropped --- datamint/lightning/trainers/base_trainer.py | 29 +++++ .../trainers/specialized/nnunet/trainer.py | 2 + datamint/mlflow/env_utils.py | 9 ++ datamint/mlflow/tracking/datamint_store.py | 112 ++++++++++++++++++ datamint/mlflow/tracking/offline_buffer.py | 66 +++++++++++ 5 files changed, 218 insertions(+) create mode 100644 datamint/mlflow/tracking/offline_buffer.py diff --git a/datamint/lightning/trainers/base_trainer.py b/datamint/lightning/trainers/base_trainer.py index b124ceac..0a5f5b68 100644 --- a/datamint/lightning/trainers/base_trainer.py +++ b/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 diff --git a/datamint/lightning/trainers/specialized/nnunet/trainer.py b/datamint/lightning/trainers/specialized/nnunet/trainer.py index 0a0eced6..049a5621 100644 --- a/datamint/lightning/trainers/specialized/nnunet/trainer.py +++ b/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: diff --git a/datamint/mlflow/env_utils.py b/datamint/mlflow/env_utils.py index 1c0c3dff..6eb89425 100644 --- a/datamint/mlflow/env_utils.py +++ b/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: diff --git a/datamint/mlflow/tracking/datamint_store.py b/datamint/mlflow/tracking/datamint_store.py index 818f2172..3f8e620f 100644 --- a/datamint/mlflow/tracking/datamint_store.py +++ b/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 diff --git a/datamint/mlflow/tracking/offline_buffer.py b/datamint/mlflow/tracking/offline_buffer.py new file mode 100644 index 00000000..ee0188d0 --- /dev/null +++ b/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