scripts/run_pdrs.py

#!/usr/bin/env python
"""Reproduce DynaSchedBench Claim 3 heuristic side.

Runs all 24 composite priority dispatch rules (PDR: <OP_RULE>:<MACHINE_RULE>)
on the 70 released DynaSched-Subset instances by importing the repo's OWN
simulator/env/agent/metrics code (faithful in-process replica of the
`dsbx-agent run` rollout loop), and records Cmax (makespan) per (instance, PDR).

Output: outputs/pdr_makespans.json  { instance_key: { "OP:MACHINE": Cmax, ... } }
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path

REPO = Path("/home/ubuntu/samuel/dynasched-repro/repo")
sys.path.insert(0, str(REPO / "src"))

# Silence the very chatty loguru logging from the framework.
from loguru import logger  # noqa: E402

logger.remove()

from dsbx.Env import DynaSchedEnv  # noqa: E402
from dsbx.Sim.Loader import load_instance_from_events  # noqa: E402
from dsbx.Eval.Metrics import evaluate_trajectory  # noqa: E402
from dsbx.Agents.PDRs import PDRAgent  # noqa: E402

OP_RULES = ["SPT", "LPT", "MWKR", "LWKR", "MOPNR", "LOPNR", "FIFO", "LIFO"]
MACHINE_RULES = ["LIT", "LWL", "SPT"]

SUBSET = Path("/home/ubuntu/samuel/dynasched-repro/repo/data/DynaSched-Subset")
OUT = Path("/home/ubuntu/samuel/dynasched-repro/outputs/pdr_makespans.json")
MAX_STEPS = 100_000


def instance_dirs():
    dirs = []
    for sub in ["fromGrid", "fromSweep"]:
        for events in sorted((SUBSET / sub).rglob("events.jsonl")):
            dirs.append(events.parent)
    return dirs


def rollout(instance_dir: Path, op_rule: str, machine_rule: str) -> float:
    """Faithful replica of dsbx-agent run core loop for a PDR agent."""
    events_file = instance_dir / "events.jsonl"
    model, events = load_instance_from_events(events_file)
    # track_trajectory=False: the makespan is computed from the FINAL snapshot's
    # per-job completion times (identical to evaluate_trajectory's
    # traj.last_snapshot path, which is what the metric uses), so accumulating
    # every step's snapshot is pure memory overhead. Disabling it drops peak RSS
    # from ~6.5 GB to ~0.2 GB per large rollout (validated bit-identical to the
    # trajectory path), which is what lets the sweep run at high parallelism.
    env = DynaSchedEnv(
        model,
        events=events,
        auto_generate_events=False,
        traj_stream_path=None,
        track_trajectory=False,
    )

    # Emergency jobs: EXACT replica of dsbx-agent run detection logic.
    from dsbx.Sim.Events import PriorityChangeEvent

    emergency = set()
    emergency_threshold = getattr(
        getattr(model, "dynamic_scenarios", None), "emergency_priority", -1
    )
    for ev in events:
        if isinstance(ev, PriorityChangeEvent):
            new_p = getattr(ev, "new_priority", 0)
            if new_p <= emergency_threshold:
                emergency.add(str(ev.job_id))
    sj = instance_dir / "static_jobs.json"
    if sj.exists():
        try:
            data = json.loads(sj.read_text(encoding="utf-8"))
            jobs_info = data.get("jobs", {}) or {}
            for jid, info in jobs_info.items():
                if info.get("arrival_type") in ("emergency", "dynamic"):
                    emergency.add(str(jid))
        except Exception:
            pass
    scenario_info = {"events_path": str(events_file)}
    if emergency:
        scenario_info["emergency_jobs"] = sorted(emergency)
    try:
        env._static_emergency_jobs = set(emergency)
    except Exception:
        pass

    # Seed the global RNG for reproducibility. The framework uses unseeded
    # random.choice() for tie-breaking, making the CLI non-deterministic
    # (verified: same instance/rule yields makespans differing at ~0.005%).
    # A fixed seed makes our reproduction bit-for-bit repeatable without
    # changing any scheduling logic.
    ag = PDRAgent(op_rule=op_rule, machine_rule=machine_rule, random_seed=0)
    obs = env.reset()
    ag.reset(scenario_info)

    done = env.done()
    steps = 0
    while (not done) and steps < MAX_STEPS:
        legal = env.legal_actions()
        if not legal:
            obs = env.advance_if_idle()
            done = env.done()
            continue
        act = ag.act(obs, legal, env)
        if act is None:
            obs = env.advance_if_idle()
            steps += 1
            done = env.done()
            continue
        obs, reward, done, info = env.step(act)
        steps += 1

    # Silent-failure guard: a makespan from a schedule that did not run to
    # completion (hit the max_steps cap) is meaningless and would corrupt the
    # gap aggregation. Fail loudly instead.
    if not done:
        env.close()
        raise RuntimeError(
            f"Rollout {instance_dir.name} {op_rule}:{machine_rule} hit "
            f"max_steps={MAX_STEPS} without env.done(); makespan is invalid."
        )
    # Makespan = max completion time over jobs in the final snapshot. This is
    # exactly what dsbx.Eval.Metrics computes from traj.last_snapshot; using the
    # live final snapshot avoids retaining the whole trajectory.
    snap = env.get_snapshot()
    completion_times = [
        float(j.completion_time) if j.completion_time is not None else float(snap.time)
        for j in snap.jobs
    ]
    makespan = max(completion_times) if completion_times else float(snap.time)
    env.close()
    return float(makespan)


def main():
    only = sys.argv[1] if len(sys.argv) > 1 else None  # optional: single instance path
    dirs = instance_dirs()
    print(f"Found {len(dirs)} instances", flush=True)
    assert len(dirs) == 70, f"expected 70 instances, got {len(dirs)}"

    results = {}
    if OUT.exists():
        results = json.loads(OUT.read_text())

    t0 = time.time()
    for i, d in enumerate(dirs):
        key = str(d.relative_to(SUBSET))
        if only and only not in str(d):
            continue
        results.setdefault(key, {})
        for op in OP_RULES:
            for mac in MACHINE_RULES:
                rk = f"{op}:{mac}"
                if rk in results[key]:
                    continue
                cmax = rollout(d, op, mac)
                results[key][rk] = cmax
        # checkpoint after each instance
        OUT.parent.mkdir(parents=True, exist_ok=True)
        OUT.write_text(json.dumps(results, indent=2))
        el = time.time() - t0
        print(f"[{i+1}/{len(dirs)}] {key} done ({el:.0f}s elapsed)", flush=True)

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


if __name__ == "__main__":
    main()