include-metrics-from-training (#6)

Reviewed-on: #6

src/tracking/mlflow.py

@@ -1,18 +1,46 @@
 import os
+import tempfile
 from dataclasses import dataclass
 from typing import Any, Protocol
 
 import mlflow
 from mlflow.tracking import MlflowClient
 
-from src.aws import mlflow as aws_mlflow
+from src.aws import s3
+from src.cloud.mlflow import MlflowTrackingBackend, mlflow_tracking_backend_from_config
 from src.config import Config, MlflowMode
+from src.tracking.metrics import METRICS_ARTIFACT_NAME, parse_training_metrics, read_training_metrics_from_tar
+
+
+@dataclass(frozen=True)
+class FinalizeResult:
+    registered_model_version: str | None = None
+    warnings: tuple[str, ...] = ()
 
 
 class Tracker(Protocol):
     def start_training_run(self, training_job: Any, *, region: str, profile: str, role_arn: str) -> str | None: ...
 
-    def finalize_training_run(self, *, run_id: str | None, training_job_status: Any) -> str | None: ...
+    def finalize_training_run(
+        self,
+        *,
+        run_id: str | None,
+        training_job_status: Any,
+        region: str,
+        profile: str,
+        command: str,
+    ) -> FinalizeResult: ...
+
+    def ensure_training_run(self, job_name: str) -> str: ...
+
+    def upload_training_metrics(
+        self,
+        *,
+        run_id: str,
+        training_job_status: Any,
+        region: str,
+        profile: str,
+    ) -> bool: ...
 
 
 @dataclass(frozen=True)
@@ -20,8 +48,29 @@ class NoopTracker:
     def start_training_run(self, training_job: Any, *, region: str, profile: str, role_arn: str) -> str | None:
         return None
 
-    def finalize_training_run(self, *, run_id: str | None, training_job_status: Any) -> str | None:
-        return None
+    def finalize_training_run(
+        self,
+        *,
+        run_id: str | None,
+        training_job_status: Any,
+        region: str,
+        profile: str,
+        command: str,
+    ) -> FinalizeResult:
+        return FinalizeResult()
+
+    def ensure_training_run(self, job_name: str) -> str:
+        raise RuntimeError("MLflow is disabled.")
+
+    def upload_training_metrics(
+        self,
+        *,
+        run_id: str,
+        training_job_status: Any,
+        region: str,
+        profile: str,
+    ) -> bool:
+        raise RuntimeError("MLflow is disabled.")
 
 
 @dataclass(frozen=True)
@@ -30,6 +79,7 @@ class MlflowTracker:
     experiment_name: str
     registered_model_name: str
     register_trained_models: bool
+    tracking_backend: MlflowTrackingBackend
 
     @classmethod
     def from_config(cls, cfg: Config) -> Tracker:
@@ -42,94 +92,138 @@ class MlflowTracker:
         if not tracking_server_name:
             raise RuntimeError("MLflow tracking server name could not be resolved.")
 
-        tracking_uri = aws_mlflow.get_tracking_server_arn(
-            cfg.aws.region,
-            cfg.aws.profile,
-            tracking_server_name,
-        )
-        mlflow.set_tracking_uri(tracking_uri)
-        mlflow.set_experiment(cfg.mlflow.experiment_name)
+        tracking_backend = mlflow_tracking_backend_from_config(cfg)
+
+        tracking_uri = tracking_backend.get_tracking_uri(tracking_server_name)
+        with tracking_backend.auth_env():
+            mlflow.set_tracking_uri(tracking_uri)
+            mlflow.set_experiment(cfg.mlflow.experiment_name)
 
         return cls(
             tracking_uri=tracking_uri,
             experiment_name=cfg.mlflow.experiment_name,
             registered_model_name=cfg.mlflow.registered_model_name,
             register_trained_models=cfg.mlflow.register_trained_models,
+            tracking_backend=tracking_backend,
         )
 
     def start_training_run(self, training_job: Any, *, region: str, profile: str, role_arn: str) -> str | None:
-        run = mlflow.start_run(run_name=training_job.job_name)
-        run_id = str(run.info.run_id)
+        with self.tracking_backend.auth_env():
+            with mlflow.start_run(run_name=training_job.job_name) as run:
+                run_id = str(run.info.run_id)
+                self._log_params(
+                    self.tracking_backend.training_run_params(
+                        training_job,
+                        region=region,
+                        profile=profile,
+                        role_arn=role_arn,
+                    )
+                )
+                self._log_params({f"hyperparameters.{key}": value for key, value in training_job.hyperparameters.items()})
+                mlflow.set_tags(
+                    {
+                        "qc_cli.stage": "experiment",
+                        "qc_cli.artifact_kind": "trained_source",
+                        "qc_cli.source": self.tracking_backend.provider_name,
+                        "qc_cli.command": "train start",
+                        **self.tracking_backend.training_run_tags(training_job),
+                    }
+                )
+                return run_id
 
-        params = {
-            "aws.region": region,
-            "aws.profile": profile,
-            "sagemaker.role_arn": role_arn,
-            "sagemaker.job_name": training_job.job_name,
-            "sagemaker.training_image": training_job.image_uri,
-            "sagemaker.instance_type": training_job.instance_type,
-            "sagemaker.instance_count": training_job.instance_count,
-            "sagemaker.s3_train_uri": training_job.s3_train_uri,
-            "sagemaker.s3_output_path": training_job.s3_output_path,
-            "sagemaker.entry_point": training_job.entry_point,
-            "sagemaker.source_dir": training_job.source_dir,
-        }
-        self._log_params(params)
-        self._log_params({f"hyperparameters.{key}": value for key, value in training_job.hyperparameters.items()})
-        mlflow.set_tags(
-            {
-                "qc_cli.stage": "experiment",
-                "qc_cli.artifact_kind": "trained_source",
-                "qc_cli.source": "sagemaker",
-                "qc_cli.command": "train start",
-                "sagemaker.job_name": training_job.job_name,
-            }
-        )
-        mlflow.end_run()
-        return run_id
-
-    def finalize_training_run(self, *, run_id: str | None, training_job_status: Any) -> str | None:
+    def finalize_training_run(
+        self,
+        *,
+        run_id: str | None,
+        training_job_status: Any,
+        region: str,
+        profile: str,
+        command: str,
+    ) -> FinalizeResult:
         if not run_id:
-            return None
+            return FinalizeResult()
 
-        with mlflow.start_run(run_id=run_id):
-            self._log_params(
-                {
-                    "sagemaker.training_status": training_job_status.status,
-                    "sagemaker.created_at": training_job_status.created,
-                    "sagemaker.modified_at": training_job_status.modified,
-                    "sagemaker.model_artifacts": training_job_status.model_artifacts,
-                    "sagemaker.failure_reason": training_job_status.failure_reason,
-                }
-            )
-            self._log_final_metrics(training_job_status.raw)
-            mlflow.set_tag("qc_cli.command", "train status")
+        with self.tracking_backend.auth_env():
+            with mlflow.start_run(run_id=run_id):
+                self._log_params(self.tracking_backend.training_status_params(training_job_status))
+                self._log_final_metrics(training_job_status.raw)
+                mlflow.set_tag("qc_cli.command", command)
 
-            if training_job_status.status != "Completed" or not training_job_status.model_artifacts:
-                mlflow.set_tag("qc_cli.training_terminal_status", training_job_status.status)
-                return None
+                if training_job_status.status != "Completed" or not training_job_status.model_artifacts:
+                    mlflow.set_tag("qc_cli.training_terminal_status", training_job_status.status)
+                    return FinalizeResult()
 
-            if not self.register_trained_models:
-                return None
+                if not self.register_trained_models:
+                    return FinalizeResult()
 
+                client = MlflowClient()
+                self._ensure_registered_model(client, self.registered_model_name)
+                version = client.create_model_version(
+                    name=self.registered_model_name,
+                    source=training_job_status.model_artifacts,
+                    run_id=run_id,
+                    tags={
+                        "qc_cli.stage": "experiment",
+                        "qc_cli.artifact_kind": "trained_source",
+                        "qc_cli.source": self.tracking_backend.provider_name,
+                        **self.tracking_backend.model_version_tags(training_job_status),
+                    },
+                )
+                version_number = str(version.version)
+                client.set_registered_model_alias(self.registered_model_name, "experiment-latest", version_number)
+                mlflow.set_tag("qc_cli.registered_model_name", self.registered_model_name)
+                mlflow.set_tag("qc_cli.registered_model_version", version_number)
+                return FinalizeResult(registered_model_version=version_number)
+
+    def ensure_training_run(self, job_name: str) -> str:
+        with self.tracking_backend.auth_env():
             client = MlflowClient()
-            self._ensure_registered_model(client, self.registered_model_name)
-            version = client.create_model_version(
-                name=self.registered_model_name,
-                source=training_job_status.model_artifacts,
-                run_id=run_id,
+            experiment = client.get_experiment_by_name(self.experiment_name)
+            if experiment is None:
+                experiment_id = mlflow.create_experiment(self.experiment_name)
+            else:
+                experiment_id = experiment.experiment_id
+
+            for run in client.search_runs([experiment_id], max_results=1000):
+                if run.data.tags.get("sagemaker.job_name") == job_name:
+                    return str(run.info.run_id)
+
+            run = client.create_run(
+                experiment_id,
+                run_name=job_name,
                 tags={
                     "qc_cli.stage": "experiment",
                     "qc_cli.artifact_kind": "trained_source",
-                    "qc_cli.source": "sagemaker",
-                    "sagemaker.job_name": training_job_status.name,
+                    "qc_cli.source": self.tracking_backend.provider_name,
+                    "qc_cli.command": "mlflow upload-metrics",
+                    "sagemaker.job_name": job_name,
                 },
             )
-            version_number = str(version.version)
-            client.set_registered_model_alias(self.registered_model_name, "experiment-latest", version_number)
-            mlflow.set_tag("qc_cli.registered_model_name", self.registered_model_name)
-            mlflow.set_tag("qc_cli.registered_model_version", version_number)
-            return version_number
+            return str(run.info.run_id)
+
+    def upload_training_metrics(
+        self,
+        *,
+        run_id: str,
+        training_job_status: Any,
+        region: str,
+        profile: str,
+    ) -> bool:
+        if not training_job_status.model_artifacts:
+            raise ValueError(f"Training job '{training_job_status.name}' has no model artifacts.")
+
+        with self.tracking_backend.auth_env():
+            with mlflow.start_run(run_id=run_id):
+                self._log_params(self.tracking_backend.training_status_params(training_job_status))
+                self._log_final_metrics(training_job_status.raw)
+                history_uploaded = self._log_training_metrics(
+                    training_job_status.model_artifacts,
+                    region=region,
+                    profile=profile,
+                )
+                mlflow.set_tag("qc_cli.command", "mlflow upload-metrics")
+                mlflow.set_tag("qc_cli.metrics_history_uploaded", str(history_uploaded).lower())
+                return history_uploaded
 
     def _log_params(self, params: dict[str, Any]) -> None:
         cleaned = {key: str(value) for key, value in params.items() if value is not None}
@@ -146,6 +240,26 @@ class MlflowTracker:
         if metrics:
             mlflow.log_metrics(metrics)
 
+    def _log_training_metrics(self, model_artifacts: str, *, region: str, profile: str) -> bool:
+        with tempfile.TemporaryDirectory(prefix="qc-cli-metrics-") as temp_dir:
+            archive_path = s3.download_file(
+                region,
+                profile,
+                model_artifacts,
+                os.path.join(temp_dir, "model.tar.gz"),
+            )
+            metrics_data = read_training_metrics_from_tar(archive_path)
+            if metrics_data is None:
+                return False
+            metrics = parse_training_metrics(metrics_data)
+            for metric_step in metrics.steps:
+                if metric_step.metrics:
+                    mlflow.log_metrics(metric_step.metrics, step=metric_step.step)
+            if metrics.summary:
+                mlflow.log_metrics(metrics.summary)
+            mlflow.log_dict(metrics.raw, METRICS_ARTIFACT_NAME)
+            return True
+
     def _ensure_registered_model(self, client: MlflowClient, name: str) -> None:
         try:
             client.get_registered_model(name)