scripts/tools/smoke_test.py

#!/usr/bin/env python3
"""
Smoke test: run every experiment in experiments.csv for 2 epochs.
Purpose: verify all code paths work, not actual training.

Usage:
    python smoke_test.py                    # run all
    python smoke_test.py --phase 1 2        # specific phases
    python smoke_test.py --resume           # skip already passed
"""

import argparse
import csv
import json
import subprocess
import sys
import time
import yaml
from pathlib import Path
from datetime import datetime

ROOT = Path(__file__).resolve().parent
EXPERIMENTS_CSV = ROOT / "experiments.csv"
SMOKE_CONFIGS_DIR = ROOT / "smoke_test_configs"
SMOKE_RESULTS = ROOT / "smoke_test_results.json"

# Force these overrides for ALL experiments
SMOKE_OVERRIDES = {
    "epochs": 2,
    "data": "sample_B",  # use existing data
}


def load_experiments(phases=None):
    exps = []
    with open(EXPERIMENTS_CSV) as f:
        for row in csv.DictReader(f):
            if row.get("id", "").startswith("##"):
                continue
            if phases and row.get("phase") not in phases:
                continue
            exps.append(row)
    return exps


def deduplicate_by_code_path(exps):
    """Many experiments differ only in hyperparams (lr, dropout, etc).
    Keep only one representative per unique code path."""
    seen = set()
    unique = []
    for exp in exps:
        # Code path is defined by: arch + objective + masking + embedding + tokenizer
        # + optimizer + forgetter + use_moe + use_attn_res + packing + use_mtp + kd_enabled
        key = (
            exp.get("arch", ""),
            exp.get("objective", ""),
            exp.get("masking", ""),
            exp.get("embedding", ""),
            exp.get("embedding_init", ""),
            exp.get("tokenizer", ""),
            exp.get("optimizer", ""),
            exp.get("forgetter", ""),
            exp.get("use_moe", ""),
            exp.get("use_attn_res", ""),
            exp.get("kd_enabled", ""),
        )
        if key not in seen:
            seen.add(key)
            unique.append(exp)
    return unique


def make_smoke_yaml(exp):
    """Create a minimal YAML config from a CSV experiment row."""
    from run_experiments import csv_row_to_yaml, DATA_PATH_MAP

    yaml_dict = csv_row_to_yaml(exp)

    # Force overrides
    yaml_dict.setdefault("training", {})["epochs"] = SMOKE_OVERRIDES["epochs"]
    yaml_dict["name"] = f"smoke_{exp['id'].replace('.', '_')}"

    # Force sample_B data
    yaml_dict.setdefault("data", {})["train_file"] = "data/8_sample_B/train.txt"

    # Disable AoA checkpoints (speed)
    yaml_dict.setdefault("checkpoint", {})["save_aoa_checkpoints"] = False
    yaml_dict["checkpoint"]["save_every_epoch"] = False

    # Skip eval
    yaml_dict["_skip_eval"] = True

    return yaml_dict


def run_one(exp, dry_run=False):
    """Run a single smoke test. Returns result dict."""
    exp_id = exp["id"]
    name = exp["name"]

    SMOKE_CONFIGS_DIR.mkdir(parents=True, exist_ok=True)
    yaml_dict = make_smoke_yaml(exp)
    yaml_path = SMOKE_CONFIGS_DIR / f"{exp_id.replace('.', '_')}.yaml"

    with open(yaml_path, "w") as f:
        yaml.dump(yaml_dict, f, default_flow_style=False)

    print(f"  [{exp_id}] {name} ...", end=" ", flush=True)

    if dry_run:
        print("DRY RUN")
        return {"id": exp_id, "name": name, "status": "dry_run"}

    # Check data exists
    data_path = ROOT / "data/8_sample_B/train.txt"
    if not data_path.exists():
        print("SKIP (no data)")
        return {"id": exp_id, "name": name, "status": "skip_no_data"}

    # Check if tokenizer exists for morfessor
    if exp.get("tokenizer") == "morfessor_bpe":
        morf_path = ROOT / "models/tokenizer_morfessor"
        if not morf_path.exists():
            print("SKIP (no morfessor tokenizer)")
            return {"id": exp_id, "name": name, "status": "skip_no_morfessor"}

    # Check KD teacher
    if exp.get("kd_enabled") in ("True", "true", True):
        print("SKIP (no KD teacher)")
        return {"id": exp_id, "name": name, "status": "skip_no_teacher"}

    start = time.time()
    try:
        result = subprocess.run(
            [sys.executable, "-m", "scripts.03_training.train",
             "--config", str(yaml_path), "--skip-eval"],
            cwd=str(ROOT),
            capture_output=True, text=True, timeout=1800,  # 30 min max
        )
        elapsed = time.time() - start

        if result.returncode == 0:
            print(f"OK ({elapsed:.0f}s)")
            return {"id": exp_id, "name": name, "status": "pass", "time": round(elapsed)}
        else:
            # Extract last few lines of error
            err_lines = result.stderr.strip().split("\n")[-5:]
            err_msg = "\n".join(err_lines)
            print(f"FAIL ({elapsed:.0f}s)")
            print(f"    Error: {err_lines[-1][:120]}")
            return {"id": exp_id, "name": name, "status": "fail",
                    "time": round(elapsed), "error": err_msg}

    except subprocess.TimeoutExpired:
        print("TIMEOUT")
        return {"id": exp_id, "name": name, "status": "timeout"}
    except Exception as e:
        print(f"ERROR: {e}")
        return {"id": exp_id, "name": name, "status": "error", "error": str(e)}


def load_results():
    if SMOKE_RESULTS.exists():
        with open(SMOKE_RESULTS) as f:
            return json.load(f)
    return {}


def save_results(results):
    with open(SMOKE_RESULTS, "w") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)


def main():
    parser = argparse.ArgumentParser(description="Smoke test all experiments")
    parser.add_argument("--phase", nargs="*", help="Filter by phase")
    parser.add_argument("--resume", action="store_true", help="Skip already passed")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--all-variants", action="store_true",
                        help="Don't deduplicate — test every single experiment")
    args = parser.parse_args()

    exps = load_experiments(phases=args.phase)
    if not args.all_variants:
        original = len(exps)
        exps = deduplicate_by_code_path(exps)
        print(f"Deduplicated {original} experiments → {len(exps)} unique code paths")

    results = load_results() if args.resume else {}
    passed = failed = skipped = 0

    print(f"\n{'='*60}")
    print(f"  Smoke Test: {len(exps)} experiments, 2 epochs each, skip eval")
    print(f"{'='*60}\n")

    for exp in exps:
        if args.resume and results.get(exp["id"], {}).get("status") == "pass":
            skipped += 1
            continue

        result = run_one(exp, dry_run=args.dry_run)
        results[result["id"]] = result
        save_results(results)

        if result["status"] == "pass":
            passed += 1
        elif result["status"] == "fail":
            failed += 1

    # Summary
    print(f"\n{'='*60}")
    print(f"  Results: {passed} passed, {failed} failed, {skipped} skipped")
    if failed > 0:
        print(f"\n  FAILED experiments:")
        for eid, r in results.items():
            if r.get("status") == "fail":
                print(f"    {eid}: {r.get('name', '?')}")
                err = r.get("error", "")
                if err:
                    print(f"      {err.split(chr(10))[-1][:100]}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()