scripts/train_models.py

"""Fit the WattGPU power and ITL models and save them for the demo.

Two things are produced:

1. `power.joblib` / `itl.joblib` — the pipelines from the paper, refitted on
   the full Watt Counts subset.
2. `meta.json` — the sets of profiled LLMs and GPUs (which decide the demo's
   certainty tier), plus the accuracy each tier can be expected to deliver.

The tier accuracies come from the paper's own validation protocols, so the
number shown next to a prediction is measured under exactly the conditions
that prediction is made in:

  green  (seen LLM, seen GPU)   -> 5-fold grouped CV
  yellow (unseen LLM, seen GPU) -> leave-one-LLM-out
  orange (seen LLM, unseen GPU) -> leave-one-GPU-out

Usage:  python scripts/train_models.py
"""

from __future__ import annotations

import argparse
import ast
import json
import os
import sys

import joblib
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.model_selection import GroupKFold, LeaveOneGroupOut
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OrdinalEncoder, StandardScaler
from xgboost import XGBRegressor

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, REPO_ROOT)

from wattgpu_demo.features import (  # noqa: E402
    ID_COLUMNS,
    ITL_FEATURES,
    ITL_TARGET,
    POWER_FEATURES,
    POWER_TARGET,
    build_training_frame,
    clean_frame,
)

def find_paper_data(start: str) -> str | None:
    """Locate the paper's `data/` directory by walking up from `start`.

    The Space lives in its own git repository nested inside the research
    repository, and how deeply is not fixed, so the measurement files are found
    by their contents rather than by a hard-coded number of parent directories.
    """
    current = os.path.abspath(start)
    while True:
        candidate = os.path.join(current, "data")
        if os.path.exists(os.path.join(candidate, "watt_counts_subset.csv")):
            return candidate
        parent = os.path.dirname(current)
        if parent == current:
            return None
        current = parent


PAPER_DATA_DIR = find_paper_data(REPO_ROOT)

# Hyperparameters as reported in the paper.
POWER_REGRESSOR = dict(max_depth=6, reg_lambda=150, n_estimators=200)
ITL_REGRESSOR = dict(max_depth=5, reg_lambda=100, n_estimators=100)


def build_pipeline(X: pd.DataFrame, regressor) -> Pipeline:
    """Preprocessing + regressor, identical to the notebook's `_build_pipeline`."""
    numeric = X.select_dtypes(include=[np.number]).columns.tolist()
    categorical = X.select_dtypes(include=["object", "category"]).columns.tolist()

    preprocessor = ColumnTransformer([
        ("num", Pipeline([
            ("imputer", SimpleImputer(strategy="mean")),
            ("scaler", StandardScaler()),
        ]), numeric),
        ("cat", Pipeline([
            ("imputer", SimpleImputer(strategy="most_frequent")),
            # Unseen categories (a new `model_type` or `memory_type`) encode to
            # -1 rather than raising, which is what lets the demo predict for
            # architectures that were never profiled.
            ("encoder", OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1)),
        ]), categorical),
    ])
    return Pipeline([("preprocessor", preprocessor), ("regressor", regressor)])


def mdape(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """Median absolute percentage error, the paper's headline metric."""
    return float(np.median(np.abs((y_true - y_pred) / y_true)) * 100)


def cross_validate(df: pd.DataFrame, group_col: str | None, target_col: str,
                   regressor_kwargs: dict, log_transform_y: bool) -> np.ndarray:
    """Out-of-fold predictions under CV, LOGO or LOLO."""
    X = df.drop(columns=[c for c in [target_col, *ID_COLUMNS] if c in df.columns])
    y = df[target_col]
    predictions = np.full(len(df), np.nan)

    if group_col is not None:
        splits = LeaveOneGroupOut().split(X, y, df[group_col])
    else:
        # Group on the configuration so replicates of one (LLM, GPU) pair never
        # straddle the train/test boundary.
        config_id = df[ID_COLUMNS].astype(str).agg("|".join, axis=1)
        splits = GroupKFold(n_splits=5).split(X, y, groups=config_id)

    for train_idx, test_idx in splits:
        pipeline = build_pipeline(X, XGBRegressor(**regressor_kwargs))
        y_train = np.log(y.iloc[train_idx]) if log_transform_y else y.iloc[train_idx]
        pipeline.fit(X.iloc[train_idx], y_train)
        y_pred = pipeline.predict(X.iloc[test_idx])
        predictions[test_idx] = np.exp(y_pred) if log_transform_y else y_pred

    return predictions


def double_holdout_predictions(df: pd.DataFrame, target_col: str,
                               regressor_kwargs: dict, log_transform_y: bool) -> np.ndarray:
    """Predictions for pairs whose LLM *and* GPU are both held out.

    The paper validates generalisation one axis at a time (LOGO and LOLO). This
    is the natural extension: for every measured (LLM, GPU) pair, train on the
    data with that GPU and that LLM both removed entirely, then predict the
    pair. It is the only honest way to attach an error to an estimate where
    neither side was measured -- without it, such an estimate would carry no
    validated accuracy at all.
    """
    X = df.drop(columns=[c for c in [target_col, *ID_COLUMNS] if c in df.columns])
    y = df[target_col]
    predictions = np.full(len(df), np.nan)

    models = df["model"].to_numpy()
    gpus = df["gpu_db_name"].to_numpy()
    pairs = df[["model", "gpu_db_name"]].drop_duplicates().itertuples(index=False)

    for n, (model, gpu) in enumerate(pairs, start=1):
        test = (models == model) & (gpus == gpu)
        train = (models != model) & (gpus != gpu)
        if not train.any() or not test.any():
            continue

        pipeline = build_pipeline(X, XGBRegressor(**regressor_kwargs))
        y_train = np.log(y[train]) if log_transform_y else y[train]
        pipeline.fit(X[train], y_train)
        y_pred = pipeline.predict(X[test])
        predictions[test] = np.exp(y_pred) if log_transform_y else y_pred
        if n % 40 == 0:
            print(f"      double holdout: {n} pairs")

    return predictions


def tier_accuracy(df: pd.DataFrame, target_col: str, regressor_kwargs: dict,
                  log_transform_y: bool, to_watts: bool) -> dict[str, dict[str, float]]:
    """MdAPE per certainty tier and scenario, each under its own protocol.

    Reported separately for offline and server operation because, as in the
    paper's Tables 2 and 3, the two regimes differ substantially -- especially
    for ITL, where offline throughput depends on batching effects that the
    features capture only partly.
    """
    scale = df["thermal_design_power_w"].to_numpy() if to_watts else 1.0
    y_true = df[target_col].to_numpy() * scale
    # The paper reports the two server load levels together.
    regime = np.where(df["scenario"].to_numpy() == "offline", "offline", "server")

    results: dict[str, dict[str, float]] = {}
    tiers = (("green", None), ("yellow", "model"), ("orange", "gpu_db_name"), ("red", "both"))
    for tier, group_col in tiers:
        if group_col == "both":
            y_pred = double_holdout_predictions(
                df, target_col, regressor_kwargs, log_transform_y) * scale
        else:
            y_pred = cross_validate(
                df, group_col, target_col, regressor_kwargs, log_transform_y) * scale
        valid = ~np.isnan(y_pred)
        results[tier] = {
            scenario: round(mdape(y_true[valid & (regime == scenario)],
                                  y_pred[valid & (regime == scenario)]), 1)
            for scenario in ("offline", "server")
        }
        print(f"    {tier:<7} MdAPE  offline {results[tier]['offline']:5.1f}%   "
              f"server {results[tier]['server']:5.1f}%")
    return results


def fit_final(df: pd.DataFrame, target_col: str, regressor_kwargs: dict,
              log_transform_y: bool) -> Pipeline:
    """Refit on every row, which is what the demo serves predictions from."""
    X = df.drop(columns=[c for c in [target_col, *ID_COLUMNS] if c in df.columns])
    y = np.log(df[target_col]) if log_transform_y else df[target_col]
    pipeline = build_pipeline(X, XGBRegressor(**regressor_kwargs))
    pipeline.fit(X, y)
    return pipeline


def _parse_architectures(value) -> list[str]:
    """`model_features.csv` stores the architecture list as a Python literal."""
    if value is None or (isinstance(value, float) and pd.isna(value)):
        return []
    if isinstance(value, list):
        return [str(v) for v in value]
    try:
        parsed = ast.literal_eval(str(value))
    except (ValueError, SyntaxError):
        return [str(value)]
    return [str(v) for v in parsed] if isinstance(parsed, (list, tuple)) else [str(parsed)]


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", default=PAPER_DATA_DIR, required=PAPER_DATA_DIR is None,
                        help="the paper's data/ directory; found automatically when the "
                             "Space repository sits inside the research repository")
    parser.add_argument("--out-dir", default=os.path.join(REPO_ROOT, "data", "models"))
    parser.add_argument("--skip-validation", action="store_true",
                        help="fit only; keep the tier accuracies from a previous run")
    args = parser.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    print(f"loading measurements from {args.data_dir}")
    df_all = build_training_frame(args.data_dir)
    print(f"  {len(df_all)} runs, {df_all.model.nunique()} LLMs, {df_all.gpu_db_name.nunique()} GPUs")

    meta: dict = {}
    accuracies: dict = {}

    # --- power draw ---------------------------------------------------------
    df_power = clean_frame(df_all, [*ID_COLUMNS, POWER_TARGET, *POWER_FEATURES])
    tdp = df_all.set_index("gpu_db_name")["thermal_design_power_w"].drop_duplicates()
    df_power["thermal_design_power_w"] = df_power["gpu_db_name"].map(tdp)
    print(f"\npower model: {len(df_power)} rows, {len(POWER_FEATURES)} features")
    if not args.skip_validation:
        accuracies["power"] = tier_accuracy(
            df_power.drop(columns=["thermal_design_power_w"]).assign(
                thermal_design_power_w=df_power["thermal_design_power_w"]),
            POWER_TARGET, POWER_REGRESSOR, log_transform_y=False, to_watts=True)

    power_pipeline = fit_final(
        df_power.drop(columns=["thermal_design_power_w"]),
        POWER_TARGET, POWER_REGRESSOR, log_transform_y=False)
    joblib.dump(power_pipeline, os.path.join(args.out_dir, "power.joblib"))

    # --- inter-token latency ------------------------------------------------
    df_itl = clean_frame(df_all, [*ID_COLUMNS, ITL_TARGET, *ITL_FEATURES])
    print(f"\nITL model: {len(df_itl)} rows, {len(ITL_FEATURES)} features")
    if not args.skip_validation:
        accuracies["itl"] = tier_accuracy(
            df_itl, ITL_TARGET, ITL_REGRESSOR, log_transform_y=True, to_watts=False)

    itl_pipeline = fit_final(df_itl, ITL_TARGET, ITL_REGRESSOR, log_transform_y=True)
    joblib.dump(itl_pipeline, os.path.join(args.out_dir, "itl.joblib"))

    # --- metadata -----------------------------------------------------------
    meta_path = os.path.join(args.out_dir, "meta.json")
    if args.skip_validation and os.path.exists(meta_path):
        with open(meta_path) as fh:
            accuracies = json.load(fh).get("tier_accuracy_mdape", accuracies)

    # A model or GPU counts as "seen" only if it survived into a training frame.
    seen_models = sorted(set(df_power["model"]) | set(df_itl["model"]))
    seen_gpus = sorted(set(df_power["gpu_db_name"]) | set(df_itl["gpu_db_name"]))

    meta = {
        "seen_models": seen_models,
        "seen_gpus": seen_gpus,
        "power_features": POWER_FEATURES,
        "itl_features": ITL_FEATURES,
        "tier_accuracy_mdape": accuracies,
        # The ITL model is fitted on log(itl); predictions must be exponentiated.
        "log_transformed_targets": ["itl"],
        "n_training_runs": {"power": len(df_power), "itl": len(df_itl)},
        "gpu_tdp": {k: float(v) for k, v in tdp.items()},
    }
    with open(meta_path, "w") as fh:
        json.dump(meta, fh, indent=2)

    # Cache the architecture of every profiled LLM. These features are already
    # in the paper's data, so a profiled model needs no Hub round-trip -- which
    # also makes licence-gated models (Llama, Gemma) work without a token.
    cache = {}
    for model_id, group in df_all.groupby("model"):
        if model_id not in seen_models:
            continue
        row = group.iloc[0]
        cache[model_id] = {
            "model_type": str(row["model_type"]),
            "num_layers": int(row["num_layers"]),
            "hidden_size": int(row["hidden_size"]),
            "num_attention_heads": int(row["num_attention_heads"]),
            "num_key_value_heads": int(row["num_key_value_heads"]),
            "total_b_params": float(row["total_b_params"]),
            "architectures": _parse_architectures(row.get("architectures")),
            "max_position_embeddings": (
                None if pd.isna(row.get("max_position_embeddings"))
                else int(row["max_position_embeddings"])
            ),
            "torch_dtype": (
                None if pd.isna(row.get("torch_dtype")) or row.get("torch_dtype") == "N/A"
                else str(row["torch_dtype"])
            ),
        }
    cache_path = os.path.join(args.out_dir, "profiled_llms.json")
    with open(cache_path, "w") as fh:
        json.dump(cache, fh, indent=2, sort_keys=True)
    print(f"  cached architecture for {len(cache)} profiled LLMs")

    print(f"\nsaved models and metadata to {args.out_dir}")
    print(f"  {len(seen_models)} profiled LLMs, {len(seen_gpus)} profiled GPUs")
    return 0


if __name__ == "__main__":
    sys.exit(main())