database.py

"""SQLite-based cache for training results and per-user FLOPs tracking."""

import json
import sqlite3
import threading
from pathlib import Path
from typing import Optional


DB_PATH = Path("/data/share/hw3-data/scaling_api.db")

_local = threading.local()


def _get_conn() -> sqlite3.Connection:
    if not hasattr(_local, "conn"):
        _local.conn = sqlite3.connect(str(DB_PATH), timeout=30)
        _local.conn.execute("PRAGMA journal_mode=WAL")
        _local.conn.execute("PRAGMA busy_timeout=30000")
        _local.conn.row_factory = sqlite3.Row
    return _local.conn


def init_db():
    conn = _get_conn()
    conn.executescript("""
        CREATE TABLE IF NOT EXISTS training_results (
            config_key TEXT PRIMARY KEY,
            d_model INTEGER NOT NULL,
            num_layers INTEGER NOT NULL,
            num_heads INTEGER NOT NULL,
            batch_size INTEGER NOT NULL,
            learning_rate REAL NOT NULL,
            train_flops INTEGER NOT NULL,
            loss REAL,
            status TEXT NOT NULL DEFAULT 'pending',
            slurm_job_id TEXT,
            submitted_at TEXT,
            completed_at TEXT
        );

        CREATE TABLE IF NOT EXISTS user_queries (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            api_key TEXT NOT NULL,
            config_key TEXT NOT NULL,
            train_flops INTEGER NOT NULL,
            flops_charged INTEGER NOT NULL DEFAULT 0,
            queried_at TEXT DEFAULT (datetime('now')),
            UNIQUE(api_key, config_key)
        );

        CREATE TABLE IF NOT EXISTS user_flops (
            api_key TEXT PRIMARY KEY,
            total_flops_used REAL NOT NULL DEFAULT 0
        );
    """)
    conn.commit()


def make_config_key(
    d_model: int,
    num_layers: int,
    num_heads: int,
    batch_size: int,
    learning_rate: float,
    train_flops: int,
) -> str:
    return json.dumps(
        {
            "d_model": d_model,
            "num_layers": num_layers,
            "num_heads": num_heads,
            "batch_size": batch_size,
            "learning_rate": learning_rate,
            "train_flops": train_flops,
        },
        sort_keys=True,
    )


def get_training_result(config_key: str) -> Optional[dict]:
    conn = _get_conn()
    row = conn.execute(
        "SELECT * FROM training_results WHERE config_key = ?", (config_key,)
    ).fetchone()
    if row is None:
        return None
    return dict(row)


def insert_training_job(
    config_key: str,
    d_model: int,
    num_layers: int,
    num_heads: int,
    batch_size: int,
    learning_rate: float,
    train_flops: int,
    slurm_job_id: str,
) -> bool:
    conn = _get_conn()
    try:
        conn.execute(
            """INSERT INTO training_results
               (config_key, d_model, num_layers, num_heads, batch_size,
                learning_rate, train_flops, status, slurm_job_id, submitted_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, 'running', ?, datetime('now'))""",
            (
                config_key,
                d_model,
                num_layers,
                num_heads,
                batch_size,
                learning_rate,
                train_flops,
                slurm_job_id,
            ),
        )
        conn.commit()
        return True
    except sqlite3.IntegrityError:
        return False


def complete_training_job(config_key: str, loss: float):
    conn = _get_conn()
    conn.execute(
        """UPDATE training_results
           SET loss = ?, status = 'completed', completed_at = datetime('now')
           WHERE config_key = ?""",
        (loss, config_key),
    )
    conn.commit()


def fail_training_job(config_key: str):
    conn = _get_conn()
    conn.execute(
        "UPDATE training_results SET status = 'failed' WHERE config_key = ?",
        (config_key,),
    )
    conn.commit()


def record_user_query(
    api_key: str, config_key: str, train_flops: int, charge_flops: bool
) -> float:
    """Record a user query and update FLOPs. Returns new total_flops_used."""
    conn = _get_conn()
    flops_to_charge = train_flops if charge_flops else 0

    conn.execute(
        """INSERT INTO user_queries (api_key, config_key, train_flops, flops_charged)
           VALUES (?, ?, ?, ?)
           ON CONFLICT(api_key, config_key) DO NOTHING""",
        (api_key, config_key, train_flops, flops_to_charge),
    )

    if charge_flops:
        already_queried = conn.execute(
            """SELECT COUNT(*) FROM user_queries
               WHERE api_key = ? AND config_key = ? AND flops_charged > 0""",
            (api_key, config_key),
        ).fetchone()[0]

        if already_queried <= 1:
            conn.execute(
                """INSERT INTO user_flops (api_key, total_flops_used)
                   VALUES (?, ?)
                   ON CONFLICT(api_key)
                   DO UPDATE SET total_flops_used = total_flops_used + ?""",
                (api_key, float(flops_to_charge), float(flops_to_charge)),
            )

    conn.commit()
    return get_user_total_flops(api_key)


def get_user_total_flops(api_key: str) -> float:
    conn = _get_conn()
    row = conn.execute(
        "SELECT total_flops_used FROM user_flops WHERE api_key = ?", (api_key,)
    ).fetchone()
    if row is None:
        return 0.0
    return row["total_flops_used"]


def has_user_queried(api_key: str) -> bool:
    conn = _get_conn()
    row = conn.execute(
        "SELECT COUNT(*) FROM user_queries WHERE api_key = ?", (api_key,)
    ).fetchone()
    return row[0] > 0


def user_already_queried_config(api_key: str, config_key: str) -> bool:
    conn = _get_conn()
    row = conn.execute(
        "SELECT COUNT(*) FROM user_queries WHERE api_key = ? AND config_key = ? AND flops_charged > 0",
        (api_key, config_key),
    ).fetchone()
    return row[0] > 0


def reset_failed_job(config_key: str, slurm_job_id: str):
    conn = _get_conn()
    conn.execute(
        """UPDATE training_results
           SET status = 'running', slurm_job_id = ?, loss = NULL,
               submitted_at = datetime('now'), completed_at = NULL
           WHERE config_key = ? AND status = 'failed'""",
        (slurm_job_id, config_key),
    )
    conn.commit()


def get_previous_runs(api_key: str) -> list[dict]:
    conn = _get_conn()
    rows = conn.execute(
        """SELECT tr.d_model, tr.num_layers, tr.num_heads, tr.batch_size,
                  tr.learning_rate, tr.train_flops, tr.loss
           FROM user_queries uq
           JOIN training_results tr ON uq.config_key = tr.config_key
           WHERE uq.api_key = ? AND tr.status = 'completed'
           ORDER BY uq.queried_at""",
        (api_key,),
    ).fetchall()
    return [dict(r) for r in rows]