update

match
move mlflow to its own command
2026-06-08 14:59:44 -04:00 · 2026-06-08 14:54:13 -04:00 · 2026-06-05 11:47:38 -04:00 · 2026-06-05 11:25:04 -04:00 · 2026-06-04 17:28:17 -04:00 · 2026-06-03 17:13:00 -04:00
34 changed files with 4581 additions and 73 deletions
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,224 @@
+# Byte-compiled / optimized / DLL files
 __pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#   Usually these files are written by a python script from a template
+#   before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+# Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+# uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+# poetry.lock
+# poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+# pdm.lock
+# pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+# pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# Redis
+*.rdb
+*.aof
+*.pid
+
+# RabbitMQ
+mnesia/
+rabbitmq/
+rabbitmq-data/
+
+# ActiveMQ
+activemq-data/
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#   JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#   be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#   and can be added to the global gitignore or merged into this file.  For a more nuclear
+#   option (not recommended) you can uncomment the following to ignore the entire idea folder.
+# .idea/
+
+# Abstra
+#   Abstra is an AI-powered process automation framework.
+#   Ignore directories containing user credentials, local state, and settings.
+#   Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#   that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#   and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#   you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+# Temporary file for partial code execution
+tempCodeRunnerFile.py
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml
+
 .venv/
-config.yaml
+config*.yaml
 cdk.out/
-.qai-cli-infra*
+.qc-cli*.json
+examples/*/data/
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
 # qc-cli

-A CLI for the Qualcomm model MLOps pipeline — browse and download models from Qualcomm AI Hub, fine-tune them on custom datasets using SageMaker, validate inference, and prepare artifacts for Qualcomm hardware deployment.
+A CLI for Qualcomm's MLOps pipeline — browse and download models from Qualcomm AI Hub, fine-tune them on custom datasets using SageMaker, validate inference, and prepare artifacts for Qualcomm hardware deployment.

 ## Requirements

@@ -30,11 +30,16 @@ qc-cli --help
 # 1. Create config.yaml in the current directory
 qc-cli init

-# 2. Edit config.yaml — at minimum set s3.bucket and sagemaker.role_name
+# 2. Edit config.yaml — at minimum set sagemaker.training.image_uri

 # 3. Provision AWS infrastructure (S3 bucket + SageMaker IAM role).
 #    This is the step that requires the AWS CDK CLI.
 qc-cli infra setup
+
+# 4. Upload training data, then submit a SageMaker training job.
+qc-cli upload ./my-dataset
+qc-cli train start
+qc-cli train status
 ```

 ## Configuration
@@ -42,25 +47,56 @@ qc-cli infra setup
 `qc-cli init` writes a `config.yaml` in the current directory. The fields you must fill in before using the tool:

 ```yaml
+infra:
+  stack_name: qc-cli-mlops-1a2b3c4d5e6f
+
 aws:
  region: us-east-1
  profile: default          # AWS CLI profile name

 s3:
-  bucket: your-unique-bucket-name
+  bucket: qc-cli-mlops-1a2b3c4d5e6f-data

 sagemaker:
-  role_name: qc-cli-sagemaker-role
+  training:
+    image_uri: ""           # ECR URI for your training container
+    instance_type: ml.m5.xlarge
+    instance_count: 1
+    entry_point: null       # Optional: script inside source_dir
+    source_dir: null        # Optional: local dir packaged and uploaded automatically
+    hyperparameters: {}
+
+aihub:
+  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
+  model_name: null          # Optional name for uploaded local ONNX models
+  compile_options: null
+  profile_options: null
+  quantize_options: null
+  output_dir: build/qai-hub
 ```

+`qc-cli init` generates the `infra.stack_name` and `s3.bucket` namespace once and writes it to `config.yaml`. Keep these values stable for a deployment; changing them points the CLI at different infrastructure.
+
+The CLI isolates both application resources and CDK bootstrap resources. The application CloudFormation stack uses `infra.stack_name`, the S3 bucket uses the same generated namespace because bucket names are globally unique, and the SageMaker IAM role uses a CloudFormation-generated physical name. CDK bootstrap resources are derived internally from `infra.stack_name`, including a bootstrap stack named `<stack_name>-bootstrap` and a matching non-default CDK asset bucket qualifier. `qc-cli infra destroy` removes the application stack but leaves the CDK bootstrap stack in place; the command prints the retained bootstrap stack name.
+
+`hyperparameters` is a flat map of values passed to the training container. Valid keys depend on the selected training image and entry point.
+
 To provision an MLflow tracking server, set:

 ```yaml
 mlflow:
  mode: create
-  tracking_server_name: your-tracking-server-name
+  experiment_name: qc-cli-training
+  registered_model_name: qc-cli-model
+  register_trained_models: true
 ```

+In `create` mode, the CLI manages the tracking server name from `infra.stack_name`; you do not need to set `tracking_server_name`.
+
 To use an existing MLflow tracking server, set:

 ```yaml
@@ -69,6 +105,16 @@ mlflow:
  tracking_server_name: your-tracking-server-name
 ```

+When MLflow is enabled, `train start` creates an MLflow run for the SageMaker job. `train status` finalizes that run once the job reaches a terminal state and registers completed model artifacts as experiment model versions using the `experiment-latest` MLflow alias. An experiment version is an immutable trained-source artifact; it records that training produced a model, not that the model is better than earlier versions or ready for release.
+
+To open the managed SageMaker MLflow UI, request a fresh presigned URL:
+
+```bash
+qc-cli mlflow open --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.
+
 ## Commands

 ### `init`
@@ -79,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`

 ```
@@ -91,6 +143,101 @@ qc-cli infra destroy --yes                 Destroy stack without confirmation
 qc-cli infra destroy --delete-bucket-data  Destroy stack and delete S3 data
 ```

+`--cloudformation-execution-policy` is a one-time CDK bootstrap option, not a `config.yaml` setting. Pass it on `infra setup` when you need the CDK bootstrap CloudFormation execution role to use a policy other than the default `AdministratorAccess`:
+
+```bash
+qc-cli infra setup --cloudformation-execution-policy arn:aws:iam::aws:policy/PowerUserAccess
+```
+
+### `upload`
+
+```
+qc-cli upload <file>                 Upload a single file to S3
+qc-cli upload <dir>                  Upload all files in a directory tree to S3
+qc-cli upload <file> --s3-key <key>  Upload a file to a custom S3 key
+```
+
+Uploads use `s3.bucket` and `s3.data_prefix` from `config.yaml`. File uploads default to `s3://<bucket>/<data_prefix>/<filename>`. Directory uploads are recursive, preserve paths relative to the uploaded directory, and place files under `s3://<bucket>/<data_prefix>/`.
+
+### `train`
+
+```
+qc-cli train start              Submit a SageMaker training job
+qc-cli train status [job-name]  Show job status; defaults to the last submitted job
+qc-cli train list               List recent training jobs
+qc-cli train list --limit 3     Show a custom number of recent jobs
+```
+
+`train start` uses `s3://<bucket>/<data_prefix>/` as the training channel and writes outputs under `s3://<bucket>/<model_prefix>/`. If `sagemaker.training.source_dir` is set, the CLI packages that directory, uploads it beside the job output prefix, and passes `sagemaker_program`/`sagemaker_submit_directory` to the SageMaker container.
+
+The expected output artifact is SageMaker’s `model.tar.gz`, normally containing the trained model file your container writes to `/opt/ml/model`.
+
+### `ai-hub`
+
+```
+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 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.
+
+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.
+
+## Model lifecycle
+
+The CLI uses neutral experiment naming for trained artifacts and reserves release terminology for an explicit promotion step.
+
+Current behavior:
+
+1. `qc-cli train start` submits a SageMaker training job.
+2. `qc-cli train status` finalizes the MLflow run after the job reaches a terminal state.
+3. If the job completed and `mlflow.register_trained_models` is enabled, the SageMaker `model.tar.gz` is registered as a new MLflow model version with:
+   - `qc_cli.stage=experiment`
+   - `qc_cli.artifact_kind=trained_source`
+   - `qc_cli.source=sagemaker`
+4. The MLflow alias `experiment-latest` points at the most recently registered experiment version.
+5. AI Hub upload commands create deployable derived artifacts from a trained-source experiment or local ONNX model.
+
+Future release aliases such as `v1` or `production` can point at a selected deployable artifact.
+
+Example future metadata:
+
+```text
+qc-cli-model version 12
+qc_cli.stage=experiment
+qc_cli.artifact_kind=trained_source
+qc_cli.source=sagemaker
+
+qc-cli-model-aihub version 3
+qc_cli.stage=ai_hub_compiled
+qc_cli.artifact_kind=deployable
+qc_cli.parent_registered_model_name=qc-cli-model
+qc_cli.parent_model_version=12
+qc_cli.runtime=tflite
+qc_cli.quantization=int8
+qc_cli.target_device=Samsung Galaxy S25
+```
+
+In that flow, `experiment-latest` remains a training convenience alias. Release selection is a separate promotion decision based on the derived artifact, not on the experiment name.
+
 ## AWS permissions required

 The IAM user or role running the CLI needs:
@@ -101,6 +248,7 @@ The IAM user or role running the CLI needs:
 | CreateRole, GetRole, DeleteRole, AttachRolePolicy, DetachRolePolicy | IAM |
 | CreateStack, UpdateStack, DeleteStack, DescribeStacks, DescribeStackEvents | CloudFormation |
 | GetCallerIdentity | STS |
+| CreateTrainingJob, DescribeTrainingJob, ListTrainingJobs | SageMaker AI |
 | CreateMlflowTrackingServer, DescribeMlflowTrackingServer, DeleteMlflowTrackingServer | SageMaker AI, when `mlflow.mode` is `create` or `existing` |

 `AdministratorAccess` covers all of the above.
--- a/app.py
+++ b/app.py
@@ -3,22 +3,24 @@ import os
 import aws_cdk as cdk

 from src.commands.utils import load_config
-from src.infra.stack import QaiStack
+from src.infra.stack import QCStack

 app = cdk.App()

 config_path = app.node.try_get_context("config") or "config.yaml"
-stack_name = app.node.try_get_context("stack_name") or "QaiCliStack"
 account_id = app.node.try_get_context("account_id") or os.getenv("CDK_DEFAULT_ACCOUNT")
 delete_bucket_data = str(app.node.try_get_context("delete_bucket_data") or "false").lower() == "true"

 cfg = load_config(config_path)
+stack_name = app.node.try_get_context("stack_name") or cfg.infra.stack_name
+bootstrap_qualifier = app.node.try_get_context("bootstrap_qualifier") or cfg.infra.effective_bootstrap_qualifier

-QaiStack(
+QCStack(
    app,
    stack_name,
    config=cfg,
    delete_bucket_data=delete_bucket_data,
+    synthesizer=cdk.DefaultStackSynthesizer(qualifier=bootstrap_qualifier),
    env=cdk.Environment(
        account=account_id,
        region=cfg.aws.region,
--- a/examples/ai-hub/README.md
+++ b/examples/ai-hub/README.md
@@ -0,0 +1,79 @@
+# 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
+examples/training/run_training.sh --wait
+```
+
+The `config.yaml` file must include AI Hub settings:
+
+```yaml
+aihub:
+  device:
+    name: Samsung Galaxy S25 (Family)
+  target_runtime: tflite
+  input_specs:
+    input: [[1, 3, 160, 160], float32]
+  output_dir: build/qai-hub
+```
+
+Finally, the user needs to authenticate with Qualcomm AI Hub using:
+
+```bash
+qai-hub configure --api_token
+```
+
+## 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.
+
+To generate calibration and validation inputs:
+
+```bash
+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
+
+## Upload Model to Qualcomm Workbench
+
+The model can be uploaded to Qualcomm Workbench using:
+
+```bash
+qc-cli ai-hub upload examples/training/data/aihub_calibration examples/training/data/inputs.npz
+```
+
+The first argument is the calibration path for the model and the second argument is the input file, both of which were created by the `prepare_inputs.py` script. For more details, add `--help` after the `upload` command.
+
+The `upload` command runs the following commands in order:
+1. `qc-cli ai-hub quantize`
+2. `qc-cli ai-hub compile`
+3. `qc-cli ai-hub validate`
+4. `qc-cli ai-hub profile`
+
+Finally the user can download the model from AI Workbench using the command
+```bash
+qc-cli ai-hub download
+```
--- a/examples/ai-hub/prepare_inputs.py
+++ b/examples/ai-hub/prepare_inputs.py
@@ -0,0 +1,74 @@
+#!/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/training/README.md
+++ b/examples/training/README.md
@@ -0,0 +1,89 @@
+# 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
@@ -0,0 +1,40 @@
+#!/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
@@ -0,0 +1,112 @@
+#!/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
@@ -0,0 +1 @@
+onnx==1.21.0
--- a/examples/training/source/train.py
+++ b/examples/training/source/train.py
@@ -0,0 +1,188 @@
+#!/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,15 +5,19 @@ 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",
    "typer==0.25.0",
    "boto3>=1.34,<1.42",
    "constructs>=10.0.0",
+    "mlflow>=3.0",
+    "numpy>=1.26",
    "pydantic>=2.13.3",
    "pyyaml>=6.0.3",
+    "qai-hub>=0.49.0",
+    "sagemaker-mlflow>=0.4.0",
 ]

 [project.scripts]
--- a/src/aws/cloudformation.py
+++ b/src/aws/cloudformation.py
@@ -3,13 +3,11 @@ from typing import Any
 import boto3
 from botocore.exceptions import ClientError

-from src.infra.provisioning import STACK_NAME

-
-def stack_status(region: str, profile: str) -> dict[str, Any] | None:
+def stack_status(region: str, profile: str, stack_name: str) -> dict[str, Any] | None:
    client = boto3.Session(profile_name=profile, region_name=region).client("cloudformation")
    try:
-        stack = client.describe_stacks(StackName=STACK_NAME)["Stacks"][0]
+        stack = client.describe_stacks(StackName=stack_name)["Stacks"][0]
    except ClientError as e:
        message = e.response.get("Error", {}).get("Message", "")
        if "does not exist" in message:
--- a/src/aws/iam.py
+++ b/src/aws/iam.py
@@ -0,0 +1,17 @@
+import boto3
+from botocore.exceptions import ClientError
+from mypy_boto3_iam import IAMClient
+
+
+def _client(profile: str) -> IAMClient:
+    return boto3.Session(profile_name=profile).client("iam")
+
+
+def get_role_arn(profile: str, role_name: str) -> str | None:
+    client = _client(profile)
+    try:
+        return client.get_role(RoleName=role_name)["Role"]["Arn"]
+    except ClientError as e:
+        if e.response.get("Error", {}).get("Code") == "NoSuchEntity":
+            return None
+        raise
--- a/src/aws/mlflow.py
+++ b/src/aws/mlflow.py
@@ -17,3 +17,20 @@ def describe_tracking_server(region: str, profile: str, name: str) -> dict[str,
        ):
            return None
        raise
+
+
+def get_tracking_server_arn(region: str, profile: str, name: str) -> str:
+    server = describe_tracking_server(region, profile, name)
+    if not server:
+        raise ValueError(f"MLflow tracking server not found: {name}")
+
+    arn = server.get("TrackingServerArn")
+    if not arn:
+        raise ValueError(f"MLflow tracking server has no ARN: {name}")
+    return str(arn)
+
+
+def create_presigned_tracking_server_url(region: str, profile: str, name: str) -> str:
+    client = boto3.Session(profile_name=profile, region_name=region).client("sagemaker")
+    response = client.create_presigned_mlflow_tracking_server_url(TrackingServerName=name)
+    return str(response["AuthorizedUrl"])
--- a/src/aws/s3.py
+++ b/src/aws/s3.py
@@ -0,0 +1,69 @@
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+import boto3
+from mypy_boto3_s3 import S3Client
+
+
+def _client(region: str, profile: str) -> S3Client:
+    return boto3.Session(profile_name=profile, region_name=region).client("s3")
+
+
+def upload_file(
+    region: str,
+    profile: str,
+    bucket: str,
+    local_path: str,
+    s3_key: str,
+) -> str:
+    _client(region, profile).upload_file(local_path, bucket, s3_key)
+    return f"s3://{bucket}/{s3_key}"
+
+
+def download_file(
+    region: str,
+    profile: str,
+    s3_uri: str,
+    local_path: str,
+) -> str:
+    if not s3_uri.startswith("s3://"):
+        raise ValueError(f"Expected S3 URI, got: {s3_uri}")
+    bucket_key = s3_uri.removeprefix("s3://")
+    bucket, _, key = bucket_key.partition("/")
+    if not bucket or not key:
+        raise ValueError(f"Expected S3 URI with bucket and key, got: {s3_uri}")
+    dest = Path(local_path)
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    _client(region, profile).download_file(bucket, key, str(dest))
+    return str(dest)
+
+
+def upload_dir(
+    region: str,
+    profile: str,
+    bucket: str,
+    local_dir: str,
+    s3_prefix: str,
+    on_progress: Callable[[], None] | None = None,
+) -> int:
+    root = Path(local_dir)
+    files = [file for file in root.rglob("*") if file.is_file()]
+    if not files:
+        return 0
+
+    client = _client(region, profile)
+    prefix = s3_prefix.rstrip("/")
+
+    def upload_one(file_path: Path) -> None:
+        key = f"{prefix}/{file_path.relative_to(root)}"
+        client.upload_file(str(file_path), bucket, key)
+        if on_progress:
+            on_progress()
+
+    with ThreadPoolExecutor(max_workers=10) as pool:
+        futures = [pool.submit(upload_one, file) for file in files]
+        for future in as_completed(futures):
+            future.result()
+
+    return len(files)
--- a/src/aws/sagemaker.py
+++ b/src/aws/sagemaker.py
@@ -0,0 +1,143 @@
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+import boto3
+from mypy_boto3_sagemaker import SageMakerClient
+from mypy_boto3_sagemaker.literals import TrainingInstanceTypeType
+from mypy_boto3_sagemaker.type_defs import (
+    CreateTrainingJobRequestTypeDef,
+    ResourceConfigTypeDef,
+    TrainingJobSummaryTypeDef,
+)
+
+from src.config import Boto3SessionKwargs
+
+
+@dataclass(frozen=True)
+class TrainingJobRequest:
+    role_arn: str
+    image_uri: str
+    instance_type: TrainingInstanceTypeType
+    instance_count: int
+    s3_train_uri: str
+    s3_output_path: str
+    job_name: str
+    hyperparameters: dict[str, Any] = field(default_factory=dict)
+    entry_point: str | None = None
+    source_dir: str | None = None
+
+
+@dataclass(frozen=True)
+class TrainingJobStatus:
+    name: str
+    status: str
+    created: datetime | None
+    modified: datetime | None
+    model_artifacts: str | None
+    failure_reason: str | None
+    raw: dict[str, Any] = field(default_factory=dict)
+
+
+def _sm(session: Boto3SessionKwargs) -> SageMakerClient:
+    return boto3.Session(**session).client("sagemaker")
+
+
+def _upload_source_dir(
+    session: Boto3SessionKwargs,
+    source_dir: str,
+    s3_output_path: str,
+    job_name: str,
+) -> str:
+    import io
+    import tarfile
+
+    buf = io.BytesIO()
+    with tarfile.open(fileobj=buf, mode="w:gz") as tar:
+        tar.add(source_dir, arcname=".")
+    buf.seek(0)
+
+    without_scheme = s3_output_path.removeprefix("s3://")
+    bucket, _, prefix = without_scheme.partition("/")
+    key = f"{prefix.rstrip('/')}/{job_name}/source/sourcedir.tar.gz".lstrip("/")
+
+    boto3.Session(**session).client("s3").upload_fileobj(buf, bucket, key)
+    return f"s3://{bucket}/{key}"
+
+
+def start_training_job(session: Boto3SessionKwargs, job: TrainingJobRequest) -> str:
+    hp = {k: str(v) for k, v in job.hyperparameters.items()}
+
+    if job.source_dir:
+        s3_code_uri = _upload_source_dir(
+            session,
+            job.source_dir,
+            job.s3_output_path,
+            job.job_name,
+        )
+        hp["sagemaker_program"] = job.entry_point or "train.py"
+        hp["sagemaker_submit_directory"] = s3_code_uri
+
+    resource_config: ResourceConfigTypeDef = {
+        "InstanceType": job.instance_type,
+        "InstanceCount": job.instance_count,
+        "VolumeSizeInGB": 30,
+    }
+    request: CreateTrainingJobRequestTypeDef = {
+        "TrainingJobName": job.job_name,
+        "AlgorithmSpecification": {"TrainingImage": job.image_uri, "TrainingInputMode": "File"},
+        "RoleArn": job.role_arn,
+        "InputDataConfig": [
+            {
+                "ChannelName": "train",
+                "DataSource": {
+                    "S3DataSource": {
+                        "S3DataType": "S3Prefix",
+                        "S3Uri": job.s3_train_uri,
+                        "S3DataDistributionType": "FullyReplicated",
+                    }
+                },
+            }
+        ],
+        "OutputDataConfig": {"S3OutputPath": job.s3_output_path},
+        "ResourceConfig": resource_config,
+        "StoppingCondition": {"MaxRuntimeInSeconds": 86400},
+        "HyperParameters": hp,
+    }
+    _sm(session).create_training_job(**request)
+    return job.job_name
+
+
+def get_training_job_status(session: Boto3SessionKwargs, job_name: str) -> TrainingJobStatus:
+    resp = _sm(session).describe_training_job(TrainingJobName=job_name)
+    return TrainingJobStatus(
+        name=resp["TrainingJobName"],
+        status=resp["TrainingJobStatus"],
+        created=resp.get("CreationTime"),
+        modified=resp.get("LastModifiedTime"),
+        model_artifacts=resp.get("ModelArtifacts", {}).get("S3ModelArtifacts"),
+        failure_reason=resp.get("FailureReason"),
+        raw=dict(resp),
+    )
+
+
+def get_model_artifacts(region: str, profile: str, job_name: str) -> str:
+    resp = boto3.Session(profile_name=profile, region_name=region).client("sagemaker").describe_training_job(
+        TrainingJobName=job_name
+    )
+    artifact = resp.get("ModelArtifacts", {}).get("S3ModelArtifacts")
+    if not artifact:
+        raise RuntimeError(f"Training job '{job_name}' does not have model artifacts yet.")
+    return str(artifact)
+
+
+def list_training_jobs(
+    session: Boto3SessionKwargs,
+    max_results: int = 10,
+) -> list[TrainingJobSummaryTypeDef]:
+    resp = _sm(session).list_training_jobs(
+        SortBy="CreationTime",
+        SortOrder="Descending",
+        MaxResults=max_results,
+    )
+    return list(resp["TrainingJobSummaries"])
--- a/src/cloud/init.py
+++ b/src/cloud/init.py
--- a/src/commands/ai_hub.py
+++ b/src/commands/ai_hub.py
@@ -0,0 +1,406 @@
+from collections.abc import Mapping, Sequence
+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
+
+app = typer.Typer(help="Quantize, compile, validate, profile, and download models with Qualcomm Workbench")
+
+_RUNTIME_EXTENSIONS = {
+    "tflite": "tflite",
+    "qnn_context_binary": "bin",
+    "onnx": "onnx",
+}
+
+
+class UploadStep(StrEnum):
+    quantize = "quantize"
+    compile = "compile"
+    validate = "validate"
+    profile = "profile"
+
+
+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:
+        CONSOLE.print("[red]aihub.input_specs must define at least one input.[/red]")
+        raise typer.Exit(1)
+    return specs
+
+
+def _load_inputs(
+    input_file: Path,
+    specs: Mapping[str, tuple[Sequence[int], str]],
+    input_name: str | None = None,
+) -> dict[str, Any]:
+    import numpy as np
+
+    if not input_file.exists():
+        raise FileNotFoundError(f"File not found: {input_file}")
+
+    if input_file.suffix == ".npz":
+        loaded = np.load(input_file)
+        missing = set(specs) - set(loaded.files)
+        if missing:
+            raise ValueError(f"Missing input(s) in NPZ: {', '.join(sorted(missing))}")
+        return {name: loaded[name] for name in specs}
+
+    if input_file.suffix == ".npy":
+        if input_name is None:
+            if len(specs) != 1:
+                raise ValueError("--input-name is required when config has multiple inputs")
+            input_name = next(iter(specs))
+        if input_name not in specs:
+            raise ValueError(f"Input name '{input_name}' is not defined in aihub.input_specs")
+        return {input_name: np.load(input_file)}
+
+    raise ValueError("Input file must be .npz or .npy")
+
+
+def _load_calibration(path: Path, specs: Mapping[str, tuple[Sequence[int], str]]) -> dict[str, Any]:
+    import numpy as np
+
+    if path.is_file():
+        return _load_inputs(path, specs)
+
+    if not path.is_dir():
+        raise FileNotFoundError(f"Calibration path not found: {path}")
+
+    if len(specs) != 1:
+        raise ValueError("Directory calibration data is supported only for single-input models.")
+    input_name = next(iter(specs))
+    samples = [np.load(p) for p in sorted(path.glob("*.npy"))]
+    if not samples:
+        raise ValueError(f"No .npy calibration samples found in {path}")
+    return {input_name: samples}
+
+
+def _job_name(cfg: Config, operation: str) -> str | None:
+    if not cfg.aihub.job_name:
+        return None
+    return f"{cfg.aihub.job_name}-{operation}"
+
+
+def _model_id_or_state(config_path: str, model_id: str | None, *, quantized: bool = False) -> str:
+    st = state_ops.store(config_path)
+    resolved = model_id or (st.get_last_quantized_model_id() if quantized else st.get_last_compiled_model_id())
+    if not resolved:
+        source = "quantized" if quantized else "compiled"
+        CONSOLE.print(f"[red]No {source} model found. Pass --model-id or run the previous AI Hub step first.[/red]")
+        raise typer.Exit(1)
+    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,
+    calibration_path: Path,
+    from_job: str | None,
+    model_s3_uri: str | None,
+    onnx_path: str | None,
+) -> str:
+    st = state_ops.store(config_path)
+    specs = _input_specs(cfg)
+    try:
+        resolved = resolve_onnx(
+            cfg=cfg,
+            output_dir=cfg.aihub.output_dir,
+            from_job=from_job,
+            model_s3_uri=model_s3_uri or st.get_last_model_artifact(),
+            onnx_path=onnx_path,
+            last_training_job=st.get_last_training_job(),
+        )
+        calibration_data = _load_calibration(calibration_path, specs)
+    except (FileNotFoundError, ValueError) as e:
+        CONSOLE.print(f"[red]{e}[/red]")
+        raise typer.Exit(1)
+
+    try:
+        result = aihub_jobs.submit_quantize_job(
+            resolved.onnx_path,
+            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(
+        last_model_artifact=resolved.model_artifact,
+        last_quantize_job_id=result["job_id"],
+        last_quantized_model_id=result["model_id"],
+    )
+    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 _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,
+) -> str:
+    st = state_ops.store(config_path)
+    _validate_device(cfg)
+    specs = _input_specs(cfg)
+
+    model: Any
+    model_artifact: str | None = None
+    has_explicit_source = bool(from_job or model_s3_uri or onnx_path)
+    if model_id:
+        model = model_id
+    elif prefer_quantized and not has_explicit_source and st.get_last_quantized_model_id():
+        model = st.get_last_quantized_model_id()
+    else:
+        try:
+            resolved = resolve_onnx(
+                cfg=cfg,
+                output_dir=cfg.aihub.output_dir,
+                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:
+        result = aihub_jobs.submit_compile_job(
+            model=model,
+            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]")
+        raise typer.Exit(1)
+
+    updates: dict[str, Any] = {
+        "last_compile_job_id": result["job_id"],
+        "last_compiled_model_id": result["model_id"],
+    }
+    if model_artifact:
+        updates["last_model_artifact"] = 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]")
+    return str(result["model_id"])
+
+
+def _validate_step(
+    cfg: Config,
+    config_path: str,
+    input_file: Path,
+    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:
+        inputs = _load_inputs(input_file, specs, input_name)
+    except (FileNotFoundError, ValueError) as e:
+        CONSOLE.print(f"[red]{e}[/red]")
+        raise typer.Exit(1)
+
+    run = datetime.now().strftime("%Y%m%d-%H%M%S")
+    out_dir = Path(cfg.aihub.output_dir) / run / "validation"
+    try:
+        result = aihub_jobs.submit_inference_job(
+            resolved_model_id,
+            cfg.aihub.device,
+            inputs,
+            out_dir,
+            job_name=_job_name(cfg, "validate"),
+        )
+    except Exception as e:
+        CONSOLE.print(f"[red]AI Hub inference failed: {e}[/red]")
+        raise typer.Exit(1)
+
+    state_ops.store(config_path).update(last_inference_job_id=result["job_id"])
+    CONSOLE.print(f"[green]✓[/green] Inference job: [bold]{result['job_id']}[/bold]")
+    outputs = result.get("outputs")
+    if isinstance(outputs, dict):
+        for name, value in outputs.items():
+            CONSOLE.print(f"  {name}: shape={getattr(value, 'shape', '?')}")
+    CONSOLE.print(f"Outputs: [cyan]{out_dir}[/cyan]")
+    return str(result["job_id"])
+
+
+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(
+            resolved_model_id,
+            cfg.aihub.device,
+            cfg.aihub.profile_options,
+            job_name=_job_name(cfg, "profile"),
+        )
+    except Exception as e:
+        CONSOLE.print(f"[red]AI Hub profile failed: {e}[/red]")
+        raise typer.Exit(1)
+    state_ops.store(config_path).update(last_profile_job_id=result["job_id"])
+    CONSOLE.print(f"[green]✓[/green] Profile job: [bold]{result['job_id']}[/bold]")
+    return str(result["job_id"])
+
+
+@app.command()
+def quantize(
+    calibration_path: Path = typer.Argument(..., help="Calibration .npz file or directory of .npy samples"),
+    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(
+        None, "--onnx-path", help="Local ONNX path or ONNX path inside extracted artifact"
+    ),
+    config: str = CONFIG_OPT,
+) -> None:
+    """Quantize an ONNX model to INT8."""
+    cfg = load_cfg(config)
+    _quantize_step(cfg, config, calibration_path, from_job, model_s3_uri, onnx_path)
+
+
+@app.command()
+def compile(
+    model_id: str | None = typer.Option(None, "--model-id", help="AI Hub model ID to compile"),
+    from_job: str | None = typer.Option(None, "--from-job", help="Training job name whose model artifact should compile"),
+    model_s3_uri: str | None = typer.Option(None, "--model-s3-uri", help="S3 URI of model.tar.gz to compile"),
+    onnx_path: str | None = typer.Option(
+        None, "--onnx-path", help="Local ONNX path or ONNX path inside extracted artifact"
+    ),
+    config: str = CONFIG_OPT,
+) -> 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)
+
+
+@app.command()
+def validate(
+    input_file: Path = typer.Argument(..., help="NumPy .npz or .npy inputs to run on device"),
+    model_id: str | None = typer.Option(None, "--model-id", help="AI Hub compiled model ID"),
+    input_name: str | None = typer.Option(None, "--input-name", help="Input name for .npy files"),
+    config: str = CONFIG_OPT,
+) -> None:
+    """Run an AI Hub inference job using sample inputs."""
+    cfg = load_cfg(config)
+    _validate_step(cfg, config, input_file, model_id, input_name)
+
+
+@app.command()
+def profile(
+    model_id: str | None = typer.Option(None, "--model-id", help="AI Hub compiled model ID"),
+    config: str = CONFIG_OPT,
+) -> None:
+    """Profile a compiled model on the configured AI Hub device."""
+    cfg = load_cfg(config)
+    _profile_step(cfg, config, model_id)
+
+
+@app.command()
+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_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(
+        None, "--onnx-path", help="Local ONNX path or ONNX path inside extracted artifact"
+    ),
+    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."""
+    cfg = load_cfg(config)
+    steps = [UploadStep.quantize, UploadStep.compile, UploadStep.validate, UploadStep.profile]
+    selected = steps[steps.index(from_step) :]
+
+    quantized_model_id: str | None = None
+    compiled_model_id: str | None = None
+    if UploadStep.quantize in selected:
+        quantized_model_id = _quantize_step(cfg, config, calibration_path, from_job, model_s3_uri, onnx_path)
+    if UploadStep.compile in selected:
+        compiled_model_id = _compile_step(
+            cfg,
+            config,
+            model_id=quantized_model_id,
+            from_job=from_job,
+            model_s3_uri=model_s3_uri,
+            onnx_path=onnx_path,
+            prefer_quantized=True,
+        )
+    if UploadStep.validate in selected:
+        _validate_step(cfg, config, input_file, compiled_model_id, input_name)
+    if UploadStep.profile in selected:
+        _profile_step(cfg, config, compiled_model_id)
+
+
+@app.command()
+def download(
+    model_id: str | None = typer.Option(None, "--model-id", help="AI Hub compiled model ID"),
+    output: Path | None = typer.Option(None, "--output", "-o", help="Destination file path"),
+    config: str = CONFIG_OPT,
+) -> None:
+    """Download the last compiled deployable artifact from AI Hub."""
+    cfg = load_cfg(config)
+    resolved_model_id = _model_id_or_state(config, model_id)
+    ext = _RUNTIME_EXTENSIONS.get(cfg.aihub.target_runtime, cfg.aihub.target_runtime)
+    dest = output or (Path(cfg.aihub.output_dir) / f"model.{ext}")
+
+    try:
+        written = aihub_jobs.download_model(resolved_model_id, dest)
+    except Exception as e:
+        CONSOLE.print(f"[red]AI Hub download failed: {e}[/red]")
+        raise typer.Exit(1)
+    state_ops.store(config).update(last_downloaded_model=written)
+    CONSOLE.print(f"[green]✓[/green] Downloaded model: [cyan]{written}[/cyan]")
--- a/src/commands/infra.py
+++ b/src/commands/infra.py
@@ -51,6 +51,8 @@ def setup(
                    profile=cfg.aws.profile,
                    account_id=account_id,
                    region=cfg.aws.region,
+                    bootstrap_qualifier=cfg.infra.effective_bootstrap_qualifier,
+                    toolkit_stack_name=cfg.infra.effective_toolkit_stack_name,
                    cloudformation_execution_policy=cloudformation_execution_policy,
                )
        with CONSOLE.status("Running cdk deploy..."):
@@ -58,6 +60,9 @@ def setup(
                profile=cfg.aws.profile,
                account_id=account_id,
                region=cfg.aws.region,
+                stack_name=cfg.infra.stack_name,
+                bootstrap_qualifier=cfg.infra.effective_bootstrap_qualifier,
+                toolkit_stack_name=cfg.infra.effective_toolkit_stack_name,
                config_path=config,
                config_dir=str(Path(config).parent),
                config_snapshot=cfg.model_dump(mode="json"),
@@ -72,7 +77,8 @@ def setup(
    if outputs.get("SageMakerRoleArn"):
        CONSOLE.print(f"[green]✓[/green] IAM role:  {outputs['SageMakerRoleArn']}")
    if cfg.mlflow.mode is MlflowMode.create and outputs.get("MlflowTrackingServerArn"):
-        CONSOLE.print(f"[green]✓[/green] MLflow:    {outputs['MlflowTrackingServerArn']}")
+        mlflow_name = outputs.get("MlflowTrackingServerName", cfg.managed_mlflow_tracking_server_name)
+        CONSOLE.print(f"[green]✓[/green] MLflow:    {mlflow_name}")
    elif cfg.mlflow.mode is MlflowMode.existing:
        CONSOLE.print(f"[green]✓[/green] MLflow:    {cfg.mlflow.tracking_server_name}")
    CONSOLE.print("\n[bold green]Infrastructure ready.[/bold green]")
@@ -82,7 +88,7 @@ def setup(
 def status(config: str = CONFIG_OPT) -> None:
    """Show current infrastructure status."""
    cfg = load_cfg(config)
-    stack = cloudformation.stack_status(cfg.aws.region, cfg.aws.profile)
+    stack = cloudformation.stack_status(cfg.aws.region, cfg.aws.profile, cfg.infra.stack_name)

    table = Table(title="Infrastructure Status")
    table.add_column("Resource", style="cyan")
@@ -91,13 +97,13 @@ def status(config: str = CONFIG_OPT) -> None:
    table.add_column("ARN / URI")

    if not stack:
-        table.add_row("CDK Stack", provisioning.STACK_NAME, "[red]missing[/red]", "-")
+        table.add_row("CDK Stack", cfg.infra.stack_name, "[red]missing[/red]", "-")
        table.add_row("S3 Bucket", cfg.s3.bucket, "[red]unknown[/red]", "-")
        table.add_row("IAM Role", cfg.sagemaker.role_name, "[red]unknown[/red]", "-")
        if cfg.mlflow.mode is not MlflowMode.disabled:
            table.add_row(
                "MLflow",
-                cfg.mlflow.tracking_server_name or "-",
+                cfg.effective_mlflow_tracking_server_name or "-",
                "[red]unknown[/red]",
                "-",
            )
@@ -114,14 +120,14 @@ def status(config: str = CONFIG_OPT) -> None:
    )
    table.add_row(
        "IAM Role",
-        cfg.sagemaker.role_name,
+        _role_name(cfg.sagemaker.role_name, outputs.get("SageMakerRoleArn", "")),
        "[green]managed[/green]",
        outputs.get("SageMakerRoleArn", "-"),
    )
    if cfg.mlflow.mode is MlflowMode.create:
        table.add_row(
            "MLflow",
-            cfg.mlflow.tracking_server_name or "-",
+            outputs.get("MlflowTrackingServerName", cfg.managed_mlflow_tracking_server_name),
            "[green]managed[/green]",
            outputs.get("MlflowTrackingServerArn", outputs.get("MlflowArtifactUri", "-")),
        )
@@ -156,10 +162,13 @@ def destroy(
 ) -> None:
    """Destroy the CDK stack."""
    cfg = _destroy_config(config)
+    stack_name = _destroy_stack_name(config, cfg)
+    bootstrap_qualifier = _destroy_bootstrap_qualifier(config, cfg)
+    toolkit_stack_name = _destroy_toolkit_stack_name(config, cfg)

    if not yes and not delete_bucket_data:
        typer.confirm(
-            f"Destroy CDK stack '{provisioning.STACK_NAME}' while retaining S3 bucket data?",
+            f"Destroy CDK stack '{stack_name}' while retaining S3 bucket data?",
            abort=True,
        )

@@ -172,13 +181,17 @@ def destroy(
                provisioning.destroy(
                    profile=cfg.aws.profile,
                    account_id=account_id,
+                    stack_name=stack_name,
+                    bootstrap_qualifier=bootstrap_qualifier,
+                    toolkit_stack_name=toolkit_stack_name,
                    config_path=str(snapshot_path),
                    delete_bucket_data=delete_bucket_data,
                )
    except RuntimeError as e:
        CONSOLE.print(f"[red]{e}[/red]")
        raise typer.Exit(1)
-    CONSOLE.print(f"[green]✓[/green] Destroyed stack: {provisioning.STACK_NAME}")
+    CONSOLE.print(f"[green]✓[/green] Destroyed stack: {stack_name}")
+    CONSOLE.print(f"[yellow]CDK bootstrap stack retained: {toolkit_stack_name}[/yellow]")


 def _destroy_config(config_path: str) -> Config:
@@ -190,6 +203,14 @@ def _destroy_config(config_path: str) -> Config:
    return load_cfg(config_path)


+def _role_name(configured_name: str, role_arn: str) -> str:
+    if configured_name:
+        return configured_name
+    if role_arn:
+        return role_arn.rsplit("/", 1)[-1]
+    return "-"
+
+
 def _destroy_account_id(config_path: str, cfg: Config) -> str:
    config_dir = str(Path(config_path).parent)
    state = read_infra_state(config_dir)
@@ -197,3 +218,30 @@ def _destroy_account_id(config_path: str, cfg: Config) -> str:
    if account_id:
        return str(account_id)
    return identity.account_id(cfg.aws.region, cfg.aws.profile)
+
+
+def _destroy_stack_name(config_path: str, cfg: Config) -> str:
+    config_dir = str(Path(config_path).parent)
+    state = read_infra_state(config_dir)
+    stack_name = state.get("stack_name")
+    if stack_name:
+        return str(stack_name)
+    return cfg.infra.stack_name
+
+
+def _destroy_bootstrap_qualifier(config_path: str, cfg: Config) -> str:
+    config_dir = str(Path(config_path).parent)
+    state = read_infra_state(config_dir)
+    bootstrap_qualifier = state.get("bootstrap_qualifier")
+    if bootstrap_qualifier:
+        return str(bootstrap_qualifier)
+    return cfg.infra.effective_bootstrap_qualifier
+
+
+def _destroy_toolkit_stack_name(config_path: str, cfg: Config) -> str:
+    config_dir = str(Path(config_path).parent)
+    state = read_infra_state(config_dir)
+    toolkit_stack_name = state.get("toolkit_stack_name")
+    if toolkit_stack_name:
+        return str(toolkit_stack_name)
+    return cfg.infra.effective_toolkit_stack_name
--- 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
@@ -0,0 +1,184 @@
+from datetime import datetime
+from pathlib import Path
+
+import typer
+from rich.table import Table
+
+from src import state as state_ops
+from src.aws import iam
+from src.aws import sagemaker as sm_ops
+from src.commands.utils import CONFIG_OPT, CONSOLE, load_cfg
+from src.config import Config, MlflowMode
+from src.infra.state import read_infra_state
+from src.tracking.mlflow import MlflowTracker
+
+app = typer.Typer(help="Manage SageMaker training jobs")
+
+_STATUS_COLOR = {
+    "Completed": "green",
+    "Failed": "red",
+    "InProgress": "yellow",
+    "Stopping": "yellow",
+    "Stopped": "dim",
+}
+
+
+def _tracker(cfg):
+    try:
+        return MlflowTracker.from_config(cfg)
+    except Exception as e:
+        CONSOLE.print(f"[red]MLflow setup failed: {e}[/red]")
+        raise typer.Exit(1)
+
+
+def _config_dir(config_path: str) -> str:
+    return str(Path(config_path).parent)
+
+
+def _sagemaker_role_arn(config_path: str, cfg: Config) -> str:
+    state = read_infra_state(_config_dir(config_path))
+    role_arn = state.get("outputs", {}).get("SageMakerRoleArn")
+    if role_arn:
+        return str(role_arn)
+    if cfg.sagemaker.role_name:
+        role_arn = iam.get_role_arn(cfg.aws.profile, cfg.sagemaker.role_name)
+        if role_arn:
+            return role_arn
+        raise RuntimeError(f"IAM role '{cfg.sagemaker.role_name}' not found. Run 'qc-cli infra setup' first.")
+    raise RuntimeError("SageMaker role not found in infra state. Run 'qc-cli infra setup' first.")
+
+
+@app.command()
+def start(config: str = CONFIG_OPT) -> None:
+    """Submit a SageMaker training job."""
+    cfg = load_cfg(config)
+
+    if not cfg.sagemaker.training.image_uri:
+        CONSOLE.print("[red]sagemaker.training.image_uri is required in config.yaml.[/red]")
+        CONSOLE.print(
+            "Find pre-built images at: "
+            "https://aws.github.io/deep-learning-containers/reference/available_images"
+        )
+        raise typer.Exit(1)
+
+    try:
+        role_arn = _sagemaker_role_arn(config, cfg)
+    except RuntimeError as e:
+        CONSOLE.print(f"[red]{e}[/red]")
+        raise typer.Exit(1)
+
+    tracker = _tracker(cfg)
+    job_name = f"qc-cli-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+    s3_train_uri = f"s3://{cfg.s3.bucket}/{cfg.s3.data_prefix}"
+    s3_output = f"s3://{cfg.s3.bucket}/{cfg.s3.model_prefix}"
+
+    CONSOLE.print(f"Submitting training job [cyan]{job_name}[/cyan]...")
+    training_job = sm_ops.TrainingJobRequest(
+        role_arn=role_arn,
+        image_uri=cfg.sagemaker.training.image_uri,
+        instance_type=cfg.sagemaker.training.instance_type,
+        instance_count=cfg.sagemaker.training.instance_count,
+        s3_train_uri=s3_train_uri,
+        s3_output_path=s3_output,
+        job_name=job_name,
+        hyperparameters=cfg.sagemaker.training.hyperparameters,
+        entry_point=cfg.sagemaker.training.entry_point,
+        source_dir=cfg.sagemaker.training.source_dir,
+    )
+    sm_ops.start_training_job(cfg.aws.boto3_session, training_job)
+
+    st = state_ops.store(config)
+    st.set_last_training_job(job_name)
+    run_id = tracker.start_training_run(
+        training_job,
+        region=cfg.aws.region,
+        profile=cfg.aws.profile,
+        role_arn=role_arn,
+    )
+    if run_id:
+        st.update_training_job(job_name, mlflow_run_id=run_id)
+
+    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("Track progress: [cyan]qc-cli train status[/cyan]")
+
+
+@app.command()
+def status(
+    job_name: str | None = typer.Argument(None, help="Training job name (default: last submitted job)"),
+    config: str = CONFIG_OPT,
+) -> None:
+    """Show training job status."""
+    cfg = load_cfg(config)
+    st = state_ops.store(config)
+
+    if not job_name:
+        job_name = st.get_last_training_job()
+        if not job_name:
+            CONSOLE.print(
+                "[red]No training job found in state. Pass a job name or run 'qc-cli train start' first.[/red]"
+            )
+            raise typer.Exit(1)
+
+    status = sm_ops.get_training_job_status(cfg.aws.boto3_session, job_name)
+    color = _STATUS_COLOR.get(status.status, "white")
+
+    CONSOLE.print(f"Job:     [cyan]{status.name}[/cyan]")
+    CONSOLE.print(f"Status:  [{color}]{status.status}[/{color}]")
+    if status.created:
+        CONSOLE.print(f"Created: {status.created}")
+    if status.model_artifacts:
+        CONSOLE.print(f"Artifacts: {status.model_artifacts}")
+    if status.failure_reason:
+        CONSOLE.print(f"[red]Failure: {status.failure_reason}[/red]")
+
+    job_state = st.get_training_job(job_name)
+    run_id = job_state.get("mlflow_run_id")
+    already_registered = job_state.get("registered_model_version")
+    if run_id and not already_registered and status.status in {"Completed", "Failed", "Stopped"}:
+        tracker = _tracker(cfg)
+        version = tracker.finalize_training_run(
+            run_id=str(run_id),
+            training_job_status=status,
+        )
+        updates = {"mlflow_finalized_status": status.status}
+        if version:
+            updates["registered_model_version"] = version
+        st.update_training_job(job_name, **updates)
+        if version:
+            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]")
+
+
+@app.command(name="list")
+def list_jobs(
+    limit: int = typer.Option(10, "--limit", "-n", help="Number of jobs to show"),
+    config: str = CONFIG_OPT,
+) -> None:
+    """List recent training jobs."""
+    cfg = load_cfg(config)
+    jobs = sm_ops.list_training_jobs(cfg.aws.boto3_session, max_results=limit)
+
+    if not jobs:
+        CONSOLE.print("[yellow]No training jobs found.[/yellow]")
+        return
+
+    table = Table(title="Training Jobs")
+    table.add_column("Name", style="cyan")
+    table.add_column("Status")
+    table.add_column("Created")
+
+    for job in jobs:
+        status_value = str(job["TrainingJobStatus"])
+        color = _STATUS_COLOR.get(status_value, "white")
+        table.add_row(
+            str(job["TrainingJobName"]),
+            f"[{color}]{status_value}[/{color}]",
+            str(job.get("CreationTime", "")),
+        )
+
+    CONSOLE.print(table)
--- 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
@@ -1,26 +1,65 @@
-from enum import Enum
-from typing import Literal
+import re
+from enum import StrEnum
+from typing import Any, Literal, TypedDict

 from mypy_boto3_s3.literals import BucketLocationConstraintType
-from pydantic import BaseModel, Field, model_validator
+from mypy_boto3_sagemaker.literals import TrainingInstanceTypeType
+from pydantic import BaseModel, Field, field_validator, model_validator
+from qai_hub.client import Device


-class MlflowMode(str, Enum):
+class MlflowMode(StrEnum):
    disabled = "disabled"
    create = "create"
    existing = "existing"


-class MlflowServerSize(str, Enum):
+class MlflowServerSize(StrEnum):
    small = "Small"
    medium = "Medium"
    large = "Large"


+class Boto3SessionKwargs(TypedDict):
+    profile_name: str
+    region_name: str
+
+
 class AwsConfig(BaseModel):
    region: BucketLocationConstraintType | Literal["us-east-1"] = "us-east-1"
    profile: str = "default"

+    @property
+    def boto3_session(self) -> Boto3SessionKwargs:
+        return {"profile_name": self.profile, "region_name": self.region}
+
+
+DEFAULT_BOOTSTRAP_QUALIFIER = "hnb659fds"
+GENERATED_STACK_PREFIX = "qc-cli-mlops-"
+
+
+class InfraConfig(BaseModel):
+    stack_name: str
+
+    @property
+    def effective_bootstrap_qualifier(self) -> str:
+        sanitized = re.sub(r"[^a-z0-9]", "", self.stack_name.lower())
+        if not sanitized:
+            return DEFAULT_BOOTSTRAP_QUALIFIER
+        if self.stack_name.startswith(GENERATED_STACK_PREFIX):
+            suffix = re.sub(r"[^a-z0-9]", "", self.stack_name.removeprefix(GENERATED_STACK_PREFIX).lower())
+            if suffix:
+                return f"q{suffix}"[:10]
+        return f"q{sanitized}"[:10]
+
+    @property
+    def effective_toolkit_stack_name(self) -> str:
+        if self.stack_name.startswith(GENERATED_STACK_PREFIX):
+            suffix = re.sub(r"[^A-Za-z0-9-]", "", self.stack_name.removeprefix(GENERATED_STACK_PREFIX))
+            if suffix:
+                return f"{self.stack_name}-bootstrap"
+        return f"{self.stack_name}-bootstrap"
+

 class S3Config(BaseModel):
    bucket: str = "my-qc-mlops-bucket"
@@ -28,13 +67,45 @@ class S3Config(BaseModel):
    model_prefix: str = "models/"


+class TrainingConfig(BaseModel):
+    instance_type: TrainingInstanceTypeType = "ml.m5.xlarge"
+    instance_count: int = 1
+    image_uri: str = ""
+    entry_point: str | None = None
+    source_dir: str | None = None
+    hyperparameters: dict[str, Any] = Field(default_factory=dict)
+
+
 class SageMakerConfig(BaseModel):
-    role_name: str = "qc-cli-sagemaker-role"
+    role_name: str = ""
+    training: TrainingConfig = Field(default_factory=TrainingConfig)
+
+
+class AIHubConfig(BaseModel):
+    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
+    model_name: str | None = None
+    compile_options: str | None = None
+    profile_options: str | None = None
+    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
    tracking_server_name: str | None = None
+    experiment_name: str = "qc-cli-training"
+    registered_model_name: str = "qc-cli-model"
+    register_trained_models: bool = True
    artifact_prefix: str = "mlflow/"
    tracking_server_size: MlflowServerSize = MlflowServerSize.small
    mlflow_version: str | None = None
@@ -43,13 +114,27 @@ class MlflowConfig(BaseModel):

    @model_validator(mode="after")
    def require_tracking_server_name(self) -> "MlflowConfig":
-        if self.mode in {MlflowMode.create, MlflowMode.existing} and not self.tracking_server_name:
-            raise ValueError("mlflow.tracking_server_name is required when mlflow.mode is create or existing")
+        if self.mode is MlflowMode.existing and not self.tracking_server_name:
+            raise ValueError("mlflow.tracking_server_name is required when mlflow.mode is existing")
        return self


 class Config(BaseModel):
+    infra: InfraConfig
    aws: AwsConfig = Field(default_factory=AwsConfig)
    s3: S3Config = Field(default_factory=S3Config)
    sagemaker: SageMakerConfig = Field(default_factory=SageMakerConfig)
+    aihub: AIHubConfig = Field(default_factory=AIHubConfig)
    mlflow: MlflowConfig = Field(default_factory=MlflowConfig)
+
+    @property
+    def managed_mlflow_tracking_server_name(self) -> str:
+        return f"{self.infra.stack_name}-mlflow"
+
+    @property
+    def effective_mlflow_tracking_server_name(self) -> str | None:
+        if self.mlflow.mode is MlflowMode.disabled:
+            return None
+        if self.mlflow.mode is MlflowMode.existing:
+            return self.mlflow.tracking_server_name
+        return self.managed_mlflow_tracking_server_name
--- a/src/infra/provisioning.py
+++ b/src/infra/provisioning.py
@@ -5,17 +5,27 @@ from typing import Any

 from src.infra.state import state_path, write_infra_state

-STACK_NAME = "QaiCliStack"
-

 def bootstrap(
    *,
    profile: str,
    account_id: str,
    region: str,
+    bootstrap_qualifier: str,
+    toolkit_stack_name: str,
    cloudformation_execution_policy: str | None = None,
 ) -> None:
-    cmd = ["cdk", "bootstrap", f"aws://{account_id}/{region}", "--profile", profile]
+    cmd = [
+        "cdk",
+        "bootstrap",
+        f"aws://{account_id}/{region}",
+        "--profile",
+        profile,
+        "--qualifier",
+        bootstrap_qualifier,
+        "--toolkit-stack-name",
+        toolkit_stack_name,
+    ]
    if cloudformation_execution_policy:
        cmd.extend(["--cloudformation-execution-policies", cloudformation_execution_policy])
    _run(cmd)
@@ -26,6 +36,9 @@ def deploy(
    profile: str,
    account_id: str,
    region: str,
+    stack_name: str,
+    bootstrap_qualifier: str,
+    toolkit_stack_name: str,
    config_path: str,
    config_dir: str,
    config_snapshot: dict[str, Any],
@@ -35,19 +48,24 @@ def deploy(
        "deploy",
        profile=profile,
        account_id=account_id,
+        stack_name=stack_name,
+        bootstrap_qualifier=bootstrap_qualifier,
+        toolkit_stack_name=toolkit_stack_name,
        config_path=config_path,
        delete_bucket_data=False,
    ) + ["--require-approval", "never", "--outputs-file", str(outputs_file)]
    _run(cmd)

-    outputs = _read_outputs(outputs_file)
+    outputs = _read_outputs(outputs_file, stack_name)
    state = {
-        "stack_name": STACK_NAME,
+        "stack_name": stack_name,
        "aws": {
            "account_id": account_id,
            "region": region,
            "profile": profile,
        },
+        "bootstrap_qualifier": bootstrap_qualifier,
+        "toolkit_stack_name": toolkit_stack_name,
        "config": config_snapshot,
        "outputs": outputs,
    }
@@ -59,6 +77,9 @@ def destroy(
    *,
    profile: str,
    account_id: str,
+    stack_name: str,
+    bootstrap_qualifier: str,
+    toolkit_stack_name: str,
    config_path: str,
    delete_bucket_data: bool,
 ) -> None:
@@ -67,6 +88,9 @@ def destroy(
            "deploy",
            profile=profile,
            account_id=account_id,
+            stack_name=stack_name,
+            bootstrap_qualifier=bootstrap_qualifier,
+            toolkit_stack_name=toolkit_stack_name,
            config_path=config_path,
            delete_bucket_data=True,
        ) + ["--require-approval", "never"]
@@ -76,6 +100,9 @@ def destroy(
        "destroy",
        profile=profile,
        account_id=account_id,
+        stack_name=stack_name,
+        bootstrap_qualifier=bootstrap_qualifier,
+        toolkit_stack_name=toolkit_stack_name,
        config_path=config_path,
        delete_bucket_data=delete_bucket_data,
    ) + ["--force"]
@@ -87,26 +114,35 @@ def _cdk_cmd(
    *,
    profile: str,
    account_id: str,
+    stack_name: str,
+    bootstrap_qualifier: str,
+    toolkit_stack_name: str,
    config_path: str,
    delete_bucket_data: bool,
 ) -> list[str]:
    cmd = [
        "cdk",
        action,
-        STACK_NAME,
+        stack_name,
        "--app",
        "python app.py",
        "--profile",
        profile,
+    ]
+    if action == "deploy":
+        cmd.extend(["--toolkit-stack-name", toolkit_stack_name])
+    cmd.extend([
        "-c",
        f"account_id={account_id}",
        "-c",
        f"config={config_path}",
        "-c",
-        f"stack_name={STACK_NAME}",
+        f"stack_name={stack_name}",
+        "-c",
+        f"bootstrap_qualifier={bootstrap_qualifier}",
        "-c",
        f"delete_bucket_data={str(delete_bucket_data).lower()}",
-    ]
+    ])
    return cmd


@@ -119,9 +155,9 @@ def _run(cmd: list[str]) -> None:
        raise RuntimeError(f"CDK command failed with exit code {e.returncode}.") from e


-def _read_outputs(path: Path) -> dict[str, str]:
+def _read_outputs(path: Path, stack_name: str) -> dict[str, str]:
    if not path.exists():
        return {}
    with open(path) as f:
        data = json.load(f)
-    return data.get(STACK_NAME, {})
+    return data.get(stack_name, {})
--- a/src/infra/stack.py
+++ b/src/infra/stack.py
@@ -9,7 +9,7 @@ from constructs import Construct
 from src.config import Config, MlflowMode


-class QaiStack(Stack):
+class QCStack(Stack):
    def __init__(
        self,
        scope: Construct,
@@ -34,7 +34,7 @@ class QaiStack(Stack):
        role = iam.CfnRole(
            self,
            "SageMakerRole",
-            role_name=config.sagemaker.role_name,
+            role_name=config.sagemaker.role_name or None,
            assume_role_policy_document=self._sagemaker_trust_policy(),
            managed_policy_arns=[
                f"arn:{self.partition}:iam::aws:policy/AmazonSageMakerFullAccess",
@@ -74,6 +74,7 @@ class QaiStack(Stack):
        CfnOutput(self, "SageMakerRoleArn", value=role.attr_arn)

        if config.mlflow.mode is MlflowMode.create:
+            tracking_server_name = config.managed_mlflow_tracking_server_name
            artifact_prefix = config.mlflow.artifact_prefix.strip("/")
            artifact_uri = (
                f"s3://{data_bucket.bucket_name}/{artifact_prefix}/"
@@ -145,14 +146,14 @@ class QaiStack(Stack):
                "MlflowTrackingServer",
                artifact_store_uri=artifact_uri,
                role_arn=mlflow_role.attr_arn,
-                tracking_server_name=config.mlflow.tracking_server_name or "",
+                tracking_server_name=tracking_server_name,
                automatic_model_registration=config.mlflow.automatic_model_registration,
                mlflow_version=config.mlflow.mlflow_version,
                tracking_server_size=config.mlflow.tracking_server_size.value,
                weekly_maintenance_window_start=config.mlflow.weekly_maintenance_window_start,
            )

-            CfnOutput(self, "MlflowTrackingServerName", value=config.mlflow.tracking_server_name or "")
+            CfnOutput(self, "MlflowTrackingServerName", value=tracking_server_name)
            CfnOutput(self, "MlflowTrackingServerArn", value=tracking_server.attr_tracking_server_arn)
            CfnOutput(self, "MlflowArtifactUri", value=artifact_uri)
            CfnOutput(self, "MlflowRoleArn", value=mlflow_role.attr_arn)
--- a/src/main.py
+++ b/src/main.py
@@ -1,36 +1,14 @@
-from pathlib import Path
-
 import typer
-import yaml
-from rich.console import Console

-from src.commands import infra
-from src.config import Config
+from src.commands import ai_hub, infra, init, mlflow, train, upload

 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")
-
-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 = Config()
-    dest.parent.mkdir(parents=True, exist_ok=True)
-    with open(dest, "w") as f:
-        yaml.safe_dump(config.model_dump(mode="json"), f, sort_keys=False)
-
-    console.print(f"[green]✓[/green] Config written to [bold]{dest}[/bold]")
-    console.print("Edit it (especially [cyan]s3.bucket[/cyan]) before running other commands.")
+app.add_typer(train.app, name="train")
+app.add_typer(ai_hub.app, name="ai-hub")
--- a/src/qualcomm/init.py
+++ b/src/qualcomm/init.py
--- a/src/qualcomm/aihub_jobs.py
+++ b/src/qualcomm/aihub_jobs.py
@@ -0,0 +1,129 @@
+from pathlib import Path
+from typing import Any, TypedDict
+
+import qai_hub.hub as hub
+from qai_hub.client import CompileJob, Device, InferenceJob, Model, ProfileJob, QuantizeDtype, QuantizeJob
+
+
+class ModelJobResult(TypedDict):
+    job: CompileJob | QuantizeJob
+    job_id: str
+    model: Model
+    model_id: str
+
+
+class InferenceJobResult(TypedDict):
+    job: InferenceJob
+    job_id: str
+    outputs: Any
+
+
+class ProfileJobResult(TypedDict):
+    job: ProfileJob
+    job_id: str
+
+
+def _dataset_entries(inputs: dict[str, Any]) -> dict[str, list[Any]]:
+    return {name: value if isinstance(value, list) else [value] for name, value in inputs.items()}
+
+
+def submit_compile_job(
+    model: Any,
+    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,
+        device=device,
+        name=job_name,
+        input_specs=input_specs,
+        options=compile_options,
+    )
+    target_model = job.get_target_model()
+    if target_model is None:
+        raise RuntimeError(f"Compile job {job.job_id} did not produce a target model.")
+    return {"job": job, "job_id": str(job.job_id), "model": target_model, "model_id": str(target_model.model_id)}
+
+
+def submit_inference_job(
+    model_id: 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),
+        device=device,
+        inputs=_dataset_entries(inputs),
+        name=job_name,
+    )
+    out = Path(output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    data = job.download_output_data(str(out))
+    return {"job": job, "job_id": str(job.job_id), "outputs": data}
+
+
+def submit_profile_job(
+    model_id: str,
+    device: Device,
+    options: str | None = None,
+    job_name: str | None = None,
+) -> ProfileJobResult:
+    job = hub.submit_profile_job(
+        model=hub.get_model(model_id),
+        device=device,
+        name=job_name,
+        options=options or "",
+    )
+    return {"job": job, "job_id": str(job.job_id)}
+
+
+def submit_quantize_job(
+    model: str | Path,
+    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,
+        calibration_data=_dataset_entries(calibration_data),
+        weights_dtype=QuantizeDtype.INT8,
+        activations_dtype=QuantizeDtype.INT8,
+        name=job_name,
+        options=options or "",
+    )
+    target_model = job.get_target_model()
+    if target_model is None:
+        raise RuntimeError(f"Quantize job {job.job_id} did not produce a target model.")
+    return {"job": job, "job_id": str(job.job_id), "model": target_model, "model_id": str(target_model.model_id)}
+
+
+def download_model(model_id: str, output_path: str | Path) -> str:
+    dest = Path(output_path)
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    model = hub.get_model(model_id)
+    result = model.download(str(dest))
+    return str(result or dest)
--- a/src/qualcomm/artifacts.py
+++ b/src/qualcomm/artifacts.py
@@ -0,0 +1,83 @@
+import tarfile
+from dataclasses import dataclass
+from pathlib import Path
+
+from src.aws import s3 as s3_ops
+from src.aws import sagemaker as sm_ops
+from src.config import Config
+
+
+@dataclass(frozen=True)
+class ResolvedOnnx:
+    onnx_path: Path
+    model_artifact: str | None
+    run_name: str
+
+
+def _safe_extract(tar: tarfile.TarFile, dest: Path) -> None:
+    dest_root = dest.resolve()
+    for member in tar.getmembers():
+        target = (dest / member.name).resolve()
+        if dest_root != target and dest_root not in target.parents:
+            raise ValueError(f"Unsafe tar member path: {member.name}")
+    tar.extractall(dest, filter="data")
+
+
+def _find_onnx(root: Path, explicit: str | None = None) -> Path:
+    if explicit:
+        p = Path(explicit)
+        if not p.is_absolute():
+            p = root / p
+        if not p.exists():
+            raise FileNotFoundError(f"ONNX file not found: {p}")
+        return p
+
+    matches = sorted(root.rglob("model.onnx"))
+    if not matches:
+        matches = sorted(root.rglob("*.onnx"))
+    if not matches:
+        raise FileNotFoundError(f"No ONNX file found under {root}")
+    if len(matches) > 1:
+        joined = ", ".join(str(p.relative_to(root)) for p in matches)
+        raise ValueError(f"Multiple ONNX files found ({joined}). Pass --onnx-path.")
+    return matches[0]
+
+
+def resolve_onnx(
+    cfg: Config,
+    output_dir: str,
+    from_job: str | None = None,
+    model_s3_uri: str | None = None,
+    onnx_path: str | None = None,
+    last_training_job: str | None = None,
+) -> ResolvedOnnx:
+    if onnx_path:
+        path = Path(onnx_path)
+        if path.exists():
+            return ResolvedOnnx(onnx_path=path, model_artifact=None, run_name=path.stem)
+
+    job = from_job or last_training_job
+    artifact = model_s3_uri
+    if not artifact:
+        if not job:
+            raise ValueError("No model source found. Pass --onnx-path, --model-s3-uri, --from-job, or run training first.")
+        artifact = sm_ops.get_model_artifacts(cfg.aws.region, cfg.aws.profile, job)
+
+    run_name = job or Path(artifact).name.removesuffix(".tar.gz").replace("/", "-")
+    root = Path(output_dir) / run_name / "source"
+    tar_path = root / "model.tar.gz"
+    s3_ops.download_file(cfg.aws.region, cfg.aws.profile, artifact, str(tar_path))
+
+    extract_dir = root / "extracted"
+    extract_dir.mkdir(parents=True, exist_ok=True)
+    try:
+        with tarfile.open(tar_path, "r:gz") as tar:
+            _safe_extract(tar, extract_dir)
+    except tarfile.TarError as e:
+        raise ValueError(f"Invalid model tarball: {tar_path}") from e
+
+    return ResolvedOnnx(
+        onnx_path=_find_onnx(extract_dir, onnx_path),
+        model_artifact=artifact,
+        run_name=run_name,
+    )
--- a/src/state.py
+++ b/src/state.py
@@ -0,0 +1,81 @@
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+STATE_FILE = ".qc-cli.json"
+
+
+@dataclass(frozen=True)
+class CliStateStore:
+    config_dir: str = "."
+
+    @property
+    def path(self) -> Path:
+        return Path(self.config_dir) / STATE_FILE
+
+    def read(self) -> dict[str, Any]:
+        if not self.path.exists():
+            return {}
+        with open(self.path) as f:
+            value = json.load(f)
+        return dict(value) if isinstance(value, dict) else {}
+
+    def update(self, **updates: Any) -> None:
+        state = self.read()
+        state.update(updates)
+        self._write(state)
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return self.read().get(key, default)
+
+    def get_last_training_job(self) -> str | None:
+        value = self.get("last_training_job")
+        return str(value) if value else None
+
+    def get_last_model_artifact(self) -> str | None:
+        value = self.get("last_model_artifact")
+        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
+
+    def get_last_compiled_model_id(self) -> str | None:
+        value = self.get("last_compiled_model_id")
+        return str(value) if value else None
+
+    def get_last_downloaded_model(self) -> str | None:
+        value = self.get("last_downloaded_model")
+        return str(value) if value else None
+
+    def set_last_training_job(self, job_name: str) -> None:
+        self.update(last_training_job=job_name)
+
+    def get_training_job(self, job_name: str) -> dict[str, Any]:
+        jobs = self._training_jobs(self.read())
+        value = jobs.get(job_name, {})
+        return dict(value) if isinstance(value, dict) else {}
+
+    def update_training_job(self, job_name: str, **updates: Any) -> None:
+        state = self.read()
+        jobs = self._training_jobs(state)
+        jobs[job_name] = {**jobs.get(job_name, {}), **updates}
+        state["training_jobs"] = jobs
+        self._write(state)
+
+    def set_latest_experiment_model_version(self, version: str) -> None:
+        self.update(latest_experiment_model_version=version)
+
+    def _write(self, state: dict[str, Any]) -> None:
+        with open(self.path, "w") as f:
+            json.dump(state, f, indent=2)
+
+    def _training_jobs(self, state: dict[str, Any]) -> dict[str, Any]:
+        value = state.get("training_jobs", {})
+        return dict(value) if isinstance(value, dict) else {}
+
+
+def store(config_path: str) -> CliStateStore:
+    config_dir = str(Path(config_path).parent)
+    return CliStateStore(config_dir)
--- a/src/tracking/init.py
+++ b/src/tracking/init.py
@@ -0,0 +1,3 @@
+from src.tracking.mlflow import MlflowTracker, NoopTracker, Tracker
+
+__all__ = ["MlflowTracker", "NoopTracker", "Tracker"]
--- a/src/tracking/mlflow.py
+++ b/src/tracking/mlflow.py
@@ -0,0 +1,153 @@
+import os
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+import mlflow
+from mlflow.tracking import MlflowClient
+
+from src.aws import mlflow as aws_mlflow
+from src.config import Config, MlflowMode
+
+
+class Tracker(Protocol):
+    def start_training_run(self, training_job: Any, *, region: str, profile: str, role_arn: str) -> str | None: ...
+
+    def finalize_training_run(self, *, run_id: str | None, training_job_status: Any) -> str | None: ...
+
+
+@dataclass(frozen=True)
+class NoopTracker:
+    def start_training_run(self, training_job: Any, *, region: str, profile: str, role_arn: str) -> str | None:
+        return None
+
+    def finalize_training_run(self, *, run_id: str | None, training_job_status: Any) -> str | None:
+        return None
+
+
+@dataclass(frozen=True)
+class MlflowTracker:
+    tracking_uri: str
+    experiment_name: str
+    registered_model_name: str
+    register_trained_models: bool
+
+    @classmethod
+    def from_config(cls, cfg: Config) -> Tracker:
+        if cfg.mlflow.mode is MlflowMode.disabled:
+            return NoopTracker()
+
+        os.environ.setdefault("MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT", "true")
+
+        tracking_server_name = cfg.effective_mlflow_tracking_server_name
+        if not tracking_server_name:
+            raise RuntimeError("MLflow tracking server name could not be resolved.")
+
+        tracking_uri = aws_mlflow.get_tracking_server_arn(
+            cfg.aws.region,
+            cfg.aws.profile,
+            tracking_server_name,
+        )
+        mlflow.set_tracking_uri(tracking_uri)
+        mlflow.set_experiment(cfg.mlflow.experiment_name)
+
+        return cls(
+            tracking_uri=tracking_uri,
+            experiment_name=cfg.mlflow.experiment_name,
+            registered_model_name=cfg.mlflow.registered_model_name,
+            register_trained_models=cfg.mlflow.register_trained_models,
+        )
+
+    def start_training_run(self, training_job: Any, *, region: str, profile: str, role_arn: str) -> str | None:
+        run = mlflow.start_run(run_name=training_job.job_name)
+        run_id = str(run.info.run_id)
+
+        params = {
+            "aws.region": region,
+            "aws.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,
+        }
+        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": "sagemaker",
+                "qc_cli.command": "train start",
+                "sagemaker.job_name": training_job.job_name,
+            }
+        )
+        mlflow.end_run()
+        return run_id
+
+    def finalize_training_run(self, *, run_id: str | None, training_job_status: Any) -> str | None:
+        if not run_id:
+            return None
+
+        with mlflow.start_run(run_id=run_id):
+            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")
+
+            if training_job_status.status != "Completed" or not training_job_status.model_artifacts:
+                mlflow.set_tag("qc_cli.training_terminal_status", training_job_status.status)
+                return None
+
+            if not self.register_trained_models:
+                return None
+
+            client = MlflowClient()
+            self._ensure_registered_model(client, self.registered_model_name)
+            version = client.create_model_version(
+                name=self.registered_model_name,
+                source=training_job_status.model_artifacts,
+                run_id=run_id,
+                tags={
+                    "qc_cli.stage": "experiment",
+                    "qc_cli.artifact_kind": "trained_source",
+                    "qc_cli.source": "sagemaker",
+                    "sagemaker.job_name": training_job_status.name,
+                },
+            )
+            version_number = str(version.version)
+            client.set_registered_model_alias(self.registered_model_name, "experiment-latest", version_number)
+            mlflow.set_tag("qc_cli.registered_model_name", self.registered_model_name)
+            mlflow.set_tag("qc_cli.registered_model_version", version_number)
+            return version_number
+
+    def _log_params(self, params: dict[str, Any]) -> None:
+        cleaned = {key: str(value) for key, value in params.items() if value is not None}
+        if cleaned:
+            mlflow.log_params(cleaned)
+
+    def _log_final_metrics(self, training_job: dict[str, Any]) -> None:
+        metrics = {}
+        for metric in training_job.get("FinalMetricDataList", []):
+            name = metric.get("MetricName")
+            value = metric.get("Value")
+            if name and value is not None:
+                metrics[str(name)] = float(value)
+        if metrics:
+            mlflow.log_metrics(metrics)
+
+    def _ensure_registered_model(self, client: MlflowClient, name: str) -> None:
+        try:
+            client.get_registered_model(name)
+        except Exception:
+            client.create_registered_model(name)
--- a/uv.lock
+++ b/uv.lock
Author	SHA1	Message	Date
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
slalom	e9ada2612f	Mlflow implementation (#2 ) Reviewed-on: #2	2026-06-02 19:04:23 +00:00
slalom	6ac9702dc5	make sure resources are set up in isolated namespaces (#1 ) Reviewed-on: #1	2026-05-27 12:51:26 +00:00
slalom	0e728cc193	command to start sagemaker training include sample training	2026-05-25 16:48:31 -04:00
slalom	62ffe163e8	enable s3 upload	2026-05-20 16:42:07 -04:00
samirodr	cfc04b473f	update	2026-05-20 15:21:45 -04:00