code/scripts/launch_alpha_grid.py

"""Run the all-block alpha grid sequentially on a configured CUDA device."""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
from pathlib import Path

GPU_DEVICES = {1: "0", 2: "1"}
MODEL_PATH_ENV = {
    "llama": "MFFLORA_LLAMA_MODEL_PATH",
    "qwen3": "MFFLORA_QWEN3_MODEL_PATH",
    "ministral": "MFFLORA_MINISTRAL_MODEL_PATH",
    "gemma4": "MFFLORA_GEMMA4_MODEL_PATH",
}

BLOCKS = [
    ("qwen3", "xnli"),
    ("ministral", "xnli"),
    ("gemma4", "xnli"),
    ("llama", "rte"),
    ("qwen3", "rte"),
    ("ministral", "rte"),
    ("gemma4", "rte"),
]
METHODS = ["fpeft_low", "fpeft_bi_high", "random_orthogonal", "peft_scaled_random"]
ALPHAS = [0.25, 0.5, 1.0, 2.0]
SEEDS = [42, 1337, 2024, 7, 123]
EXPECTED_CELLS_PER_JOB = len(METHODS) * len(SEEDS)


def _block_name(model: str, task: str) -> str:
    return f"{model}_{task}"


def _alpha_tag(alpha: float) -> str:
    return f"alpha{alpha:g}"


def _commands(
    repo: Path,
    output_root: Path,
    seeds: list[int] = SEEDS,
    alphas: list[float] = ALPHAS,
) -> dict[str, list[str]]:
    commands = {}
    for alpha in alphas:
        for model, task in BLOCKS:
            block = _block_name(model, task)
            job = f"{_alpha_tag(alpha)}__{block}"
            output_dir = output_root / _alpha_tag(alpha) / "comparison" / block
            command = [
                sys.executable,
                str(repo / "experiments/broad_sweep/run.py"),
                "--model-id", model,
                "--dataset", task,
                "--methods", *METHODS,
                "--ranks", "32",
                "--seeds", *(str(seed) for seed in seeds),
                "--alpha-fpeft", str(alpha),
                "--alpha-baseline", str(alpha),
                "--output-dir", str(output_dir),
                "--no-wandb",
            ]
            if model == "llama" and task == "rte":
                command.extend([
                    "--batch-size", "1", "--grad-accum", "16", "--max-length", "128",
                    "--gradient-checkpointing",
                ])
            model_path = os.environ.get(MODEL_PATH_ENV[model])
            if model_path and Path(model_path).is_dir():
                command[command.index("--dataset"):command.index("--dataset")] = [
                    "--model-path", model_path
                ]
            commands[job] = command
    return commands


def _inspect(path: Path, seeds: list[int] = SEEDS) -> dict[str, int | bool]:
    result = {
        "complete": 0,
        "errors": 0,
        "retryable_errors": 0,
        "non_retryable_errors": 0,
        "constant": 0,
        "total": 0,
        "finished": False,
    }
    if not path.exists():
        return result
    data = json.loads(path.read_text())
    result["total"] = len(data)
    for method in METHODS:
        for seed in seeds:
            value = data.get(f"{method}__rank32__seed{seed}", {})
            if value.get("status") == "complete":
                result["complete"] += 1
                distribution = value.get("prediction_distribution")
                if distribution and max(distribution) == sum(distribution):
                    result["constant"] += 1
            elif value.get("status") == "error":
                result["errors"] += 1
                if value.get("retryable") is True:
                    result["retryable_errors"] += 1
                else:
                    result["non_retryable_errors"] += 1
    result["finished"] = result["complete"] == len(METHODS) * len(seeds)
    return result


def _should_retry(state: dict[str, int | bool]) -> bool:
    return bool(state["errors"]) and state["retryable_errors"] == state["errors"]


def _can_resume(state: dict[str, int | bool]) -> bool:
    return state["non_retryable_errors"] == 0


def _write_json(path: Path, value: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(value, indent=2) + "\n")
    temporary.replace(path)


def _gpu_state() -> str:
    result = subprocess.run(
        ["nvidia-smi", "--query-gpu=index,memory.used,utilization.gpu", "--format=csv,noheader"],
        capture_output=True,
        text=True,
        check=False,
    )
    return result.stdout.strip()


def _heartbeat(
    path: Path,
    *,
    alpha: float | None,
    block: str | None,
    pid: int | None,
    pending: list[str],
    done: list[str],
    attempts: dict[str, int],
    state: dict | None = None,
    gpu: int = 1,
    gpu_device: str = "0",
) -> None:
    record = {
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "gpu": gpu,
        "gpu_device": gpu_device,
        "alpha": alpha,
        "block": block,
        "pid": pid,
        "pending": pending,
        "done": done,
        "attempts": attempts,
        "gpu_state": _gpu_state(),
    }
    if state is not None:
        record["cell_state"] = state
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record) + "\n")


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo", default=".")
    parser.add_argument("--gpu", type=int, default=1)
    parser.add_argument("--output-root", default="outputs/paper_upgrade_alpha_grid_20260720")
    parser.add_argument("--heartbeat-seconds", type=int, default=600)
    parser.add_argument("--poll-seconds", type=int, default=30)
    parser.add_argument("--max-attempts", type=int, default=3)
    parser.add_argument("--budget-hours", type=float, default=120.0)
    parser.add_argument("--seeds", nargs="+", type=int, default=SEEDS)
    parser.add_argument("--alphas", nargs="+", type=float, default=ALPHAS)
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = _parse_args()
    if args.gpu not in GPU_DEVICES:
        raise ValueError("alpha-grid launcher supports configured CUDA devices 1 or 2")
    gpu_device = GPU_DEVICES[args.gpu]
    repo = Path(args.repo).resolve()
    output_root = repo / args.output_root
    output_root.mkdir(parents=True, exist_ok=True)
    commands = _commands(output_root=output_root, repo=repo, seeds=args.seeds, alphas=args.alphas)
    if args.dry_run:
        for job, command in commands.items():
            print(job, " ".join(command))
        return

    manifest_path = output_root / "manifest.json"
    watchdog_path = output_root / "watchdog.jsonl"
    manifest = json.loads(manifest_path.read_text()) if manifest_path.exists() else {
        "gpu": args.gpu,
        "gpu_device": gpu_device,
        "mode": "sequential",
        "methods": METHODS,
        "alphas": args.alphas,
        "seeds": args.seeds,
        "jobs": {},
    }
    manifest["gpu"] = args.gpu
    manifest["gpu_device"] = gpu_device
    _write_json(manifest_path, manifest)
    started = time.monotonic()
    done: list[str] = []

    for job, command in commands.items():
        output_dir = Path(command[command.index("--output-dir") + 1])
        results_path = output_dir / "results.json"
        state = _inspect(results_path, args.seeds)
        if state["finished"]:
            done.append(job)
            manifest["jobs"][job] = {"status": "complete", **state}
            _write_json(manifest_path, manifest)
            continue
        if results_path.exists() and not _can_resume(state):
            raise RuntimeError(f"job has non-retryable cell errors: {job}")

        for attempt in range(1, args.max_attempts + 1):
            if time.monotonic() - started >= args.budget_hours * 3600:
                raise TimeoutError("alpha-grid GPU-hour budget exhausted")
            log_path = output_root / "logs" / f"{job}.attempt{attempt}.log"
            log_path.parent.mkdir(parents=True, exist_ok=True)
            env = os.environ.copy()
            env["CUDA_VISIBLE_DEVICES"] = gpu_device
            env["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
            manifest["jobs"][job] = {"status": "running", "attempt": attempt}
            _write_json(manifest_path, manifest)
            with log_path.open("w", encoding="utf-8") as log:
                process = subprocess.Popen(command, cwd=repo, env=env, stdout=log, stderr=subprocess.STDOUT)
                last_heartbeat = 0.0
                while process.poll() is None:
                    state = _inspect(results_path, args.seeds)
                    if state["constant"] >= 3:
                        process.terminate()
                        process.wait()
                        raise RuntimeError(f"constant-class stop condition in {job}")
                    now = time.monotonic()
                    if now - started >= args.budget_hours * 3600:
                        process.terminate()
                        process.wait()
                        raise TimeoutError("alpha-grid GPU-hour budget exhausted")
                    if now - last_heartbeat >= args.heartbeat_seconds:
                        alpha = float(job.split("__", 1)[0].removeprefix("alpha"))
                        block = job.split("__", 1)[1]
                        _heartbeat(
                            watchdog_path,
                            alpha=alpha,
                            block=block,
                            pid=process.pid,
                            pending=[name for name in commands if name not in done and name != job],
                            done=done,
                            attempts={name: manifest["jobs"].get(name, {}).get("attempt", 0) for name in commands},
                            state=state,
                            gpu=args.gpu,
                            gpu_device=gpu_device,
                        )
                        last_heartbeat = now
                    time.sleep(args.poll_seconds)
            state = _inspect(results_path, args.seeds)
            if state["constant"] >= 3:
                raise RuntimeError(f"constant-class stop condition in {job}")
            if state["finished"]:
                done.append(job)
                manifest["jobs"][job] = {"status": "complete", "attempt": attempt, **state}
                _write_json(manifest_path, manifest)
                break
            if not _can_resume(state):
                raise RuntimeError(f"job failed without retryable cell errors: {job}")
            status = "retrying" if _should_retry(state) else "resuming"
            manifest["jobs"][job] = {"status": status, "attempt": attempt, **state}
            _write_json(manifest_path, manifest)
            if attempt < args.max_attempts:
                time.sleep(2 ** (attempt - 1))
        else:
            raise RuntimeError(f"job failed after {args.max_attempts} attempts: {job}")

    _heartbeat(
        watchdog_path,
        alpha=None,
        block=None,
        pid=None,
        pending=[],
        done=done,
        attempts={name: manifest["jobs"].get(name, {}).get("attempt", 0) for name in commands},
        gpu=args.gpu,
        gpu_device=gpu_device,
    )
    print(json.dumps({"done": done, "cells": len(done) * EXPECTED_CELLS_PER_JOB}, indent=2))


if __name__ == "__main__":
    main()