| """Run matrix tasks concurrently across a fixed local GPU set.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from dataclasses import dataclass |
| import csv |
| import json |
| import os |
| from pathlib import Path |
| from queue import Queue |
| import re |
| import subprocess |
| import sys |
| import threading |
| from typing import Sequence |
|
|
| from dlb.conditional_matrix import ( |
| CONDITIONAL_MATRIX_COLUMNS, |
| CONDITIONAL_MATRIX_SCHEMA, |
| ConditionalMatrixTask, |
| ) |
| from dlb.matrix import MatrixTask, read_matrix |
|
|
|
|
| DEFAULT_GPUS = "0,1,2,3" |
| DEFAULT_MAX_JOBS = 4 |
| DEFAULT_METRICS = "gen_ppl,entropy,self_bleu" |
| SAFE_LOG_COMPONENT = re.compile(r"[^A-Za-z0-9_.-]+") |
|
|
|
|
| @dataclass(frozen=True) |
| class TaskResult: |
| task: MatrixTask | ConditionalMatrixTask |
| gpu: str |
| exit_code: int |
| stdout_log: Path |
| stderr_log: Path |
|
|
|
|
| def parse_gpus(value: str) -> tuple[str, ...]: |
| """Parse comma/space separated GPU IDs while rejecting shell syntax.""" |
|
|
| parts = tuple(part for part in re.split(r"[,\s]+", value.strip()) if part) |
| if not parts: |
| raise ValueError("at least one GPU must be supplied") |
| allowed = re.compile(r"^[A-Za-z0-9_.:-]+$") |
| invalid = [part for part in parts if not allowed.fullmatch(part)] |
| if invalid: |
| raise ValueError("GPU IDs may contain only letters, numbers, '.', '_', ':', or '-'") |
| if len(set(parts)) != len(parts): |
| raise ValueError("GPU IDs must not be duplicated") |
| return parts |
|
|
|
|
| def active_gpus(gpus: Sequence[str], max_jobs: int) -> tuple[str, ...]: |
| if max_jobs <= 0: |
| raise ValueError("max jobs must be positive") |
| return tuple(gpus[index % len(gpus)] for index in range(max_jobs)) |
|
|
|
|
| def _script(root: Path, name: str) -> str: |
| path = root / "scripts" / name |
| return str(path) |
|
|
|
|
| def command_for_task( |
| task: MatrixTask | ConditionalMatrixTask, |
| *, |
| root: Path, |
| stage: str, |
| metrics: str = DEFAULT_METRICS, |
| precision: str = "author", |
| ) -> list[str]: |
| """Return the argv for one matrix row and stage.""" |
|
|
| conditional = getattr(task, "protocol", None) == "c64_zs_v1" |
| if stage in {"smoke", "generate"}: |
| script = "run_conditional_one.sh" if conditional else "run_one.sh" |
| command = [ |
| "bash", |
| _script(root, script), |
| "--model", |
| task.model, |
| "--dataset", |
| task.dataset, |
| "--steps", |
| str(task.steps), |
| ] |
| if conditional: |
| command.extend(["--seed", str(task.seed)]) |
| else: |
| command.extend(["--num-samples", str(task.sample_count), "--seed", str(task.seed)]) |
| if stage == "smoke": |
| smoke_root = root / "results" / ("conditional/smoke" if conditional else "smoke") |
| command.extend(["--results-root", str(smoke_root)]) |
| return command |
| if stage == "evaluate": |
| eval_python = os.environ.get("DLB_EVAL_PYTHON") |
| if eval_python: |
| command = [eval_python] |
| else: |
| command = [ |
| os.environ.get("DLB_CONDA", "conda"), |
| "run", |
| "-n", |
| os.environ.get("DLB_EVAL_ENV", "dlb-eval"), |
| "python", |
| ] |
| return [ |
| *command, |
| "-m", |
| "evaluation.conditional_evaluate" if conditional else "evaluation.evaluate", |
| "--root", |
| str(root), |
| "--samples", |
| str(Path(task.sample_dir) / "samples.jsonl"), |
| "--metrics", |
| metrics, |
| "--dataset", |
| task.dataset, |
| "--output", |
| task.metrics_path, |
| ] |
| if stage == "benchmark": |
| return [ |
| "bash", |
| _script(root, "benchmark_conditional_one.sh" if conditional else "benchmark_one.sh"), |
| "--model", |
| task.model, |
| "--dataset", |
| task.dataset, |
| "--steps", |
| str(task.steps), |
| "--seed", |
| str(task.seed), |
| "--precision", |
| precision, |
| ] |
| raise ValueError(f"unsupported stage: {stage}") |
|
|
|
|
| def _safe_log_name(task_id: str) -> str: |
| return SAFE_LOG_COMPONENT.sub("_", task_id) |
|
|
|
|
| def _task_environment(root: Path, gpu: str) -> dict[str, str]: |
| env = os.environ.copy() |
| env["CUDA_VISIBLE_DEVICES"] = gpu |
| env["DLB_ROOT"] = str(root) |
| source_path = str(root / "src") |
| current = env.get("PYTHONPATH") |
| env["PYTHONPATH"] = source_path if not current else source_path + os.pathsep + current |
| return env |
|
|
|
|
| def _run_task( |
| task: MatrixTask, |
| *, |
| root: Path, |
| stage: str, |
| gpu: str, |
| metrics: str, |
| precision: str, |
| print_lock: threading.Lock, |
| ) -> TaskResult: |
| logs = root / "results" / "logs" |
| logs.mkdir(parents=True, exist_ok=True) |
| name = _safe_log_name(task.task_id) |
| stdout_log = logs / f"4gpu-{stage}-{name}.out" |
| stderr_log = logs / f"4gpu-{stage}-{name}.err" |
| command = command_for_task( |
| task, root=root, stage=stage, metrics=metrics, precision=precision |
| ) |
| with print_lock: |
| print(f"{stage.upper()} {task.task_id} gpu={gpu}", flush=True) |
| with stdout_log.open("w", encoding="utf-8") as stdout, stderr_log.open( |
| "w", encoding="utf-8" |
| ) as stderr: |
| try: |
| completed = subprocess.run( |
| command, |
| cwd=root, |
| env=_task_environment(root, gpu), |
| stdout=stdout, |
| stderr=stderr, |
| check=False, |
| ) |
| exit_code = completed.returncode |
| except OSError as error: |
| stderr.write(f"failed to start command: {error}\n") |
| exit_code = 127 |
| return TaskResult( |
| task=task, |
| gpu=gpu, |
| exit_code=exit_code, |
| stdout_log=stdout_log, |
| stderr_log=stderr_log, |
| ) |
|
|
|
|
| def _failure_path(root: Path, stage: str) -> Path: |
| names = {"generate": "generation", "evaluate": "evaluation"} |
| name = names.get(stage, stage) |
| return root / "results" / "logs" / f"{name}_4gpu_failures.tsv" |
|
|
|
|
| def _write_failures(root: Path, stage: str, results: Sequence[TaskResult]) -> None: |
| failures = [result for result in results if result.exit_code != 0] |
| path = _failure_path(root, stage) |
| if not failures: |
| path.unlink(missing_ok=True) |
| return |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_name(f".{path.name}.{os.getpid()}.partial") |
| try: |
| with temporary.open("w", encoding="utf-8", newline="") as handle: |
| writer = csv.DictWriter( |
| handle, |
| fieldnames=( |
| "task_id", |
| "model", |
| "dataset", |
| "steps", |
| "gpu", |
| "exit_code", |
| "stdout_log", |
| "stderr_log", |
| ), |
| delimiter="\t", |
| lineterminator="\n", |
| ) |
| writer.writeheader() |
| for result in failures: |
| writer.writerow( |
| { |
| "task_id": result.task.task_id, |
| "model": result.task.model, |
| "dataset": result.task.dataset, |
| "steps": result.task.steps, |
| "gpu": result.gpu, |
| "exit_code": result.exit_code, |
| "stdout_log": str(result.stdout_log), |
| "stderr_log": str(result.stderr_log), |
| } |
| ) |
| os.replace(temporary, path) |
| except BaseException: |
| temporary.unlink(missing_ok=True) |
| raise |
|
|
|
|
| def dry_run( |
| tasks: Sequence[MatrixTask | ConditionalMatrixTask], |
| *, |
| root: Path, |
| stage: str, |
| gpus: Sequence[str], |
| metrics: str, |
| precision: str, |
| ) -> int: |
| for index, task in enumerate(tasks): |
| gpu = gpus[index % len(gpus)] |
| command = command_for_task( |
| task, root=root, stage=stage, metrics=metrics, precision=precision |
| ) |
| print( |
| json.dumps( |
| { |
| "stage": stage, |
| "task_id": task.task_id, |
| "gpu": gpu, |
| "cuda_visible_devices": gpu, |
| "command": command, |
| }, |
| sort_keys=True, |
| ) |
| ) |
| return 0 |
|
|
|
|
| def run_tasks( |
| tasks: Sequence[MatrixTask | ConditionalMatrixTask], |
| *, |
| root: Path, |
| stage: str, |
| gpus: Sequence[str], |
| metrics: str, |
| precision: str, |
| ) -> int: |
| gpu_queue: Queue[str] = Queue() |
| for gpu in gpus: |
| gpu_queue.put(gpu) |
| print_lock = threading.Lock() |
| results: list[TaskResult] = [] |
|
|
| def run_with_gpu(task: MatrixTask | ConditionalMatrixTask) -> TaskResult: |
| gpu = gpu_queue.get() |
| try: |
| return _run_task( |
| task, |
| root=root, |
| stage=stage, |
| gpu=gpu, |
| metrics=metrics, |
| precision=precision, |
| print_lock=print_lock, |
| ) |
| finally: |
| gpu_queue.put(gpu) |
|
|
| with ThreadPoolExecutor(max_workers=len(gpus)) as executor: |
| futures = [executor.submit(run_with_gpu, task) for task in tasks] |
| for future in as_completed(futures): |
| result = future.result() |
| results.append(result) |
| if result.exit_code != 0: |
| with print_lock: |
| print( |
| f"FAILED {result.task.task_id} gpu={result.gpu} " |
| f"exit={result.exit_code}", |
| file=sys.stderr, |
| flush=True, |
| ) |
|
|
| _write_failures(root, stage, results) |
| return 1 if any(result.exit_code != 0 for result in results) else 0 |
|
|
|
|
| def _read_conditional_matrix(path: Path) -> list[ConditionalMatrixTask]: |
| if path.is_symlink() or not path.is_file(): |
| raise ValueError(f"conditional matrix TSV is missing or unsafe: {path}") |
| with path.open("r", encoding="utf-8", newline="") as handle: |
| schema = handle.readline().rstrip("\r\n") |
| if schema != f"# schema={CONDITIONAL_MATRIX_SCHEMA}": |
| raise ValueError("conditional matrix TSV has an unsupported schema header") |
| reader = csv.DictReader(handle, delimiter="\t") |
| if tuple(reader.fieldnames or ()) != CONDITIONAL_MATRIX_COLUMNS: |
| raise ValueError("conditional matrix TSV columns do not match the canonical schema") |
| tasks: list[ConditionalMatrixTask] = [] |
| for row_number, row in enumerate(reader, start=3): |
| if None in row or any(value is None for value in row.values()): |
| raise ValueError(f"conditional matrix TSV row {row_number} has missing fields") |
| try: |
| tasks.append( |
| ConditionalMatrixTask( |
| task_id=row["task_id"], |
| category=row["category"], |
| model=row["model"], |
| dataset=row["dataset"], |
| steps=int(row["steps"]), |
| sample_count=2048, |
| seed=int(row["seed"]), |
| environment=row["environment"], |
| adapter=row["adapter"], |
| source=row["source"], |
| provenance=row["provenance"], |
| protocol="c64_zs_v1", |
| conditioning_manifest=row["conditioning_manifest"], |
| conditioning_manifest_sha256=row["conditioning_manifest_sha256"], |
| sample_dir=row["sample_dir"], |
| metrics_path=row["metrics_path"], |
| timing_path=row["timing_path"], |
| ) |
| ) |
| except (KeyError, TypeError, ValueError) as error: |
| raise ValueError(f"invalid conditional matrix TSV row {row_number}") from error |
| return tasks |
|
|
|
|
| def read_any_matrix(path: Path) -> list[MatrixTask | ConditionalMatrixTask]: |
| with path.open("r", encoding="utf-8", newline="") as handle: |
| schema = handle.readline().rstrip("\r\n") |
| if schema == f"# schema={CONDITIONAL_MATRIX_SCHEMA}": |
| return _read_conditional_matrix(path) |
| return read_matrix(path) |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--root", type=Path, required=True) |
| parser.add_argument("--stage", choices=("smoke", "generate", "evaluate", "benchmark"), required=True) |
| parser.add_argument("--matrix", type=Path, required=True) |
| parser.add_argument("--gpus", default=os.environ.get("DLB_GPUS", DEFAULT_GPUS)) |
| parser.add_argument("--max-jobs", type=int, default=int(os.environ.get("DLB_MAX_JOBS", DEFAULT_MAX_JOBS))) |
| parser.add_argument("--metrics", default=DEFAULT_METRICS) |
| parser.add_argument("--precision", choices=("author",), default="author") |
| parser.add_argument("--dry-run", action="store_true") |
| arguments = parser.parse_args(argv) |
|
|
| root = arguments.root.resolve() |
| try: |
| gpu_list = active_gpus(parse_gpus(arguments.gpus), arguments.max_jobs) |
| except ValueError as error: |
| parser.error(str(error)) |
| tasks = read_any_matrix(arguments.matrix) |
| if arguments.dry_run: |
| return dry_run( |
| tasks, |
| root=root, |
| stage=arguments.stage, |
| gpus=gpu_list, |
| metrics=arguments.metrics, |
| precision=arguments.precision, |
| ) |
| return run_tasks( |
| tasks, |
| root=root, |
| stage=arguments.stage, |
| gpus=gpu_list, |
| metrics=arguments.metrics, |
| precision=arguments.precision, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
