Compare commits
7 Commits
2f77190ea7
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5360a482fc | ||
|
|
6a560a8610 | ||
| d244150d98 | |||
| d7c7158464 | |||
| 6bc25dc183 | |||
|
|
71a95aa3a7 | ||
| a3f3060e13 |
25
README.md
25
README.md
@@ -67,7 +67,8 @@ sagemaker:
|
||||
hyperparameters: {}
|
||||
|
||||
aihub:
|
||||
device: Samsung Galaxy S25 (Family)
|
||||
device:
|
||||
name: Samsung Galaxy S25 (Family)
|
||||
target_runtime: tflite
|
||||
input_specs: {} # Required before running qc-cli ai-hub commands
|
||||
job_name: null # Optional prefix for AI Hub Workbench jobs
|
||||
@@ -109,10 +110,10 @@ When MLflow is enabled, `train start` creates an MLflow run for the SageMaker jo
|
||||
To open the managed SageMaker MLflow UI, request a fresh presigned URL:
|
||||
|
||||
```bash
|
||||
qc-cli infra mlflow-url --config config.yaml
|
||||
qc-cli mlflow open --config config.yaml
|
||||
```
|
||||
|
||||
This works for `mode: create` and for `mode: existing` when the existing server is managed by Amazon SageMaker. In `create` mode, the command uses the CLI-managed tracking server name. In `existing` mode, it uses `mlflow.tracking_server_name`. If the existing MLflow server is external to SageMaker, open it with that server's own URL instead.
|
||||
This opens a browser to a fresh presigned URL. It works for `mode: create` and for `mode: existing` when the existing server is managed by Amazon SageMaker. In `create` mode, the command uses the CLI-managed tracking server name. In `existing` mode, it uses `mlflow.tracking_server_name`. If the existing MLflow server is external to SageMaker, open it with that server's own URL instead.
|
||||
|
||||
## Commands
|
||||
|
||||
@@ -124,6 +125,12 @@ qc-cli init --output <path> Write config to a custom path
|
||||
qc-cli init --force Overwrite an existing config file
|
||||
```
|
||||
|
||||
### `mlflow`
|
||||
|
||||
```
|
||||
qc-cli mlflow open Open a presigned MLflow UI URL in a browser
|
||||
```
|
||||
|
||||
### `infra`
|
||||
|
||||
```
|
||||
@@ -131,7 +138,6 @@ qc-cli infra setup Deploy the CDK stack
|
||||
qc-cli infra setup --no-bootstrap Deploy without running CDK bootstrap
|
||||
qc-cli infra setup --cloudformation-execution-policy <arn> Set CDK bootstrap execution policy ARN
|
||||
qc-cli infra status Show CDK stack/resource status
|
||||
qc-cli infra mlflow-url Print a presigned MLflow UI URL
|
||||
qc-cli infra destroy Destroy stack, retaining S3 data
|
||||
qc-cli infra destroy --yes Destroy stack without confirmation
|
||||
qc-cli infra destroy --delete-bucket-data Destroy stack and delete S3 data
|
||||
@@ -180,6 +186,17 @@ 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.
|
||||
|
||||
@@ -13,38 +13,35 @@ This example takes the ONNX model produced by the SageMaker training example and
|
||||
Run the training example first and wait for it to complete:
|
||||
|
||||
```bash
|
||||
bash examples/training/run_training.sh --config config.yaml --wait
|
||||
examples/training/run_training.sh --wait
|
||||
```
|
||||
|
||||
If the dataset is already uploaded to S3, use:
|
||||
|
||||
```bash
|
||||
bash examples/training/run_training.sh --config config.yaml --skip-upload --wait
|
||||
```
|
||||
|
||||
The training artifact must contain a static-shape `model.onnx`. The training example exports an input named `input` with shape `1x3x160x160`.
|
||||
|
||||
Your `config.yaml` must include AI Hub settings:
|
||||
The `config.yaml` file must include AI Hub settings:
|
||||
|
||||
```yaml
|
||||
aihub:
|
||||
device: Samsung Galaxy S25 (Family)
|
||||
device:
|
||||
name: Samsung Galaxy S25 (Family)
|
||||
target_runtime: tflite
|
||||
input_specs:
|
||||
input: [[1, 3, 160, 160], float32]
|
||||
output_dir: build/qai-hub
|
||||
```
|
||||
|
||||
You also need local Qualcomm AI Hub SDK authentication configured.
|
||||
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.
|
||||
|
||||
Generate calibration and validation inputs:
|
||||
To generate calibration and validation inputs:
|
||||
|
||||
```bash
|
||||
uv run python examples/ai-hub/prepare_inputs.py
|
||||
python examples/ai-hub/prepare_inputs.py
|
||||
```
|
||||
|
||||
This writes:
|
||||
@@ -60,58 +57,23 @@ The script applies the same image preprocessing used by the training example:
|
||||
- convert to channel-first `1x3x160x160`
|
||||
- normalize with ImageNet mean and standard deviation
|
||||
|
||||
Useful options:
|
||||
## Upload Model to Qualcomm Workbench
|
||||
|
||||
The model can be uploaded to Qualcomm Workbench using:
|
||||
|
||||
```bash
|
||||
uv run python examples/ai-hub/prepare_inputs.py \
|
||||
--dataset-dir examples/training/data/flower_photos_sagemaker \
|
||||
--calibration-dir examples/training/data/aihub_calibration \
|
||||
--input-file examples/training/data/inputs.npz \
|
||||
--samples 16
|
||||
qc-cli ai-hub upload examples/training/data/aihub_calibration examples/training/data/inputs.npz
|
||||
```
|
||||
|
||||
## Run AI Hub
|
||||
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.
|
||||
|
||||
After training completes and inputs are prepared:
|
||||
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
|
||||
bash examples/ai-hub/run_ai_hub.sh --config config.yaml
|
||||
qc-cli ai-hub download
|
||||
```
|
||||
|
||||
By default, the script uses the last SageMaker training job recorded in `.qc-cli.json`. It downloads that job's `model.tar.gz`, extracts `model.onnx`, runs the AI Hub workflow, and downloads the compiled artifact.
|
||||
|
||||
To use a specific training job:
|
||||
|
||||
```bash
|
||||
bash examples/ai-hub/run_ai_hub.sh \
|
||||
--config config.yaml \
|
||||
--from-job qc-cli-YYYYMMDD-HHMMSS
|
||||
```
|
||||
|
||||
To resume from a later Workbench step:
|
||||
|
||||
```bash
|
||||
bash examples/ai-hub/run_ai_hub.sh \
|
||||
--config config.yaml \
|
||||
--from-step validate
|
||||
```
|
||||
|
||||
To skip downloading the compiled artifact:
|
||||
|
||||
```bash
|
||||
bash examples/ai-hub/run_ai_hub.sh \
|
||||
--config config.yaml \
|
||||
--skip-download
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
If AI Hub reports dynamic input shapes, rerun training with the current training source. AI Hub quantization requires the exported ONNX model to use static input shapes.
|
||||
|
||||
If `run_ai_hub.sh` reports missing calibration or input files, run:
|
||||
|
||||
```bash
|
||||
uv run python examples/ai-hub/prepare_inputs.py
|
||||
```
|
||||
|
||||
If validation fails with a missing input name, make sure `config.yaml` and the generated `.npz` both use `input` as the input name.
|
||||
|
||||
0
examples/ai-hub/prepare_inputs.py
Executable file → Normal file
0
examples/ai-hub/prepare_inputs.py
Executable file → Normal file
@@ -1,156 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
CONFIG_PATH="config.yaml"
|
||||
CALIBRATION_PATH="examples/training/data/aihub_calibration"
|
||||
INPUT_FILE="examples/training/data/inputs.npz"
|
||||
FROM_STEP="quantize"
|
||||
FROM_JOB=""
|
||||
MODEL_S3_URI=""
|
||||
ONNX_PATH=""
|
||||
INPUT_NAME=""
|
||||
DOWNLOAD=true
|
||||
OUTPUT_PATH=""
|
||||
|
||||
usage() {
|
||||
cat <<EOF
|
||||
Usage: $0 [options]
|
||||
|
||||
Options:
|
||||
--config PATH Path to qc-cli config file. Default: config.yaml
|
||||
--calibration PATH Calibration .npz file or directory of .npy samples.
|
||||
Default: ${CALIBRATION_PATH}
|
||||
--input-file PATH Validation .npz or .npy inputs. Default: ${INPUT_FILE}
|
||||
--from-step STEP Resume upload from: quantize, compile, validate, profile.
|
||||
Default: ${FROM_STEP}
|
||||
--from-job NAME SageMaker training job whose model artifact should upload.
|
||||
Defaults to the last training job in local qc-cli state.
|
||||
--model-s3-uri URI S3 URI of model.tar.gz to upload.
|
||||
--onnx-path PATH Local ONNX path or ONNX path inside extracted artifact.
|
||||
--input-name NAME Input name for .npy validation files.
|
||||
--skip-download Do not download the compiled AI Hub artifact after upload.
|
||||
--output PATH Destination file for ai-hub download.
|
||||
-h, --help Show this help.
|
||||
EOF
|
||||
}
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--config)
|
||||
CONFIG_PATH="$2"
|
||||
shift 2
|
||||
;;
|
||||
--calibration)
|
||||
CALIBRATION_PATH="$2"
|
||||
shift 2
|
||||
;;
|
||||
--input-file)
|
||||
INPUT_FILE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--from-step)
|
||||
FROM_STEP="$2"
|
||||
shift 2
|
||||
;;
|
||||
--from-job)
|
||||
FROM_JOB="$2"
|
||||
shift 2
|
||||
;;
|
||||
--model-s3-uri)
|
||||
MODEL_S3_URI="$2"
|
||||
shift 2
|
||||
;;
|
||||
--onnx-path)
|
||||
ONNX_PATH="$2"
|
||||
shift 2
|
||||
;;
|
||||
--input-name)
|
||||
INPUT_NAME="$2"
|
||||
shift 2
|
||||
;;
|
||||
--skip-download)
|
||||
DOWNLOAD=false
|
||||
shift
|
||||
;;
|
||||
--output)
|
||||
OUTPUT_PATH="$2"
|
||||
shift 2
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1" >&2
|
||||
usage >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
if [[ ! -f "${CONFIG_PATH}" ]]; then
|
||||
echo "Config not found: ${CONFIG_PATH}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
case "${FROM_STEP}" in
|
||||
quantize|compile|validate|profile)
|
||||
;;
|
||||
*)
|
||||
echo "--from-step must be one of: quantize, compile, validate, profile" >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
|
||||
if [[ ! -e "${CALIBRATION_PATH}" ]]; then
|
||||
echo "Calibration path not found: ${CALIBRATION_PATH}" >&2
|
||||
echo "Pass --calibration with a .npz file or directory of .npy samples." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ ! -f "${INPUT_FILE}" ]]; then
|
||||
echo "Input file not found: ${INPUT_FILE}" >&2
|
||||
echo "Pass --input-file with a validation .npz or .npy file." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
run() {
|
||||
echo "+ $*"
|
||||
"$@"
|
||||
}
|
||||
|
||||
UPLOAD_ARGS=(
|
||||
"${CALIBRATION_PATH}"
|
||||
"${INPUT_FILE}"
|
||||
--from-step "${FROM_STEP}"
|
||||
--config "${CONFIG_PATH}"
|
||||
)
|
||||
|
||||
if [[ -n "${FROM_JOB}" ]]; then
|
||||
UPLOAD_ARGS+=(--from-job "${FROM_JOB}")
|
||||
fi
|
||||
|
||||
if [[ -n "${MODEL_S3_URI}" ]]; then
|
||||
UPLOAD_ARGS+=(--model-s3-uri "${MODEL_S3_URI}")
|
||||
fi
|
||||
|
||||
if [[ -n "${ONNX_PATH}" ]]; then
|
||||
UPLOAD_ARGS+=(--onnx-path "${ONNX_PATH}")
|
||||
fi
|
||||
|
||||
if [[ -n "${INPUT_NAME}" ]]; then
|
||||
UPLOAD_ARGS+=(--input-name "${INPUT_NAME}")
|
||||
fi
|
||||
|
||||
run uv run qc-cli ai-hub upload "${UPLOAD_ARGS[@]}"
|
||||
|
||||
if [[ "${DOWNLOAD}" == false ]]; then
|
||||
exit 0
|
||||
fi
|
||||
|
||||
DOWNLOAD_ARGS=(--config "${CONFIG_PATH}")
|
||||
if [[ -n "${OUTPUT_PATH}" ]]; then
|
||||
DOWNLOAD_ARGS+=(--output "${OUTPUT_PATH}")
|
||||
fi
|
||||
|
||||
run uv run qc-cli ai-hub download "${DOWNLOAD_ARGS[@]}"
|
||||
@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
|
||||
[project]
|
||||
name = "qc-cli"
|
||||
version = "0.1.0"
|
||||
description = "CLI for SageMaker ONNX training and Qualcomm AI Hub optimization"
|
||||
description = "CLI for training and deploying models for Qualcomm AI Hub"
|
||||
requires-python = ">=3.13"
|
||||
dependencies = [
|
||||
"aws-cdk-lib>=2.180.0",
|
||||
@@ -29,8 +29,6 @@ packages = ["src"]
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"boto3-stubs[iam,s3,sagemaker]",
|
||||
"pytest>=8.0",
|
||||
"pytest-mock>=3.12",
|
||||
"pyright>=1.1.409",
|
||||
"types-PyYAML",
|
||||
"ruff>=0.4",
|
||||
|
||||
0
src/cloud/__init__.py
Normal file
0
src/cloud/__init__.py
Normal file
@@ -4,7 +4,9 @@ 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
|
||||
@@ -12,7 +14,7 @@ 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 AI Hub")
|
||||
app = typer.Typer(help="Quantize, compile, validate, profile, and download models with Qualcomm Workbench")
|
||||
|
||||
_RUNTIME_EXTENSIONS = {
|
||||
"tflite": "tflite",
|
||||
@@ -99,6 +101,33 @@ def _model_id_or_state(config_path: str, model_id: str | None, *, quantized: boo
|
||||
return resolved
|
||||
|
||||
|
||||
def _device_selector(device: Device) -> str:
|
||||
parts: list[str] = []
|
||||
if device.name:
|
||||
parts.append(f"name={device.name!r}")
|
||||
if device.os:
|
||||
parts.append(f"os={device.os!r}")
|
||||
if device.attributes:
|
||||
parts.append(f"attributes={device.attributes!r}")
|
||||
return ", ".join(parts) if parts else "empty selector"
|
||||
|
||||
|
||||
def _validate_device(cfg: Config) -> None:
|
||||
device = cfg.aihub.device
|
||||
try:
|
||||
matches = hub.get_devices(name=device.name, os=device.os, attributes=device.attributes)
|
||||
except Exception as e:
|
||||
CONSOLE.print(f"[red]Unable to validate AI Hub device {_device_selector(device)}: {e}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
if matches:
|
||||
return
|
||||
|
||||
CONSOLE.print(f"[red]AI Hub device not found: {_device_selector(device)}[/red]")
|
||||
CONSOLE.print("Run [bold]qai-hub list-devices[/bold] to see valid device names.")
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
def _quantize_step(
|
||||
cfg: Config,
|
||||
config_path: str,
|
||||
@@ -156,6 +185,7 @@ def _compile_step(
|
||||
prefer_quantized: bool,
|
||||
) -> str:
|
||||
st = state_ops.store(config_path)
|
||||
_validate_device(cfg)
|
||||
specs = _input_specs(cfg)
|
||||
|
||||
model: Any
|
||||
@@ -184,7 +214,7 @@ def _compile_step(
|
||||
try:
|
||||
result = aihub_jobs.submit_compile_job(
|
||||
model=model,
|
||||
device_name=cfg.aihub.device,
|
||||
device=cfg.aihub.device,
|
||||
input_specs=specs,
|
||||
target_runtime=cfg.aihub.target_runtime,
|
||||
options=cfg.aihub.compile_options,
|
||||
@@ -214,6 +244,7 @@ def _validate_step(
|
||||
model_id: str | None,
|
||||
input_name: str | None,
|
||||
) -> str:
|
||||
_validate_device(cfg)
|
||||
specs = _input_specs(cfg)
|
||||
resolved_model_id = _model_id_or_state(config_path, model_id)
|
||||
try:
|
||||
@@ -247,6 +278,7 @@ def _validate_step(
|
||||
|
||||
|
||||
def _profile_step(cfg: Config, config_path: str, model_id: str | None) -> str:
|
||||
_validate_device(cfg)
|
||||
resolved_model_id = _model_id_or_state(config_path, model_id)
|
||||
try:
|
||||
result = aihub_jobs.submit_profile_job(
|
||||
|
||||
@@ -150,35 +150,6 @@ def status(config: str = CONFIG_OPT) -> None:
|
||||
CONSOLE.print(table)
|
||||
|
||||
|
||||
@app.command(name="mlflow-url")
|
||||
def mlflow_url(config: str = CONFIG_OPT) -> None:
|
||||
"""Print a presigned URL for the configured MLflow tracking server."""
|
||||
cfg = load_cfg(config)
|
||||
tracking_server_name = cfg.effective_mlflow_tracking_server_name
|
||||
if not tracking_server_name:
|
||||
CONSOLE.print("[red]MLflow is disabled in config.yaml.[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
try:
|
||||
url = mlflow.create_presigned_tracking_server_url(
|
||||
cfg.aws.region,
|
||||
cfg.aws.profile,
|
||||
tracking_server_name,
|
||||
)
|
||||
except Exception as e:
|
||||
CONSOLE.print("[yellow]Could not create a SageMaker MLflow UI URL.[/yellow]")
|
||||
CONSOLE.print(f"Tracking server: [cyan]{tracking_server_name}[/cyan]")
|
||||
CONSOLE.print(f"Reason: {e}")
|
||||
CONSOLE.print(
|
||||
"This command can create presigned URLs only for MLflow tracking servers managed by "
|
||||
"Amazon SageMaker. If this is an external MLflow server, open it with that server's own URL."
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
CONSOLE.print(f"MLflow tracking server: [cyan]{tracking_server_name}[/cyan]")
|
||||
CONSOLE.print(f"MLflow UI: {url}")
|
||||
|
||||
|
||||
@app.command()
|
||||
def destroy(
|
||||
config: str = CONFIG_OPT,
|
||||
|
||||
40
src/commands/init.py
Normal file
40
src/commands/init.py
Normal file
@@ -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
|
||||
41
src/commands/mlflow.py
Normal file
41
src/commands/mlflow.py
Normal file
@@ -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]")
|
||||
@@ -101,7 +101,7 @@ def start(config: str = CONFIG_OPT) -> None:
|
||||
CONSOLE.print(f"[green]✓[/green] Job submitted: [bold]{job_name}[/bold]")
|
||||
if run_id:
|
||||
CONSOLE.print(f"MLflow run: [cyan]{run_id}[/cyan]")
|
||||
CONSOLE.print("Open MLflow: [cyan]qc-cli infra mlflow-url[/cyan]")
|
||||
CONSOLE.print("Open MLflow: [cyan]qc-cli mlflow open[/cyan]")
|
||||
CONSOLE.print("Track progress: [cyan]qc-cli train status[/cyan]")
|
||||
|
||||
|
||||
@@ -151,7 +151,7 @@ def status(
|
||||
st.set_latest_experiment_model_version(version)
|
||||
CONSOLE.print(f"MLflow model version: [cyan]{version}[/cyan] ([cyan]experiment-latest[/cyan])")
|
||||
if run_id and cfg.mlflow.mode is not MlflowMode.disabled:
|
||||
CONSOLE.print("Open MLflow: [cyan]qc-cli infra mlflow-url[/cyan]")
|
||||
CONSOLE.print("Open MLflow: [cyan]qc-cli mlflow open[/cyan]")
|
||||
|
||||
|
||||
@app.command(name="list")
|
||||
|
||||
70
src/commands/upload.py
Normal file
70
src/commands/upload.py
Normal file
@@ -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)
|
||||
@@ -4,7 +4,8 @@ from typing import Any, Literal, TypedDict
|
||||
|
||||
from mypy_boto3_s3.literals import BucketLocationConstraintType
|
||||
from mypy_boto3_sagemaker.literals import TrainingInstanceTypeType
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
from qai_hub.client import Device
|
||||
|
||||
|
||||
class MlflowMode(StrEnum):
|
||||
@@ -81,7 +82,7 @@ class SageMakerConfig(BaseModel):
|
||||
|
||||
|
||||
class AIHubConfig(BaseModel):
|
||||
device: str = "Samsung Galaxy S25 (Family)"
|
||||
device: Device = Field(default_factory=lambda: Device("Samsung Galaxy S25 (Family)"))
|
||||
target_runtime: str = "tflite"
|
||||
input_specs: dict[str, tuple[list[int], str]] = Field(default_factory=dict)
|
||||
job_name: str | None = None
|
||||
@@ -91,6 +92,13 @@ class AIHubConfig(BaseModel):
|
||||
quantize_options: str | None = None
|
||||
output_dir: str = "build/qai-hub"
|
||||
|
||||
@field_validator("device", mode="before")
|
||||
@classmethod
|
||||
def parse_device(cls, value: Any) -> Any:
|
||||
if isinstance(value, str):
|
||||
return Device(value)
|
||||
return value
|
||||
|
||||
|
||||
class MlflowConfig(BaseModel):
|
||||
mode: MlflowMode = MlflowMode.disabled
|
||||
|
||||
109
src/main.py
109
src/main.py
@@ -1,115 +1,14 @@
|
||||
import secrets
|
||||
from pathlib import Path
|
||||
|
||||
import typer
|
||||
import yaml
|
||||
from rich.console import Console
|
||||
from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
|
||||
|
||||
from src.aws import s3 as s3_ops
|
||||
from src.commands import ai_hub, infra, train
|
||||
from src.commands.utils import CONFIG_OPT, load_cfg
|
||||
from src.config import GENERATED_STACK_PREFIX, Config, InfraConfig, S3Config
|
||||
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")
|
||||
app.add_typer(train.app, name="train")
|
||||
app.add_typer(ai_hub.app, name="ai-hub")
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
@app.command()
|
||||
def init(
|
||||
output: str = typer.Option("config.yaml", help="Destination path for the config file"),
|
||||
force: bool = typer.Option(False, "--force", "-f", help="Overwrite an existing config file"),
|
||||
) -> None:
|
||||
"""Write a starter config.yaml to the current directory."""
|
||||
dest = Path(output)
|
||||
if dest.exists() and not force:
|
||||
console.print(f"[yellow]{dest} already exists.[/yellow] Use --force to overwrite.")
|
||||
raise typer.Exit(1)
|
||||
|
||||
config = _new_isolated_config()
|
||||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
config_data = config.model_dump(mode="json")
|
||||
config_data["sagemaker"].pop("role_name", None)
|
||||
with open(dest, "w") as f:
|
||||
yaml.safe_dump(config_data, f, sort_keys=False)
|
||||
|
||||
console.print(f"[green]✓[/green] Config written to [bold]{dest}[/bold]")
|
||||
console.print(
|
||||
"Edit [cyan]sagemaker.training.image_uri[/cyan] before running training commands."
|
||||
)
|
||||
|
||||
|
||||
def _new_isolated_config() -> Config:
|
||||
suffix = secrets.token_hex(6)
|
||||
namespace = f"{GENERATED_STACK_PREFIX}{suffix}"
|
||||
config = Config(infra=InfraConfig(stack_name=namespace))
|
||||
config.s3 = S3Config(bucket=f"{namespace}-data")
|
||||
return config
|
||||
|
||||
|
||||
@app.command()
|
||||
def upload(
|
||||
path: Path = typer.Argument(..., help="Local file or directory to upload"),
|
||||
s3_key: str | None = typer.Option(None, "--s3-key", help="S3 key for file uploads"),
|
||||
config: str = CONFIG_OPT,
|
||||
) -> None:
|
||||
"""Upload a local file or directory to S3."""
|
||||
cfg = load_cfg(config)
|
||||
|
||||
if path.is_file():
|
||||
key = s3_key or f"{cfg.s3.data_prefix.rstrip('/')}/{path.name}"
|
||||
try:
|
||||
with console.status(f"Uploading {path.name}..."):
|
||||
uri = s3_ops.upload_file(cfg.aws.region, cfg.aws.profile, cfg.s3.bucket, str(path), key)
|
||||
except Exception as e:
|
||||
console.print(f"[red]Upload failed: {e}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
console.print(f"[green]✓[/green] {path.name} -> {uri}")
|
||||
return
|
||||
|
||||
if path.is_dir():
|
||||
if s3_key is not None:
|
||||
console.print("[red]--s3-key can only be used when uploading a single file.[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
files = [file for file in path.rglob("*") if file.is_file()]
|
||||
if not files:
|
||||
console.print("[yellow]No files found in directory.[/yellow]")
|
||||
raise typer.Exit(0)
|
||||
|
||||
prefix = cfg.s3.data_prefix
|
||||
console.print(f"Uploading {len(files)} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/")
|
||||
try:
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
task = progress.add_task("Uploading...", total=len(files))
|
||||
count = s3_ops.upload_dir(
|
||||
cfg.aws.region,
|
||||
cfg.aws.profile,
|
||||
cfg.s3.bucket,
|
||||
str(path),
|
||||
prefix,
|
||||
on_progress=lambda: progress.advance(task),
|
||||
)
|
||||
except Exception as e:
|
||||
console.print(f"[red]Upload failed: {e}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
console.print(f"[green]✓[/green] Uploaded {count} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/")
|
||||
return
|
||||
|
||||
console.print(f"[red]Path not found: {path}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
|
||||
@@ -29,7 +29,7 @@ def _dataset_entries(inputs: dict[str, Any]) -> dict[str, list[Any]]:
|
||||
|
||||
def submit_compile_job(
|
||||
model: Any,
|
||||
device_name: str,
|
||||
device: Device,
|
||||
input_specs: dict[str, tuple[tuple[int, ...], str]],
|
||||
target_runtime: str,
|
||||
options: str | None = None,
|
||||
@@ -52,7 +52,7 @@ def submit_compile_job(
|
||||
|
||||
job = hub.submit_compile_job(
|
||||
model=model_arg,
|
||||
device=Device(device_name),
|
||||
device=device,
|
||||
name=job_name,
|
||||
input_specs=input_specs,
|
||||
options=compile_options,
|
||||
@@ -65,14 +65,14 @@ def submit_compile_job(
|
||||
|
||||
def submit_inference_job(
|
||||
model_id: str,
|
||||
device_name: str,
|
||||
device: Device,
|
||||
inputs: dict[str, Any],
|
||||
output_dir: str | Path,
|
||||
job_name: str | None = None,
|
||||
) -> InferenceJobResult:
|
||||
job = hub.submit_inference_job(
|
||||
model=hub.get_model(model_id),
|
||||
device=Device(device_name),
|
||||
device=device,
|
||||
inputs=_dataset_entries(inputs),
|
||||
name=job_name,
|
||||
)
|
||||
@@ -84,13 +84,13 @@ def submit_inference_job(
|
||||
|
||||
def submit_profile_job(
|
||||
model_id: str,
|
||||
device_name: str,
|
||||
device: Device,
|
||||
options: str | None = None,
|
||||
job_name: str | None = None,
|
||||
) -> ProfileJobResult:
|
||||
job = hub.submit_profile_job(
|
||||
model=hub.get_model(model_id),
|
||||
device=Device(device_name),
|
||||
device=device,
|
||||
name=job_name,
|
||||
options=options or "",
|
||||
)
|
||||
|
||||
50
uv.lock
generated
50
uv.lock
generated
@@ -1003,15 +1003,6 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/8a/db/55a262f3606bebcae07cc14095338471ad7c0bbcaa37707e6f0ee49725b7/importlib_resources-7.1.0-py3-none-any.whl", hash = "sha256:1bd7b48b4088eddb2cd16382150bb515af0bd2c70128194392725f82ad2c96a1", size = 37232, upload-time = "2026-04-12T16:36:08.219Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iniconfig"
|
||||
version = "2.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itsdangerous"
|
||||
version = "2.2.0"
|
||||
@@ -1674,15 +1665,6 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ff/6e/cf826fae916b8658848d7b9f38d88da6396895c676e8086fc0988073aaf8/pillow-12.2.0-cp314-cp314t-win_arm64.whl", hash = "sha256:aa88ccfe4e32d362816319ed727a004423aab09c5cea43c01a4b435643fa34eb", size = 2556579, upload-time = "2026-04-01T14:45:52.529Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pluggy"
|
||||
version = "1.6.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "prettytable"
|
||||
version = "3.17.0"
|
||||
@@ -1963,34 +1945,6 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/16/6b/330d8ebae582b30c2959a1ef4c3bc344ebde48c2ff0c3f113c4710735e11/pyright-1.1.409-py3-none-any.whl", hash = "sha256:aa3ea228cab90c845c7a60d28db7a844c04315356392aa09fafcee98c8c22fb3", size = 6438161, upload-time = "2026-04-23T11:02:01.309Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest"
|
||||
version = "9.0.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "colorama", marker = "sys_platform == 'win32'" },
|
||||
{ name = "iniconfig" },
|
||||
{ name = "packaging" },
|
||||
{ name = "pluggy" },
|
||||
{ name = "pygments" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/7d/0d/549bd94f1a0a402dc8cf64563a117c0f3765662e2e668477624baeec44d5/pytest-9.0.3.tar.gz", hash = "sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c", size = 1572165, upload-time = "2026-04-07T17:16:18.027Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d4/24/a372aaf5c9b7208e7112038812994107bc65a84cd00e0354a88c2c77a617/pytest-9.0.3-py3-none-any.whl", hash = "sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9", size = 375249, upload-time = "2026-04-07T17:16:16.13Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest-mock"
|
||||
version = "3.15.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pytest" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/68/14/eb014d26be205d38ad5ad20d9a80f7d201472e08167f0bb4361e251084a9/pytest_mock-3.15.1.tar.gz", hash = "sha256:1849a238f6f396da19762269de72cb1814ab44416fa73a8686deac10b0d87a0f", size = 34036, upload-time = "2025-09-16T16:37:27.081Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl", hash = "sha256:0a25e2eb88fe5168d535041d09a4529a188176ae608a6d249ee65abc0949630d", size = 10095, upload-time = "2025-09-16T16:37:25.734Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "python-dateutil"
|
||||
version = "2.9.0.post0"
|
||||
@@ -2114,8 +2068,6 @@ dependencies = [
|
||||
dev = [
|
||||
{ name = "boto3-stubs", extra = ["iam", "s3", "sagemaker"] },
|
||||
{ name = "pyright" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-mock" },
|
||||
{ name = "ruff" },
|
||||
{ name = "types-pyyaml" },
|
||||
]
|
||||
@@ -2138,8 +2090,6 @@ requires-dist = [
|
||||
dev = [
|
||||
{ name = "boto3-stubs", extras = ["iam", "s3", "sagemaker"] },
|
||||
{ name = "pyright", specifier = ">=1.1.409" },
|
||||
{ name = "pytest", specifier = ">=8.0" },
|
||||
{ name = "pytest-mock", specifier = ">=3.12" },
|
||||
{ name = "ruff", specifier = ">=0.4" },
|
||||
{ name = "types-pyyaml" },
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user