store.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.
 
"""On-disk layout under ./aquin_run/ — revisions, runs, checkpoints, gates."""
 
from __future__ import annotations
 
import hashlib
import json
import shutil
from pathlib import Path
from typing import Any
 
from aquin.recipe.schema import DataRevision, RecipeError, RunRecord
 
RUN_ROOT_NAME = "aquin_run"
 
 
def run_root(cwd: Path | None = None) -> Path:
    return (cwd or Path.cwd()) / RUN_ROOT_NAME
 
 
def revision_dir(rev_id: str, *, cwd: Path | None = None) -> Path:
    return run_root(cwd) / "revisions" / rev_id
 
 
def run_dir(run_id: str, *, cwd: Path | None = None) -> Path:
    return run_root(cwd) / "runs" / run_id
 
 
def hash_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as fh:
        while True:
            chunk = fh.read(1024 * 1024)
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()
 
 
def _count_rows(path: Path, fmt: str) -> int | None:
    try:
        if fmt == "csv":
            with path.open(encoding="utf-8", errors="replace") as fh:
                n = sum(1 for _ in fh)
            return max(0, n - 1)
        if fmt == "jsonl":
            with path.open(encoding="utf-8", errors="replace") as fh:
                return sum(1 for line in fh if line.strip())
    except OSError:
        return None
    return None
 
 
def _csv_columns(path: Path) -> list[str] | None:
    try:
        with path.open(encoding="utf-8", errors="replace") as fh:
            header = fh.readline()
        return [c.strip() for c in header.strip().split(",") if c.strip()]
    except OSError:
        return None
 
 
def capture_revision(
    path: str | Path,
    *,
    snapshot: bool = False,
    cwd: Path | None = None,
) -> DataRevision:
    src = Path(path).expanduser().resolve()
    if not src.exists():
        raise RecipeError(f"Data path not found: {src}")
    if src.is_dir():
        raise RecipeError(f"Data path must be a file (csv/jsonl), got directory: {src}")
 
    suffix = src.suffix.lower()
    fmt = {".csv": "csv", ".jsonl": "jsonl", ".json": "json"}.get(suffix, suffix.lstrip(".") or "file")
    sha = hash_file(src)
    rev_id = sha[:12]
    dest = revision_dir(rev_id, cwd=cwd)
    dest.mkdir(parents=True, exist_ok=True)
 
    snapshot_dir = None
    data_path = src
    if snapshot:
        copied = dest / src.name
        shutil.copy2(src, copied)
        snapshot_dir = str(copied)
        data_path = copied
 
    rev = DataRevision(
        id=rev_id,
        path=str(data_path),
        sha256=sha,
        format=fmt,
        n_rows=_count_rows(src, fmt),
        columns=_csv_columns(src) if fmt == "csv" else None,
        snapshot_dir=snapshot_dir,
    )
    (dest / "manifest.json").write_text(json.dumps(rev.to_dict(), indent=2), encoding="utf-8")
    return rev
 
 
def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")
 
 
def save_run(record: RunRecord) -> Path:
    d = Path(record.run_dir)
    d.mkdir(parents=True, exist_ok=True)
    write_json(d / "run.json", record.to_dict())
    if record.checkpoint:
        write_json(d / "checkpoint.json", record.checkpoint)
    if record.gate:
        write_json(d / "gate.json", record.gate)
    if record.metrics:
        write_json(d / "metrics.json", record.metrics)
    if record.inspect:
        write_json(d / "inspect.json", record.inspect)
    return d / "run.json"
 
 
def load_run(run_id: str, *, cwd: Path | None = None) -> RunRecord:
    path = run_dir(run_id, cwd=cwd) / "run.json"
    if not path.is_file():
        raise RecipeError(f"No run {run_id!r} under {run_root(cwd) / 'runs'}")
    data = json.loads(path.read_text(encoding="utf-8"))
    return RunRecord(
        run_id=data["run_id"],
        name=data.get("name") or "",
        family=data.get("family") or "",
        status=data.get("status") or "unknown",
        run_dir=data.get("run_dir") or str(path.parent),
        recipe=dict(data.get("recipe") or {}),
        revision=dict(data.get("revision") or {}),
        checkpoint=data.get("checkpoint"),
        gate=data.get("gate"),
        metrics=dict(data.get("metrics") or {}),
        inspect=dict(data.get("inspect") or {}),
        error=data.get("error"),
    )