sweep.py

"""soup sweep — hyperparameter search over training configs."""

import itertools
import random
from pathlib import Path
from typing import Optional

import typer
from rich.console import Console
from rich.panel import Panel
from rich.table import Table

from soup_cli.config.loader import load_config

console = Console()


def sweep(
    config: str = typer.Option(
        "soup.yaml",
        "--config",
        "-c",
        help="Path to base soup.yaml config file",
    ),
    param: list[str] = typer.Option(
        ...,
        "--param",
        "-p",
        help="Parameter to sweep: key=val1,val2,val3 (e.g., lr=1e-5,2e-5,5e-5)",
    ),
    strategy: str = typer.Option(
        "grid",
        "--strategy",
        "-s",
        help="Search strategy: grid, random",
    ),
    max_runs: Optional[int] = typer.Option(
        None,
        "--max-runs",
        help="Max number of runs (useful for random strategy)",
    ),
    name: Optional[str] = typer.Option(
        None,
        "--name",
        "-n",
        help="Sweep experiment name prefix",
    ),
    dry_run: bool = typer.Option(
        False,
        "--dry-run",
        help="Show planned runs without executing",
    ),
    early_stop: float = typer.Option(
        None,
        "--early-stop",
        help="Stop early if run's loss exceeds best loss by this factor (e.g. 1.5 = 50% worse)",
    ),
    yes: bool = typer.Option(
        False,
        "--yes",
        "-y",
        help="Skip confirmation prompt",
    ),
):
    """Run hyperparameter sweep: grid or random search over training parameters."""
    config_path = Path(config)
    if not config_path.exists():
        console.print(f"[red]Config not found: {config_path}[/]")
        raise typer.Exit(1)

    if strategy not in ("grid", "random"):
        console.print(f"[red]Invalid strategy: {strategy}. Must be grid or random.[/]")
        raise typer.Exit(1)

    # Parse sweep parameters
    sweep_params = _parse_sweep_params(param)
    if not sweep_params:
        console.print("[red]No valid sweep parameters provided.[/]")
        raise typer.Exit(1)

    # Generate parameter combinations
    combinations = _generate_combinations(sweep_params, strategy, max_runs)

    # Validate before anything is printed (#642). --dry-run used to return
    # below without ever loading the config, so neither the loader's
    # unknown-key warning (#627) nor the sweep-parameter pre-flight (#628)
    # was reachable under it — a dry run whose job is catching mistakes
    # before a long run caught neither. Loading a config file executes
    # nothing, so both paths now validate at the same point and share the
    # single load.
    base_cfg = load_config(config_path)

    # Refuse the whole sweep before any arm starts — and before printing a
    # grid that can never run (#627, #642). The arm loop below wraps each run
    # in `except Exception`, so the guard inside `_run_single` would be
    # caught, recorded as a per-arm failure, and the command would still
    # exit 0 — an entirely invalid sweep that nothing downstream can detect.
    # Every combination carries the same parameter names, so one probe built
    # the way `_run_single` builds its config settles it for the grid.
    if combinations:
        probe = base_cfg.model_dump()
        for key, val in combinations[0].items():
            _set_nested_param(probe, key, val)
        try:
            _reject_unknown_sweep_params(probe)
        except ValueError as exc:
            console.print(f"[red]{exc}[/]")
            raise typer.Exit(1) from exc

    console.print(
        Panel(
            f"Config:   [bold]{config_path}[/]\n"
            f"Strategy: [bold]{strategy}[/]\n"
            f"Params:   [bold]{', '.join(sweep_params.keys())}[/]\n"
            f"Runs:     [bold]{len(combinations)}[/]",
            title="Sweep Plan",
        )
    )

    # Show parameter table
    param_table = Table(title="Parameter Grid")
    param_table.add_column("Run", style="bold")
    for key in sweep_params:
        param_table.add_column(key)

    for idx, combo in enumerate(combinations):
        row_values = [str(combo[key]) for key in sweep_params]
        param_table.add_row(f"#{idx + 1}", *row_values)

    console.print(param_table)

    if dry_run:
        console.print("[yellow]Dry run - no training will be executed.[/]")
        raise typer.Exit()

    if not yes:
        if not typer.confirm(f"Start {len(combinations)} training run(s)?", default=True):
            console.print("[yellow]Cancelled.[/]")
            raise typer.Exit()

    results = []
    best_loss = float("inf")
    skipped = 0

    for idx, combo in enumerate(combinations):
        run_name = f"{name or 'sweep'}_{idx + 1}"
        console.print(f"\n[bold]--- Run {idx + 1}/{len(combinations)}: {run_name} ---[/]")

        for key, val in combo.items():
            console.print(f"  {key} = {val}")

        try:
            result = _run_single(base_cfg, combo, run_name, config_path)
            final_loss = result.get("final_loss", 0)
            results.append({
                "name": run_name,
                "params": combo,
                "run_id": result.get("run_id", ""),
                "final_loss": final_loss,
                "duration": result.get("duration", ""),
                "status": "completed",
            })

            # Update best loss and check early stopping for remaining runs
            if final_loss and final_loss < best_loss:
                best_loss = final_loss

            if early_stop and final_loss and best_loss < float("inf"):
                if final_loss > best_loss * early_stop:
                    console.print(
                        f"[yellow]Loss {final_loss:.4f} exceeds threshold "
                        f"({best_loss:.4f} x {early_stop} = {best_loss * early_stop:.4f})[/]"
                    )
        except Exception as exc:
            console.print(f"[red]Run {run_name} failed: {exc}[/]")
            results.append({
                "name": run_name,
                "params": combo,
                "run_id": "",
                "final_loss": 0,
                "duration": "",
                "status": "failed",
            })

        # Early stopping: skip remaining runs if too many are poor
        if early_stop and len(results) >= 2:
            completed = [r for r in results if r["status"] == "completed" and r["final_loss"]]
            if completed:
                recent = completed[-1]
                if recent["final_loss"] > best_loss * early_stop:
                    remaining = len(combinations) - idx - 1
                    if remaining > 0:
                        skipped = remaining
                        console.print(
                            f"[yellow]Early stopping: skipping {remaining} remaining run(s). "
                            f"Last loss {recent['final_loss']:.4f} exceeded threshold.[/]"
                        )
                        break

    # Summary table
    _display_summary(results, sweep_params)

    if skipped:
        console.print(f"\n[yellow]Early stopping: {skipped} run(s) skipped.[/]")


def _parse_sweep_params(params: list[str]) -> dict[str, list]:
    """Parse sweep parameter strings into a dict of {key: [values]}."""
    result = {}
    for param_str in params:
        if "=" not in param_str:
            console.print(f"[yellow]Skipping invalid param: {param_str} (missing '=')[/]")
            continue

        key, values_str = param_str.split("=", 1)
        key = key.strip()
        values = []

        for val in values_str.split(","):
            val = val.strip()
            values.append(_parse_value(val))

        if values:
            result[key] = values

    return result


def _parse_value(val: str):
    """Parse a string value into the appropriate Python type."""
    # Bool
    if val.lower() in ("true", "false"):
        return val.lower() == "true"
    # None
    if val.lower() == "none":
        return None
    # Int
    try:
        return int(val)
    except ValueError:
        pass
    # Float (including scientific notation)
    try:
        return float(val)
    except ValueError:
        pass
    # String
    return val


def _generate_combinations(
    sweep_params: dict[str, list],
    strategy: str,
    max_runs: Optional[int],
) -> list[dict]:
    """Generate parameter combinations based on strategy."""
    keys = list(sweep_params.keys())
    value_lists = [sweep_params[k] for k in keys]

    if strategy == "grid":
        combos = [dict(zip(keys, vals)) for vals in itertools.product(*value_lists)]
    elif strategy == "random":
        total_possible = 1
        for vals in value_lists:
            total_possible *= len(vals)

        num_runs = max_runs or min(total_possible, 10)
        num_runs = min(num_runs, total_possible)

        if num_runs >= total_possible:
            # Just do all of them
            combos = [dict(zip(keys, vals)) for vals in itertools.product(*value_lists)]
        else:
            seen = set()
            combos = []
            while len(combos) < num_runs:
                vals = tuple(random.choice(vals_list) for vals_list in value_lists)
                if vals not in seen:
                    seen.add(vals)
                    combos.append(dict(zip(keys, vals)))
    else:
        combos = []

    if max_runs and len(combos) > max_runs:
        combos = combos[:max_runs]

    return combos


def _set_nested_param(config_dict: dict, key: str, value) -> dict:
    """Set a nested parameter in a config dict using dot notation.

    Supports keys like: lr, lora.r, training.epochs, etc.
    Maps common short names to their full paths.
    """
    # Short name mappings
    shortcuts = {
        "lr": "training.lr",
        "epochs": "training.epochs",
        "batch_size": "training.batch_size",
        "lora_r": "training.lora.r",
        "lora_alpha": "training.lora.alpha",
        "lora_dropout": "training.lora.dropout",
        "quantization": "training.quantization",
        "warmup_ratio": "training.warmup_ratio",
        "weight_decay": "training.weight_decay",
        "gradient_accumulation_steps": "training.gradient_accumulation_steps",
        "max_grad_norm": "training.max_grad_norm",
        "optimizer": "training.optimizer",
        "scheduler": "training.scheduler",
        "val_split": "data.val_split",
        "max_length": "data.max_length",
        "dpo_beta": "training.dpo_beta",
        "kto_beta": "training.kto_beta",
        "grpo_beta": "training.grpo_beta",
        "num_generations": "training.num_generations",
        "reward_fn": "training.reward_fn",
        "ppo_epochs": "training.ppo_epochs",
        "ppo_clip_ratio": "training.ppo_clip_ratio",
        "ppo_kl_penalty": "training.ppo_kl_penalty",
        "reward_model": "training.reward_model",
        "orpo_beta": "training.orpo_beta",
        "simpo_gamma": "training.simpo_gamma",
        "cpo_alpha": "training.cpo_alpha",
        "ipo_tau": "training.ipo_tau",
        "bco_beta": "training.bco_beta",
        "loraplus_lr_ratio": "training.loraplus_lr_ratio",
        "use_dora": "training.lora.use_dora",
        "use_galore": "training.use_galore",
        "galore_rank": "training.galore_rank",
        "moe_lora": "training.moe_lora",
        "moe_aux_loss_coeff": "training.moe_aux_loss_coeff",
        "embedding_loss": "training.embedding_loss",
        "embedding_margin": "training.embedding_margin",
        "embedding_pooling": "training.embedding_pooling",
        "embedding_temperature": "training.embedding_temperature",
        "neftune_alpha": "training.neftune_alpha",
        "use_rslora": "training.lora.use_rslora",
        "backend": "backend",
    }

    full_key = shortcuts.get(key, key)
    parts = full_key.split(".")

    obj = config_dict
    for part in parts[:-1]:
        if part not in obj:
            obj[part] = {}
        obj = obj[part]
    obj[parts[-1]] = value

    return config_dict


def _reject_unknown_sweep_params(config_dict: dict) -> None:
    """Refuse a sweep whose parameter names no config field (#627).

    ``config_dict`` starts from a validated ``model_dump()``, so anything the
    schema cannot place got there from a ``--param`` name. Dropping it silently
    would run the whole grid with the swept knob never applied, producing arms
    that are all identical and a winner that means nothing -- so this raises
    regardless of the loader's severity switch, and carries no deadline: there
    is no partially-useful result to preserve by continuing.

    Kept out of :func:`_run_single` so it is reachable without importing the
    training stack, and so removing it fails a test rather than a review.
    """
    from soup_cli.config.unknown_keys import find_unknown_config_keys, format_unknown_keys

    unknown = find_unknown_config_keys(config_dict)
    if unknown:
        detail = format_unknown_keys(unknown, include_deadline=False)
        raise ValueError(f"sweep parameter does not match any config field: {detail}")


def _run_single(base_cfg, params: dict, run_name: str, config_path: Path) -> dict:
    """Run a single training with modified parameters."""
    # Deep copy and modify config
    config_dict = base_cfg.model_dump()
    for key, val in params.items():
        _set_nested_param(config_dict, key, val)

    # Override experiment name
    config_dict["experiment_name"] = run_name

    # Before the heavy imports, so this refusal is reachable -- and testable --
    # without the training stack. `sweep()` pre-checks the grid too; this stays
    # so a direct caller cannot bypass it.
    _reject_unknown_sweep_params(config_dict)

    from soup_cli.config.schema import SoupConfig
    from soup_cli.data.loader import load_dataset
    from soup_cli.experiment.tracker import ExperimentTracker
    from soup_cli.monitoring.display import TrainingDisplay
    from soup_cli.trainer.sft import SFTTrainerWrapper
    from soup_cli.utils.gpu import detect_device, get_gpu_info

    cfg = SoupConfig(**config_dict)

    # Detect hardware
    device, device_name = detect_device()
    gpu_info = get_gpu_info()

    # Load data
    dataset = load_dataset(
        cfg.data,
        preserve_source_columns=cfg.task == "grpo",
    )
    console.print(f"[dim]Loaded {len(dataset['train'])} train samples[/]")

    # Start tracking
    tracker = ExperimentTracker()
    run_id = tracker.start_run(
        config_dict=cfg.model_dump(),
        device=device,
        device_name=device_name,
        gpu_info=gpu_info,
        experiment_name=run_name,
    )

    # Build trainer
    if cfg.task == "dpo":
        from soup_cli.trainer.dpo import DPOTrainerWrapper

        trainer_wrapper = DPOTrainerWrapper(cfg, device=device)
    elif cfg.task == "kto":
        from soup_cli.trainer.kto import KTOTrainerWrapper

        trainer_wrapper = KTOTrainerWrapper(cfg, device=device)
    elif cfg.task == "grpo":
        from soup_cli.trainer.grpo import GRPOTrainerWrapper

        trainer_wrapper = GRPOTrainerWrapper(cfg, device=device)
    elif cfg.task == "ppo":
        from soup_cli.trainer.ppo import PPOTrainerWrapper

        trainer_wrapper = PPOTrainerWrapper(cfg, device=device)
    elif cfg.task == "orpo":
        from soup_cli.trainer.orpo import ORPOTrainerWrapper

        trainer_wrapper = ORPOTrainerWrapper(cfg, device=device)
    elif cfg.task == "simpo":
        from soup_cli.trainer.simpo import SimPOTrainerWrapper

        trainer_wrapper = SimPOTrainerWrapper(cfg, device=device)
    elif cfg.task == "ipo":
        from soup_cli.trainer.ipo import IPOTrainerWrapper

        trainer_wrapper = IPOTrainerWrapper(cfg, device=device)
    elif cfg.task == "bco":
        from soup_cli.trainer.bco import BCOTrainerWrapper

        trainer_wrapper = BCOTrainerWrapper(cfg, device=device)
    elif cfg.task == "preference":
        from soup_cli.trainer.preference import PreferenceTrainerWrapper

        trainer_wrapper = PreferenceTrainerWrapper(cfg, device=device)
    elif cfg.task == "reward_model":
        from soup_cli.trainer.reward_model import RewardModelTrainerWrapper

        trainer_wrapper = RewardModelTrainerWrapper(cfg, device=device)
    elif cfg.task == "pretrain":
        from soup_cli.trainer.pretrain import PretrainTrainerWrapper

        trainer_wrapper = PretrainTrainerWrapper(cfg, device=device)
    elif cfg.task == "embedding":
        from soup_cli.trainer.embedding import EmbeddingTrainerWrapper

        trainer_wrapper = EmbeddingTrainerWrapper(cfg, device=device)
    else:
        trainer_wrapper = SFTTrainerWrapper(cfg, device=device)
    trainer_wrapper.setup(dataset)

    # Train
    display = TrainingDisplay(cfg, device_name=device_name)
    try:
        result = trainer_wrapper.train(display=display, tracker=tracker, run_id=run_id)
        tracker.finish_run(
            run_id=run_id,
            initial_loss=result["initial_loss"],
            final_loss=result["final_loss"],
            total_steps=result["total_steps"],
            duration_secs=result["duration_secs"],
            output_dir=result["output_dir"],
        )
        result["run_id"] = run_id
        return result
    except Exception:
        tracker.fail_run(run_id)
        raise


def _display_summary(results: list[dict], sweep_params: dict[str, list]):
    """Display sweep results summary table."""
    table = Table(title="Sweep Results")
    table.add_column("Run", style="bold")
    for key in sweep_params:
        table.add_column(key)
    table.add_column("Final Loss", justify="right", style="green")
    table.add_column("Duration", justify="right")
    table.add_column("Status")

    # Sort by final loss (best first)
    sorted_results = sorted(results, key=lambda r: r.get("final_loss", float("inf")))

    for idx, res in enumerate(sorted_results):
        status_style = "green" if res["status"] == "completed" else "red"
        param_vals = [str(res["params"].get(k, "")) for k in sweep_params]
        loss_str = f"{res['final_loss']:.4f}" if res["final_loss"] else "-"
        best_marker = " [bold yellow]*[/]" if idx == 0 and res["status"] == "completed" else ""
        table.add_row(
            res["name"],
            *param_vals,
            f"{loss_str}{best_marker}",
            res.get("duration", "-"),
            f"[{status_style}]{res['status']}[/]",
        )

    console.print(table)

    # Best run
    completed = [r for r in sorted_results if r["status"] == "completed"]
    if completed:
        best = completed[0]
        console.print(
            f"\n[bold green]Best run:[/] {best['name']} "
            f"(loss: {best['final_loss']:.4f})"
        )
        for key, val in best["params"].items():
            console.print(f"  {key} = {val}")
        if best.get("run_id"):
            console.print(f"\n[dim]View details: soup runs show {best['run_id']}[/]")