trajectory_analysis_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""aquin check trajectory — multi-checkpoint weight trajectory vs base."""
 
import sys
 
from aquin.cli_flags import reject_legacy_output_flags
from pathlib import Path
from typing import Any
 
 
def _parse_flag(args: list[str], name: str) -> str | None:
    for i, a in enumerate(args):
        if a == name and i + 1 < len(args):
            return args[i + 1]
    return None
 
 
def _has_flag(args: list[str], name: str) -> bool:
    return name in args
 
 
def _ensure_compute_env() -> None:
    from aquin.compute.loader_shim import apply as _shim_apply
    from aquin.engine.local_server import start as _start_local_server
 
    _shim_apply()
    _start_local_server()
 
 
def _resolve_model_id() -> str:
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
 
    active = (get_active_model_id() or "").strip()
    if not active:
        print("Error: no model loaded. Run: aquin load --model <id>")
        sys.exit(1)
    try:
        return resolve_model_id(active)
    except ValueError as e:
        print(f"Error: {e}")
        sys.exit(1)
 
 
def _print_help() -> None:
    print("Training trajectory: weight-diff summary for each checkpoint vs base.")
    print("")
    print("Prerequisite: aquin load --model <id>")
    print("")
    print("Usage: aquin check trajectory --checkpoints <glob>")
    print("       aquin check trajectory --dir <checkpoint-dir>")
    print("       [--name <prefix>] [--save <path>]")
    print("")
    print("  --checkpoints  Glob of .pt checkpoints (sorted by training step when present).")
    print("  --dir          Scan directory recursively for *.pt files.")
    print("  --name         Optional prefix for step labels.")
    print("  --save         Write schema_version=1 JSON export.")
    print("")
    print("Example:")
    print("  aquin check trajectory --dir ~/runs/checkpoints --name my-run")
    print("")
    print("Docs: https://aquin.app/docs/checkpoint-sae")
 
 
def cmd_trajectory_analysis(args: list[str]) -> None:
    if _has_flag(args, "--help") or _has_flag(args, "-h"):
        _print_help()
        return
 
    reject_legacy_output_flags(args)
 
    pattern = _parse_flag(args, "--checkpoints")
    directory = _parse_flag(args, "--dir")
    if not pattern and not directory:
        print("Error: --checkpoints <glob> or --dir <path> is required.")
        print("")
        _print_help()
        sys.exit(1)
 
    if _parse_flag(args, "--model") is not None:
        print("Error: check trajectory uses the loaded session model only.")
        sys.exit(1)
 
    _ensure_compute_env()
 
    from aquin.cli import _build_tool_ctx
    from aquin.engine.sync_dispatch import dispatch_with_sync, require_active_session
 
    mid = _resolve_model_id()
    ctx = _build_tool_ctx(model_id=mid)
    require_active_session(ctx, label="aquin check trajectory")
 
    tool_args: dict[str, Any] = {
        "model_id": mid,
        "checkpoints": pattern,
        "dir": directory,
        "name": _parse_flag(args, "--name"),
        "save": _parse_flag(args, "--save"),
    }
    try:
        label = pattern or directory or ""
        print(f"[check trajectory] model={mid}  source={label}")
        result = dispatch_with_sync("run_trajectory_analysis", tool_args, ctx)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    from aquin.cli_output import print_tool_result
 
    print_tool_result("check trajectory", result)
 
    if isinstance(result, dict) and result.get("error"):
        sys.exit(1)