374
src/commands/ai_hub.py
Normal file
374
src/commands/ai_hub.py
Normal file
@@ -0,0 +1,374 @@
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from enum import StrEnum
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import typer
|
||||
|
||||
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 AI Hub")
|
||||
|
||||
_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 _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)
|
||||
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_name=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:
|
||||
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:
|
||||
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]")
|
||||
Reference in New Issue
Block a user