run.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
from __future__ import annotations
 
import json
import os
import threading
import time
import uuid
from pathlib import Path
from typing import Any
 
 
def init(
    base_model: str | None = None,
    run_name: str | None = None,
    config: dict[str, Any] | None = None,
) -> "Run":
    """Start recording a training run. Call this before your training loop."""
    return Run(base_model=base_model, run_name=run_name, config=config)
 
 
class Run:
    """
    Records metrics and checkpoints locally during training.
    Call aquin.init() to create one.
 
    Writes under ``./aquin_run/`` (metrics, config, checkpoint).
 
    Example
    -------
    run = aquin.init(
        base_model="meta-llama/Llama-3.2-1B-Instruct",
        run_name="my-run",
        config={"lr": 2e-4, "epochs": 3, "rank": 16, "lora_alpha": 32,
                "method": "qlora", "per_device_train_batch_size": 2,
                "gradient_accumulation_steps": 8, "dataset": "data.jsonl"},
    )
    for step, batch in enumerate(dataloader):
        loss = train_step(batch)
        run.log(step, loss=loss.item(), learning_rate=scheduler.get_last_lr()[0],
                grad_norm=grad_norm, epoch=epoch)
    run.checkpoint(model, step=step)
    run.finish()
    """
 
    def __init__(
        self,
        base_model: str | None = None,
        run_name: str | None = None,
        config: dict[str, Any] | None = None,
        *,
        quant: str = "none",
        mode: str = "external",
        run_dir: str | Path | None = None,
    ) -> None:
        _ = quant, mode  # reserved for future recipe / runtime tags
        self._run_id = str(uuid.uuid4())
        self._run_name = run_name or f"run-{self._run_id[:8]}"
        self._base_model = base_model or os.environ.get("AQUIN_BASE_MODEL") or ""
        self._started_at = time.time()
 
        self._run_dir = Path(run_dir) if run_dir is not None else Path.cwd() / "aquin_run"
        self._run_dir.mkdir(parents=True, exist_ok=True)
        (self._run_dir / "checkpoints").mkdir(exist_ok=True)
 
        self._metrics: list[dict[str, Any]] = []
        self._obs: list[dict[str, Any]] = []
        self._signals: list[dict[str, Any]] = []
        self._metrics_lock = threading.Lock()
 
        self._write_meta()
        if config:
            self._write_config(config)
 
        print(f"[aquin] Recording run '{self._run_name}' (id: {self._run_id})")
        print(f"[aquin] Run data: {self._run_dir}")
 
    def log(
        self,
        step: int,
        *,
        loss: float,
        learning_rate: float | None = None,
        grad_norm: float | None = None,
        momentum_norm: float | None = None,
        epoch: int | None = None,
        batch: int | None = None,
        total_batches: int | None = None,
        step_ms: float | None = None,
        **extra: float,
    ) -> None:
        """Record metrics for one training step. Call every step inside your loop."""
        entry: dict[str, Any] = {"step": step, "metrics": {"loss": loss}}
        if learning_rate is not None:
            entry["metrics"]["learning_rate"] = learning_rate
        if grad_norm is not None:
            entry["metrics"]["grad_norm"] = grad_norm
        if momentum_norm is not None:
            entry["metrics"]["momentum_norm"] = momentum_norm
        if step_ms is not None:
            entry["metrics"]["step_ms"] = step_ms
        if extra:
            entry["metrics"].update(extra)
        if epoch is not None:
            entry["epoch"] = epoch
        if batch is not None:
            entry["batch"] = batch
        if total_batches is not None:
            entry["total_batches"] = total_batches
 
        with self._metrics_lock:
            self._metrics.append(entry)
        self._write_metrics()
 
    def observe(self, step: int, channel: str, value: float) -> None:
        """Emit a custom scalar observation (eval metrics, etc.)."""
        row = {"step": step, "channel": channel, "value": float(value), "ts": time.time()}
        with self._metrics_lock:
            self._obs.append(row)
        self._write_obs()
 
    def signal(self, step: int, message: str, *, severity: str = "info") -> None:
        """Emit a non-scalar training signal (checkpoint saved, OOM, etc.)."""
        row = {
            "step": step,
            "message": message,
            "severity": severity,
            "ts": time.time(),
        }
        with self._metrics_lock:
            self._signals.append(row)
        self._write_signals()
 
    def checkpoint(self, model: Any, step: int) -> None:
        """Save the final model checkpoint. One per run — replaces the previous save."""
        try:
            import torch
        except ImportError:
            raise ImportError("torch is required to save checkpoints.")
 
        ckpt_path = self._run_dir / "checkpoints" / "checkpoint.pt"
        print(f"[aquin] Saving checkpoint (step {step})...")
        torch.save({"step": step, "state_dict": model.state_dict()}, ckpt_path)
        self._write_metrics()
        self.signal(step, f"checkpoint saved (step {step})", severity="info")
        print(f"[aquin] Checkpoint saved (step {step}).")
 
    def finish(self, config: dict[str, Any] | None = None) -> None:
        """Finalise the run. Flushes metrics and optional config to disk."""
        self._write_metrics()
        if config:
            self._write_config(config)
        elapsed = round(time.time() - self._started_at)
        print(f"[aquin] Run finished. {len(self._metrics)} steps recorded in {elapsed}s.")
        print(f"[aquin] Data: {self._run_dir}")
 
    @property
    def run_id(self) -> str:
        return self._run_id
 
    def _write_meta(self) -> None:
        meta = {
            "run_id": self._run_id,
            "run_name": self._run_name,
            "base_model": self._base_model,
            "started_at": self._started_at,
        }
        with open(self._run_dir / "meta.json", "w") as f:
            json.dump(meta, f, indent=2)
 
    def _write_metrics(self) -> None:
        with self._metrics_lock:
            data = list(self._metrics)
        with open(self._run_dir / "metrics.json", "w") as f:
            json.dump(data, f)
 
    def _write_obs(self) -> None:
        with self._metrics_lock:
            data = list(self._obs)
        with open(self._run_dir / "observations.json", "w") as f:
            json.dump(data, f)
 
    def _write_signals(self) -> None:
        with self._metrics_lock:
            data = list(self._signals)
        with open(self._run_dir / "signals.json", "w") as f:
            json.dump(data, f)
 
    def _write_config(self, config: dict[str, Any]) -> None:
        with open(self._run_dir / "config.json", "w") as f:
            json.dump(config, f, indent=2)