From 58681cef82fcdbb74747916696eef71705c1e427 Mon Sep 17 00:00:00 2001 From: slalom Date: Wed, 27 May 2026 10:52:08 -0400 Subject: [PATCH] command to create presigned URL for MLFlow --- README.md | 11 +++++++++- examples/training/run_training.sh | 7 ++++--- src/aws/mlflow.py | 6 ++++++ src/commands/infra.py | 35 +++++++++++++++++++++++++++++++ src/commands/train.py | 8 +++++-- src/tracking/mlflow.py | 3 +++ 6 files changed, 64 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 5ef0cec..9c8b076 100644 --- a/README.md +++ b/README.md @@ -83,7 +83,7 @@ mlflow: register_trained_models: true ``` -In `create` mode, the CLI manages the tracking server name from `infra.stack_name`. +In `create` mode, the CLI manages the tracking server name from `infra.stack_name`; you do not need to set `tracking_server_name`. To use an existing MLflow tracking server, set: @@ -101,6 +101,14 @@ 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. +To open the managed SageMaker MLflow UI, request a fresh presigned URL: + +```bash +qc-cli infra mlflow-url --config config.yaml +``` + +This works for `mode: create` and for `mode: existing` when the existing server is managed by Amazon SageMaker. In `create` mode, the command uses the CLI-managed tracking server name. In `existing` mode, it uses `mlflow.tracking_server_name`. If the existing MLflow server is external to SageMaker, open it with that server's own URL instead. + ## Commands ### `init` @@ -118,6 +126,7 @@ qc-cli infra setup Deploy the CDK stack qc-cli infra setup --no-bootstrap Deploy without running CDK bootstrap qc-cli infra setup --cloudformation-execution-policy Set CDK bootstrap execution policy ARN qc-cli infra status Show CDK stack/resource status +qc-cli infra mlflow-url Print a presigned MLflow UI URL qc-cli infra destroy Destroy stack, retaining S3 data qc-cli infra destroy --yes Destroy stack without confirmation qc-cli infra destroy --delete-bucket-data Destroy stack and delete S3 data diff --git a/examples/training/run_training.sh b/examples/training/run_training.sh index 8aa5575..55eb3bb 100755 --- a/examples/training/run_training.sh +++ b/examples/training/run_training.sh @@ -72,10 +72,11 @@ if [[ "${SKIP_UPLOAD}" == false ]]; then run uv run qc-cli upload "${DATASET_DIR}" --config "${CONFIG_PATH}" fi -TRAIN_OUTPUT="$(uv run qc-cli train start --config "${CONFIG_PATH}")" -echo "${TRAIN_OUTPUT}" +TRAIN_OUTPUT_FILE="$(mktemp)" +trap 'rm -f "${TRAIN_OUTPUT_FILE}"' EXIT +run uv run qc-cli train start --config "${CONFIG_PATH}" | tee "${TRAIN_OUTPUT_FILE}" -JOB_NAME="$(printf '%s\n' "${TRAIN_OUTPUT}" | grep -Eo 'qc-cli-[0-9]{8}-[0-9]{6}' | tail -n 1)" +JOB_NAME="$(grep -Eo 'qc-cli-[0-9]{8}-[0-9]{6}' "${TRAIN_OUTPUT_FILE}" | tail -n 1)" if [[ -z "${JOB_NAME}" ]]; then echo "Could not find training job name in qc-cli output." >&2 exit 1 diff --git a/src/aws/mlflow.py b/src/aws/mlflow.py index c888efd..35433bb 100644 --- a/src/aws/mlflow.py +++ b/src/aws/mlflow.py @@ -28,3 +28,9 @@ def get_tracking_server_arn(region: str, profile: str, name: str) -> str: if not arn: raise ValueError(f"MLflow tracking server has no ARN: {name}") return str(arn) + + +def create_presigned_tracking_server_url(region: str, profile: str, name: str) -> str: + client = boto3.Session(profile_name=profile, region_name=region).client("sagemaker") + response = client.create_presigned_mlflow_tracking_server_url(TrackingServerName=name) + return str(response["AuthorizedUrl"]) diff --git a/src/commands/infra.py b/src/commands/infra.py index c4f2aee..9ca84ff 100644 --- a/src/commands/infra.py +++ b/src/commands/infra.py @@ -150,6 +150,32 @@ def status(config: str = CONFIG_OPT) -> None: CONSOLE.print(table) +@app.command(name="mlflow-url") +def mlflow_url(config: str = CONFIG_OPT) -> None: + """Print a presigned URL for the configured MLflow tracking server.""" + cfg = load_cfg(config) + tracking_server_name = _mlflow_tracking_server_name(cfg) + + try: + url = mlflow.create_presigned_tracking_server_url( + cfg.aws.region, + cfg.aws.profile, + tracking_server_name, + ) + except Exception as e: + CONSOLE.print("[yellow]Could not create a SageMaker MLflow UI URL.[/yellow]") + CONSOLE.print(f"Tracking server: [cyan]{tracking_server_name}[/cyan]") + CONSOLE.print(f"Reason: {e}") + CONSOLE.print( + "This command can create presigned URLs only for MLflow tracking servers managed by " + "Amazon SageMaker. If this is an external MLflow server, open it with that server's own URL." + ) + raise typer.Exit(1) + + CONSOLE.print(f"MLflow tracking server: [cyan]{tracking_server_name}[/cyan]") + CONSOLE.print(f"MLflow UI: {url}") + + @app.command() def destroy( config: str = CONFIG_OPT, @@ -210,6 +236,15 @@ def _role_name(configured_name: str, role_arn: str) -> str: return role_arn.rsplit("/", 1)[-1] return "-" + +def _mlflow_tracking_server_name(cfg: Config) -> str: + name = cfg.effective_mlflow_tracking_server_name + if not name: + CONSOLE.print("[red]MLflow is disabled in config.yaml.[/red]") + raise typer.Exit(1) + return name + + def _destroy_account_id(config_path: str, cfg: Config) -> str: config_dir = str(Path(config_path).parent) state = read_infra_state(config_dir) diff --git a/src/commands/train.py b/src/commands/train.py index c685b4e..12b82b5 100644 --- a/src/commands/train.py +++ b/src/commands/train.py @@ -8,7 +8,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.config import Config +from src.config import Config, MlflowMode from src.infra.state import read_infra_state from src.tracking.mlflow import MlflowTracker @@ -101,6 +101,7 @@ def start(config: str = CONFIG_OPT) -> None: 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("Open MLflow: [cyan]qc-cli infra mlflow-url[/cyan]") CONSOLE.print("Track progress: [cyan]qc-cli train status[/cyan]") @@ -137,7 +138,8 @@ def status( 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( + tracker = _tracker(cfg) + version = tracker.finalize_training_run( run_id=str(run_id), training_job_status=status, ) @@ -148,6 +150,8 @@ def status( if version: st.set_latest_prerelease_model_version(version) CONSOLE.print(f"MLflow model version: [cyan]{version}[/cyan] ([cyan]prerelease-latest[/cyan])") + if run_id and cfg.mlflow.mode is not MlflowMode.disabled: + CONSOLE.print("Open MLflow: [cyan]qc-cli infra mlflow-url[/cyan]") @app.command(name="list") diff --git a/src/tracking/mlflow.py b/src/tracking/mlflow.py index ac8dd87..7019b65 100644 --- a/src/tracking/mlflow.py +++ b/src/tracking/mlflow.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from dataclasses import dataclass from typing import Any, Protocol @@ -40,6 +41,8 @@ class MlflowTracker: if cfg.mlflow.mode is MlflowMode.disabled: return NoopTracker() + os.environ.setdefault("MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT", "true") + try: import mlflow except ImportError as e: