scripts/run_experiments.py

"""
Experiment runner (Person E).

Generates one merged config per ablation run under a dedicated output tree, optionally
invokes train.py and evaluate.py.

Examples:
    # Emit configs only (no training), useful while train.py is still being wired up
    python scripts/run_experiments.py --dry-run

    # Run train + eval for each experiment (requires a working scripts/train.py)
    python scripts/run_experiments.py --execute

    # Single experiment
    python scripts/run_experiments.py --execute --only exp1_baseline_transformer
"""

from __future__ import annotations

import argparse
import json
import logging
import subprocess
import sys
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))

from omegaconf import DictConfig, open_dict

from easytranslate.utils.config import config_from_cli, overrides_to_cli_args, save_config

REPO_ROOT = Path(__file__).resolve().parent.parent

# Six ablations aligned with TASK_ASSIGNMENT.md; turn off unrelated toggles for single-factor runs.
EXPERIMENTS: list[dict[str, Any]] = [
    {
        "name": "exp1_baseline_transformer",
        "description": "Baseline: Transformer from scratch (6 layers, d=512), sinusoidal PE, standard attention",
        "overrides": {
            "model.type": "transformer_scratch",
            "model.transformer.num_encoder_layers": 6,
            "model.transformer.num_decoder_layers": 6,
            "model.transformer.d_model": 512,
            "model.transformer.use_flash_attention": False,
            "model.transformer.use_rotary_embedding": False,
        },
    },
    {
        "name": "exp2_transformer_rope",
        "description": "Ablation: RoPE only (Flash off to isolate RoPE)",
        "overrides": {
            "model.type": "transformer_scratch",
            "model.transformer.use_flash_attention": False,
            "model.transformer.use_rotary_embedding": True,
        },
    },
    {
        "name": "exp3_transformer_flash_attn",
        "description": "Ablation: Flash attention only (RoPE off)",
        "overrides": {
            "model.type": "transformer_scratch",
            "model.transformer.use_flash_attention": True,
            "model.transformer.use_rotary_embedding": False,
        },
    },
    {
        "name": "exp4_transformer_full",
        "description": "Full stack: Transformer + RoPE + Flash attention",
        "overrides": {
            "model.type": "transformer_scratch",
            "model.transformer.use_flash_attention": True,
            "model.transformer.use_rotary_embedding": True,
        },
    },
    {
        "name": "exp5_nllb_lora",
        "description": "Pretrained finetune: NLLB-600M + LoRA",
        "overrides": {
            "model.type": "finetune_nllb",
            "model.pretrained.use_lora": True,
            "model.pretrained.lora.r": 16,
        },
    },
    {
        "name": "exp6_nllb_full_finetune",
        "description": "Pretrained full finetune: NLLB-600M",
        "overrides": {
            "model.type": "finetune_nllb",
            "model.pretrained.use_lora": False,
        },
    },
]


def _experiment_paths(output_root: Path, exp_name: str) -> dict[str, Path]:
    root = output_root / exp_name
    return {
        "root": root,
        "checkpoints": root / "checkpoints",
        "logs": root / "logs",
        "config": root / "config.yaml",
        "eval_json": root / "evaluation_results.json",
        "meta_json": root / "experiment_meta.json",
    }


def build_experiment_config(
    base_config_path: Path,
    exp: dict[str, Any],
    output_root: Path,
) -> DictConfig:
    """Merge base YAML with overrides; point checkpoint/log dirs at the experiment folder."""
    cli_args = overrides_to_cli_args(exp["overrides"])
    cfg = config_from_cli(str(base_config_path), cli_args)

    paths = _experiment_paths(output_root, exp["name"])
    paths["root"].mkdir(parents=True, exist_ok=True)
    paths["checkpoints"].mkdir(parents=True, exist_ok=True)
    paths["logs"].mkdir(parents=True, exist_ok=True)

    with open_dict(cfg):
        cfg.experiment.name = exp["name"]
        cfg.experiment.output_dir = str(paths["root"])
        if "training" in cfg and "checkpoint" in cfg.training:
            cfg.training.checkpoint.save_dir = str(paths["checkpoints"])
        if "logging" in cfg:
            cfg.logging.log_dir = str(paths["logs"])

    return cfg


def run_single_experiment(
    exp: dict[str, Any],
    *,
    base_config_path: Path,
    output_root: Path,
    execute: bool,
    skip_train: bool,
    skip_eval: bool,
    extra_train_args: list[str],
    extra_eval_args: list[str],
) -> dict[str, Any]:
    paths = _experiment_paths(output_root, exp["name"])
    cfg = build_experiment_config(base_config_path, exp, output_root)
    save_config(cfg, paths["config"])

    meta: dict[str, Any] = {
        "name": exp["name"],
        "description": exp["description"],
        "config_path": str(paths["config"]),
        "output_root": str(paths["root"]),
        "train_returncode": None,
        "eval_returncode": None,
        "status": "prepared",
    }

    if not execute:
        with open(paths["meta_json"], "w", encoding="utf-8") as f:
            json.dump(meta, f, indent=2, ensure_ascii=False)
        return meta

    train_script = REPO_ROOT / "scripts" / "train.py"
    eval_script = REPO_ROOT / "scripts" / "evaluate.py"
    best_ckpt = paths["checkpoints"] / "best_model.pt"

    if not skip_train:
        cmd = [
            sys.executable,
            str(train_script),
            "--config",
            str(paths["config"]),
            *extra_train_args,
        ]
        logging.info("Running: %s", " ".join(cmd))
        proc = subprocess.run(cmd, cwd=str(REPO_ROOT))
        meta["train_returncode"] = proc.returncode
        if proc.returncode != 0:
            meta["status"] = "train_failed"
            with open(paths["meta_json"], "w", encoding="utf-8") as f:
                json.dump(meta, f, indent=2, ensure_ascii=False)
            return meta
    else:
        meta["train_returncode"] = None

    if not skip_eval and best_ckpt.exists():
        cmd = [
            sys.executable,
            str(eval_script),
            "--config",
            str(paths["config"]),
            "--checkpoint",
            str(best_ckpt),
            "--output",
            str(paths["eval_json"]),
            *extra_eval_args,
        ]
        logging.info("Running: %s", " ".join(cmd))
        proc = subprocess.run(cmd, cwd=str(REPO_ROOT))
        meta["eval_returncode"] = proc.returncode
        if proc.returncode != 0:
            meta["status"] = "eval_failed"
        else:
            meta["status"] = "ok"
    elif not skip_eval:
        meta["eval_returncode"] = None
        meta["status"] = "eval_skipped_no_checkpoint"
        logging.warning("Checkpoint not found at %s; skipping evaluation", best_ckpt)
    else:
        meta["status"] = "train_only"

    with open(paths["meta_json"], "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2, ensure_ascii=False)

    return meta


def _load_eval_metrics(path: Path) -> dict[str, float]:
    if not path.exists():
        return {}
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    out: dict[str, float] = {}
    for k, v in data.items():
        if isinstance(v, (int, float)) and not isinstance(v, bool):
            out[k] = float(v)
    return out


def write_experiments_summary(output_root: Path, records: list[dict[str, Any]]) -> Path:
    """Write experiments_summary.json for visualize.py comparison plots."""
    rows: list[dict[str, Any]] = []
    for rec in records:
        name = rec.get("name")
        paths = _experiment_paths(output_root, name) if name else None
        metrics = _load_eval_metrics(paths["eval_json"]) if paths else {}
        rows.append({**rec, "metrics": metrics})

    out_path = output_root / "experiments_summary.json"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(rows, f, indent=2, ensure_ascii=False)
    return out_path


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="EasyTranslate experiment runner")
    p.add_argument("--config", type=str, default="configs/default_config.yaml", help="Base YAML config")
    p.add_argument("--output-root", type=str, default="outputs/experiments", help="Root directory for all runs")
    p.add_argument("--dry-run", action="store_true", help="Only write per-experiment config.yaml")
    p.add_argument("--execute", action="store_true", help="Run train.py then evaluate.py per experiment")
    p.add_argument("--skip-train", action="store_true", help="Evaluate only (expects best_model.pt)")
    p.add_argument("--skip-eval", action="store_true", help="Train only, no evaluation")
    p.add_argument("--only", type=str, default=None, help="Run a single experiment id (see EXPERIMENTS names)")
    p.add_argument(
        "extra",
        nargs="*",
        default=[],
        help="Extra key=value args forwarded to train.py / evaluate.py",
    )
    return p.parse_args()


def main() -> None:
    args = parse_args()
    logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")

    base = (REPO_ROOT / args.config).resolve()
    if not base.exists():
        raise FileNotFoundError(f"Base config not found: {base}")

    output_root = (REPO_ROOT / args.output_root).resolve()
    output_root.mkdir(parents=True, exist_ok=True)

    execute = bool(args.execute) and not bool(args.dry_run)
    if not args.dry_run and not args.execute and not args.skip_train:
        logging.info("Neither --execute nor --dry-run set; defaulting to dry-run (configs only).")
        args.dry_run = True

    selected = EXPERIMENTS
    if args.only:
        selected = [e for e in EXPERIMENTS if e["name"] == args.only]
        if not selected:
            raise ValueError(f"Unknown experiment {args.only!r}; choose one of {[e['name'] for e in EXPERIMENTS]}")

    print("=" * 60)
    print("  EasyTranslate - Experiment Runner")
    print("=" * 60)
    print(f"  Base config: {base}")
    print(f"  Output root: {output_root}")
    print(f"  Mode: {'dry-run' if args.dry_run else 'execute' if execute else 'custom'}")
    print("=" * 60)

    records: list[dict[str, Any]] = []
    extra = list(args.extra)

    for exp in selected:
        print(f"\n>>> {exp['name']}: {exp['description']}")
        meta = run_single_experiment(
            exp,
            base_config_path=base,
            output_root=output_root,
            execute=execute,
            skip_train=args.skip_train,
            skip_eval=args.skip_eval,
            extra_train_args=extra,
            extra_eval_args=extra,
        )
        records.append(meta)
        print(f"    status: {meta.get('status')}, meta: {meta.get('output_root')}/experiment_meta.json")

    summary_path = write_experiments_summary(output_root, records)
    print("\n" + "=" * 60)
    print(f"  Summary: {summary_path}")
    print(f"  Compare: python scripts/visualize.py --task comparison --results-dir {output_root}")
    print("=" * 60)


if __name__ == "__main__":
    main()