experiments/run_weekend_experiments.py

import argparse
import copy
import csv
import json
import logging
import subprocess
import sys
from datetime import datetime
from pathlib import Path

import yaml

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from train.train import compute_split_baselines, list_split_batches


logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")


def _load_yaml(path):
    with open(path, "r", encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def _write_yaml(path, payload):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as handle:
        yaml.safe_dump(payload, handle, sort_keys=False)


def _deep_update(target, updates):
    for key, value in updates.items():
        if isinstance(value, dict) and isinstance(target.get(key), dict):
            _deep_update(target[key], value)
        else:
            target[key] = value


def _slugify(name):
    return "".join(ch if ch.isalnum() or ch in "-_" else "_" for ch in name)


def _to_jsonable(value):
    if isinstance(value, dict):
        return {key: _to_jsonable(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [_to_jsonable(item) for item in value]
    if hasattr(value, "tolist"):
        return value.tolist()
    return value


def _torchrun_path():
    candidate = Path(sys.executable).with_name("torchrun")
    if candidate.exists():
        return str(candidate)
    return "torchrun"


def _run_command(command, log_path, cwd):
    log_path.parent.mkdir(parents=True, exist_ok=True)
    logging.info("Running: %s", " ".join(command))
    with open(log_path, "w", encoding="utf-8") as handle:
        process = subprocess.run(command, cwd=cwd, stdout=handle, stderr=subprocess.STDOUT, text=True)
    return process.returncode


def _find_run_dir(prefix_path):
    parent = prefix_path.parent
    pattern = prefix_path.name + "_*"
    candidates = [path for path in parent.glob(pattern) if path.is_dir()]
    if not candidates:
        raise FileNotFoundError(f"No run directory found for prefix {str(prefix_path)}")
    candidates.sort(key=lambda path: path.stat().st_mtime, reverse=True)
    return candidates[0]


def _write_summary_csv(path, rows):
    path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "experiment",
        "seed",
        "status",
        "run_dir",
        "weights_path",
        "model_miou",
        "model_mf1",
        "majority_miou",
        "majority_mf1",
        "anyview_miou",
        "anyview_mf1",
        "metrics_path",
        "log_train",
        "log_test",
    ]
    with open(path, "w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow({key: row.get(key, "") for key in fieldnames})


def build_experiment_config(base_config, experiment, seed, output_root, defaults=None):
    cfg = copy.deepcopy(base_config)
    _deep_update(cfg, experiment.get("overrides", {}))
    cfg["training"]["seed"] = int(seed)
    cfg.setdefault("data", {})
    cfg["data"].setdefault("split_roots", {})
    defaults = defaults or {}
    for split_name in ("train", "val", "test"):
        split_key = f"{split_name}_batches_root"
        split_root = experiment.get(split_key, defaults.get(split_key))
        if split_root is None:
            split_root = base_config.get("data", {}).get("split_roots", {}).get(split_name)
        if split_root:
            cfg["data"]["split_roots"][split_name] = str(split_root)
    # Weekend sweeps already compute split baselines once globally and again at test time.
    # Disabling per-train baseline recomputation avoids redundant DDP startup work.
    cfg.setdefault("training", {})
    cfg["training"]["compute_validation_baselines"] = bool(
        experiment.get("compute_validation_baselines", False)
    )
    run_prefix = output_root / "runs" / f"{_slugify(experiment['name'])}_seed{seed}"
    cfg["training"]["output_dir"] = str(run_prefix)
    return cfg, run_prefix


def build_train_command(config_path, nproc_per_node):
    if int(nproc_per_node) > 1:
        return [
            _torchrun_path(),
            "--standalone",
            f"--nproc_per_node={int(nproc_per_node)}",
            "main.py",
            "--config",
            str(config_path),
            "--mode",
            "train",
        ]
    return [sys.executable, "main.py", "--config", str(config_path), "--mode", "train"]


def build_test_command(config_path, weights_path, split_name, metrics_output, limit_files=None):
    command = [
        sys.executable,
        "main.py",
        "--config",
        str(config_path),
        "--mode",
        "test",
        "--weights_path",
        str(weights_path),
        "--split",
        split_name,
        "--metrics_output",
        str(metrics_output),
    ]
    if limit_files is not None:
        command.extend(["--limit_files", str(int(limit_files))])
    return command


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-config", default="configs/config_deepchoice_base.yaml")
    parser.add_argument("--plan", default="experiments/weekend_plan.yaml")
    parser.add_argument("--output-root", default=None)
    parser.add_argument("--continue-on-error", action="store_true")
    args = parser.parse_args()

    repo_root = REPO_ROOT
    base_config = _load_yaml(repo_root / args.base_config) if not Path(args.base_config).is_absolute() else _load_yaml(args.base_config)
    plan = _load_yaml(repo_root / args.plan) if not Path(args.plan).is_absolute() else _load_yaml(args.plan)

    output_root = Path(args.output_root) if args.output_root else repo_root / "artifacts" / "experiments" / datetime.now().strftime("weekend_%Y%m%d_%H%M%S")
    output_root.mkdir(parents=True, exist_ok=True)

    defaults = plan.get("defaults", {})
    summary_rows = []
    baselines = None

    baseline_cfg = copy.deepcopy(base_config)
    baseline_eval = plan.get("baseline_evaluation", {})
    if baseline_eval.get("enabled", False):
        split_name = baseline_eval.get("split", defaults.get("split", "test"))
        try:
            _deep_update(baseline_cfg, baseline_eval.get("overrides", {}))
            baseline_cfg.setdefault("data", {})
            baseline_cfg["data"].setdefault("split_roots", {})
            for split_key in ("train", "val", "test"):
                default_root = defaults.get(f"{split_key}_batches_root")
                if default_root:
                    baseline_cfg["data"]["split_roots"][split_key] = str(default_root)
            baseline_limit = baseline_eval.get("limit_files")
            baseline_paths = list_split_batches(baseline_cfg, split_name, limit=baseline_limit)
            baselines = compute_split_baselines(
                baseline_cfg,
                paths=baseline_paths,
                file_batch_size=baseline_cfg.get("test", {}).get("file_batch_size", baseline_cfg["training"].get("eval_file_batch_size", 1)),
                desc=f"Computing {split_name} baselines",
            )
            baseline_path = output_root / "metrics" / f"baselines_{split_name}.json"
            baseline_path.parent.mkdir(parents=True, exist_ok=True)
            baseline_path.write_text(json.dumps(_to_jsonable(baselines), indent=2), encoding="utf-8")
            summary_rows.append(
                {
                    "experiment": f"baseline_{split_name}",
                    "seed": "",
                    "status": "ok",
                    "run_dir": "",
                    "weights_path": "",
                    "model_miou": "",
                    "model_mf1": "",
                    "majority_miou": baselines["majority"]["miou"],
                    "majority_mf1": baselines["majority"]["mf1"],
                    "anyview_miou": baselines["anyview"]["miou"],
                    "anyview_mf1": baselines["anyview"]["mf1"],
                    "metrics_path": str(baseline_path),
                    "log_train": "",
                    "log_test": "",
                }
            )
        except Exception as exc:
            logging.exception("Baseline evaluation failed on split %s", split_name)
            summary_rows.append(
                {
                    "experiment": f"baseline_{split_name}",
                    "seed": "",
                    "status": f"failed: {exc}",
                }
            )
            if not args.continue_on_error:
                raise

    for experiment in plan.get("experiments", []):
        exp_name = experiment["name"]
        split_name = experiment.get("split", defaults.get("split", "test"))
        test_limit_files = experiment.get("limit_files", defaults.get("limit_files"))
        nproc = experiment.get("nproc_per_node", defaults.get("nproc_per_node", 1))
        run_test_after_train = experiment.get("run_test_after_train", defaults.get("run_test_after_train", True))

        for seed in experiment.get("seeds", [base_config["training"].get("seed", 42)]):
            cfg, run_prefix = build_experiment_config(base_config, experiment, seed, output_root, defaults=defaults)
            config_path = output_root / "configs" / f"{_slugify(exp_name)}__seed{seed}.yaml"
            train_log = output_root / "logs" / f"{_slugify(exp_name)}__seed{seed}__train.log"
            test_log = output_root / "logs" / f"{_slugify(exp_name)}__seed{seed}__test.log"
            metrics_path = output_root / "metrics" / f"{_slugify(exp_name)}__seed{seed}__test.json"
            _write_yaml(config_path, cfg)

            row = {
                "experiment": exp_name,
                "seed": seed,
                "status": "pending",
                "run_dir": "",
                "weights_path": "",
                "metrics_path": str(metrics_path),
                "log_train": str(train_log),
                "log_test": str(test_log),
            }

            try:
                if (
                    baselines is not None
                    and baseline_eval.get("enabled", False)
                    and split_name == baseline_eval.get("split", defaults.get("split", "test"))
                ):
                    cfg.setdefault("training", {})
                    cfg["training"]["precomputed_validation_baselines"] = _to_jsonable(baselines)
                    _write_yaml(config_path, cfg)
                train_cmd = build_train_command(config_path, nproc)
                train_rc = _run_command(train_cmd, train_log, repo_root)
                if train_rc != 0:
                    raise RuntimeError(f"Training failed with return code {train_rc}")

                run_dir = _find_run_dir(run_prefix)
                weights_path = run_dir / "best_model.pt"
                if not weights_path.exists():
                    raise FileNotFoundError(f"Missing checkpoint {str(weights_path)}")

                row["run_dir"] = str(run_dir)
                row["weights_path"] = str(weights_path)

                if run_test_after_train:
                    test_cmd = build_test_command(config_path, weights_path, split_name, metrics_path, limit_files=test_limit_files)
                    test_rc = _run_command(test_cmd, test_log, repo_root)
                    if test_rc != 0:
                        raise RuntimeError(f"Test failed with return code {test_rc}")
                    metrics = json.loads(metrics_path.read_text(encoding="utf-8"))
                    row["model_miou"] = metrics["miou"]
                    row["model_mf1"] = metrics["mf1"]
                    row["majority_miou"] = metrics["baselines"]["majority"]["miou"]
                    row["majority_mf1"] = metrics["baselines"]["majority"]["mf1"]
                    row["anyview_miou"] = metrics["baselines"]["anyview"]["miou"]
                    row["anyview_mf1"] = metrics["baselines"]["anyview"]["mf1"]

                row["status"] = "ok"
            except Exception as exc:
                logging.exception("Experiment failed: %s seed=%s", exp_name, seed)
                row["status"] = f"failed: {exc}"
                summary_rows.append(row)
                _write_summary_csv(output_root / "summary.csv", summary_rows)
                (output_root / "summary.json").write_text(json.dumps(_to_jsonable(summary_rows), indent=2), encoding="utf-8")
                if not args.continue_on_error:
                    raise
                continue

            summary_rows.append(row)
            _write_summary_csv(output_root / "summary.csv", summary_rows)
            (output_root / "summary.json").write_text(json.dumps(_to_jsonable(summary_rows), indent=2), encoding="utf-8")

    logging.info("Weekend experiment run complete. Summary written under %s", output_root)


if __name__ == "__main__":
    main()