run-adasearch.py

from __future__ import annotations

import argparse
import json
import csv
import os
import pathlib
import time
from copy import deepcopy
from typing import Any, Dict, List, Tuple

import torch
import yaml  
from concurrent.futures import ProcessPoolExecutor, as_completed
from termcolor import cprint
from utils.prompts import build_prompt


def _format_duration(seconds: float) -> str:
    """Format a duration in seconds into a human-readable H:M:S string."""
    if seconds < 0:
        return "0s"
    m, s = divmod(seconds, 60)
    h, m = divmod(m, 60)
    return f"{int(h)}h {int(m)}m {s:.2f}s"


def _flatten_dataset(ds: Any, format_identifier: str) -> List[str]:
    """Take a loaded dataset and return a flat list of strings based on a format identifier."""
    prompts: List[str] = []
    if format_identifier == "hh-rlhf":
        for bucket in ds.values():
            for item in bucket:
                if isinstance(item, dict) and "question" in item:
                    prompts.append(item["question"].strip())
    elif format_identifier == "imdb":
        import re

        for row in ds:
            raw_review = row.get("review", "").strip()
            if not raw_review:
                continue

            cleaned = re.sub(r"<br\s*/?>", " ", raw_review, flags=re.IGNORECASE)

            first_words = " ".join(cleaned.split()[:25])

            prompt_text = (
                f"The following is a movie review. Continue writing the review: {first_words} "
            )

            prompts.append(prompt_text)

    elif format_identifier == "openai-gsm":
        for row in ds:
            q = row.get("question", "").strip()
            if q:
                prompts.append(q)

    elif format_identifier == "harmfulqa":
        if isinstance(ds, dict) and "toxic-conversations" in ds:
            for item in ds["toxic-conversations"]:
                if isinstance(item, dict) and "question" in item:
                    prompts.append(item["question"].strip())
        elif isinstance(ds, list):
            for item in ds:
                if isinstance(item, dict) and "question" in item:
                    prompts.append(item["question"].strip())
    else:
        raise ValueError(f"Unsupported dataset format identifier: {format_identifier}")
    return prompts


def _ensure_parent_dir(path: pathlib.Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)


def _run_llm_group(group_idx: int, cfgs: List[Dict[str, Any]], dataset: List[str], n_gpus: int) -> List[Tuple[str, str]]:
    if n_gpus > 0:
        cuda_idx = group_idx % n_gpus 
        visible_str = os.environ.get("CUDA_VISIBLE_DEVICES")
        if visible_str and visible_str.strip():
            visible_ids = [tok.strip() for tok in visible_str.split(",") if tok.strip()]

            chosen_id = visible_ids[cuda_idx % len(visible_ids)]
        else:
            chosen_id = str(cuda_idx)

        os.environ["CUDA_VISIBLE_DEVICES"] = chosen_id

    from adasearch import (
        LLMGenerator,
        RewardScorer,
        RejectionSampler,
    )

    first_cfg = cfgs[0]

    generator_kwargs = {
        "model_name": first_cfg["llm_model_name"],
        "temperature": first_cfg["temperature"],
        "top_p": first_cfg["top_p"],
        "top_k": first_cfg["top_k"],
        "repetition_penalty": first_cfg["repetition_penalty"],
    }

    generator = LLMGenerator(**generator_kwargs)
    scorer = RewardScorer(model_name=first_cfg["rm_model_name"])

    results: List[Tuple[str, str]] = []
    for cfg in cfgs:
        try:
            sampler = RejectionSampler(
                generator=generator,
                scorer=scorer,
                segment_size=cfg["segment_size"],
                attempt_schedule=cfg["attempt_schedule"],
            )

            output_path = pathlib.Path(cfg["output_json_path"])
            _ensure_parent_dir(output_path)

            generations: List[Dict[str, Any]] = []

            if "dataset_format_identifier" not in cfg or cfg["dataset_format_identifier"] is None:
                raise KeyError("'dataset_format_identifier' missing in job config while running _run_llm_group")
            format_identifier = cfg["dataset_format_identifier"]
            for idx, question in enumerate(dataset, start=1):
                # Build prompt via shared helper so wording is consistent across scripts
                prompt = build_prompt(question, format_identifier)
                answer = sampler.run(prompt, question_idx=idx)

                generations.append({
                    "index": idx,
                    "question": question,
                    "prompt": prompt,
                    "answer": answer,
                })
                
                try:
                    with output_path.open("w", encoding="utf-8") as f:
                        json.dump(generations, f, ensure_ascii=False, indent=2)
                except Exception as write_err:
                    print(f"[Warning] Failed to write interim results to {output_path}: {write_err}")

            results.append((str(output_path), "ok"))

            # Progress logging: show how many schedules are done for this LLM
            completed = len(results)
            total = len(cfgs)
            print(f"[Progress] {first_cfg['llm_model_name']}: {completed}/{total} schedules completed")

        except Exception as e:
            results.append((cfg.get("output_json_path", "<unknown>"), f"error: {e}"))

    return results


def _build_all_jobs(yaml_cfg: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Expand the Cartesian product of LLM × attempt-schedule into a list of job configs."""
    defaults: Dict[str, Any] = deepcopy(yaml_cfg["defaults"])
    llm_model_names: List[str] = yaml_cfg["llm_model_names"]
    attempt_schedules: Dict[str, List[int]] = yaml_cfg["attempt_schedules"]

    jobs: List[Dict[str, Any]] = []
    for model_name in llm_model_names:
        model_tag = model_name.split("/")[-1]
        for sched_tag, sched_vals in attempt_schedules.items():
            cfg = deepcopy(defaults)
            if "dataset_format_identifier" not in defaults or defaults["dataset_format_identifier"] is None:
                raise KeyError("'dataset_format_identifier' missing from YAML defaults; please specify explicitly.")
            dataset_id = defaults["dataset_format_identifier"]

            cfg.update({
                "llm_model_name": model_name,
                "attempt_schedule": sched_vals,
                "output_json_path": defaults["output_json_path"].format(
                    model_tag=model_tag,
                    schedule_tag=sched_tag,
                    dataset_format_identifier=dataset_id,
                ),

            })
            jobs.append(cfg)
    return jobs


def main() -> None:
    start_time = time.monotonic()

    parser = argparse.ArgumentParser(description="Run greedy block-wise alignment across multiple configs.")
    parser.add_argument("--yaml-path", default="configs/config.yaml",
                        help="Path to the unified YAML config file.")
    parser.add_argument("--max-workers", type=int, default=0,
                        help="Maximum number of parallel workers. 0 = auto (one per GPU or per distinct LLM, whichever is smaller).")
    parser.add_argument("--model_name", default=None,
                        help="If set, run only this LLM (used by SLURM job-array task).")
    parser.add_argument("--all_schedules", action="store_true",
                        help="If set, run all schedules for the selected model (requires --model_name).")
    parser.add_argument("--schedule-workers", type=int, default=1,
                        help="Number of parallel workers when running all schedules for a single model (duplicates the model on GPU).")
    parser.add_argument("--cuda-visible-devices", type=str, default='7',
                        help="Comma-separated list of GPU IDs to use . Overrides the environment variable.")
    parser.add_argument("--num-datapoints", type=int, default=50,
                        help="Number of data points to use from the dataset. 0 means all.")
    args = parser.parse_args()

    if args.cuda_visible_devices:
        os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices

    import multiprocessing as _mp
    try:
        _mp.set_start_method("spawn", force=True)
    except RuntimeError:
        pass

    with open(args.yaml_path, "r", encoding="utf-8") as f:
        yaml_cfg: Dict[str, Any] = yaml.safe_load(f)

    n_gpus_available = torch.cuda.device_count()

    if "dataset_format_identifier" not in yaml_cfg["defaults"] or yaml_cfg["defaults"]["dataset_format_identifier"] is None:
        raise KeyError("'dataset_format_identifier' missing from YAML defaults; please specify explicitly.")
    format_identifier = yaml_cfg["defaults"]["dataset_format_identifier"]
    raw_dataset: Any
    input_path: str = yaml_cfg["defaults"]["input_json_path"]
    
    if input_path.endswith(".json"):
        with open(input_path, "r", encoding="utf-8") as f:
            raw_dataset = json.load(f)
    elif input_path.endswith(".csv"):
        with open(input_path, "r", encoding="utf-8") as f:
            raw_dataset = list(csv.DictReader(f))
    else:
        raise ValueError(f"Unsupported dataset file type: {input_path}")

    dataset_questions = _flatten_dataset(raw_dataset, format_identifier)


    print('--------------------------------')
    print(f"Loaded {len(dataset_questions)} prompts from {input_path} (format: {format_identifier})")
    print('--------------------------------')

    if args.num_datapoints > 0:
        dataset_questions = dataset_questions[:args.num_datapoints]
        print(f"Limiting run to the first {len(dataset_questions)} data points.")

    all_jobs = _build_all_jobs(yaml_cfg)

    if args.model_name and args.all_schedules:
        filtered_jobs = [j for j in all_jobs if j["llm_model_name"] == args.model_name]
        print(f"[Model mode] Running all schedules for model: {args.model_name}")

        if args.schedule_workers <= 1:
            _run_llm_group(0, filtered_jobs, dataset_questions, n_gpus=n_gpus_available)
            return

        print(f"Launching {min(args.schedule_workers, len(filtered_jobs))} parallel schedule workers …")

        max_workers = min(args.schedule_workers, len(filtered_jobs))
        aggregated_results: List[Tuple[str, str]] = []
        with ProcessPoolExecutor(max_workers=max_workers) as executor:
            futures = {
                executor.submit(_run_llm_group, idx, [cfg], dataset_questions, n_gpus_available): cfg["output_json_path"]
                for idx, cfg in enumerate(filtered_jobs)
            }
            for future in as_completed(futures):
                out_path = futures[future]
                try:
                    aggregated_results.extend(future.result())
                    print(f"[✓] Completed {out_path}")
                except Exception as e:
                    print(f"[✗] Job {out_path} failed: {e}")

        print("\n==== Summary (parallel schedules) ====")
        for out, status in aggregated_results:
            print(f"{out}: {status}")
        return

    groups: Dict[str, List[Dict[str, Any]]] = {}
    for cfg in all_jobs:
        groups.setdefault(cfg["llm_model_name"], []).append(cfg)

    n_llms = len(groups)
    max_workers = args.max_workers if args.max_workers > 0 else n_gpus_available
    max_workers = max(1, max_workers)

    work_packages: List[List[Dict[str, Any]]] = []

    if max_workers <= n_llms:
        work_packages.extend(groups.values())
        max_workers = min(max_workers, len(work_packages))
    else:
        slots_per_llm = {llm: max_workers // n_llms for llm in groups}
        for i, llm in enumerate(list(groups.keys())[:max_workers % n_llms]):
            slots_per_llm[llm] += 1

        for llm_name, jobs in groups.items():
            num_chunks = slots_per_llm[llm_name]
            if num_chunks > 1 and len(jobs) > 1:
                base_size, rem = divmod(len(jobs), num_chunks)
                chunks, start_idx = [], 0
                for i in range(num_chunks):
                    end_idx = start_idx + base_size + (1 if i < rem else 0)
                    chunk = jobs[start_idx:end_idx]
                    if chunk:
                        chunks.append(chunk)
                    start_idx = end_idx
                work_packages.extend(chunks)
            elif jobs:
                work_packages.append(jobs)

        max_workers = min(max_workers, len(work_packages))
    

    print(f"Detected {n_gpus_available} GPU(s). Using {max_workers} parallel worker(s).")
    print(f"Total LLMs: {len(groups)}; Total job configs: {len(all_jobs)}.")
    if max_workers > n_llms:
        print("Distributing LLM schedules across workers to utilize all available GPUs.")
    print()

    aggregated_results: List[Tuple[str, str]] = []
    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(_run_llm_group, idx, pkg, dataset_questions, n_gpus_available): pkg[0]["llm_model_name"]
            for idx, pkg in enumerate(work_packages) if pkg
        }
        for future in as_completed(futures):
            llm_name = futures[future]
            try:
                group_results = future.result()
                aggregated_results.extend(group_results)
                # Custom message for split-schedule jobs
                num_completed = len([r for r in group_results if r[1] == "ok"])
                if len(group_results) < len(groups.get(llm_name, [])):
                    print(f"[✓] Completed {num_completed}/{len(group_results)} schedules for {llm_name} (worker finished)")
                else:
                    print(f"[✓] Completed all jobs for {llm_name}")
            except Exception as e:
                print(f"[✗] Jobs for {llm_name} failed with: {e}")

    print("\n==== Summary ====")
    for out, status in aggregated_results:
        print(f"{out}: {status}")

    end_time = time.monotonic()
    duration = end_time - start_time
    cprint(f"\nTotal execution time: {_format_duration(duration)}", "green")


if __name__ == "__main__":
    main()