simplify jobs script

README.md

@@ -67,8 +67,7 @@ sagemaker:
     hyperparameters: {}
 
 aihub:
-  device:
+  device: Samsung Galaxy S25 (Family)
-    name: Samsung Galaxy S25 (Family)
   target_runtime: tflite
   input_specs: {}           # Required before running qc-cli ai-hub commands
   job_name: null            # Optional prefix for AI Hub Workbench jobs
@@ -110,10 +109,10 @@ When MLflow is enabled, `train start` creates an MLflow run for the SageMaker jo
 To open the managed SageMaker MLflow UI, request a fresh presigned URL:
 
 ```bash
-qc-cli mlflow open --config config.yaml
+qc-cli infra mlflow-url --config config.yaml
 ```
 
-This opens a browser to a fresh presigned URL. It 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.
+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
 
@@ -125,12 +124,6 @@ qc-cli init --output <path>  Write config to a custom path
 qc-cli init --force          Overwrite an existing config file
 ```
 
-### `mlflow`
-
-```
-qc-cli mlflow open  Open a presigned MLflow UI URL in a browser
-```
-
 ### `infra`
 
 ```
@@ -138,6 +131,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 <arn> 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
@@ -186,17 +180,6 @@ qc-cli ai-hub download [--model-id ID] [--output PATH]
 
 `ai-hub upload` runs the four Workbench upload steps in order: quantize, compile, validate, and profile. Use `--from-step compile`, `--from-step validate`, or `--from-step profile` to resume from saved local state after a completed earlier step.
 
-Resume behavior:
-
-```text
---from-step quantize  Run quantize, compile, validate, and profile.
---from-step compile   Skip quantize; compile the last quantized model unless an explicit source is passed.
---from-step validate  Skip quantize and compile; validate the last compiled model.
---from-step profile   Skip quantize, compile, and validate; profile the last compiled model.
-```
-
-When a step runs in the current command, `upload` passes its returned model ID directly to the next step. When a step is skipped, the next step resolves the needed model ID from `.qc-cli.json`. This avoids re-running earlier AI Hub jobs when you only need to continue from a later step.
-
 `ai-hub compile` resolves model sources in this order: `--model-id`, explicit source options (`--onnx-path`, `--model-s3-uri`, `--from-job`), last quantized model from state, then the last training job from local state. `ai-hub download` is separate because downloading the optimized artifact is outside the four-step Workbench upload loop.
 
 AI Hub authentication currently uses the local `qai-hub` SDK configuration. A planned follow-up is to support AWS Systems Manager Parameter Store `SecureString` for team-managed tokens, where `config.yaml` stores only a parameter name such as `/qc-cli/aihub/token`, AWS KMS encrypts the token at rest, and the CLI retrieves it at runtime with `ssm:GetParameter` plus `kms:Decrypt` permissions.

examples/ai-hub/README.md

@@ -28,8 +28,7 @@ Your `config.yaml` must include AI Hub settings:
 
 ```yaml
 aihub:
-  device:
+  device: Samsung Galaxy S25 (Family)
-    name: Samsung Galaxy S25 (Family)
   target_runtime: tflite
   input_specs:
     input: [[1, 3, 160, 160], float32]

examples/ai-hub/prepare_inputs.py

@@ -9,6 +9,7 @@ from pathlib import Path
 import numpy as np
 from PIL import Image
 
+
 IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"}
 
 

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 [project]
 name = "qc-cli"
 version = "0.1.0"
-description = "CLI for training and deploying models for Qualcomm AI Hub"
+description = "CLI for SageMaker ONNX training and Qualcomm AI Hub optimization"
 requires-python = ">=3.13"
 dependencies = [
     "aws-cdk-lib>=2.180.0",
@@ -29,6 +29,8 @@ packages = ["src"]
 [dependency-groups]
 dev = [
     "boto3-stubs[iam,s3,sagemaker]",
+    "pytest>=8.0",
+    "pytest-mock>=3.12",
     "pyright>=1.1.409",
     "types-PyYAML",
     "ruff>=0.4",

src/aws/mlflow.py

@@ -1,6 +1,3 @@
-import os
-from collections.abc import Generator
-from contextlib import contextmanager
 from typing import Any, cast
 
 import boto3
@@ -37,38 +34,3 @@ def create_presigned_tracking_server_url(region: str, profile: str, name: 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"])
-
-
-@contextmanager
-def tracking_auth_env(profile: str, region: str) -> Generator[None]:
-    credentials = boto3.Session(profile_name=profile, region_name=region).get_credentials()
-    if credentials is None:
-        raise RuntimeError(f"AWS credentials could not be resolved for profile '{profile}'.")
-
-    frozen_credentials = credentials.get_frozen_credentials()
-    if not frozen_credentials.access_key or not frozen_credentials.secret_key:
-        raise RuntimeError(f"AWS credentials are incomplete for profile '{profile}'.")
-
-    env_updates = {
-        "AWS_PROFILE": profile,
-        "AWS_DEFAULT_REGION": region,
-        "AWS_REGION": region,
-        "AWS_ACCESS_KEY_ID": frozen_credentials.access_key,
-        "AWS_SECRET_ACCESS_KEY": frozen_credentials.secret_key,
-    }
-    if frozen_credentials.token:
-        env_updates["AWS_SESSION_TOKEN"] = frozen_credentials.token
-
-    restore_keys = set(env_updates) | {"AWS_SESSION_TOKEN"}
-    previous_env = {key: os.environ.get(key) for key in restore_keys}
-    try:
-        os.environ.update(env_updates)
-        if not frozen_credentials.token:
-            os.environ.pop("AWS_SESSION_TOKEN", None)
-        yield
-    finally:
-        for key, value in previous_env.items():
-            if value is None:
-                os.environ.pop(key, None)
-            else:
-                os.environ[key] = value

src/cloud/__init__.py

@@ -1 +0,0 @@
-"""Cloud provider adapters."""

src/cloud/mlflow.py

@@ -1,77 +0,0 @@
-from contextlib import AbstractContextManager
-from dataclasses import dataclass
-from typing import Any, Protocol
-
-from src.aws import mlflow as aws_mlflow
-from src.config import Config
-
-
-class MlflowTrackingBackend(Protocol):
-    @property
-    def provider_name(self) -> str: ...
-
-    @property
-    def profile(self) -> str: ...
-
-    @property
-    def region(self) -> str: ...
-
-    def get_tracking_uri(self, tracking_server_name: str) -> str: ...
-
-    def auth_env(self) -> AbstractContextManager[None]: ...
-
-    def training_run_params(self, training_job: Any, *, region: str, profile: str, role_arn: str) -> dict[str, Any]: ...
-
-    def training_run_tags(self, training_job: Any) -> dict[str, Any]: ...
-
-    def training_status_params(self, training_job_status: Any) -> dict[str, Any]: ...
-
-    def model_version_tags(self, training_job_status: Any) -> dict[str, Any]: ...
-
-
-@dataclass(frozen=True)
-class AwsMlflowTrackingBackend:
-    profile: str
-    region: str
-    provider_name: str = "aws"
-
-    def get_tracking_uri(self, tracking_server_name: str) -> str:
-        return aws_mlflow.get_tracking_server_arn(self.region, self.profile, tracking_server_name)
-
-    def auth_env(self) -> AbstractContextManager[None]:
-        return aws_mlflow.tracking_auth_env(self.profile, self.region)
-
-    def training_run_params(self, training_job: Any, *, region: str, profile: str, role_arn: str) -> dict[str, Any]:
-        return {
-            "provider.name": self.provider_name,
-            "provider.region": region,
-            "provider.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,
-        }
-
-    def training_run_tags(self, training_job: Any) -> dict[str, Any]:
-        return {"sagemaker.job_name": training_job.job_name}
-
-    def training_status_params(self, training_job_status: Any) -> dict[str, Any]:
-        return {
-            "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,
-        }
-
-    def model_version_tags(self, training_job_status: Any) -> dict[str, Any]:
-        return {"sagemaker.job_name": training_job_status.name}
-
-
-def mlflow_tracking_backend_from_config(cfg: Config) -> MlflowTrackingBackend:
-    return AwsMlflowTrackingBackend(profile=cfg.aws.profile, region=cfg.aws.region)

src/commands/ai_hub.py

@@ -4,9 +4,7 @@ from enum import StrEnum
 from pathlib import Path
 from typing import Any
 
-import qai_hub.hub as hub
 import typer
-from qai_hub.client import Device
 
 from src import state as state_ops
 from src.commands.utils import CONFIG_OPT, CONSOLE, load_cfg
@@ -101,33 +99,6 @@ def _model_id_or_state(config_path: str, model_id: str | None, *, quantized: boo
     return resolved
 
 
-def _device_selector(device: Device) -> str:
-    parts: list[str] = []
-    if device.name:
-        parts.append(f"name={device.name!r}")
-    if device.os:
-        parts.append(f"os={device.os!r}")
-    if device.attributes:
-        parts.append(f"attributes={device.attributes!r}")
-    return ", ".join(parts) if parts else "empty selector"
-
-
-def _validate_device(cfg: Config) -> None:
-    device = cfg.aihub.device
-    try:
-        matches = hub.get_devices(name=device.name, os=device.os, attributes=device.attributes)
-    except Exception as e:
-        CONSOLE.print(f"[red]Unable to validate AI Hub device {_device_selector(device)}: {e}[/red]")
-        raise typer.Exit(1)
-
-    if matches:
-        return
-
-    CONSOLE.print(f"[red]AI Hub device not found: {_device_selector(device)}[/red]")
-    CONSOLE.print("Run [bold]qai-hub list-devices[/bold] to see valid device names.")
-    raise typer.Exit(1)
-
-
 def _quantize_step(
     cfg: Config,
     config_path: str,
@@ -185,7 +156,6 @@ def _compile_step(
     prefer_quantized: bool,
 ) -> str:
     st = state_ops.store(config_path)
-    _validate_device(cfg)
     specs = _input_specs(cfg)
 
     model: Any
@@ -214,7 +184,7 @@ def _compile_step(
     try:
         result = aihub_jobs.submit_compile_job(
             model=model,
-            device=cfg.aihub.device,
+            device_name=cfg.aihub.device,
             input_specs=specs,
             target_runtime=cfg.aihub.target_runtime,
             options=cfg.aihub.compile_options,
@@ -244,7 +214,6 @@ def _validate_step(
     model_id: str | None,
     input_name: str | None,
 ) -> str:
-    _validate_device(cfg)
     specs = _input_specs(cfg)
     resolved_model_id = _model_id_or_state(config_path, model_id)
     try:
@@ -278,7 +247,6 @@ def _validate_step(
 
 
 def _profile_step(cfg: Config, config_path: str, model_id: str | None) -> str:
-    _validate_device(cfg)
     resolved_model_id = _model_id_or_state(config_path, model_id)
     try:
         result = aihub_jobs.submit_profile_job(

src/commands/infra.py

@@ -150,6 +150,35 @@ 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 = cfg.effective_mlflow_tracking_server_name
+    if not tracking_server_name:
+        CONSOLE.print("[red]MLflow is disabled in config.yaml.[/red]")
+        raise typer.Exit(1)
+
+    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,

src/commands/init.py

@@ -1,40 +0,0 @@
-import secrets
-from pathlib import Path
-
-import typer
-import yaml
-
-from src.commands.utils import CONSOLE
-from src.config import GENERATED_STACK_PREFIX, Config, InfraConfig, S3Config
-
-app = typer.Typer()
-
-
-@app.command()
-def init(
-    output: str = typer.Option("config.yaml", help="Destination path for the config file"),
-    force: bool = typer.Option(False, "--force", "-f", help="Overwrite an existing config file"),
-) -> None:
-    """Write a starter config.yaml to the current directory."""
-    dest = Path(output)
-    if dest.exists() and not force:
-        CONSOLE.print(f"[yellow]{dest} already exists.[/yellow] Use --force to overwrite.")
-        raise typer.Exit(1)
-
-    config = _new_isolated_config()
-    dest.parent.mkdir(parents=True, exist_ok=True)
-    config_data = config.model_dump(mode="json")
-    config_data["sagemaker"].pop("role_name", None)
-    with open(dest, "w") as f:
-        yaml.safe_dump(config_data, f, sort_keys=False)
-
-    CONSOLE.print(f"[green]✓[/green] Config written to [bold]{dest}[/bold]")
-    CONSOLE.print("Edit [cyan]sagemaker.training.image_uri[/cyan] before running training commands.")
-
-
-def _new_isolated_config() -> Config:
-    suffix = secrets.token_hex(6)
-    namespace = f"{GENERATED_STACK_PREFIX}{suffix}"
-    config = Config(infra=InfraConfig(stack_name=namespace))
-    config.s3 = S3Config(bucket=f"{namespace}-data")
-    return config

src/commands/mlflow.py

@@ -1,41 +0,0 @@
-import webbrowser
-
-import typer
-
-from src.aws import mlflow as aws_mlflow
-from src.commands.utils import CONFIG_OPT, CONSOLE, load_cfg
-
-app = typer.Typer(help="Manage MLflow tracking server access")
-
-
-@app.command(name="open")
-def open_mlflow(config: str = CONFIG_OPT) -> None:
-    """Open a presigned URL for the configured MLflow tracking server."""
-    cfg = load_cfg(config)
-    tracking_server_name = cfg.effective_mlflow_tracking_server_name
-    if not tracking_server_name:
-        CONSOLE.print("[red]MLflow is disabled in config.yaml.[/red]")
-        raise typer.Exit(1)
-
-    try:
-        url = aws_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}")
-    if webbrowser.open(url):
-        CONSOLE.print("[green]✓[/green] Opened MLflow UI in your browser.")
-    else:
-        CONSOLE.print("[yellow]Could not open a browser automatically. Open the URL above manually.[/yellow]")

src/commands/train.py

@@ -101,7 +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 mlflow open[/cyan]")
+        CONSOLE.print("Open MLflow: [cyan]qc-cli infra mlflow-url[/cyan]")
     CONSOLE.print("Track progress: [cyan]qc-cli train status[/cyan]")
 
 
@@ -151,7 +151,7 @@ def status(
             st.set_latest_experiment_model_version(version)
             CONSOLE.print(f"MLflow model version: [cyan]{version}[/cyan] ([cyan]experiment-latest[/cyan])")
     if run_id and cfg.mlflow.mode is not MlflowMode.disabled:
-        CONSOLE.print("Open MLflow: [cyan]qc-cli mlflow open[/cyan]")
+        CONSOLE.print("Open MLflow: [cyan]qc-cli infra mlflow-url[/cyan]")
 
 
 @app.command(name="list")

src/commands/upload.py

@@ -1,70 +0,0 @@
-from pathlib import Path
-
-import typer
-from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
-
-from src.aws import s3 as s3_ops
-from src.commands.utils import CONFIG_OPT, CONSOLE, load_cfg
-
-app = typer.Typer()
-
-
-@app.command()
-def upload(
-    path: Path = typer.Argument(..., help="Local file or directory to upload"),
-    s3_key: str | None = typer.Option(None, "--s3-key", help="S3 key for file uploads"),
-    config: str = CONFIG_OPT,
-) -> None:
-    """Upload a local file or directory to S3."""
-    cfg = load_cfg(config)
-
-    if path.is_file():
-        key = s3_key or f"{cfg.s3.data_prefix.rstrip('/')}/{path.name}"
-        try:
-            with CONSOLE.status(f"Uploading {path.name}..."):
-                uri = s3_ops.upload_file(cfg.aws.region, cfg.aws.profile, cfg.s3.bucket, str(path), key)
-        except Exception as e:
-            CONSOLE.print(f"[red]Upload failed: {e}[/red]")
-            raise typer.Exit(1)
-
-        CONSOLE.print(f"[green]✓[/green] {path.name} -> {uri}")
-        return
-
-    if path.is_dir():
-        if s3_key is not None:
-            CONSOLE.print("[red]--s3-key can only be used when uploading a single file.[/red]")
-            raise typer.Exit(1)
-
-        files = [file for file in path.rglob("*") if file.is_file()]
-        if not files:
-            CONSOLE.print("[yellow]No files found in directory.[/yellow]")
-            raise typer.Exit(0)
-
-        prefix = cfg.s3.data_prefix
-        CONSOLE.print(f"Uploading {len(files)} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/")
-        try:
-            with Progress(
-                SpinnerColumn(),
-                TextColumn("[progress.description]{task.description}"),
-                BarColumn(),
-                TaskProgressColumn(),
-                console=CONSOLE,
-            ) as progress:
-                task = progress.add_task("Uploading...", total=len(files))
-                count = s3_ops.upload_dir(
-                    cfg.aws.region,
-                    cfg.aws.profile,
-                    cfg.s3.bucket,
-                    str(path),
-                    prefix,
-                    on_progress=lambda: progress.advance(task),
-                )
-        except Exception as e:
-            CONSOLE.print(f"[red]Upload failed: {e}[/red]")
-            raise typer.Exit(1)
-
-        CONSOLE.print(f"[green]✓[/green] Uploaded {count} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/")
-        return
-
-    CONSOLE.print(f"[red]Path not found: {path}[/red]")
-    raise typer.Exit(1)

src/config.py

@@ -4,8 +4,7 @@ from typing import Any, Literal, TypedDict
 
 from mypy_boto3_s3.literals import BucketLocationConstraintType
 from mypy_boto3_sagemaker.literals import TrainingInstanceTypeType
-from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic import BaseModel, Field, model_validator
-from qai_hub.client import Device
 
 
 class MlflowMode(StrEnum):
@@ -82,7 +81,7 @@ class SageMakerConfig(BaseModel):
 
 
 class AIHubConfig(BaseModel):
-    device: Device = Field(default_factory=lambda: Device("Samsung Galaxy S25 (Family)"))
+    device: str = "Samsung Galaxy S25 (Family)"
     target_runtime: str = "tflite"
     input_specs: dict[str, tuple[list[int], str]] = Field(default_factory=dict)
     job_name: str | None = None
@@ -92,13 +91,6 @@ class AIHubConfig(BaseModel):
     quantize_options: str | None = None
     output_dir: str = "build/qai-hub"
 
-    @field_validator("device", mode="before")
-    @classmethod
-    def parse_device(cls, value: Any) -> Any:
-        if isinstance(value, str):
-            return Device(value)
-        return value
-
 
 class MlflowConfig(BaseModel):
     mode: MlflowMode = MlflowMode.disabled

src/main.py

@@ -1,14 +1,115 @@
-import typer
+import secrets
+from pathlib import Path
 
-from src.commands import ai_hub, infra, init, mlflow, train, upload
+import typer
+import yaml
+from rich.console import Console
+from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
+
+from src.aws import s3 as s3_ops
+from src.commands import ai_hub, infra, train
+from src.commands.utils import CONFIG_OPT, load_cfg
+from src.config import GENERATED_STACK_PREFIX, Config, InfraConfig, S3Config
 
 app = typer.Typer(
     help="qc-cli: End-to-end model managment for Qualcomm AI Hub.",
     no_args_is_help=True,
 )
-app.add_typer(init.app)
-app.add_typer(upload.app)
-app.add_typer(mlflow.app, name="mlflow")
 app.add_typer(infra.app, name="infra")
 app.add_typer(train.app, name="train")
 app.add_typer(ai_hub.app, name="ai-hub")
+
+console = Console()
+
+
+@app.command()
+def init(
+    output: str = typer.Option("config.yaml", help="Destination path for the config file"),
+    force: bool = typer.Option(False, "--force", "-f", help="Overwrite an existing config file"),
+) -> None:
+    """Write a starter config.yaml to the current directory."""
+    dest = Path(output)
+    if dest.exists() and not force:
+        console.print(f"[yellow]{dest} already exists.[/yellow] Use --force to overwrite.")
+        raise typer.Exit(1)
+
+    config = _new_isolated_config()
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    config_data = config.model_dump(mode="json")
+    config_data["sagemaker"].pop("role_name", None)
+    with open(dest, "w") as f:
+        yaml.safe_dump(config_data, f, sort_keys=False)
+
+    console.print(f"[green]✓[/green] Config written to [bold]{dest}[/bold]")
+    console.print(
+        "Edit [cyan]sagemaker.training.image_uri[/cyan] before running training commands."
+    )
+
+
+def _new_isolated_config() -> Config:
+    suffix = secrets.token_hex(6)
+    namespace = f"{GENERATED_STACK_PREFIX}{suffix}"
+    config = Config(infra=InfraConfig(stack_name=namespace))
+    config.s3 = S3Config(bucket=f"{namespace}-data")
+    return config
+
+
+@app.command()
+def upload(
+    path: Path = typer.Argument(..., help="Local file or directory to upload"),
+    s3_key: str | None = typer.Option(None, "--s3-key", help="S3 key for file uploads"),
+    config: str = CONFIG_OPT,
+) -> None:
+    """Upload a local file or directory to S3."""
+    cfg = load_cfg(config)
+
+    if path.is_file():
+        key = s3_key or f"{cfg.s3.data_prefix.rstrip('/')}/{path.name}"
+        try:
+            with console.status(f"Uploading {path.name}..."):
+                uri = s3_ops.upload_file(cfg.aws.region, cfg.aws.profile, cfg.s3.bucket, str(path), key)
+        except Exception as e:
+            console.print(f"[red]Upload failed: {e}[/red]")
+            raise typer.Exit(1)
+
+        console.print(f"[green]✓[/green] {path.name} -> {uri}")
+        return
+
+    if path.is_dir():
+        if s3_key is not None:
+            console.print("[red]--s3-key can only be used when uploading a single file.[/red]")
+            raise typer.Exit(1)
+
+        files = [file for file in path.rglob("*") if file.is_file()]
+        if not files:
+            console.print("[yellow]No files found in directory.[/yellow]")
+            raise typer.Exit(0)
+
+        prefix = cfg.s3.data_prefix
+        console.print(f"Uploading {len(files)} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/")
+        try:
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("[progress.description]{task.description}"),
+                BarColumn(),
+                TaskProgressColumn(),
+                console=console,
+            ) as progress:
+                task = progress.add_task("Uploading...", total=len(files))
+                count = s3_ops.upload_dir(
+                    cfg.aws.region,
+                    cfg.aws.profile,
+                    cfg.s3.bucket,
+                    str(path),
+                    prefix,
+                    on_progress=lambda: progress.advance(task),
+                )
+        except Exception as e:
+            console.print(f"[red]Upload failed: {e}[/red]")
+            raise typer.Exit(1)
+
+        console.print(f"[green]✓[/green] Uploaded {count} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/")
+        return
+
+    console.print(f"[red]Path not found: {path}[/red]")
+    raise typer.Exit(1)

src/qualcomm/aihub_jobs.py

@@ -29,7 +29,7 @@ def _dataset_entries(inputs: dict[str, Any]) -> dict[str, list[Any]]:
 
 def submit_compile_job(
     model: Any,
-    device: Device,
+    device_name: str,
     input_specs: dict[str, tuple[tuple[int, ...], str]],
     target_runtime: str,
     options: str | None = None,
@@ -52,7 +52,7 @@ def submit_compile_job(
 
     job = hub.submit_compile_job(
         model=model_arg,
-        device=device,
+        device=Device(device_name),
         name=job_name,
         input_specs=input_specs,
         options=compile_options,
@@ -65,14 +65,14 @@ def submit_compile_job(
 
 def submit_inference_job(
     model_id: str,
-    device: Device,
+    device_name: str,
     inputs: dict[str, Any],
     output_dir: str | Path,
     job_name: str | None = None,
 ) -> InferenceJobResult:
     job = hub.submit_inference_job(
         model=hub.get_model(model_id),
-        device=device,
+        device=Device(device_name),
         inputs=_dataset_entries(inputs),
         name=job_name,
     )
@@ -84,13 +84,13 @@ def submit_inference_job(
 
 def submit_profile_job(
     model_id: str,
-    device: Device,
+    device_name: str,
     options: str | None = None,
     job_name: str | None = None,
 ) -> ProfileJobResult:
     job = hub.submit_profile_job(
         model=hub.get_model(model_id),
-        device=device,
+        device=Device(device_name),
         name=job_name,
         options=options or "",
     )

src/tracking/mlflow.py

@@ -5,7 +5,7 @@ from typing import Any, Protocol
 import mlflow
 from mlflow.tracking import MlflowClient
 
-from src.cloud.mlflow import MlflowTrackingBackend, mlflow_tracking_backend_from_config
+from src.aws import mlflow as aws_mlflow
 from src.config import Config, MlflowMode
 
 
@@ -30,7 +30,6 @@ class MlflowTracker:
     experiment_name: str
     registered_model_name: str
     register_trained_models: bool
-    tracking_backend: MlflowTrackingBackend
 
     @classmethod
     def from_config(cls, cfg: Config) -> Tracker:
@@ -43,10 +42,11 @@ class MlflowTracker:
         if not tracking_server_name:
             raise RuntimeError("MLflow tracking server name could not be resolved.")
 
-        tracking_backend = mlflow_tracking_backend_from_config(cfg)
+        tracking_uri = aws_mlflow.get_tracking_server_arn(
-
+            cfg.aws.region,
-        tracking_uri = tracking_backend.get_tracking_uri(tracking_server_name)
+            cfg.aws.profile,
-        with tracking_backend.auth_env():
+            tracking_server_name,
+        )
         mlflow.set_tracking_uri(tracking_uri)
         mlflow.set_experiment(cfg.mlflow.experiment_name)
 
@@ -55,30 +55,34 @@ class MlflowTracker:
             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:
-        with self.tracking_backend.auth_env():
         run = mlflow.start_run(run_name=training_job.job_name)
         run_id = str(run.info.run_id)
 
-            self._log_params(
+        params = {
-                self.tracking_backend.training_run_params(
+            "aws.region": region,
-                    training_job,
+            "aws.profile": profile,
-                    region=region,
+            "sagemaker.role_arn": role_arn,
-                    profile=profile,
+            "sagemaker.job_name": training_job.job_name,
-                    role_arn=role_arn,
+            "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": self.tracking_backend.provider_name,
+                "qc_cli.source": "sagemaker",
                 "qc_cli.command": "train start",
-                    **self.tracking_backend.training_run_tags(training_job),
+                "sagemaker.job_name": training_job.job_name,
             }
         )
         mlflow.end_run()
@@ -88,9 +92,16 @@ class MlflowTracker:
         if not run_id:
             return None
 
-        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_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")
 
@@ -110,8 +121,8 @@ class MlflowTracker:
                 tags={
                     "qc_cli.stage": "experiment",
                     "qc_cli.artifact_kind": "trained_source",
-                        "qc_cli.source": self.tracking_backend.provider_name,
+                    "qc_cli.source": "sagemaker",
-                        **self.tracking_backend.model_version_tags(training_job_status),
+                    "sagemaker.job_name": training_job_status.name,
                 },
             )
             version_number = str(version.version)

uv.lock

@@ -1003,6 +1003,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/8a/db/55a262f3606bebcae07cc14095338471ad7c0bbcaa37707e6f0ee49725b7/importlib_resources-7.1.0-py3-none-any.whl", hash = "sha256:1bd7b48b4088eddb2cd16382150bb515af0bd2c70128194392725f82ad2c96a1", size = 37232, upload-time = "2026-04-12T16:36:08.219Z" },
 ]
 
+[[package]]
+name = "iniconfig"
+version = "2.3.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
+]
+
 [[package]]
 name = "itsdangerous"
 version = "2.2.0"
@@ -1665,6 +1674,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ff/6e/cf826fae916b8658848d7b9f38d88da6396895c676e8086fc0988073aaf8/pillow-12.2.0-cp314-cp314t-win_arm64.whl", hash = "sha256:aa88ccfe4e32d362816319ed727a004423aab09c5cea43c01a4b435643fa34eb", size = 2556579, upload-time = "2026-04-01T14:45:52.529Z" },
 ]
 
+[[package]]
+name = "pluggy"
+version = "1.6.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
+]
+
 [[package]]
 name = "prettytable"
 version = "3.17.0"
@@ -1945,6 +1963,34 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/16/6b/330d8ebae582b30c2959a1ef4c3bc344ebde48c2ff0c3f113c4710735e11/pyright-1.1.409-py3-none-any.whl", hash = "sha256:aa3ea228cab90c845c7a60d28db7a844c04315356392aa09fafcee98c8c22fb3", size = 6438161, upload-time = "2026-04-23T11:02:01.309Z" },
 ]
 
+[[package]]
+name = "pytest"
+version = "9.0.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "iniconfig" },
+    { name = "packaging" },
+    { name = "pluggy" },
+    { name = "pygments" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/7d/0d/549bd94f1a0a402dc8cf64563a117c0f3765662e2e668477624baeec44d5/pytest-9.0.3.tar.gz", hash = "sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c", size = 1572165, upload-time = "2026-04-07T17:16:18.027Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d4/24/a372aaf5c9b7208e7112038812994107bc65a84cd00e0354a88c2c77a617/pytest-9.0.3-py3-none-any.whl", hash = "sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9", size = 375249, upload-time = "2026-04-07T17:16:16.13Z" },
+]
+
+[[package]]
+name = "pytest-mock"
+version = "3.15.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pytest" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/68/14/eb014d26be205d38ad5ad20d9a80f7d201472e08167f0bb4361e251084a9/pytest_mock-3.15.1.tar.gz", hash = "sha256:1849a238f6f396da19762269de72cb1814ab44416fa73a8686deac10b0d87a0f", size = 34036, upload-time = "2025-09-16T16:37:27.081Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl", hash = "sha256:0a25e2eb88fe5168d535041d09a4529a188176ae608a6d249ee65abc0949630d", size = 10095, upload-time = "2025-09-16T16:37:25.734Z" },
+]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"
@@ -2068,6 +2114,8 @@ dependencies = [
 dev = [
     { name = "boto3-stubs", extra = ["iam", "s3", "sagemaker"] },
     { name = "pyright" },
+    { name = "pytest" },
+    { name = "pytest-mock" },
     { name = "ruff" },
     { name = "types-pyyaml" },
 ]
@@ -2090,6 +2138,8 @@ requires-dist = [
 dev = [
     { name = "boto3-stubs", extras = ["iam", "s3", "sagemaker"] },
     { name = "pyright", specifier = ">=1.1.409" },
+    { name = "pytest", specifier = ">=8.0" },
+    { name = "pytest-mock", specifier = ">=3.12" },
     { name = "ruff", specifier = ">=0.4" },
     { name = "types-pyyaml" },
 ]