mirror of https://github.com/razor-ai/soup.git
275 lines
8.7 KiB
Python
275 lines
8.7 KiB
Python
"""Experiment tracking — stores runs in local SQLite.
|
|
|
|
Usage:
|
|
tracker = ExperimentTracker()
|
|
run_id = tracker.start_run(config_dict, device, device_name, gpu_info)
|
|
tracker.log_metrics(run_id, step=10, loss=2.3, lr=1e-5)
|
|
tracker.finish_run(run_id, initial_loss=2.5, final_loss=0.8, ...)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import secrets
|
|
import sqlite3
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from soup_cli.utils.constants import EXPERIMENTS_DB, SOUP_DIR
|
|
|
|
_SCHEMA_SQL = """
|
|
CREATE TABLE IF NOT EXISTS runs (
|
|
run_id TEXT PRIMARY KEY,
|
|
experiment_name TEXT,
|
|
created_at TEXT NOT NULL,
|
|
status TEXT NOT NULL DEFAULT 'running',
|
|
config_json TEXT NOT NULL,
|
|
device TEXT,
|
|
device_name TEXT,
|
|
gpu_memory TEXT,
|
|
initial_loss REAL,
|
|
final_loss REAL,
|
|
total_steps INTEGER,
|
|
duration_secs REAL,
|
|
output_dir TEXT,
|
|
base_model TEXT,
|
|
task TEXT
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS metrics (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
run_id TEXT NOT NULL REFERENCES runs(run_id),
|
|
step INTEGER NOT NULL,
|
|
epoch REAL,
|
|
loss REAL,
|
|
lr REAL,
|
|
grad_norm REAL,
|
|
speed REAL,
|
|
gpu_mem TEXT,
|
|
timestamp TEXT NOT NULL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS eval_results (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
run_id TEXT REFERENCES runs(run_id),
|
|
model_path TEXT NOT NULL,
|
|
benchmark TEXT NOT NULL,
|
|
score REAL NOT NULL,
|
|
details_json TEXT,
|
|
created_at TEXT NOT NULL
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON metrics(run_id);
|
|
CREATE INDEX IF NOT EXISTS idx_eval_run_id ON eval_results(run_id);
|
|
"""
|
|
|
|
|
|
def _get_db_path() -> Path:
|
|
"""Return path to experiments DB, creating parent dir if needed."""
|
|
import os
|
|
|
|
# Allow override via env var (useful for tests and CI)
|
|
env_path = os.environ.get("SOUP_DB_PATH")
|
|
if env_path:
|
|
return Path(env_path)
|
|
|
|
soup_dir = Path.home() / SOUP_DIR
|
|
soup_dir.mkdir(parents=True, exist_ok=True)
|
|
return soup_dir / EXPERIMENTS_DB
|
|
|
|
|
|
def generate_run_id() -> str:
|
|
"""Generate a unique, sortable run ID: run_YYYYMMDD_HHMMSS_xxxxxxxx."""
|
|
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
suffix = secrets.token_hex(4)
|
|
return f"run_{ts}_{suffix}"
|
|
|
|
|
|
class ExperimentTracker:
|
|
"""SQLite-backed experiment tracker for training runs and evaluations."""
|
|
|
|
def __init__(self, db_path: Optional[Path] = None):
|
|
self.db_path = db_path or _get_db_path()
|
|
self._conn: Optional[sqlite3.Connection] = None
|
|
self._ensure_schema()
|
|
|
|
def _get_conn(self) -> sqlite3.Connection:
|
|
if self._conn is None:
|
|
self._conn = sqlite3.connect(str(self.db_path))
|
|
self._conn.row_factory = sqlite3.Row
|
|
self._conn.execute("PRAGMA journal_mode=WAL")
|
|
return self._conn
|
|
|
|
def _ensure_schema(self) -> None:
|
|
"""Create tables if they don't exist."""
|
|
conn = self._get_conn()
|
|
conn.executescript(_SCHEMA_SQL)
|
|
conn.commit()
|
|
|
|
def start_run(
|
|
self,
|
|
config_dict: dict,
|
|
device: str,
|
|
device_name: str,
|
|
gpu_info: dict,
|
|
experiment_name: Optional[str] = None,
|
|
) -> str:
|
|
"""Insert a new run and return its run_id."""
|
|
run_id = generate_run_id()
|
|
now = datetime.now().isoformat()
|
|
config_json = json.dumps(config_dict, default=str)
|
|
|
|
base_model = config_dict.get("base", "")
|
|
task = config_dict.get("task", "sft")
|
|
gpu_memory = gpu_info.get("memory_total", "")
|
|
|
|
conn = self._get_conn()
|
|
conn.execute(
|
|
"""INSERT INTO runs
|
|
(run_id, experiment_name, created_at, status, config_json,
|
|
device, device_name, gpu_memory, base_model, task)
|
|
VALUES (?, ?, ?, 'running', ?, ?, ?, ?, ?, ?)""",
|
|
(
|
|
run_id, experiment_name, now, config_json,
|
|
device, device_name, gpu_memory, base_model, task,
|
|
),
|
|
)
|
|
conn.commit()
|
|
return run_id
|
|
|
|
def log_metrics(
|
|
self,
|
|
run_id: str,
|
|
step: int,
|
|
epoch: float = 0.0,
|
|
loss: float = 0.0,
|
|
lr: float = 0.0,
|
|
grad_norm: float = 0.0,
|
|
speed: float = 0.0,
|
|
gpu_mem: str = "",
|
|
) -> None:
|
|
"""Log a single metrics row for the given run."""
|
|
now = datetime.now().isoformat()
|
|
conn = self._get_conn()
|
|
conn.execute(
|
|
"""INSERT INTO metrics
|
|
(run_id, step, epoch, loss, lr, grad_norm, speed, gpu_mem, timestamp)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
|
(run_id, step, epoch, loss, lr, grad_norm, speed, gpu_mem, now),
|
|
)
|
|
conn.commit()
|
|
|
|
def finish_run(
|
|
self,
|
|
run_id: str,
|
|
initial_loss: float,
|
|
final_loss: float,
|
|
total_steps: int,
|
|
duration_secs: float,
|
|
output_dir: str,
|
|
) -> None:
|
|
"""Mark run as completed and fill summary fields."""
|
|
conn = self._get_conn()
|
|
conn.execute(
|
|
"""UPDATE runs SET
|
|
status = 'completed',
|
|
initial_loss = ?, final_loss = ?,
|
|
total_steps = ?, duration_secs = ?, output_dir = ?
|
|
WHERE run_id = ?""",
|
|
(initial_loss, final_loss, total_steps, duration_secs, output_dir, run_id),
|
|
)
|
|
conn.commit()
|
|
|
|
def fail_run(self, run_id: str) -> None:
|
|
"""Mark run as failed."""
|
|
conn = self._get_conn()
|
|
conn.execute("UPDATE runs SET status = 'failed' WHERE run_id = ?", (run_id,))
|
|
conn.commit()
|
|
|
|
def list_runs(self, limit: int = 50) -> list[dict]:
|
|
"""Return list of runs ordered by created_at desc."""
|
|
conn = self._get_conn()
|
|
rows = conn.execute(
|
|
"SELECT * FROM runs ORDER BY created_at DESC, rowid DESC LIMIT ?", (limit,)
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
def get_run(self, run_id: str) -> Optional[dict]:
|
|
"""Get full details of a single run. Supports prefix matching."""
|
|
conn = self._get_conn()
|
|
# Try exact match first
|
|
row = conn.execute(
|
|
"SELECT * FROM runs WHERE run_id = ?", (run_id,)
|
|
).fetchone()
|
|
|
|
if row is None:
|
|
# Try prefix match
|
|
rows = conn.execute(
|
|
"SELECT * FROM runs WHERE run_id LIKE ? ORDER BY created_at DESC",
|
|
(f"{run_id}%",),
|
|
).fetchall()
|
|
if len(rows) == 1:
|
|
row = rows[0]
|
|
elif len(rows) > 1:
|
|
return None # ambiguous prefix
|
|
|
|
return dict(row) if row else None
|
|
|
|
def get_metrics(self, run_id: str) -> list[dict]:
|
|
"""Get all metric rows for a run, ordered by step."""
|
|
conn = self._get_conn()
|
|
rows = conn.execute(
|
|
"SELECT * FROM metrics WHERE run_id = ? ORDER BY step", (run_id,)
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
def save_eval_result(
|
|
self,
|
|
model_path: str,
|
|
benchmark: str,
|
|
score: float,
|
|
details: dict,
|
|
run_id: Optional[str] = None,
|
|
) -> None:
|
|
"""Save an evaluation result."""
|
|
now = datetime.now().isoformat()
|
|
details_json = json.dumps(details, default=str)
|
|
conn = self._get_conn()
|
|
conn.execute(
|
|
"""INSERT INTO eval_results
|
|
(run_id, model_path, benchmark, score, details_json, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?)""",
|
|
(run_id, model_path, benchmark, score, details_json, now),
|
|
)
|
|
conn.commit()
|
|
|
|
def get_eval_results(self, run_id: Optional[str] = None) -> list[dict]:
|
|
"""Get eval results, optionally filtered by run_id."""
|
|
conn = self._get_conn()
|
|
if run_id:
|
|
rows = conn.execute(
|
|
"SELECT * FROM eval_results WHERE run_id = ? ORDER BY created_at DESC",
|
|
(run_id,),
|
|
).fetchall()
|
|
else:
|
|
rows = conn.execute(
|
|
"SELECT * FROM eval_results ORDER BY created_at DESC"
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
def delete_run(self, run_id: str) -> bool:
|
|
"""Delete a run and its metrics. Returns True if found."""
|
|
conn = self._get_conn()
|
|
conn.execute("DELETE FROM metrics WHERE run_id = ?", (run_id,))
|
|
conn.execute("DELETE FROM eval_results WHERE run_id = ?", (run_id,))
|
|
cursor = conn.execute("DELETE FROM runs WHERE run_id = ?", (run_id,))
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def close(self) -> None:
|
|
"""Close the database connection."""
|
|
if self._conn:
|
|
self._conn.close()
|
|
self._conn = None
|