train.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.
 
"""Run a Recipe: capture DataRevision → train → Checkpoint → EvalGate."""
 
from __future__ import annotations
 
import time
import uuid
from pathlib import Path
from typing import Any
 
from aquin.recipe.load import load_recipe
from aquin.recipe.schema import Recipe, RecipeError, RunRecord
from aquin.recipe.store import capture_revision, run_dir, save_run
from aquin.run import Run
 
 
def train_recipe(
    recipe: str | Path | Recipe,
    *,
    cwd: Path | None = None,
    snapshot: bool | None = None,
    dry_run: bool = False,
) -> RunRecord:
    cwd = cwd or Path.cwd()
    rec = recipe if isinstance(recipe, Recipe) else load_recipe(recipe)
    snap = bool(rec.data.get("snapshot")) if snapshot is None else snapshot
    revision = capture_revision(rec.data["path"], snapshot=snap, cwd=cwd)
    rec.data["path"] = revision.path
 
    run_id = uuid.uuid4().hex[:12]
    dest = run_dir(run_id, cwd=cwd)
    dest.mkdir(parents=True, exist_ok=True)
 
    record = RunRecord(
        run_id=run_id,
        name=rec.name,
        family=rec.family,
        status="planned" if dry_run else "running",
        run_dir=str(dest),
        recipe=rec.to_dict(),
        revision=revision.to_dict(),
    )
    if dry_run:
        record.status = "dry_run"
        save_run(record)
        return record
 
    recorder = Run(
        base_model=str(rec.train.get("base") or rec.family),
        run_name=rec.name,
        config=rec.to_dict(),
        run_dir=dest,
    )
    t0 = time.time()
    try:
        if rec.family == "tabular":
            from aquin.recipe.tabular import train_tabular
 
            result = train_tabular(rec, recorder, dest)
        elif rec.family == "llm":
            from aquin.recipe.llm_lora import train_llm
 
            result = train_llm(rec, recorder, dest)
        else:
            raise RecipeError(f"Unsupported family {rec.family!r}")
    except Exception as exc:
        record.status = "failed"
        record.error = str(exc)
        save_run(record)
        recorder.finish()
        raise
 
    record.checkpoint = result["checkpoint"].to_dict()
    record.gate = result["gate"].to_dict()
    record.metrics = dict(result.get("metrics") or {})
    record.inspect = dict(result.get("inspect") or {})
    record.status = "passed" if result["gate"].passed else "failed_gate"
    record.metrics["elapsed_s"] = round(time.time() - t0, 3)
    save_run(record)
    recorder.finish()
    return record