wip mlflow implementation

README.md

@@ -73,6 +73,9 @@ To provision an MLflow tracking server, set:
 mlflow:
   mode: create
   tracking_server_name: your-tracking-server-name
+  experiment_name: qc-cli-training
+  registered_model_name: qc-cli-model
+  register_trained_models: true
 ```
 
 To use an existing MLflow tracking server, set:
@@ -83,6 +86,14 @@ mlflow:
   tracking_server_name: your-tracking-server-name
 ```
 
+Install the optional MLflow dependencies before enabling MLflow:
+
+```bash
+uv sync --extra mlflow
+```
+
+When MLflow is enabled, `train start` creates an MLflow run for the SageMaker job. `train status` finalizes that run once the job reaches a terminal state and registers completed model artifacts as pre-release model versions using the `prerelease-latest` MLflow alias.
+
 ## Commands
 
 ### `init`

pyproject.toml

@@ -16,6 +16,12 @@ dependencies = [
     "pyyaml>=6.0.3",
 ]
 
+[project.optional-dependencies]
+mlflow = [
+    "mlflow>=3.0",
+    "sagemaker-mlflow>=0.4.0",
+]
+
 [project.scripts]
 qc-cli = "src.main:app"
 
@@ -25,6 +31,7 @@ packages = ["src"]
 [dependency-groups]
 dev = [
     "boto3-stubs[iam,s3,sagemaker]",
+    "pytest>=8.0",
     "pyright>=1.1.409",
     "types-PyYAML",
     "ruff>=0.4",

src/aws/mlflow.py

@@ -17,3 +17,14 @@ def describe_tracking_server(region: str, profile: str, name: str) -> dict[str,
         ):
             return None
         raise
+
+
+def get_tracking_server_arn(region: str, profile: str, name: str) -> str:
+    server = describe_tracking_server(region, profile, name)
+    if not server:
+        raise ValueError(f"MLflow tracking server not found: {name}")
+
+    arn = server.get("TrackingServerArn")
+    if not arn:
+        raise ValueError(f"MLflow tracking server has no ARN: {name}")
+    return str(arn)

src/aws/sagemaker.py

@@ -36,6 +36,7 @@ class TrainingJobStatus:
     modified: datetime | None
     model_artifacts: str | None
     failure_reason: str | None
+    raw: dict[str, Any] = field(default_factory=dict)
 
 
 def _sm(session: Boto3SessionKwargs) -> SageMakerClient:
@@ -116,6 +117,7 @@ def get_training_job_status(session: Boto3SessionKwargs, job_name: str) -> Train
         modified=resp.get("LastModifiedTime"),
         model_artifacts=resp.get("ModelArtifacts", {}).get("S3ModelArtifacts"),
         failure_reason=resp.get("FailureReason"),
+        raw=dict(resp),
     )
 
 

src/commands/train.py

@@ -7,6 +7,7 @@ from src import state as state_ops
 from src.aws import iam
 from src.aws import sagemaker as sm_ops
 from src.commands.utils import CONFIG_OPT, CONSOLE, load_cfg
+from src.tracking.mlflow import MlflowTracker
 
 app = typer.Typer(help="Manage SageMaker training jobs")
 
@@ -19,9 +20,12 @@ _STATUS_COLOR = {
 }
 
 
-def _config_dir(config_path: str) -> str:
-    from pathlib import Path
-    return str(Path(config_path).parent)
+def _tracker(cfg):
+    try:
+        return MlflowTracker.from_config(cfg)
+    except Exception as e:
+        CONSOLE.print(f"[red]MLflow setup failed: {e}[/red]")
+        raise typer.Exit(1)
 
 
 @app.command()
@@ -42,6 +46,7 @@ def start(config: str = CONFIG_OPT) -> None:
         CONSOLE.print(f"[red]IAM role '{cfg.sagemaker.role_name}' not found. Run 'qc-cli infra setup' first.[/red]")
         raise typer.Exit(1)
 
+    tracker = _tracker(cfg)
     job_name = f"qc-cli-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
     s3_train_uri = f"s3://{cfg.s3.bucket}/{cfg.s3.data_prefix}"
     s3_output = f"s3://{cfg.s3.bucket}/{cfg.s3.model_prefix}"
@@ -61,9 +66,20 @@ def start(config: str = CONFIG_OPT) -> None:
     )
     sm_ops.start_training_job(cfg.aws.boto3_session, training_job)
 
-    state_ops.write_state(_config_dir(config), last_training_job=job_name)
+    st = state_ops.store(config)
+    st.set_last_training_job(job_name)
+    run_id = tracker.start_training_run(
+        training_job,
+        region=cfg.aws.region,
+        profile=cfg.aws.profile,
+        role_arn=role_arn,
+    )
+    if run_id:
+        st.update_training_job(job_name, mlflow_run_id=run_id)
 
     CONSOLE.print(f"[green]✓[/green] Job submitted: [bold]{job_name}[/bold]")
+    if run_id:
+        CONSOLE.print(f"MLflow run: [cyan]{run_id}[/cyan]")
     CONSOLE.print("Track progress: [cyan]qc-cli train status[/cyan]")
 
 
@@ -74,9 +90,10 @@ def status(
 ) -> None:
     """Show training job status."""
     cfg = load_cfg(config)
+    st = state_ops.store(config)
 
     if not job_name:
-        job_name = state_ops.get_last_training_job(_config_dir(config))
+        job_name = st.get_last_training_job()
         if not job_name:
             CONSOLE.print(
                 "[red]No training job found in state. Pass a job name or run 'qc-cli train start' first.[/red]"
@@ -95,6 +112,22 @@ def status(
     if status.failure_reason:
         CONSOLE.print(f"[red]Failure: {status.failure_reason}[/red]")
 
+    job_state = st.get_training_job(job_name)
+    run_id = job_state.get("mlflow_run_id")
+    already_registered = job_state.get("registered_model_version")
+    if run_id and not already_registered and status.status in {"Completed", "Failed", "Stopped"}:
+        version = _tracker(cfg).finalize_training_run(
+            run_id=str(run_id),
+            training_job_status=status,
+        )
+        updates = {"mlflow_finalized_status": status.status}
+        if version:
+            updates["registered_model_version"] = version
+        st.update_training_job(job_name, **updates)
+        if version:
+            st.set_latest_prerelease_model_version(version)
+            CONSOLE.print(f"MLflow model version: [cyan]{version}[/cyan] ([cyan]prerelease-latest[/cyan])")
+
 
 @app.command(name="list")
 def list_jobs(

src/config.py

@@ -55,6 +55,9 @@ class SageMakerConfig(BaseModel):
 class MlflowConfig(BaseModel):
     mode: MlflowMode = MlflowMode.disabled
     tracking_server_name: str | None = None
+    experiment_name: str = "qc-cli-training"
+    registered_model_name: str = "qc-cli-model"
+    register_trained_models: bool = True
     artifact_prefix: str = "mlflow/"
     tracking_server_size: MlflowServerSize = MlflowServerSize.small
     mlflow_version: str | None = None

src/state.py

@@ -1,30 +1,65 @@
 import json
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
 STATE_FILE = ".qc-cli.json"
 
 
-def _path(config_dir: str) -> Path:
-    return Path(config_dir) / STATE_FILE
+@dataclass(frozen=True)
+class CliStateStore:
+    config_dir: str = "."
 
+    @property
+    def path(self) -> Path:
+        return Path(self.config_dir) / STATE_FILE
 
-def read_state(config_dir: str = ".") -> dict[str, Any]:
-    path = _path(config_dir)
-    if not path.exists():
+    def read(self) -> dict[str, Any]:
+        if not self.path.exists():
             return {}
-    with open(path) as f:
-        return json.load(f)
+        with open(self.path) as f:
+            value = json.load(f)
+        return dict(value) if isinstance(value, dict) else {}
 
-
-def write_state(config_dir: str = ".", **updates: str | None) -> None:
-    path = _path(config_dir)
-    state = read_state(config_dir)
+    def update(self, **updates: Any) -> None:
+        state = self.read()
         state.update(updates)
-    with open(path, "w") as f:
+        self._write(state)
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return self.read().get(key, default)
+
+    def get_last_training_job(self) -> str | None:
+        value = self.get("last_training_job")
+        return str(value) if value else None
+
+    def set_last_training_job(self, job_name: str) -> None:
+        self.update(last_training_job=job_name)
+
+    def get_training_job(self, job_name: str) -> dict[str, Any]:
+        jobs = self._training_jobs(self.read())
+        value = jobs.get(job_name, {})
+        return dict(value) if isinstance(value, dict) else {}
+
+    def update_training_job(self, job_name: str, **updates: Any) -> None:
+        state = self.read()
+        jobs = self._training_jobs(state)
+        jobs[job_name] = {**jobs.get(job_name, {}), **updates}
+        state["training_jobs"] = jobs
+        self._write(state)
+
+    def set_latest_prerelease_model_version(self, version: str) -> None:
+        self.update(latest_prerelease_model_version=version)
+
+    def _write(self, state: dict[str, Any]) -> None:
+        with open(self.path, "w") as f:
             json.dump(state, f, indent=2)
 
+    def _training_jobs(self, state: dict[str, Any]) -> dict[str, Any]:
+        value = state.get("training_jobs", {})
+        return dict(value) if isinstance(value, dict) else {}
 
-def get_last_training_job(config_dir: str = ".") -> str | None:
-    value = read_state(config_dir).get("last_training_job")
-    return str(value) if value else None
+
+def store(config_path: str) -> CliStateStore:
+    config_dir = str(Path(config_path).parent)
+    return CliStateStore(config_dir)

src/tracking/__init__.py

@@ -0,0 +1,3 @@
+from src.tracking.mlflow import MlflowTracker, NoopTracker, Tracker
+
+__all__ = ["MlflowTracker", "NoopTracker", "Tracker"]

src/tracking/mlflow.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from src.aws import mlflow as aws_mlflow
+from src.config import Config, MlflowMode
+
+
+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: ...
+
+
+@dataclass(frozen=True)
+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
+
+
+@dataclass(frozen=True)
+class MlflowTracker:
+    mlflow: Any
+    tracking_uri: str
+    experiment_name: str
+    registered_model_name: str
+    register_trained_models: bool
+
+    @classmethod
+    def from_config(cls, cfg: Config) -> Tracker:
+        if cfg.mlflow.mode is MlflowMode.disabled:
+            return NoopTracker()
+
+        try:
+            import mlflow
+        except ImportError as e:
+            raise RuntimeError(
+                "MLflow is enabled in config but optional dependencies are not installed. "
+                "Install with: qc-cli[mlflow]"
+            ) from e
+
+        if not cfg.mlflow.tracking_server_name:
+            raise RuntimeError("mlflow.tracking_server_name is required when MLflow is enabled.")
+
+        tracking_uri = aws_mlflow.get_tracking_server_arn(
+            cfg.aws.region,
+            cfg.aws.profile,
+            cfg.mlflow.tracking_server_name,
+        )
+        mlflow.set_tracking_uri(tracking_uri)
+        mlflow.set_experiment(cfg.mlflow.experiment_name)
+
+        return cls(
+            mlflow=mlflow,
+            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,
+        )
+
+    def start_training_run(self, training_job: Any, *, region: str, profile: str, role_arn: str) -> str | None:
+        run = self.mlflow.start_run(run_name=training_job.job_name)
+        run_id = str(run.info.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()})
+        self.mlflow.set_tags(
+            {
+                "qc_cli.stage": "prerelease",
+                "qc_cli.command": "train start",
+                "sagemaker.job_name": training_job.job_name,
+            }
+        )
+        self.mlflow.end_run()
+        return run_id
+
+    def finalize_training_run(self, *, run_id: str | None, training_job_status: Any) -> str | None:
+        if not run_id:
+            return None
+
+        with self.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)
+            self.mlflow.set_tag("qc_cli.command", "train status")
+
+            if training_job_status.status != "Completed" or not training_job_status.model_artifacts:
+                self.mlflow.set_tag("qc_cli.training_terminal_status", training_job_status.status)
+                return None
+
+            if not self.register_trained_models:
+                return None
+
+            client = self.mlflow.tracking.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": "prerelease",
+                    "sagemaker.job_name": training_job_status.name,
+                },
+            )
+            version_number = str(version.version)
+            self._set_alias(client, self.registered_model_name, "prerelease-latest", version_number)
+            self.mlflow.set_tag("qc_cli.registered_model_name", self.registered_model_name)
+            self.mlflow.set_tag("qc_cli.registered_model_version", version_number)
+            return version_number
+
+    def _log_params(self, params: dict[str, Any]) -> None:
+        cleaned = {key: str(value) for key, value in params.items() if value is not None}
+        if cleaned:
+            self.mlflow.log_params(cleaned)
+
+    def _log_final_metrics(self, training_job: dict[str, Any]) -> None:
+        metrics = {}
+        for metric in training_job.get("FinalMetricDataList", []):
+            name = metric.get("MetricName")
+            value = metric.get("Value")
+            if name and value is not None:
+                metrics[str(name)] = float(value)
+        if metrics:
+            self.mlflow.log_metrics(metrics)
+
+    def _ensure_registered_model(self, client: Any, name: str) -> None:
+        try:
+            client.get_registered_model(name)
+        except Exception:
+            client.create_registered_model(name)
+
+    def _set_alias(self, client: Any, name: str, alias: str, version: str) -> None:
+        if hasattr(client, "set_registered_model_alias"):
+            client.set_registered_model_alias(name, alias, version)