update ai-hub to first optimize model for Workbench

Remove old examples
clean
2026-06-09 14:55:26 -04:00 · 2026-06-09 12:21:25 -04:00 · 2026-06-09 12:19:17 -04:00 · 2026-06-09 12:18:41 -04:00 · 2026-06-09 12:06:49 -04:00 · 2026-06-09 11:55:03 -04:00
28 changed files with 1001 additions and 1070 deletions
--- a/README.md
+++ b/README.md
@@ -67,7 +67,8 @@ sagemaker:
    hyperparameters: {}
 aihub:
-  device: Samsung Galaxy S25 (Family)
+  device:
    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
@@ -109,10 +110,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 infra mlflow-url --config config.yaml
+qc-cli mlflow open --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.
+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.
 ## Commands
@@ -124,6 +125,12 @@ 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`
 ```
@@ -131,7 +138,6 @@ 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
@@ -171,16 +177,35 @@ The expected output artifact is SageMaker’s `model.tar.gz`, normally containin
 ```
 qc-cli ai-hub upload <calibration.npz|calibration-dir> <inputs.npz|inputs.npy>
 qc-cli ai-hub upload <calibration> <inputs> --from-step validate
-qc-cli ai-hub quantize <calibration.npz|calibration-dir> [--onnx-path PATH] [--model-s3-uri URI] [--from-job NAME]
+qc-cli ai-hub optimize [--onnx-path PATH] [--model-s3-uri URI] [--from-job NAME]
 qc-cli ai-hub quantize <calibration.npz|calibration-dir> [--model-id ID] [--onnx-path PATH] [--model-s3-uri URI] [--from-job NAME]
 qc-cli ai-hub compile [--model-id ID] [--onnx-path PATH] [--model-s3-uri URI] [--from-job NAME]
 qc-cli ai-hub validate <inputs.npz|inputs.npy> [--model-id ID] [--input-name NAME]
 qc-cli ai-hub profile [--model-id ID]
 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.
+`ai-hub upload` optimizes to ONNX, quantizes, validates, and profiles. When `aihub.target_runtime` is not `onnx`, it
 also compiles the quantized model to that deployment runtime. The initial ONNX optimization gives external models
 Workbench provenance and applies compiler optimization passes before quantization.
-`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.
+Resume behavior:
 ```text
 --from-step optimize  Run optimize, quantize, optional final compile, validate, and profile.
 --from-step quantize  Quantize the last optimized ONNX, then optionally compile, validate, and profile.
 --from-step compile   Skip optimize and quantize; finalize the last quantized model for the target runtime.
 --from-step validate  Skip optimize, quantize, and compile; validate the last compiled model.
 --from-step profile   Skip optimize, 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 optimize` compiles an external model with `--target_runtime onnx`. `ai-hub quantize` uses an explicit
 `--model-id`, the last optimized ONNX model, or an explicit/local model source in that order. `ai-hub compile` resolves
 model sources in this order: `--model-id`, explicit source options, last quantized model, then the last training job.
 For `target_runtime: onnx`, upload treats the quantized ONNX as the final model and skips a redundant second compile.
 `ai-hub download` remains separate because downloading is outside the Workbench processing 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.
--- a/examples/ai-hub/README.md
+++ b/examples/ai-hub/README.md
@@ -1,117 +0,0 @@
 # Qualcomm AI Hub Example
 This example takes the ONNX model produced by the SageMaker training example and runs the Qualcomm AI Hub upload workflow:
 1. Quantize
 2. Compile
 3. Validate
 4. Profile
 5. Download the compiled artifact
 ## Prerequisites
 Run the training example first and wait for it to complete:
 ```bash
 bash examples/training/run_training.sh --config config.yaml --wait
 ```
 If the dataset is already uploaded to S3, use:
 ```bash
 bash examples/training/run_training.sh --config config.yaml --skip-upload --wait
 ```
 The training artifact must contain a static-shape `model.onnx`. The training example exports an input named `input` with shape `1x3x160x160`.
 Your `config.yaml` must include AI Hub settings:
 ```yaml
 aihub:
  device: Samsung Galaxy S25 (Family)
  target_runtime: tflite
  input_specs:
    input: [[1, 3, 160, 160], float32]
  output_dir: build/qai-hub
 ```
 You also need local Qualcomm AI Hub SDK authentication configured.
 ## Prepare Inputs
 AI Hub does not consume the raw JPG training images directly. It needs NumPy tensors that match the ONNX model input shape and preprocessing.
 Generate calibration and validation inputs:
 ```bash
 uv run python examples/ai-hub/prepare_inputs.py
 ```
 This writes:
 ```text
 examples/training/data/aihub_calibration/*.npy
 examples/training/data/inputs.npz
 ```
 The script applies the same image preprocessing used by the training example:
 - resize to `160x160`
 - convert to channel-first `1x3x160x160`
 - normalize with ImageNet mean and standard deviation
 Useful options:
 ```bash
 uv run python examples/ai-hub/prepare_inputs.py \
  --dataset-dir examples/training/data/flower_photos_sagemaker \
  --calibration-dir examples/training/data/aihub_calibration \
  --input-file examples/training/data/inputs.npz \
  --samples 16
 ```
 ## Run AI Hub
 After training completes and inputs are prepared:
 ```bash
 bash examples/ai-hub/run_ai_hub.sh --config config.yaml
 ```
 By default, the script uses the last SageMaker training job recorded in `.qc-cli.json`. It downloads that job's `model.tar.gz`, extracts `model.onnx`, runs the AI Hub workflow, and downloads the compiled artifact.
 To use a specific training job:
 ```bash
 bash examples/ai-hub/run_ai_hub.sh \
  --config config.yaml \
  --from-job qc-cli-YYYYMMDD-HHMMSS
 ```
 To resume from a later Workbench step:
 ```bash
 bash examples/ai-hub/run_ai_hub.sh \
  --config config.yaml \
  --from-step validate
 ```
 To skip downloading the compiled artifact:
 ```bash
 bash examples/ai-hub/run_ai_hub.sh \
  --config config.yaml \
  --skip-download
 ```
 ## Troubleshooting
 If AI Hub reports dynamic input shapes, rerun training with the current training source. AI Hub quantization requires the exported ONNX model to use static input shapes.
 If `run_ai_hub.sh` reports missing calibration or input files, run:
 ```bash
 uv run python examples/ai-hub/prepare_inputs.py
 ```
 If validation fails with a missing input name, make sure `config.yaml` and the generated `.npz` both use `input` as the input name.
--- a/examples/ai-hub/prepare_inputs.py
+++ b/examples/ai-hub/prepare_inputs.py
@@ -1,74 +0,0 @@
 #!/usr/bin/env python3
 """Prepare Qualcomm AI Hub calibration and validation inputs for the training example."""
 from __future__ import annotations
 import argparse
 from pathlib import Path
 import numpy as np
 from PIL import Image
 IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"}
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dataset-dir",
        type=Path,
        default=Path("examples/training/data/flower_photos_sagemaker"),
        help="ImageFolder-style dataset used for training.",
    )
    parser.add_argument(
        "--calibration-dir",
        type=Path,
        default=Path("examples/training/data/aihub_calibration"),
        help="Directory where .npy calibration samples will be written.",
    )
    parser.add_argument(
        "--input-file",
        type=Path,
        default=Path("examples/training/data/inputs.npz"),
        help="Validation .npz input file for qc-cli ai-hub validate.",
    )
    parser.add_argument("--input-name", default="input", help="ONNX input name.")
    parser.add_argument("--image-size", type=int, default=160, help="Square image size used by training.")
    parser.add_argument("--samples", type=int, default=16, help="Number of calibration samples to write.")
    return parser.parse_args()
 def preprocess_image(path: Path, image_size: int) -> np.ndarray:
    image = Image.open(path).convert("RGB").resize((image_size, image_size), Image.Resampling.BILINEAR)
    array = np.asarray(image, dtype=np.float32) / 255.0
    array = np.transpose(array, (2, 0, 1))
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)[:, None, None]
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)[:, None, None]
    return ((array - mean) / std)[None, ...].astype("float32")
 def main() -> None:
    args = parse_args()
    images = sorted(p for p in args.dataset_dir.rglob("*") if p.suffix.lower() in IMAGE_EXTENSIONS)
    if not images:
        raise SystemExit(f"No images found under {args.dataset_dir}")
    if args.samples < 1:
        raise SystemExit("--samples must be at least 1")
    args.calibration_dir.mkdir(parents=True, exist_ok=True)
    args.input_file.parent.mkdir(parents=True, exist_ok=True)
    sample_count = min(args.samples, len(images))
    prepared = []
    for index, image_path in enumerate(images[:sample_count]):
        sample = preprocess_image(image_path, args.image_size)
        np.save(args.calibration_dir / f"sample_{index:03d}.npy", sample)
        prepared.append(sample)
    np.savez(args.input_file, **{args.input_name: prepared[0]})
    print(f"Wrote {sample_count} calibration samples to {args.calibration_dir}")
    print(f"Wrote validation input to {args.input_file}")
 if __name__ == "__main__":
    main()
--- a/examples/ai-hub/run_ai_hub.sh
+++ b/examples/ai-hub/run_ai_hub.sh
@@ -1,156 +0,0 @@
 #!/usr/bin/env bash
 set -euo pipefail
 CONFIG_PATH="config.yaml"
 CALIBRATION_PATH="examples/training/data/aihub_calibration"
 INPUT_FILE="examples/training/data/inputs.npz"
 FROM_STEP="quantize"
 FROM_JOB=""
 MODEL_S3_URI=""
 ONNX_PATH=""
 INPUT_NAME=""
 DOWNLOAD=true
 OUTPUT_PATH=""
 usage() {
  cat <<EOF
 Usage: $0 [options]
 Options:
  --config PATH              Path to qc-cli config file. Default: config.yaml
  --calibration PATH         Calibration .npz file or directory of .npy samples.
                              Default: ${CALIBRATION_PATH}
  --input-file PATH          Validation .npz or .npy inputs. Default: ${INPUT_FILE}
  --from-step STEP           Resume upload from: quantize, compile, validate, profile.
                              Default: ${FROM_STEP}
  --from-job NAME            SageMaker training job whose model artifact should upload.
                              Defaults to the last training job in local qc-cli state.
  --model-s3-uri URI         S3 URI of model.tar.gz to upload.
  --onnx-path PATH           Local ONNX path or ONNX path inside extracted artifact.
  --input-name NAME          Input name for .npy validation files.
  --skip-download            Do not download the compiled AI Hub artifact after upload.
  --output PATH              Destination file for ai-hub download.
  -h, --help                 Show this help.
 EOF
 }
 while [[ $# -gt 0 ]]; do
  case "$1" in
    --config)
      CONFIG_PATH="$2"
      shift 2
      ;;
    --calibration)
      CALIBRATION_PATH="$2"
      shift 2
      ;;
    --input-file)
      INPUT_FILE="$2"
      shift 2
      ;;
    --from-step)
      FROM_STEP="$2"
      shift 2
      ;;
    --from-job)
      FROM_JOB="$2"
      shift 2
      ;;
    --model-s3-uri)
      MODEL_S3_URI="$2"
      shift 2
      ;;
    --onnx-path)
      ONNX_PATH="$2"
      shift 2
      ;;
    --input-name)
      INPUT_NAME="$2"
      shift 2
      ;;
    --skip-download)
      DOWNLOAD=false
      shift
      ;;
    --output)
      OUTPUT_PATH="$2"
      shift 2
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown option: $1" >&2
      usage >&2
      exit 1
      ;;
  esac
 done
 if [[ ! -f "${CONFIG_PATH}" ]]; then
  echo "Config not found: ${CONFIG_PATH}" >&2
  exit 1
 fi
 case "${FROM_STEP}" in
  quantize|compile|validate|profile)
    ;;
  *)
    echo "--from-step must be one of: quantize, compile, validate, profile" >&2
    exit 1
    ;;
 esac
 if [[ ! -e "${CALIBRATION_PATH}" ]]; then
  echo "Calibration path not found: ${CALIBRATION_PATH}" >&2
  echo "Pass --calibration with a .npz file or directory of .npy samples." >&2
  exit 1
 fi
 if [[ ! -f "${INPUT_FILE}" ]]; then
  echo "Input file not found: ${INPUT_FILE}" >&2
  echo "Pass --input-file with a validation .npz or .npy file." >&2
  exit 1
 fi
 run() {
  echo "+ $*"
  "$@"
 }
 UPLOAD_ARGS=(
  "${CALIBRATION_PATH}"
  "${INPUT_FILE}"
  --from-step "${FROM_STEP}"
  --config "${CONFIG_PATH}"
 )
 if [[ -n "${FROM_JOB}" ]]; then
  UPLOAD_ARGS+=(--from-job "${FROM_JOB}")
 fi
 if [[ -n "${MODEL_S3_URI}" ]]; then
  UPLOAD_ARGS+=(--model-s3-uri "${MODEL_S3_URI}")
 fi
 if [[ -n "${ONNX_PATH}" ]]; then
  UPLOAD_ARGS+=(--onnx-path "${ONNX_PATH}")
 fi
 if [[ -n "${INPUT_NAME}" ]]; then
  UPLOAD_ARGS+=(--input-name "${INPUT_NAME}")
 fi
 run uv run qc-cli ai-hub upload "${UPLOAD_ARGS[@]}"
 if [[ "${DOWNLOAD}" == false ]]; then
  exit 0
 fi
 DOWNLOAD_ARGS=(--config "${CONFIG_PATH}")
 if [[ -n "${OUTPUT_PATH}" ]]; then
  DOWNLOAD_ARGS+=(--output "${OUTPUT_PATH}")
 fi
 run uv run qc-cli ai-hub download "${DOWNLOAD_ARGS[@]}"
--- a/examples/meter-detection/README.md
+++ b/examples/meter-detection/README.md
@@ -0,0 +1,266 @@
 # YOLO26 Electric Meter Detection Example
 This example trains a YOLO26 object detection model on the Roboflow Universe electric meter dataset using the existing `qc-cli` SageMaker training flow.
 The workflow is intentionally command driven. Run each step yourself so you can inspect the dataset, update `config.yaml`, and decide when to submit the SageMaker job.
 Dataset:
 ```text
 https://universe.roboflow.com/kemals-workspace-kbc8l/electric-meter-detection-o4tfi/dataset/1
 ```
 ## Prerequisites
 - Install or sync the project dependencies: `uv sync`
 - The virtual environment is activated.
 - AWS credentials configured for the profile in `config.yaml`
 - Infrastructure already deployed with `qc-cli infra setup`
 ## 1. Download The Dataset
 Register or sign in to Roboflow, then open the dataset page:
 ```text
 https://universe.roboflow.com/kemals-workspace-kbc8l/electric-meter-detection-o4tfi/dataset/1
 ```
 Download the dataset in YOLOv26 format from the Roboflow UI, then extract the downloaded archive into:
 ```text
 examples/meter-detection/data/electric-meter-detection
 ```
 The `data.yaml` file should be directly under that folder:
 ```text
 examples/meter-detection/data/electric-meter-detection/data.yaml
 ```
 Do not move `data.yaml` into the `train/` split folder.
 After extracting, confirm the dataset has a YOLO data file and image splits:
 ```bash
 find examples/meter-detection/data/electric-meter-detection -maxdepth 2 -type d | sort
 find examples/meter-detection/data/electric-meter-detection -name data.yaml -print
 ```
 Open `examples/meter-detection/data/electric-meter-detection/data.yaml` and make sure the split paths are relative to that folder:
 ```yaml
 path: .
 train: train/images
 val: valid/images
 test: test/images
 ```
 If your downloaded dataset does not include a `test/` folder, remove the `test:` line.
 The expected layout is similar to:
 ```text
 examples/meter-detection/data/electric-meter-detection/
  data.yaml
  train/
  valid/
  test/
 ```
 ## 2. Configure SageMaker Training
 Update `config.yaml` so the training section points at this example's source directory:
 ```yaml
 sagemaker:
  training:
    image_uri: 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.6-cpu-py312-ubuntu22.04-sagemaker-v1
    instance_type: ml.g4dn.xlarge
    instance_count: 1
    source_dir: examples/meter-detection/source
    entry_point: train.py
    hyperparameters:
      model: yolo26n.pt
      epochs: 25
      imgsz: 640
      batch: 16
      workers: 2
 ```
 Use `yolo26n.pt` for a lightweight first YOLO26 run. If those weights are unavailable in the installed Ultralytics package, use `yolo11n.pt` as the established fallback:
 ```yaml
      model: yolo11n.pt
 ```
 The `source/requirements.txt` file is installed by the SageMaker PyTorch container before running `train.py`.
 For a CPU smoke test, use a CPU instance and reduce the workload:
 ```yaml
 sagemaker:
  training:
    image_uri: 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.6-cpu-py312-ubuntu22.04-sagemaker-v1
    instance_type: ml.m4.xlarge
    instance_count: 1
    source_dir: examples/meter-detection/source
    entry_point: train.py
    hyperparameters:
      model: yolo26n.pt
      epochs: 1
      imgsz: 320
      batch: 4
      workers: 2
 ```
 ## 3. Check Infrastructure
 Confirm the CLI can see the configured SageMaker role and S3 bucket:
 ```bash
 qc-cli infra status
 ```
 ## 4. Upload The Dataset
 Upload the downloaded Roboflow dataset to the `s3.data_prefix` configured in `config.yaml`:
 ```bash
 qc-cli upload examples/meter-detection/data/electric-meter-detection
 ```
 Directory uploads preserve paths relative to the uploaded directory, so SageMaker receives the dataset root with `data.yaml` plus the split directories.
 In SageMaker, this uploaded dataset root is mounted at `/opt/ml/input/data/train`. That `train` path is the SageMaker channel name, not the YOLO `train/` split folder.
 ## 5. Start Training
 Submit the SageMaker training job:
 ```bash
 qc-cli train start
 ```
 The command prints the submitted SageMaker job name. Check progress with:
 ```bash
 qc-cli train status
 ```
 Or pass the job name explicitly:
 ```bash
 qc-cli train status qc-cli-YYYYMMDD-HHMMSS
 ```
 ## SageMaker Outputs
 When the job completes, SageMaker packages the files written under `/opt/ml/model` into `model.tar.gz`.
 This example writes:
 ```text
 best.pt
 model.onnx
 metrics.json
 ```
 The archive is stored under the configured `s3.model_prefix`.
 ## 6. Configure Qualcomm AI Hub
 Authenticate with Qualcomm AI Hub:
 ```bash
 qai-hub configure --api_token
 ```
 Add AI Hub settings to `config.yaml`. The input name and image size must match the ONNX model exported by this example:
 ```yaml
 aihub:
  device:
    name: Dragonwing IQ-9075 EVK
  target_runtime: onnx
  input_specs:
    images: [[1, 3, 640, 640], float32]
  job_name: meter-detection
  model_name: meter-detection
  output_dir: build/qai-hub/meter-detection
 ```
 The ONNX graph is the source of truth. The export normally uses the same value as `sagemaker.training.hyperparameters.imgsz`, but changing `config.yaml` after training does not resize an existing model. For example, a model exported with `imgsz: 320` requires `images: [[1, 3, 320, 320], float32]`.
 ## 7. Prepare AI Hub Inputs
 Generate calibration samples and a validation input from the downloaded dataset:
 ```bash
 uv run python examples/meter-detection/prepare_aihub_inputs.py --image-size 640
 ```
 This writes:
 ```text
 examples/meter-detection/data/aihub_calibration/*.npy
 examples/meter-detection/data/inputs.npz
 ```
 The script applies the preprocessing expected by the exported YOLO model: aspect-ratio-preserving letterboxing, RGB channel order, channel-first layout, and pixel values normalized to `[0, 1]`.
 ## 8. Upload To Qualcomm AI Hub
 Use the SageMaker job name printed by `qc-cli train start`:
 ```bash
 qc-cli ai-hub upload \
  examples/meter-detection/data/aihub_calibration \
  examples/meter-detection/data/inputs.npz \
  --from-job qc-cli-YYYYMMDD-HHMMSS
 ```
 The command downloads the job's `model.tar.gz`, finds `model.onnx`, and runs the following AI Hub workflow:
 1. Compile the external ONNX to a Workbench-optimized ONNX model.
 2. Quantize the optimized ONNX model.
 3. Compile the quantized model when the configured deployment runtime is not `onnx`.
 4. Validate and profile the final model.
 The training example sanitizes the Ultralytics ONNX export before saving `model.onnx`. This removes graph input or output names, such as `output0`, that are duplicated in the ONNX `value_info` metadata and rejected by AI Hub.
 For a model already downloaded by a failed upload attempt, sanitize the extracted ONNX file and retry using the local model. Replace the job name in both paths:
 ```bash
 uv run --with onnx python examples/meter-detection/source/sanitize_onnx.py \
  build/qai-hub/meter-detection/qc-cli-YYYYMMDD-HHMMSS/source/extracted/model.onnx \
  --output build/qai-hub/meter-detection/model.aihub.onnx
 qc-cli ai-hub upload \
  examples/meter-detection/data/aihub_calibration \
  examples/meter-detection/data/inputs.npz \
  --onnx-path build/qai-hub/meter-detection/model.aihub.onnx
 ```
 Download the compiled artifact after the workflow completes:
 ```bash
 qc-cli ai-hub download --output build/qai-hub/meter-detection/model.tflite
 ```
 ## Training Hyperparameters
 Values under `sagemaker.training.hyperparameters` are passed to `source/train.py` as command-line arguments.
 | Name | Type | Default | Description |
 |---|---:|---:|---|
 | `model` | string | `yolo26n.pt` | Ultralytics model weights or model YAML. |
 | `epochs` | int | `25` | Number of training epochs. |
 | `imgsz` | int | `640` | Square training image size. |
 | `batch` | int | `16` | Images per training batch. |
 | `workers` | int | `2` | DataLoader worker count. |
 | `patience` | int | `20` | Early stopping patience. |
 | `device` | string | auto | Optional Ultralytics device value such as `0` or `cpu`. |
 | `data-yaml` | string | auto | Optional path to `data.yaml`; normally discovered from the uploaded dataset root. |
 | `dataset-dir` | string | `SM_CHANNEL_TRAIN` | Uploaded dataset root mounted by SageMaker. |
 Do not set `dataset-dir` or `model-dir` in normal SageMaker runs. SageMaker sets those automatically through `SM_CHANNEL_TRAIN` and `SM_MODEL_DIR`.
--- a/examples/meter-detection/prepare_aihub_inputs.py
+++ b/examples/meter-detection/prepare_aihub_inputs.py
@@ -0,0 +1,92 @@
 #!/usr/bin/env python3
 """Prepare Qualcomm AI Hub calibration and validation inputs for the meter detector."""
 from __future__ import annotations
 import argparse
 from pathlib import Path
 import numpy as np
 from PIL import Image
 IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"}
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dataset-dir",
        type=Path,
        default=Path("examples/meter-detection/data/electric-meter-detection"),
        help="Root of the extracted Roboflow dataset.",
    )
    parser.add_argument(
        "--calibration-dir",
        type=Path,
        default=Path("examples/meter-detection/data/aihub_calibration"),
        help="Directory where .npy calibration samples will be written.",
    )
    parser.add_argument(
        "--input-file",
        type=Path,
        default=Path("examples/meter-detection/data/inputs.npz"),
        help="Validation .npz input file for qc-cli ai-hub validate.",
    )
    parser.add_argument("--input-name", default="images", help="ONNX input name.")
    parser.add_argument("--image-size", type=int, default=640, help="Square image size used for ONNX export.")
    parser.add_argument("--samples", type=int, default=16, help="Number of calibration samples to write.")
    return parser.parse_args()
 def preprocess_image(path: Path, image_size: int) -> np.ndarray:
    """Apply Ultralytics-style letterboxing and produce an NCHW float32 tensor."""
    with Image.open(path) as source:
        image = source.convert("RGB")
    scale = min(image_size / image.width, image_size / image.height)
    resized_width = round(image.width * scale)
    resized_height = round(image.height * scale)
    image = image.resize((resized_width, resized_height), Image.Resampling.BILINEAR)
    canvas = Image.new("RGB", (image_size, image_size), (114, 114, 114))
    left = round((image_size - resized_width) / 2 - 0.1)
    top = round((image_size - resized_height) / 2 - 0.1)
    canvas.paste(image, (left, top))
    array = np.asarray(canvas, dtype=np.float32) / 255.0
    return np.transpose(array, (2, 0, 1))[None, ...].astype(np.float32)
 def main() -> None:
    args = parse_args()
    if args.image_size < 1:
        raise SystemExit("--image-size must be at least 1")
    if args.samples < 1:
        raise SystemExit("--samples must be at least 1")
    images = sorted(
        path
        for path in args.dataset_dir.rglob("*")
        if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS and path.parent.name == "images"
    )
    if not images:
        raise SystemExit(f"No images found under {args.dataset_dir}")
    args.calibration_dir.mkdir(parents=True, exist_ok=True)
    args.input_file.parent.mkdir(parents=True, exist_ok=True)
    for stale_sample in args.calibration_dir.glob("sample_*.npy"):
        stale_sample.unlink()
    prepared: list[np.ndarray] = []
    for index, image_path in enumerate(images[: args.samples]):
        sample = preprocess_image(image_path, args.image_size)
        np.save(args.calibration_dir / f"sample_{index:03d}.npy", sample)
        prepared.append(sample)
    np.savez(args.input_file, **{args.input_name: prepared[0]})  # pyright: ignore[reportArgumentType]
    print(f"Wrote {len(prepared)} calibration samples to {args.calibration_dir}")
    print(f"Wrote validation input to {args.input_file}")
 if __name__ == "__main__":
    main()
--- a/examples/meter-detection/source/requirements.txt
+++ b/examples/meter-detection/source/requirements.txt
@@ -0,0 +1,3 @@
 ultralytics>=8.3.0
 pyyaml>=6.0.3
 onnx>=1.16.0
--- a/examples/meter-detection/source/sanitize_onnx.py
+++ b/examples/meter-detection/source/sanitize_onnx.py
@@ -0,0 +1,38 @@
 #!/usr/bin/env python3
 """Remove ONNX value_info entries that duplicate graph inputs or outputs."""
 from __future__ import annotations
 import argparse
 from pathlib import Path
 import onnx  # type: ignore[reportMissingImports]
 def sanitize_onnx(path: Path, output_path: Path | None = None) -> Path:
    model = onnx.load(path)
    io_names = {value.name for value in (*model.graph.input, *model.graph.output)}
    retained_value_info = [value for value in model.graph.value_info if value.name not in io_names]
    destination = output_path or path
    if len(retained_value_info) != len(model.graph.value_info):
        del model.graph.value_info[:]
        model.graph.value_info.extend(retained_value_info)
    destination.parent.mkdir(parents=True, exist_ok=True)
    onnx.save(model, destination)
    return destination
 def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("onnx_path", type=Path)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    written = sanitize_onnx(args.onnx_path, args.output)
    print(f"Saved sanitized ONNX model to {written}")
 if __name__ == "__main__":
    main()
--- a/examples/meter-detection/source/train.py
+++ b/examples/meter-detection/source/train.py
@@ -0,0 +1,126 @@
 #!/usr/bin/env python3
 """SageMaker entry point for YOLO electric meter detection training."""
 from __future__ import annotations
 import argparse
 import json
 import os
 import shutil
 from pathlib import Path
 from typing import Any
 import yaml
 from sanitize_onnx import sanitize_onnx
 from ultralytics import YOLO  # type: ignore[reportMissingImports]
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="yolo26n.pt")
    parser.add_argument("--epochs", type=int, default=25)
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--batch", type=int, default=16)
    parser.add_argument("--workers", type=int, default=2)
    parser.add_argument("--patience", type=int, default=20)
    parser.add_argument("--device", default=None)
    parser.add_argument("--data-yaml", default=None)
    parser.add_argument("--dataset-dir", default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
    parser.add_argument("--train-dir", dest="dataset_dir", help=argparse.SUPPRESS)
    parser.add_argument("--model-dir", default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
    return parser.parse_args()
 def find_data_yaml(dataset_dir: Path, explicit_path: str | None) -> Path:
    if explicit_path:
        data_yaml = Path(explicit_path)
        if data_yaml.is_file():
            return data_yaml
        raise FileNotFoundError(f"Configured data.yaml does not exist: {data_yaml}")
    matches = sorted(dataset_dir.rglob("data.yaml"))
    if not matches:
        raise FileNotFoundError(f"Could not find data.yaml under {dataset_dir}")
    if len(matches) > 1:
        print(f"Found multiple data.yaml files; using {matches[0]}")
    return matches[0]
 def prepare_data_yaml(data_yaml: Path) -> Path:
    """Write a SageMaker-local data file rooted at the uploaded dataset."""
    dataset_root = data_yaml.parent
    data = yaml.safe_load(data_yaml.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise ValueError(f"Expected a mapping in {data_yaml}")
    normalized = dict(data)
    normalized["path"] = str(dataset_root)
    if "val" not in normalized and "valid" in normalized:
        normalized["val"] = normalized.pop("valid")
    prepared_path = dataset_root / "data.sagemaker.yaml"
    prepared_path.write_text(yaml.safe_dump(normalized, sort_keys=False), encoding="utf-8")
    print(f"Prepared dataset config: {prepared_path}")
    return prepared_path
 def copy_if_exists(source: Path, destination: Path) -> None:
    if source.exists():
        shutil.copy2(source, destination)
        print(f"Saved {destination}")
 def main() -> None:
    args = parse_args()
    dataset_dir = Path(args.dataset_dir)
    model_dir = Path(args.model_dir)
    model_dir.mkdir(parents=True, exist_ok=True)
    data_yaml = prepare_data_yaml(find_data_yaml(dataset_dir, args.data_yaml))
    model = YOLO(args.model)
    train_kwargs: dict[str, Any] = {
        "data": str(data_yaml),
        "epochs": args.epochs,
        "imgsz": args.imgsz,
        "batch": args.batch,
        "workers": args.workers,
        "patience": args.patience,
        "project": str(model_dir / "runs"),
        "name": "train",
        "exist_ok": True,
    }
    if args.device:
        train_kwargs["device"] = args.device
    results = model.train(**train_kwargs)
    save_dir = Path(results.save_dir)
    best_pt = save_dir / "weights" / "best.pt"
    last_pt = save_dir / "weights" / "last.pt"
    trained_weights = best_pt if best_pt.exists() else last_pt
    if not trained_weights.exists():
        raise FileNotFoundError(f"Could not find trained weights in {save_dir / 'weights'}")
    copy_if_exists(trained_weights, model_dir / "best.pt")
    trained_model = YOLO(str(trained_weights))
    onnx_path = Path(trained_model.export(format="onnx", imgsz=args.imgsz))
    saved_onnx_path = sanitize_onnx(onnx_path, model_dir / "model.onnx")
    print(f"Saved {saved_onnx_path}")
    metrics = {
        "model": args.model,
        "epochs": args.epochs,
        "imgsz": args.imgsz,
        "batch": args.batch,
        "workers": args.workers,
        "patience": args.patience,
        "data_yaml": str(data_yaml),
        "weights": str(trained_weights),
        "onnx": str(saved_onnx_path),
    }
    (model_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    print(f"Saved model artifacts to {model_dir}")
 if __name__ == "__main__":
    main()
--- a/examples/training/README.md
+++ b/examples/training/README.md
@@ -1,89 +0,0 @@
 # SageMaker Training Example
 This example downloads a small image-classification dataset, uploads it through `qc-cli`, and submits a live SageMaker training job.
 ## Prerequisites
 - AWS credentials configured for the profile in `config.yaml`
 - Infrastructure already deployed with `qc-cli infra setup`
 - `config.yaml` updated with:
 ```yaml
 s3:
  bucket: your-bucket-name
 sagemaker:
  training:
    image_uri: 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.6-cpu-py312-ubuntu22.04-sagemaker-v1
    instance_type: ml.m4.xlarge
    instance_count: 1
    source_dir: examples/training/source
    entry_point: train.py
    hyperparameters:
      epochs: 1
      batch-size: 32
      learning-rate: 0.001
      image-size: 160
      validation-split: 0.2
 ```
 ## Training Hyperparameters
 Values under `sagemaker.training.hyperparameters` are passed to the training entry point as command-line arguments. For this example, they map to arguments defined in [source/train.py](source/train.py).
 Supported by this example:
 | Name | Type | Default | Description |
 |---|---:|---:|---|
 | `epochs` | int | `1` | Number of training epochs. |
 | `batch-size` | int | `32` | Images per training batch. |
 | `learning-rate` | float | `0.001` | Adam optimizer learning rate. |
 | `image-size` | int | `160` | Resize images to square `image-size x image-size`. |
 | `validation-split` | float | `0.2` | Fraction of data used for validation. |
 | `max-samples` | int | `0` | Optional cap for smoke tests; `0` means use all images. |
 | `seed` | int | `13` | Random seed for reproducible splitting. |
 | `num-workers` | int | `2` | DataLoader worker count. |
 Do not set `train-dir` or `model-dir` in normal SageMaker runs. SageMaker sets those automatically through `SM_CHANNEL_TRAIN` and `SM_MODEL_DIR`.
 ## 1. Download The Dataset
 ```bash
 bash examples/training/download_flower_photos.sh
 ```
 This creates:
 ```text
 examples/training/data/flower_photos_sagemaker/
  daisy/
  dandelion/
  roses/
  sunflowers/
  tulips/
 ```
 ## 2. Run Training
 Run the training script and wait until it finishes:
 ```bash
 bash examples/training/run_training.sh --config config.yaml --wait
 ```
 Use a dataset that is already uploaded to `s3.data_prefix`:
 ```bash
 bash examples/training/run_training.sh \
  --config config.yaml \
  --skip-upload \
  --wait
 ```
 ## Notes
 - The default dataset path is `examples/training/data/flower_photos_sagemaker`.
 - Uploaded data uses the `s3.bucket` and `s3.data_prefix` values from `config.yaml`.
 - Training artifacts are written under `s3://<bucket>/<model_prefix>/`.
 - The SageMaker `model.tar.gz` contains `model.onnx`, `model.pt`, `class_to_idx.json`, and `metrics.json`.
 - SageMaker packages `examples/training/source`, installs `requirements.txt`, and runs `train.py`.
--- a/examples/training/download_flower_photos.sh
+++ b/examples/training/download_flower_photos.sh
@@ -1,40 +0,0 @@
 #!/usr/bin/env bash
 set -euo pipefail
 DATASET_URL="https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz"
 DEST_DIR="${1:-examples/training/data}"
 ARCHIVE_PATH="${DEST_DIR}/flower_photos.tgz"
 RAW_DATASET_DIR="${DEST_DIR}/flower_photos"
 DATASET_DIR="${DEST_DIR}/flower_photos_sagemaker"
 CLASS_NAMES=("daisy" "dandelion" "roses" "sunflowers" "tulips")
 mkdir -p "${DEST_DIR}"
 if [[ -d "${DATASET_DIR}" ]]; then
  echo "Dataset already exists: ${DATASET_DIR}"
  echo "Use this path with run_training.py:"
  echo "  ${DATASET_DIR}"
  exit 0
 fi
 echo "Downloading TensorFlow flower_photos dataset..."
 if command -v curl >/dev/null 2>&1; then
  curl -L "${DATASET_URL}" -o "${ARCHIVE_PATH}"
 elif command -v wget >/dev/null 2>&1; then
  wget -O "${ARCHIVE_PATH}" "${DATASET_URL}"
 else
  echo "Either curl or wget is required." >&2
  exit 1
 fi
 echo "Extracting dataset..."
 tar -xzf "${ARCHIVE_PATH}" -C "${DEST_DIR}"
 echo "Preparing SageMaker directory layout..."
 mkdir -p "${DATASET_DIR}"
 for class_name in "${CLASS_NAMES[@]}"; do
  cp -R "${RAW_DATASET_DIR}/${class_name}" "${DATASET_DIR}/${class_name}"
 done
 echo "Dataset ready: ${DATASET_DIR}"
 find "${DATASET_DIR}" -mindepth 1 -maxdepth 1 -type d -print | sort
--- a/examples/training/run_training.sh
+++ b/examples/training/run_training.sh
@@ -1,112 +0,0 @@
 #!/usr/bin/env bash
 set -euo pipefail
 CONFIG_PATH="config.yaml"
 DATASET_DIR="examples/training/data/flower_photos_sagemaker"
 WAIT=false
 SKIP_UPLOAD=false
 POLL_SECONDS=60
 usage() {
  cat <<EOF
 Usage: $0 [options]
 Options:
  --config PATH          Path to qc-cli config file. Default: config.yaml
  --dataset-dir PATH     Dataset directory to upload. Default: ${DATASET_DIR}
  --skip-upload          Train against data already uploaded to s3.data_prefix.
  --wait                 Poll until training completes.
  -h, --help             Show this help.
 EOF
 }
 while [[ $# -gt 0 ]]; do
  case "$1" in
    --config)
      CONFIG_PATH="$2"
      shift 2
      ;;
    --dataset-dir)
      DATASET_DIR="$2"
      shift 2
      ;;
    --skip-upload)
      SKIP_UPLOAD=true
      shift
      ;;
    --wait)
      WAIT=true
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown option: $1" >&2
      usage >&2
      exit 1
      ;;
  esac
 done
 if [[ ! -f "${CONFIG_PATH}" ]]; then
  echo "Config not found: ${CONFIG_PATH}" >&2
  exit 1
 fi
 if [[ "${SKIP_UPLOAD}" == false && ! -d "${DATASET_DIR}" ]]; then
  echo "Dataset not found: ${DATASET_DIR}" >&2
  echo "Run: bash examples/training/download_flower_photos.sh" >&2
  exit 1
 fi
 run() {
  echo "+ $*"
  "$@"
 }
 run uv run qc-cli infra status --config "${CONFIG_PATH}"
 if [[ "${SKIP_UPLOAD}" == false ]]; then
  run uv run qc-cli upload "${DATASET_DIR}" --config "${CONFIG_PATH}"
 fi
 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="$(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
 fi
 echo "Submitted SageMaker training job: ${JOB_NAME}"
 if [[ "${WAIT}" == false ]]; then
  run uv run qc-cli train status "${JOB_NAME}" --config "${CONFIG_PATH}"
  exit 0
 fi
 while true; do
  STATUS_OUTPUT="$(uv run qc-cli train status "${JOB_NAME}" --config "${CONFIG_PATH}")"
  echo "${STATUS_OUTPUT}"
  if printf '%s\n' "${STATUS_OUTPUT}" | grep -q 'Status:.*Completed'; then
    echo "Training completed successfully."
    exit 0
  fi
  if printf '%s\n' "${STATUS_OUTPUT}" | grep -q 'Status:.*Failed'; then
    echo "Training failed." >&2
    exit 1
  fi
  if printf '%s\n' "${STATUS_OUTPUT}" | grep -q 'Status:.*Stopped'; then
    echo "Training stopped." >&2
    exit 1
  fi
  sleep "${POLL_SECONDS}"
 done
--- a/examples/training/source/requirements.txt
+++ b/examples/training/source/requirements.txt
@@ -1 +0,0 @@
 onnx==1.21.0
--- a/examples/training/source/train.py
+++ b/examples/training/source/train.py
@@ -1,188 +0,0 @@
 #!/usr/bin/env python3
 """SageMaker entry point for CPU image-classification training."""
 from __future__ import annotations
 import argparse
 import json
 import os
 import random
 from pathlib import Path
 import torch
 from torch import nn
 from torch.utils.data import DataLoader, Subset, random_split
 from torchvision import datasets, transforms
 class SmallImageClassifier(nn.Module):
    def __init__(self, class_count: int) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.classifier = nn.Linear(64, class_count)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--learning-rate", type=float, default=0.001)
    parser.add_argument("--image-size", type=int, default=160)
    parser.add_argument("--validation-split", type=float, default=0.2)
    parser.add_argument("--max-samples", type=int, default=0)
    parser.add_argument("--seed", type=int, default=13)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--train-dir", default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
    parser.add_argument("--model-dir", default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
    return parser.parse_args()
 def build_datasets(args: argparse.Namespace) -> tuple[Subset, Subset, dict[str, int]]:
    transform = transforms.Compose(
        [
            transforms.Resize((args.image_size, args.image_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ]
    )
    dataset = datasets.ImageFolder(args.train_dir, transform=transform)
    if len(dataset.classes) < 2:
        raise ValueError(f"Expected at least two classes in {args.train_dir}. Found: {dataset.classes}")
    if args.max_samples > 0 and args.max_samples < len(dataset):
        indices = list(range(len(dataset)))
        random.Random(args.seed).shuffle(indices)
        dataset = Subset(dataset, indices[: args.max_samples])
    validation_size = max(1, int(len(dataset) * args.validation_split))
    train_size = len(dataset) - validation_size
    if train_size < 1:
        raise ValueError("Not enough images to create a train/validation split.")
    generator = torch.Generator().manual_seed(args.seed)
    train_dataset, validation_dataset = random_split(dataset, [train_size, validation_size], generator=generator)
    return train_dataset, validation_dataset, getattr(dataset, "dataset", dataset).class_to_idx
 def run_epoch(
    model: nn.Module,
    data_loader: DataLoader,
    criterion: nn.Module,
    optimizer: torch.optim.Optimizer | None,
    device: torch.device,
 ) -> tuple[float, float]:
    training = optimizer is not None
    model.train(training)
    total_loss = 0.0
    total_correct = 0
    total_examples = 0
    for images, labels in data_loader:
        images = images.to(device)
        labels = labels.to(device)
        with torch.set_grad_enabled(training):
            logits = model(images)
            loss = criterion(logits, labels)
            if training:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
        total_loss += loss.item() * images.size(0)
        total_correct += (logits.argmax(dim=1) == labels).sum().item()
        total_examples += images.size(0)
    return total_loss / total_examples, total_correct / total_examples
 def export_onnx(model: nn.Module, model_dir: Path, image_size: int) -> None:
    model.eval()
    dummy_input = torch.randn(1, 3, image_size, image_size)
    torch.onnx.export(
        model,
        dummy_input,
        model_dir / "model.onnx",
        export_params=True,
        opset_version=17,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["logits"],
    )
 def main() -> None:
    args = parse_args()
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    train_dataset, validation_dataset, class_to_idx = build_datasets(args)
    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
    )
    validation_loader = DataLoader(
        validation_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
    )
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SmallImageClassifier(class_count=len(class_to_idx)).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
    print(f"Training on {device}. Classes: {sorted(class_to_idx)}")
    metrics = []
    for epoch in range(1, args.epochs + 1):
        train_loss, train_accuracy = run_epoch(model, train_loader, criterion, optimizer, device)
        validation_loss, validation_accuracy = run_epoch(model, validation_loader, criterion, None, device)
        epoch_metrics = {
            "epoch": epoch,
            "train_loss": train_loss,
            "train_accuracy": train_accuracy,
            "validation_loss": validation_loss,
            "validation_accuracy": validation_accuracy,
        }
        metrics.append(epoch_metrics)
        print(json.dumps(epoch_metrics, sort_keys=True))
    model_dir = Path(args.model_dir)
    model_dir.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "model_state_dict": model.cpu().state_dict(),
            "class_to_idx": class_to_idx,
            "image_size": args.image_size,
        },
        model_dir / "model.pt",
    )
    export_onnx(model, model_dir, args.image_size)
    (model_dir / "class_to_idx.json").write_text(json.dumps(class_to_idx, indent=2), encoding="utf-8")
    (model_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    print(f"Saved model artifacts to {model_dir}")
 if __name__ == "__main__":
    main()
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 [project]
 name = "qc-cli"
 version = "0.1.0"
-description = "CLI for SageMaker ONNX training and Qualcomm AI Hub optimization"
+description = "CLI for training and deploying models for Qualcomm AI Hub"
 requires-python = ">=3.13"
 dependencies = [
    "aws-cdk-lib>=2.180.0",
@@ -29,8 +29,6 @@ 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",
--- a/src/cloud/init.py
+++ b/src/cloud/init.py
--- a/src/commands/ai_hub.py
+++ b/src/commands/ai_hub.py
@@ -1,18 +1,21 @@
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from datetime import datetime
 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
 from src.config import Config
 from src.qualcomm import aihub_jobs
-from src.qualcomm.artifacts import resolve_onnx
+from src.qualcomm.artifacts import ResolvedOnnx, resolve_onnx
-app = typer.Typer(help="Quantize, compile, validate, profile, and download models with Qualcomm AI Hub")
+app = typer.Typer(help="Optimize, quantize, compile, validate, profile, and download models with Qualcomm Workbench")
 _RUNTIME_EXTENSIONS = {
    "tflite": "tflite",
@@ -22,12 +25,19 @@ _RUNTIME_EXTENSIONS = {
 class UploadStep(StrEnum):
    optimize = "optimize"
    quantize = "quantize"
    compile = "compile"
    validate = "validate"
    profile = "profile"
@dataclass(frozen=True)
 class ResolvedModelSource:
    model: str | Path
    model_artifact: str | None = None
 def _input_specs(cfg: Config) -> dict[str, tuple[tuple[int, ...], str]]:
    specs = {name: (tuple(shape), dtype) for name, (shape, dtype) in cfg.aihub.input_specs.items()}
    if not specs:
@@ -99,24 +109,105 @@ def _model_id_or_state(config_path: str, model_id: str | None, *, quantized: boo
    return resolved
 def _resolve_model_source(
    cfg: Config,
    config_path: str,
    *,
    model_id: str | None = None,
    previous_model_id: str | None = None,
    from_job: str | None = None,
    model_s3_uri: str | None = None,
    onnx_path: str | None = None,
 ) -> ResolvedModelSource:
    if model_id:
        return ResolvedModelSource(model_id)
    has_explicit_source = bool(from_job or model_s3_uri or onnx_path)
    if previous_model_id and not has_explicit_source:
        return ResolvedModelSource(previous_model_id)
    resolved = _resolve_onnx_source(
        cfg,
        config_path,
        from_job=from_job,
        model_s3_uri=model_s3_uri,
        onnx_path=onnx_path,
    )
    return ResolvedModelSource(resolved.onnx_path, resolved.model_artifact)
 def _resolve_onnx_source(
    cfg: Config,
    config_path: str,
    *,
    from_job: str | None = None,
    model_s3_uri: str | None = None,
    onnx_path: str | None = None,
 ) -> ResolvedOnnx:
    st = state_ops.store(config_path)
    last_training_job = st.get_last_training_job()
    saved_model_artifact = None
    if not from_job and not model_s3_uri and not onnx_path and not last_training_job:
        saved_model_artifact = st.get_last_model_artifact()
    return resolve_onnx(
        cfg=cfg,
        output_dir=cfg.aihub.output_dir,
        from_job=from_job,
        model_s3_uri=model_s3_uri or saved_model_artifact,
        onnx_path=onnx_path,
        last_training_job=last_training_job,
    )
 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,
    calibration_path: Path,
-    from_job: str | None,
+    *,
-    model_s3_uri: str | None,
+    model_id: str | None = None,
-    onnx_path: str | None,
+    from_job: str | None = None,
    model_s3_uri: str | None = None,
    onnx_path: str | None = None,
 ) -> str:
    st = state_ops.store(config_path)
    specs = _input_specs(cfg)
    try:
-        resolved = resolve_onnx(
+        source = _resolve_model_source(
-            cfg=cfg,
+            cfg,
-            output_dir=cfg.aihub.output_dir,
+            config_path,
            model_id=model_id,
            previous_model_id=st.get_last_optimized_model_id(),
            from_job=from_job,
-            model_s3_uri=model_s3_uri or st.get_last_model_artifact(),
+            model_s3_uri=model_s3_uri,
            onnx_path=onnx_path,
            last_training_job=st.get_last_training_job(),
        )
        calibration_data = _load_calibration(calibration_path, specs)
    except (FileNotFoundError, ValueError) as e:
@@ -124,72 +215,117 @@ def _quantize_step(
        raise typer.Exit(1)
    try:
        hub_model = (
            hub.upload_model(str(source.model), name=cfg.aihub.model_name)
            if isinstance(source.model, Path)
            else hub.get_model(source.model)
        )
        result = aihub_jobs.submit_quantize_job(
-            resolved.onnx_path,
+            hub_model,
            calibration_data,
            cfg.aihub.quantize_options,
            job_name=_job_name(cfg, "quantize"),
            model_name=cfg.aihub.model_name,
        )
    except Exception as e:
        CONSOLE.print(f"[red]AI Hub quantize failed: {e}[/red]")
        raise typer.Exit(1)
-    st.update(
+    updates: dict[str, Any] = {
-        last_model_artifact=resolved.model_artifact,
+        "last_quantize_job_id": result["job_id"],
-        last_quantize_job_id=result["job_id"],
+        "last_quantized_model_id": result["model_id"],
-        last_quantized_model_id=result["model_id"],
+    }
-    )
+    if source.model_artifact:
        updates["last_model_artifact"] = source.model_artifact
    st.update(**updates)
    CONSOLE.print(f"[green]✓[/green] Quantize job: [bold]{result['job_id']}[/bold]")
    CONSOLE.print(f"[green]✓[/green] Quantized model: [bold]{result['model_id']}[/bold]")
    return str(result["model_id"])
 def _optimize_step(
    cfg: Config,
    config_path: str,
    from_job: str | None,
    model_s3_uri: str | None,
    onnx_path: str | None,
 ) -> str:
    st = state_ops.store(config_path)
    _validate_device(cfg)
    specs = _input_specs(cfg)
    try:
        source = _resolve_onnx_source(
            cfg,
            config_path,
            from_job=from_job,
            model_s3_uri=model_s3_uri,
            onnx_path=onnx_path,
        )
    except (FileNotFoundError, ValueError) as e:
        CONSOLE.print(f"[red]{e}[/red]")
        raise typer.Exit(1)
    try:
        hub_model = hub.upload_model(str(source.onnx_path), name=cfg.aihub.model_name)
        result = aihub_jobs.submit_compile_job(
            model=hub_model,
            device=cfg.aihub.device,
            input_specs=specs,
            target_runtime="onnx",
            job_name=_job_name(cfg, "optimize"),
        )
    except Exception as e:
        CONSOLE.print(f"[red]AI Hub ONNX optimization failed: {e}[/red]")
        raise typer.Exit(1)
    st.update(
        last_model_artifact=source.model_artifact,
        last_optimize_job_id=result["job_id"],
        last_optimized_model_id=result["model_id"],
    )
    CONSOLE.print(f"[green]✓[/green] ONNX optimization job: [bold]{result['job_id']}[/bold]")
    CONSOLE.print(f"[green]✓[/green] Optimized ONNX model: [bold]{result['model_id']}[/bold]")
    return str(result["model_id"])
 def _compile_step(
    cfg: Config,
    config_path: str,
    model_id: str | None,
    from_job: str | None,
    model_s3_uri: str | None,
    onnx_path: str | None,
    *,
-    prefer_quantized: bool,
+    model_id: str | None = None,
    from_job: str | None = None,
    model_s3_uri: str | None = None,
    onnx_path: str | None = None,
 ) -> str:
    st = state_ops.store(config_path)
    _validate_device(cfg)
    specs = _input_specs(cfg)
-
+    try:
-    model: Any
+        source = _resolve_model_source(
-    model_artifact: str | None = None
+            cfg,
-    has_explicit_source = bool(from_job or model_s3_uri or onnx_path)
+            config_path,
-    if model_id:
+            model_id=model_id,
-        model = model_id
+            previous_model_id=st.get_last_quantized_model_id(),
-    elif prefer_quantized and not has_explicit_source and st.get_last_quantized_model_id():
+            from_job=from_job,
-        model = st.get_last_quantized_model_id()
+            model_s3_uri=model_s3_uri,
-    else:
+            onnx_path=onnx_path,
-        try:
+        )
-            resolved = resolve_onnx(
+    except (FileNotFoundError, ValueError) as e:
-                cfg=cfg,
+        CONSOLE.print(f"[red]{e}[/red]")
-                output_dir=cfg.aihub.output_dir,
+        raise typer.Exit(1)
                from_job=from_job,
                model_s3_uri=model_s3_uri,
                onnx_path=onnx_path,
                last_training_job=st.get_last_training_job(),
            )
        except (FileNotFoundError, ValueError) as e:
            CONSOLE.print(f"[red]{e}[/red]")
            raise typer.Exit(1)
        model = resolved.onnx_path
        model_artifact = resolved.model_artifact
    try:
        hub_model = (
            hub.upload_model(str(source.model), name=cfg.aihub.model_name)
            if isinstance(source.model, Path)
            else hub.get_model(source.model)
        )
        result = aihub_jobs.submit_compile_job(
-            model=model,
+            model=hub_model,
-            device_name=cfg.aihub.device,
+            device=cfg.aihub.device,
            input_specs=specs,
            target_runtime=cfg.aihub.target_runtime,
            options=cfg.aihub.compile_options,
            job_name=_job_name(cfg, "compile"),
            model_name=cfg.aihub.model_name if isinstance(model, Path) else None,
        )
    except Exception as e:
        CONSOLE.print(f"[red]AI Hub compile failed: {e}[/red]")
@@ -199,8 +335,8 @@ def _compile_step(
        "last_compile_job_id": result["job_id"],
        "last_compiled_model_id": result["model_id"],
    }
-    if model_artifact:
+    if source.model_artifact:
-        updates["last_model_artifact"] = model_artifact
+        updates["last_model_artifact"] = source.model_artifact
    st.update(**updates)
    CONSOLE.print(f"[green]✓[/green] Compile job: [bold]{result['job_id']}[/bold]")
    CONSOLE.print(f"[green]✓[/green] Compiled model: [bold]{result['model_id']}[/bold]")
@@ -214,6 +350,7 @@ 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:
@@ -225,8 +362,9 @@ def _validate_step(
    run = datetime.now().strftime("%Y%m%d-%H%M%S")
    out_dir = Path(cfg.aihub.output_dir) / run / "validation"
    try:
        hub_model = hub.get_model(resolved_model_id)
        result = aihub_jobs.submit_inference_job(
-            resolved_model_id,
+            hub_model,
            cfg.aihub.device,
            inputs,
            out_dir,
@@ -247,10 +385,12 @@ 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:
        hub_model = hub.get_model(resolved_model_id)
        result = aihub_jobs.submit_profile_job(
-            resolved_model_id,
+            hub_model,
            cfg.aihub.device,
            cfg.aihub.profile_options,
            job_name=_job_name(cfg, "profile"),
@@ -263,9 +403,24 @@ def _profile_step(cfg: Config, config_path: str, model_id: str | None) -> str:
    return str(result["job_id"])
@app.command()
 def optimize(
    from_job: str | None = typer.Option(None, "--from-job", help="Training job name whose model artifact should optimize"),
    model_s3_uri: str | None = typer.Option(None, "--model-s3-uri", help="S3 URI of model.tar.gz to optimize"),
    onnx_path: str | None = typer.Option(
        None, "--onnx-path", help="Local ONNX path or ONNX path inside extracted artifact"
    ),
    config: str = CONFIG_OPT,
 ) -> None:
    """Optimize an external model into a Workbench-produced ONNX model."""
    cfg = load_cfg(config)
    _optimize_step(cfg, config, from_job, model_s3_uri, onnx_path)
@app.command()
 def quantize(
    calibration_path: Path = typer.Argument(..., help="Calibration .npz file or directory of .npy samples"),
    model_id: str | None = typer.Option(None, "--model-id", help="AI Hub optimized ONNX model ID"),
    from_job: str | None = typer.Option(None, "--from-job", help="Training job name whose model artifact should quantize"),
    model_s3_uri: str | None = typer.Option(None, "--model-s3-uri", help="S3 URI of model.tar.gz to quantize"),
    onnx_path: str | None = typer.Option(
@@ -275,7 +430,15 @@ def quantize(
 ) -> None:
    """Quantize an ONNX model to INT8."""
    cfg = load_cfg(config)
-    _quantize_step(cfg, config, calibration_path, from_job, model_s3_uri, onnx_path)
+    _quantize_step(
        cfg,
        config,
        calibration_path,
        model_id=model_id,
        from_job=from_job,
        model_s3_uri=model_s3_uri,
        onnx_path=onnx_path,
    )
@app.command()
@@ -290,7 +453,14 @@ def compile(
 ) -> None:
    """Compile a model for the configured Qualcomm AI Hub target."""
    cfg = load_cfg(config)
-    _compile_step(cfg, config, model_id, from_job, model_s3_uri, onnx_path, prefer_quantized=True)
+    _compile_step(
        cfg,
        config,
        model_id=model_id,
        from_job=from_job,
        model_s3_uri=model_s3_uri,
        onnx_path=onnx_path,
    )
@app.command()
@@ -319,7 +489,7 @@ def profile(
 def upload(
    calibration_path: Path = typer.Argument(..., help="Calibration .npz file or directory of .npy samples"),
    input_file: Path = typer.Argument(..., help="Validation .npz or .npy inputs to run on device"),
-    from_step: UploadStep = typer.Option(UploadStep.quantize, "--from-step", help="Resume from this Workbench step"),
+    from_step: UploadStep = typer.Option(UploadStep.optimize, "--from-step", help="Resume from this Workbench step"),
    from_job: str | None = typer.Option(None, "--from-job", help="Training job name whose model artifact should upload"),
    model_s3_uri: str | None = typer.Option(None, "--model-s3-uri", help="S3 URI of model.tar.gz to upload"),
    onnx_path: str | None = typer.Option(
@@ -328,25 +498,48 @@ def upload(
    input_name: str | None = typer.Option(None, "--input-name", help="Input name for .npy validation files"),
    config: str = CONFIG_OPT,
 ) -> None:
-    """Run the four Workbench upload steps: quantize, compile, validate, and profile."""
+    """Optimize, quantize, optionally compile, validate, and profile a model."""
    cfg = load_cfg(config)
-    steps = [UploadStep.quantize, UploadStep.compile, UploadStep.validate, UploadStep.profile]
+    steps = [UploadStep.optimize, UploadStep.quantize, UploadStep.compile, UploadStep.validate, UploadStep.profile]
    selected = steps[steps.index(from_step) :]
    optimized_model_id: str | None = None
    quantized_model_id: str | None = None
    compiled_model_id: str | None = None
    if UploadStep.optimize in selected:
        optimized_model_id = _optimize_step(cfg, config, from_job, model_s3_uri, onnx_path)
    if UploadStep.quantize in selected:
-        quantized_model_id = _quantize_step(cfg, config, calibration_path, from_job, model_s3_uri, onnx_path)
+        if UploadStep.optimize not in selected:
-    if UploadStep.compile in selected:
+            optimized_model_id = state_ops.store(config).get_last_optimized_model_id()
-        compiled_model_id = _compile_step(
+            if not optimized_model_id:
                CONSOLE.print(
                    "[red]No optimized ONNX model found. Resume from --from-step optimize or run "
                    "'qc-cli ai-hub optimize' first.[/red]"
                )
                raise typer.Exit(1)
        quantized_model_id = _quantize_step(
            cfg,
            config,
-            model_id=quantized_model_id,
+            calibration_path,
-            from_job=from_job,
+            model_id=optimized_model_id,
            model_s3_uri=model_s3_uri,
            onnx_path=onnx_path,
            prefer_quantized=True,
        )
    if UploadStep.compile in selected:
        if cfg.aihub.target_runtime == "onnx":
            compiled_model_id = quantized_model_id or state_ops.store(config).get_last_quantized_model_id()
            if not compiled_model_id:
                CONSOLE.print(
                    "[red]No quantized ONNX model found. Resume from --from-step quantize or run "
                    "'qc-cli ai-hub quantize' first.[/red]"
                )
                raise typer.Exit(1)
            state_ops.store(config).update(last_compiled_model_id=compiled_model_id)
            CONSOLE.print("[green]✓[/green] Target runtime is ONNX; skipping final compile.")
        else:
            compiled_model_id = _compile_step(
                cfg,
                config,
                model_id=quantized_model_id,
            )
    if UploadStep.validate in selected:
        _validate_step(cfg, config, input_file, compiled_model_id, input_name)
    if UploadStep.profile in selected:
--- a/src/commands/infra.py
+++ b/src/commands/infra.py
@@ -150,35 +150,6 @@ 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,
--- a/src/commands/init.py
+++ b/src/commands/init.py
@@ -0,0 +1,40 @@
 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
--- a/src/commands/mlflow.py
+++ b/src/commands/mlflow.py
@@ -0,0 +1,41 @@
 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]")
--- a/src/commands/train.py
+++ b/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 infra mlflow-url[/cyan]")
+        CONSOLE.print("Open MLflow: [cyan]qc-cli mlflow open[/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 infra mlflow-url[/cyan]")
+        CONSOLE.print("Open MLflow: [cyan]qc-cli mlflow open[/cyan]")
@app.command(name="list")
--- a/src/commands/upload.py
+++ b/src/commands/upload.py
@@ -0,0 +1,70 @@
 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)
--- a/src/config.py
+++ b/src/config.py
@@ -4,7 +4,8 @@ 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, model_validator
+from pydantic import BaseModel, Field, field_validator, model_validator
 from qai_hub.client import Device
 class MlflowMode(StrEnum):
@@ -81,7 +82,7 @@ class SageMakerConfig(BaseModel):
 class AIHubConfig(BaseModel):
-    device: str = "Samsung Galaxy S25 (Family)"
+    device: Device = Field(default_factory=lambda: Device("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
@@ -91,6 +92,13 @@ 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
--- a/src/main.py
+++ b/src/main.py
@@ -1,115 +1,14 @@
 import secrets
 from pathlib import Path
 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, init, mlflow, train, upload
 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)
--- a/src/qualcomm/init.py
+++ b/src/qualcomm/init.py
@@ -1 +0,0 @@
--- a/src/qualcomm/aihub_jobs.py
+++ b/src/qualcomm/aihub_jobs.py
@@ -28,31 +28,20 @@ def _dataset_entries(inputs: dict[str, Any]) -> dict[str, list[Any]]:
 def submit_compile_job(
-    model: Any,
+    model: Model,
-    device_name: str,
+    device: Device,
    input_specs: dict[str, tuple[tuple[int, ...], str]],
    target_runtime: str,
    options: str | None = None,
    job_name: str | None = None,
    model_name: str | None = None,
 ) -> ModelJobResult:
    compile_options = f"--target_runtime {target_runtime}"
    if options:
        compile_options = f"{compile_options} {options}"
    model_arg = model
    if isinstance(model, Path):
        model_arg = str(model)
    elif isinstance(model, str):
        candidate = Path(model)
        model_arg = model if candidate.exists() or candidate.suffix else hub.get_model(model)
    if model_name and isinstance(model_arg, str) and Path(model_arg).exists():
        model_arg = hub.upload_model(model_arg, name=model_name)
    job = hub.submit_compile_job(
-        model=model_arg,
+        model=model,
-        device=Device(device_name),
+        device=device,
        name=job_name,
        input_specs=input_specs,
        options=compile_options,
@@ -64,15 +53,15 @@ def submit_compile_job(
 def submit_inference_job(
-    model_id: str,
+    model: Model,
-    device_name: str,
+    device: Device,
    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),
+        model=model,
-        device=Device(device_name),
+        device=device,
        inputs=_dataset_entries(inputs),
        name=job_name,
    )
@@ -83,14 +72,14 @@ def submit_inference_job(
 def submit_profile_job(
-    model_id: str,
+    model: Model,
-    device_name: str,
+    device: Device,
    options: str | None = None,
    job_name: str | None = None,
 ) -> ProfileJobResult:
    job = hub.submit_profile_job(
-        model=hub.get_model(model_id),
+        model=model,
-        device=Device(device_name),
+        device=device,
        name=job_name,
        options=options or "",
    )
@@ -98,17 +87,13 @@ def submit_profile_job(
 def submit_quantize_job(
-    model: str | Path,
+    model: Model,
    calibration_data: dict[str, Any],
    options: str | None = None,
    job_name: str | None = None,
    model_name: str | None = None,
 ) -> ModelJobResult:
    model_arg = str(model)
    if model_name and Path(model_arg).exists():
        model_arg = hub.upload_model(model_arg, name=model_name)
    job = hub.submit_quantize_job(
-        model=model_arg,
+        model=model,
        calibration_data=_dataset_entries(calibration_data),
        weights_dtype=QuantizeDtype.INT8,
        activations_dtype=QuantizeDtype.INT8,
--- a/src/state.py
+++ b/src/state.py
@@ -37,6 +37,10 @@ class CliStateStore:
        value = self.get("last_model_artifact")
        return str(value) if value else None
    def get_last_optimized_model_id(self) -> str | None:
        value = self.get("last_optimized_model_id")
        return str(value) if value else None
    def get_last_quantized_model_id(self) -> str | None:
        value = self.get("last_quantized_model_id")
        return str(value) if value else None
--- a/uv.lock
+++ b/uv.lock
@@ -1003,15 +1003,6 @@ 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"
@@ -1674,15 +1665,6 @@ 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"
@@ -1963,34 +1945,6 @@ 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"
@@ -2114,8 +2068,6 @@ dependencies = [
 dev = [
    { name = "boto3-stubs", extra = ["iam", "s3", "sagemaker"] },
    { name = "pyright" },
    { name = "pytest" },
    { name = "pytest-mock" },
    { name = "ruff" },
    { name = "types-pyyaml" },
 ]
@@ -2138,8 +2090,6 @@ 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" },
 ]
Author	SHA1	Message	Date
slalom	f26e8256f0	update ai-hub to first optimize model for Workbench Remove old examples	2026-06-09 14:55:26 -04:00
slalom	6c9f30d290	clean	2026-06-09 12:21:25 -04:00
slalom	c6d4da66ae	remove	2026-06-09 12:19:17 -04:00
slalom	9dc6f478bd	add ONNX sanitize required for Ultralytics	2026-06-09 12:18:41 -04:00
samirodr	c2d3f44498	add	2026-06-09 12:06:49 -04:00
slalom	46cf2d5afe	include steps for ai-hub	2026-06-09 11:55:03 -04:00
slalom	98b4d0d200	another one	2026-06-09 10:14:49 -04:00
slalom	f1f5dcbed7	update	2026-06-09 10:01:09 -04:00
slalom	75f66f81c1	initial version to train yolo model	2026-06-09 09:15:35 -04:00
samirodr	5360a482fc	update	2026-06-08 14:59:44 -04:00
samirodr	6a560a8610	match	2026-06-08 14:54:13 -04:00
slalom	d244150d98	move mlflow to its own command	2026-06-05 11:47:38 -04:00
slalom	d7c7158464	clean main file	2026-06-05 11:25:04 -04:00
slalom	6bc25dc183	restructure config to use Device class directly Also include device validation	2026-06-04 17:28:17 -04:00
samirodr	71a95aa3a7	update description	2026-06-03 17:13:00 -04:00
slalom	a3f3060e13	ai-hub (#3 ) Reviewed-on: #3	2026-06-03 21:06:06 +00:00