command to create presigned URL for MLFlow

This commit is contained in:
2026-05-27 10:52:08 -04:00
parent e1c8d6574f
commit 58681cef82
6 changed files with 64 additions and 6 deletions

View File

@@ -83,7 +83,7 @@ mlflow:
register_trained_models: true 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: 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. 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 ## Commands
### `init` ### `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 --no-bootstrap Deploy without running CDK bootstrap
qc-cli infra setup --cloudformation-execution-policy <arn> Set CDK bootstrap execution policy ARN qc-cli infra setup --cloudformation-execution-policy <arn> Set CDK bootstrap execution policy ARN
qc-cli infra status Show CDK stack/resource status 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 Destroy stack, retaining S3 data
qc-cli infra destroy --yes Destroy stack without confirmation qc-cli infra destroy --yes Destroy stack without confirmation
qc-cli infra destroy --delete-bucket-data Destroy stack and delete S3 data qc-cli infra destroy --delete-bucket-data Destroy stack and delete S3 data

View File

@@ -72,10 +72,11 @@ if [[ "${SKIP_UPLOAD}" == false ]]; then
run uv run qc-cli upload "${DATASET_DIR}" --config "${CONFIG_PATH}" run uv run qc-cli upload "${DATASET_DIR}" --config "${CONFIG_PATH}"
fi fi
TRAIN_OUTPUT="$(uv run qc-cli train start --config "${CONFIG_PATH}")" TRAIN_OUTPUT_FILE="$(mktemp)"
echo "${TRAIN_OUTPUT}" 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 if [[ -z "${JOB_NAME}" ]]; then
echo "Could not find training job name in qc-cli output." >&2 echo "Could not find training job name in qc-cli output." >&2
exit 1 exit 1

View File

@@ -28,3 +28,9 @@ def get_tracking_server_arn(region: str, profile: str, name: str) -> str:
if not arn: if not arn:
raise ValueError(f"MLflow tracking server has no ARN: {name}") raise ValueError(f"MLflow tracking server has no ARN: {name}")
return str(arn) 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"])

View File

@@ -150,6 +150,32 @@ def status(config: str = CONFIG_OPT) -> None:
CONSOLE.print(table) 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() @app.command()
def destroy( def destroy(
config: str = CONFIG_OPT, 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 role_arn.rsplit("/", 1)[-1]
return "-" 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: def _destroy_account_id(config_path: str, cfg: Config) -> str:
config_dir = str(Path(config_path).parent) config_dir = str(Path(config_path).parent)
state = read_infra_state(config_dir) state = read_infra_state(config_dir)

View File

@@ -8,7 +8,7 @@ from src import state as state_ops
from src.aws import iam from src.aws import iam
from src.aws import sagemaker as sm_ops from src.aws import sagemaker as sm_ops
from src.commands.utils import CONFIG_OPT, CONSOLE, load_cfg 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.infra.state import read_infra_state
from src.tracking.mlflow import MlflowTracker 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]") CONSOLE.print(f"[green]✓[/green] Job submitted: [bold]{job_name}[/bold]")
if run_id: if run_id:
CONSOLE.print(f"MLflow run: [cyan]{run_id}[/cyan]") 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]") CONSOLE.print("Track progress: [cyan]qc-cli train status[/cyan]")
@@ -137,7 +138,8 @@ def status(
run_id = job_state.get("mlflow_run_id") run_id = job_state.get("mlflow_run_id")
already_registered = job_state.get("registered_model_version") already_registered = job_state.get("registered_model_version")
if run_id and not already_registered and status.status in {"Completed", "Failed", "Stopped"}: 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), run_id=str(run_id),
training_job_status=status, training_job_status=status,
) )
@@ -148,6 +150,8 @@ def status(
if version: if version:
st.set_latest_prerelease_model_version(version) st.set_latest_prerelease_model_version(version)
CONSOLE.print(f"MLflow model version: [cyan]{version}[/cyan] ([cyan]prerelease-latest[/cyan])") 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") @app.command(name="list")

View File

@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import os
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, Protocol from typing import Any, Protocol
@@ -40,6 +41,8 @@ class MlflowTracker:
if cfg.mlflow.mode is MlflowMode.disabled: if cfg.mlflow.mode is MlflowMode.disabled:
return NoopTracker() return NoopTracker()
os.environ.setdefault("MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT", "true")
try: try:
import mlflow import mlflow
except ImportError as e: except ImportError as e: