scripts/run_pdrs_parallel.py

#!/usr/bin/env python
"""Parallel driver for the 24-PDR x 70-instance sweep (Claim 3 heuristic side).

Resumes from outputs/pdr_makespans.json: only (instance, rule) pairs not already
present are recomputed. Parallelism is at the (instance, rule) task granularity
(up to 24 x 70 = 1680 tasks) so heavy "large" instances do not serialize a whole
worker. Each task calls rollout() from run_pdrs.py (the faithful in-process
replica of `dsbx-agent run`). Results are checkpointed to disk after every
completion, so a broken pool can simply be re-run to continue.
"""
from __future__ import annotations

import json
import os
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from concurrent.futures.process import BrokenProcessPool
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
from run_pdrs import (  # noqa: E402
    rollout,
    instance_dirs,
    OP_RULES,
    MACHINE_RULES,
    SUBSET,
    OUT,
)

MAX_WORKERS = int(os.environ.get("PDR_WORKERS", "48"))
RULES = [f"{op}:{mac}" for op in OP_RULES for mac in MACHINE_RULES]


def one_task(dir_str: str, op: str, mac: str):
    return dir_str, f"{op}:{mac}", rollout(Path(dir_str), op, mac)


def _instance_weight(d: Path) -> int:
    """Cheap size proxy: number of event lines (drives rollout time & memory)."""
    ev = d / "events.jsonl"
    try:
        with ev.open("rb") as fh:
            return sum(1 for _ in fh)
    except OSError:
        return 0


def _remaining_tasks(results, dirs, key_of):
    # Smallest instances first: gives fast early progress (confirms pool
    # stability quickly) and defers the memory-heavy large instances so fewer
    # of them are ever resident at once.
    ordered = sorted(dirs, key=_instance_weight)
    tasks = []
    for d in ordered:
        key = key_of[str(d)]
        for op in OP_RULES:
            for mac in MACHINE_RULES:
                if f"{op}:{mac}" not in results[key]:
                    tasks.append((str(d), op, mac))
    return tasks


def _run_pool(tasks, results, key_of, t0):
    """One executor attempt. Returns (#done, broke) and checkpoints as it goes.

    If the pool dies abruptly (BrokenProcessPool), the completed tasks are
    already on disk; we report broke=True so the caller can rebuild a fresh
    executor and continue from the checkpoint.
    """
    done = 0
    broke = False
    with ProcessPoolExecutor(max_workers=MAX_WORKERS) as ex:
        futs = {ex.submit(one_task, *t): t for t in tasks}
        try:
            for fut in as_completed(futs):
                src = futs[fut]
                try:
                    dir_str, rule, cmax = fut.result()
                except BrokenProcessPool:
                    raise
                except Exception as e:  # noqa: BLE001 - single bad task
                    print(
                        f"  TASK ERROR {src[0]} {src[1]}:{src[2]} -> {e!r}", flush=True
                    )
                    continue
                results[key_of[dir_str]][rule] = cmax
                done += 1
                OUT.parent.mkdir(parents=True, exist_ok=True)
                OUT.write_text(json.dumps(results, indent=2))
                if done % 10 == 0:
                    el = time.time() - t0
                    print(f"  +{done} this attempt, {el:.0f}s elapsed", flush=True)
        except BrokenProcessPool:
            broke = True
            print("  POOL BROKE; will rebuild and continue from checkpoint", flush=True)
    return done, broke


def main():
    dirs = instance_dirs()
    assert len(dirs) == 70, f"expected 70, got {len(dirs)}"

    results = {}
    if OUT.exists():
        results = json.loads(OUT.read_text())
    key_of = {str(d): str(d.relative_to(SUBSET)) for d in dirs}
    for d in dirs:
        results.setdefault(key_of[str(d)], {})

    total_needed = 70 * 24
    t0 = time.time()
    attempt = 0
    while True:
        tasks = _remaining_tasks(results, dirs, key_of)
        have = total_needed - len(tasks)
        if not tasks:
            print(f"All {total_needed} rollouts present.", flush=True)
            break
        attempt += 1
        print(
            f"[attempt {attempt}] {have}/{total_needed} done; "
            f"{len(tasks)} remaining on {MAX_WORKERS} workers",
            flush=True,
        )
        done, broke = _run_pool(tasks, results, key_of, t0)
        if not broke and done == 0 and tasks:
            print(
                "No progress and pool did not break; aborting to avoid a loop.",
                flush=True,
            )
            break
        if attempt > 40:
            print("Too many attempts; aborting.", flush=True)
            break

    print(f"DONE in {time.time()-t0:.0f}s -> {OUT}", flush=True)
    _verify(results)


def _verify(results):
    n_full = sum(1 for k in results if set(results[k].keys()) >= set(RULES))
    print(f"Instances with all 24 rules: {n_full}/{len(results)}", flush=True)


if __name__ == "__main__":
    main()