Compare commits
2 Commits
35d25d8967
...
58681cef82
| Author | SHA1 | Date | |
|---|---|---|---|
| 58681cef82 | |||
| e1c8d6574f |
12
README.md
12
README.md
@@ -78,12 +78,13 @@ To provision an MLflow tracking server, set:
|
|||||||
```yaml
|
```yaml
|
||||||
mlflow:
|
mlflow:
|
||||||
mode: create
|
mode: create
|
||||||
tracking_server_name: your-tracking-server-name
|
|
||||||
experiment_name: qc-cli-training
|
experiment_name: qc-cli-training
|
||||||
registered_model_name: qc-cli-model
|
registered_model_name: qc-cli-model
|
||||||
register_trained_models: true
|
register_trained_models: true
|
||||||
```
|
```
|
||||||
|
|
||||||
|
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:
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
@@ -100,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`
|
||||||
@@ -117,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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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"])
|
||||||
|
|||||||
@@ -77,7 +77,8 @@ def setup(
|
|||||||
if outputs.get("SageMakerRoleArn"):
|
if outputs.get("SageMakerRoleArn"):
|
||||||
CONSOLE.print(f"[green]✓[/green] IAM role: {outputs['SageMakerRoleArn']}")
|
CONSOLE.print(f"[green]✓[/green] IAM role: {outputs['SageMakerRoleArn']}")
|
||||||
if cfg.mlflow.mode is MlflowMode.create and outputs.get("MlflowTrackingServerArn"):
|
if cfg.mlflow.mode is MlflowMode.create and outputs.get("MlflowTrackingServerArn"):
|
||||||
CONSOLE.print(f"[green]✓[/green] MLflow: {outputs['MlflowTrackingServerArn']}")
|
mlflow_name = outputs.get("MlflowTrackingServerName", cfg.managed_mlflow_tracking_server_name)
|
||||||
|
CONSOLE.print(f"[green]✓[/green] MLflow: {mlflow_name}")
|
||||||
elif cfg.mlflow.mode is MlflowMode.existing:
|
elif cfg.mlflow.mode is MlflowMode.existing:
|
||||||
CONSOLE.print(f"[green]✓[/green] MLflow: {cfg.mlflow.tracking_server_name}")
|
CONSOLE.print(f"[green]✓[/green] MLflow: {cfg.mlflow.tracking_server_name}")
|
||||||
CONSOLE.print("\n[bold green]Infrastructure ready.[/bold green]")
|
CONSOLE.print("\n[bold green]Infrastructure ready.[/bold green]")
|
||||||
@@ -102,7 +103,7 @@ def status(config: str = CONFIG_OPT) -> None:
|
|||||||
if cfg.mlflow.mode is not MlflowMode.disabled:
|
if cfg.mlflow.mode is not MlflowMode.disabled:
|
||||||
table.add_row(
|
table.add_row(
|
||||||
"MLflow",
|
"MLflow",
|
||||||
cfg.mlflow.tracking_server_name or "-",
|
cfg.effective_mlflow_tracking_server_name or "-",
|
||||||
"[red]unknown[/red]",
|
"[red]unknown[/red]",
|
||||||
"-",
|
"-",
|
||||||
)
|
)
|
||||||
@@ -126,7 +127,7 @@ def status(config: str = CONFIG_OPT) -> None:
|
|||||||
if cfg.mlflow.mode is MlflowMode.create:
|
if cfg.mlflow.mode is MlflowMode.create:
|
||||||
table.add_row(
|
table.add_row(
|
||||||
"MLflow",
|
"MLflow",
|
||||||
cfg.mlflow.tracking_server_name or "-",
|
outputs.get("MlflowTrackingServerName", cfg.managed_mlflow_tracking_server_name),
|
||||||
"[green]managed[/green]",
|
"[green]managed[/green]",
|
||||||
outputs.get("MlflowTrackingServerArn", outputs.get("MlflowArtifactUri", "-")),
|
outputs.get("MlflowTrackingServerArn", outputs.get("MlflowArtifactUri", "-")),
|
||||||
)
|
)
|
||||||
@@ -149,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,
|
||||||
@@ -209,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)
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
import re
|
import re
|
||||||
from enum import Enum
|
from enum import StrEnum
|
||||||
from typing import Any, Literal, TypedDict
|
from typing import Any, Literal, TypedDict
|
||||||
|
|
||||||
from mypy_boto3_s3.literals import BucketLocationConstraintType
|
from mypy_boto3_s3.literals import BucketLocationConstraintType
|
||||||
@@ -7,13 +7,13 @@ from mypy_boto3_sagemaker.literals import TrainingInstanceTypeType
|
|||||||
from pydantic import BaseModel, Field, model_validator
|
from pydantic import BaseModel, Field, model_validator
|
||||||
|
|
||||||
|
|
||||||
class MlflowMode(str, Enum):
|
class MlflowMode(StrEnum):
|
||||||
disabled = "disabled"
|
disabled = "disabled"
|
||||||
create = "create"
|
create = "create"
|
||||||
existing = "existing"
|
existing = "existing"
|
||||||
|
|
||||||
|
|
||||||
class MlflowServerSize(str, Enum):
|
class MlflowServerSize(StrEnum):
|
||||||
small = "Small"
|
small = "Small"
|
||||||
medium = "Medium"
|
medium = "Medium"
|
||||||
large = "Large"
|
large = "Large"
|
||||||
@@ -94,8 +94,8 @@ class MlflowConfig(BaseModel):
|
|||||||
|
|
||||||
@model_validator(mode="after")
|
@model_validator(mode="after")
|
||||||
def require_tracking_server_name(self) -> "MlflowConfig":
|
def require_tracking_server_name(self) -> "MlflowConfig":
|
||||||
if self.mode in {MlflowMode.create, MlflowMode.existing} and not self.tracking_server_name:
|
if self.mode is MlflowMode.existing and not self.tracking_server_name:
|
||||||
raise ValueError("mlflow.tracking_server_name is required when mlflow.mode is create or existing")
|
raise ValueError("mlflow.tracking_server_name is required when mlflow.mode is existing")
|
||||||
return self
|
return self
|
||||||
|
|
||||||
|
|
||||||
@@ -105,3 +105,15 @@ class Config(BaseModel):
|
|||||||
s3: S3Config = Field(default_factory=S3Config)
|
s3: S3Config = Field(default_factory=S3Config)
|
||||||
sagemaker: SageMakerConfig = Field(default_factory=SageMakerConfig)
|
sagemaker: SageMakerConfig = Field(default_factory=SageMakerConfig)
|
||||||
mlflow: MlflowConfig = Field(default_factory=MlflowConfig)
|
mlflow: MlflowConfig = Field(default_factory=MlflowConfig)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def managed_mlflow_tracking_server_name(self) -> str:
|
||||||
|
return f"{self.infra.stack_name}-mlflow"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def effective_mlflow_tracking_server_name(self) -> str | None:
|
||||||
|
if self.mlflow.mode is MlflowMode.disabled:
|
||||||
|
return None
|
||||||
|
if self.mlflow.mode is MlflowMode.existing:
|
||||||
|
return self.mlflow.tracking_server_name
|
||||||
|
return self.managed_mlflow_tracking_server_name
|
||||||
|
|||||||
@@ -74,6 +74,7 @@ class QCStack(Stack):
|
|||||||
CfnOutput(self, "SageMakerRoleArn", value=role.attr_arn)
|
CfnOutput(self, "SageMakerRoleArn", value=role.attr_arn)
|
||||||
|
|
||||||
if config.mlflow.mode is MlflowMode.create:
|
if config.mlflow.mode is MlflowMode.create:
|
||||||
|
tracking_server_name = config.managed_mlflow_tracking_server_name
|
||||||
artifact_prefix = config.mlflow.artifact_prefix.strip("/")
|
artifact_prefix = config.mlflow.artifact_prefix.strip("/")
|
||||||
artifact_uri = (
|
artifact_uri = (
|
||||||
f"s3://{data_bucket.bucket_name}/{artifact_prefix}/"
|
f"s3://{data_bucket.bucket_name}/{artifact_prefix}/"
|
||||||
@@ -145,14 +146,14 @@ class QCStack(Stack):
|
|||||||
"MlflowTrackingServer",
|
"MlflowTrackingServer",
|
||||||
artifact_store_uri=artifact_uri,
|
artifact_store_uri=artifact_uri,
|
||||||
role_arn=mlflow_role.attr_arn,
|
role_arn=mlflow_role.attr_arn,
|
||||||
tracking_server_name=config.mlflow.tracking_server_name or "",
|
tracking_server_name=tracking_server_name,
|
||||||
automatic_model_registration=config.mlflow.automatic_model_registration,
|
automatic_model_registration=config.mlflow.automatic_model_registration,
|
||||||
mlflow_version=config.mlflow.mlflow_version,
|
mlflow_version=config.mlflow.mlflow_version,
|
||||||
tracking_server_size=config.mlflow.tracking_server_size.value,
|
tracking_server_size=config.mlflow.tracking_server_size.value,
|
||||||
weekly_maintenance_window_start=config.mlflow.weekly_maintenance_window_start,
|
weekly_maintenance_window_start=config.mlflow.weekly_maintenance_window_start,
|
||||||
)
|
)
|
||||||
|
|
||||||
CfnOutput(self, "MlflowTrackingServerName", value=config.mlflow.tracking_server_name or "")
|
CfnOutput(self, "MlflowTrackingServerName", value=tracking_server_name)
|
||||||
CfnOutput(self, "MlflowTrackingServerArn", value=tracking_server.attr_tracking_server_arn)
|
CfnOutput(self, "MlflowTrackingServerArn", value=tracking_server.attr_tracking_server_arn)
|
||||||
CfnOutput(self, "MlflowArtifactUri", value=artifact_uri)
|
CfnOutput(self, "MlflowArtifactUri", value=artifact_uri)
|
||||||
CfnOutput(self, "MlflowRoleArn", value=mlflow_role.attr_arn)
|
CfnOutput(self, "MlflowRoleArn", value=mlflow_role.attr_arn)
|
||||||
|
|||||||
@@ -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:
|
||||||
@@ -48,13 +51,14 @@ class MlflowTracker:
|
|||||||
"Install with: qc-cli[mlflow]"
|
"Install with: qc-cli[mlflow]"
|
||||||
) from e
|
) from e
|
||||||
|
|
||||||
if not cfg.mlflow.tracking_server_name:
|
tracking_server_name = cfg.effective_mlflow_tracking_server_name
|
||||||
raise RuntimeError("mlflow.tracking_server_name is required when MLflow is enabled.")
|
if not tracking_server_name:
|
||||||
|
raise RuntimeError("MLflow tracking server name could not be resolved.")
|
||||||
|
|
||||||
tracking_uri = aws_mlflow.get_tracking_server_arn(
|
tracking_uri = aws_mlflow.get_tracking_server_arn(
|
||||||
cfg.aws.region,
|
cfg.aws.region,
|
||||||
cfg.aws.profile,
|
cfg.aws.profile,
|
||||||
cfg.mlflow.tracking_server_name,
|
tracking_server_name,
|
||||||
)
|
)
|
||||||
mlflow.set_tracking_uri(tracking_uri)
|
mlflow.set_tracking_uri(tracking_uri)
|
||||||
mlflow.set_experiment(cfg.mlflow.experiment_name)
|
mlflow.set_experiment(cfg.mlflow.experiment_name)
|
||||||
|
|||||||
Reference in New Issue
Block a user