src/ahdcma/cli/run_task.py

"""Single LoRA-task run via any optimizer.

Usage::

    python -m ahdcma.cli.run_task \\
        --algo ahdcma --task cifar100_vit --seed 0 \\
        --pop 8 --gens 5 --num-steps 50 --output outputs/runs/lora/

Writes ``{output}/{run_id}/result.json`` with the best LoRA config,
final fitness, history, and config snapshot. Each fitness call is a
short LoRA fine-tune (``num_steps`` optimiser steps) — the wall-time
per call dominates total runtime.
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path
from typing import Any

import numpy as np
import yaml

from ahdcma.algorithms.base import SearchSpace
from ahdcma.cli.run_benchmark import ALGO_REGISTRY
from ahdcma.fitness.lora_finetune import TaskConfig, lora_fitness
from ahdcma.search_space.encoder import DIM, DIM_NAMES
from ahdcma.utils.logging import make_run_id, setup_logging
from ahdcma.utils.seed import set_global_seed


def _load_task_config(task_name: str) -> TaskConfig:
    cfg_path = Path(__file__).resolve().parents[3] / "configs" / "tasks" / f"{task_name}.yaml"
    with cfg_path.open() as f:
        raw = yaml.safe_load(f)
    return TaskConfig(**raw)


def _load_algo_config(algo: str, *, seed: int, pop: int, gens: int) -> dict[str, Any]:
    cfg_path = Path(__file__).resolve().parents[3] / "configs" / "algo" / f"{algo}.yaml"
    with cfg_path.open() as f:
        cfg = yaml.safe_load(f)
    cfg["seed"] = seed
    cfg["population_size"] = pop
    cfg["max_generations"] = gens
    return dict(cfg)


def run_lora_task(
    algo: str,
    task_name: str,
    seed: int,
    *,
    pop: int = 8,
    gens: int = 10,
    num_steps: int = 100,
    output_dir: Path | str = "outputs/runs/lora",
) -> dict[str, Any]:
    """Run one (algo, task, seed) LoRA tuning job and write result.json."""
    if algo not in ALGO_REGISTRY:
        raise KeyError(f"unknown algorithm {algo!r}")

    set_global_seed(seed)
    run_id = make_run_id(algo, task_name, seed)
    out_root = Path(output_dir) / run_id
    out_root.mkdir(parents=True, exist_ok=True)
    setup_logging(run_id, log_dir=out_root / "logs")

    task = _load_task_config(task_name)
    algo_cfg = _load_algo_config(algo, seed=seed, pop=pop, gens=gens)

    search_space = SearchSpace.unit_cube(DIM, names=DIM_NAMES)

    eval_log: list[dict[str, Any]] = []

    def fitness(x: np.ndarray[Any, np.dtype[np.float64]]) -> float:
        f = lora_fitness(x, task, num_steps=num_steps, seed=seed)
        eval_log.append({"x": x.tolist(), "f": f})
        return f

    cls = ALGO_REGISTRY[algo]
    opt = cls(algo_cfg, fitness, search_space, run_id=run_id)

    t0 = time.time()
    result = opt.optimize()
    wall = time.time() - t0

    summary: dict[str, Any] = {
        "run_id": run_id,
        "algo": algo,
        "task": task_name,
        "seed": seed,
        "pop": pop,
        "gens": gens,
        "num_steps": num_steps,
        "best_f": float(result.best_f),
        "best_x": result.best_x.tolist(),
        "best_accuracy": float(-result.best_f),
        "wall_time": wall,
        "n_generations": len(result.history),
        "n_evals": len(eval_log),
        "config": algo_cfg,
        "best_fitness_curve": list(result.history.best_fitness),
        "mode_curve": list(result.history.mode_per_gen),
    }
    (out_root / "result.json").write_text(json.dumps(summary, indent=2))
    (out_root / "evaluations.json").write_text(json.dumps(eval_log, indent=2))
    return summary


def main() -> None:
    p = argparse.ArgumentParser(description="Run a single LoRA tuning job.")
    p.add_argument("--algo", required=True, choices=sorted(ALGO_REGISTRY))
    p.add_argument("--task", required=True)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--pop", type=int, default=8)
    p.add_argument("--gens", type=int, default=10)
    p.add_argument("--num-steps", type=int, default=100)
    p.add_argument("--output", default="outputs/runs/lora")
    args = p.parse_args()
    summary = run_lora_task(
        args.algo,
        args.task,
        args.seed,
        pop=args.pop,
        gens=args.gens,
        num_steps=args.num_steps,
        output_dir=args.output,
    )
    print(
        json.dumps(
            {
                "run_id": summary["run_id"],
                "best_accuracy": summary["best_accuracy"],
                "wall_time": summary["wall_time"],
            }
        )
    )


if __name__ == "__main__":
    main()