From 4b97eb48eb77d57b7e738b59162e631aaabac90b Mon Sep 17 00:00:00 2001
From: mmshad <31811010+mmshad@users.noreply.github.com>
Date: Mon, 2 Nov 2026 22:03:48 -0500
Subject: [PATCH 1/5] Add a partition setting and SLURM checks for the round
 jobs

---
 src/pruning_example/finish.py     | 16 ++++++---
 src/pruning_example/loop.py       | 24 +++++++++++--
 src/pruning_example/multi_loop.py | 20 +++++++----
 src/pruning_example/slurm.py      | 36 ++++++++++++++++---
 tests/test_array.py               | 46 +++++++++++++++++++++++-
 tests/test_loop.py                | 56 ++++++++++++++++++++++++++---
 tests/test_multi_loop.py          | 59 +++++++++++++++++++++++++++----
 7 files changed, 227 insertions(+), 30 deletions(-)

diff --git a/src/pruning_example/finish.py b/src/pruning_example/finish.py
index 58ffd06..735156d 100644
--- a/src/pruning_example/finish.py
+++ b/src/pruning_example/finish.py
@@ -91,10 +91,11 @@ def refuse_after_final_steps(ctx) -> None:
 
 
 def round_job_options(ctx) -> tuple:
-    """Return the sbatch options of each round job. An optional round.constraint in the experiment file
-    replaces the --constraint line of the job scripts, so that all designs of a comparison use the same GPUs."""
-    constraint = ctx.experiment["round"].get("constraint")
-    return (f"--constraint={constraint}",) if constraint else ()
+    """Return the sbatch options of each round job. An optional round.partition and round.constraint in the
+    experiment file replace the --partition and --constraint lines of the job scripts, for example to run the
+    round jobs on kempner_interactive, or so that all designs of a comparison use the same GPUs."""
+    settings = ctx.experiment["round"]
+    return tuple(f"--{key}={settings[key]}" for key in ("partition", "constraint") if settings.get(key))
 
 
 def note_pending_job(ctx, job: str, kind: str, round_no: int, tasks: int = 0) -> None:
@@ -163,6 +164,8 @@ def confirmation_round(ctx, round_no: int) -> dict | None:
     ctx.log(f"confirmation round: submitted job {job}")
     outcome = ctx.slurm.wait(job, 60 * ctx.budget["queue_limit_minutes"])
     if outcome == "queue_limit":
+        ctx.log(f"confirmation round: job {job} waited longer than the queue limit "
+                f"(SLURM reason: {ctx.slurm.last_reasons})")
         ctx.slurm.cancel(job)
     usage = job_accounting(ctx, job)
     (state / "current_job").unlink(missing_ok=True)
@@ -198,7 +201,10 @@ def confirmation_rounds(ctx, round_no: int, seeds: list[int]) -> dict[int, dict
             (state / "current_job").write_text(job + "\n")
             note_pending_job(ctx, job, "confirm", round_no, len(missing))
             ctx.log(f"confirmation rounds: submitted array job {job} for seeds {', '.join(map(str, missing))}")
-            outcome = ctx.slurm.wait_array(job, 60 * ctx.budget["queue_limit_minutes"])
+            outcome = ctx.slurm.wait_array(job, 60 * ctx.budget["queue_limit_minutes"], max_parallel)
+            if outcome == "queue_limit":
+                ctx.log(f"confirmation rounds: array job {job} waited longer than the queue limit "
+                        f"(SLURM reason: {ctx.slurm.last_reasons})")
             usage = ctx.slurm.array_accounting(job, len(missing))
             (state / "current_job").unlink(missing_ok=True)
             for index, seed in enumerate(missing):
diff --git a/src/pruning_example/loop.py b/src/pruning_example/loop.py
index 8042438..9a153a6 100644
--- a/src/pruning_example/loop.py
+++ b/src/pruning_example/loop.py
@@ -239,7 +239,8 @@ def run_attempt(ctx: Context, round_no: int) -> None:
     rows = [row for row in records.read_history(state) if int(row["round"]) == round_no]
     if outcome == "queue_limit":
         if not rows:
-            fail(ctx, round_no, "queue_limit", f"job {job} waited longer than the queue limit", round_config)
+            fail(ctx, round_no, "queue_limit", f"job {job} waited longer than the queue limit "
+                 f"(SLURM reason: {ctx.slurm.last_reasons})", round_config)
         raise StopLoop("queue_limit")
     if not rows:
         fail(ctx, round_no, "job_failed", f"job {job} ended as {usage.state} without metrics", round_config)
@@ -283,6 +284,23 @@ def check_time_lines(ctx: Context, scripts: tuple[str, ...]) -> None:
                              f"has round.time_minutes {ctx.experiment['round']['time_minutes']}.")
 
 
+def check_round_jobs(ctx: Context, scripts: tuple[str, ...]) -> None:
+    """Refuse to start if SLURM would reject the round jobs, for example for a constraint that no node of the
+    partition has. sbatch --test-only checks each job script with the options from the experiment file."""
+    options = finish.round_job_options(ctx)
+    for script in scripts:
+        try:
+            error = ctx.slurm.test_only(script, ctx.paths.project, options)
+        except (SlurmError, subprocess.TimeoutExpired, OSError) as problem:
+            ctx.log(f"could not check {script} with sbatch --test-only: {problem}")
+            continue
+        if error:
+            experiment_file = config.experiment_path(ctx.paths.project).relative_to(ctx.paths.project)
+            raise SystemExit(f"SLURM would reject the round jobs of {script} with the options "
+                             f"{' '.join(options) or '(none)'}: {error}\nCheck round.partition and "
+                             f"round.constraint in {experiment_file}.")
+
+
 def check_start(ctx: Context) -> None:
     """Refuse to start without the baseline and round 1, with settings that do not agree, or after the final steps."""
     project = ctx.paths.project
@@ -292,8 +310,10 @@ def check_start(ctx: Context) -> None:
     if not (ctx.paths.state / "rounds" / "round_01" / "metrics.json").exists():
         raise SystemExit("There is no round 1. Run the supervised round first.")
     several_seeds = len(finish.confirmation_seeds(ctx.experiment)) > 1  # their confirmation rounds run as an array job
-    check_time_lines(ctx, ("slurm/round.sbatch",) + (("slurm/round_array.sbatch",) if several_seeds else ()))
+    scripts = ("slurm/round.sbatch",) + (("slurm/round_array.sbatch",) if several_seeds else ())
+    check_time_lines(ctx, scripts)
     records.check_integrity(project)
+    check_round_jobs(ctx, scripts)
     notes = ctx.paths.state / "notes.md"
     if not notes.exists():
         first = records.read_json(ctx.paths.state / "rounds" / "round_01" / "config.json")
diff --git a/src/pruning_example/multi_loop.py b/src/pruning_example/multi_loop.py
index 4c2832a..26177a1 100644
--- a/src/pruning_example/multi_loop.py
+++ b/src/pruning_example/multi_loop.py
@@ -31,8 +31,8 @@
 from .agent import Agent, AgentError, AgentFormatError, EFFORT_LEVELS, ModelCheckError, check_model, tokens_used
 from .config import ConfigError
 from .loop import (
-    LEDGER_FIELDS, Context, Ledger, StopLoop, Terminated, check_time_lines, describe_best, job_end_time, log,
-    on_sigterm, render,
+    LEDGER_FIELDS, Context, Ledger, StopLoop, Terminated, check_round_jobs, check_time_lines, describe_best,
+    job_end_time, log, on_sigterm, render,
 )
 from .records import IntegrityError
 from .slurm import Slurm, SlurmError
@@ -143,9 +143,11 @@ def plan_minutes(ctx) -> float:
     return minutes
 
 
-def array_minutes(ctx) -> float:
-    """The time that the loop job needs to run one array job and the final steps."""
-    return (ctx.experiment["round"]["time_minutes"] + ctx.budget["queue_limit_minutes"]
+def array_minutes(ctx, tasks: int) -> float:
+    """The time that the loop job needs to run one array job of tasks rounds, and the final steps. With
+    wave.max_parallel below tasks, the rounds run in turns."""
+    turns = math.ceil(tasks / ctx.experiment["wave"]["max_parallel"])
+    return (turns * ctx.experiment["round"]["time_minutes"] + ctx.budget["queue_limit_minutes"]
             + ctx.budget["final_reserve_minutes"])
 
 
@@ -586,13 +588,15 @@ def submit_wave(ctx, wave: int, accepted: list[dict]) -> int:
     (folder / "array_job").write_text(job + "\n")
     (state / "current_job").write_text(job + "\n")
     ctx.log(f"wave {wave}: submitted array job {job}, rounds {first} to {first + len(lines) - 1}")
-    outcome = ctx.slurm.wait_array(job, 60 * ctx.budget["queue_limit_minutes"])
+    outcome = ctx.slurm.wait_array(job, 60 * ctx.budget["queue_limit_minutes"], ctx.experiment["wave"]["max_parallel"])
     complete = add_round_records(ctx, wave, rounds, job, outcome)
     (state / "current_job").unlink(missing_ok=True)
     write_outcome(ctx, wave, complete, outcome, array_job=job)
     write_hypotheses(ctx)
     log_wave(ctx, wave, [round_no for round_no, _ in rounds])
     if outcome == "queue_limit":
+        ctx.log(f"wave {wave}: array job {job} waited longer than the queue limit "
+                f"(SLURM reason: {ctx.slurm.last_reasons})")
         raise StopLoop("queue_limit")
     return complete
 
@@ -866,6 +870,8 @@ def check_start(ctx, plan_only: bool = False) -> None:
     if not plan_only and not (accepted.exists() and records.read_json(accepted)):
         raise SystemExit("Wave 1 has no accepted configuration. Run 'sbatch --time=01:30:00 slurm/multi_loop.sbatch "
                          "--plan-only', and read the plan first.")
+    one_seed = finish.confirmation_seeds(ctx.experiment) == [1]  # then the confirmation round is one round job
+    check_round_jobs(ctx, ("slurm/round_array.sbatch",) + (("slurm/round.sbatch",) if one_seed else ()))
     clean_up_after_a_stop(ctx)
 
 
@@ -969,7 +975,7 @@ def run_waves(ctx) -> str:
             if not accepted:
                 write_outcome(ctx, wave, 0, "no accepted configuration" if accepted == [] else "a role step failed")
                 continue
-            if not time_for(ctx, array_minutes(ctx)):
+            if not time_for(ctx, array_minutes(ctx, len(accepted))):
                 write_outcome(ctx, wave, 0, "not run: not enough time for the array job and the final steps")
                 return "time_left"
             submit_wave(ctx, wave, accepted)
diff --git a/src/pruning_example/slurm.py b/src/pruning_example/slurm.py
index 6a39ace..e970ff6 100644
--- a/src/pruning_example/slurm.py
+++ b/src/pruning_example/slurm.py
@@ -130,6 +130,7 @@ def __init__(self, poll_seconds: int = 60, sleep=time.sleep, clock=time.monotoni
         self.poll_seconds = poll_seconds
         self.sleep = sleep
         self.clock = clock
+        self.last_reasons = ""  # SLURM's reasons for the wait of the last job that reached the queue limit
 
     def run(self, args: list[str], timeout: int = 120, **kwargs) -> subprocess.CompletedProcess:
         return subprocess.run(args, capture_output=True, text=True, timeout=timeout, **kwargs)
@@ -151,6 +152,17 @@ def submit_array(self, script: str, tasks_file: str, count: int, max_parallel: i
         """Submit an array job with one task for each line of tasks_file, at most max_parallel at a time."""
         return self.submit(script, [tasks_file], cwd, options=(f"--array=0-{count - 1}%{max_parallel}", *options))
 
+    def test_only(self, script: str, cwd, options: tuple = ()) -> str | None:
+        """Ask SLURM if it would accept the job script with the options, without a submit (sbatch --test-only).
+        Return None if it would, or SLURM's error. Raise SlurmError if the controller is too busy to answer."""
+        result = self.run(["sbatch", "--test-only", *options, script], cwd=cwd, env=clean_env())
+        if result.returncode == 0:
+            return None
+        error = result.stderr.strip() or result.stdout.strip()
+        if any(text in error for text in TRANSIENT_ERRORS):
+            raise SlurmError(f"sbatch --test-only failed: {error}")
+        return error or "sbatch --test-only failed"
+
     def queue_state(self, job_id: str) -> str | None:
         """Return the state of the job in the queue, or None if the job left the queue."""
         result = self.run(["squeue", "-h", "-j", job_id, "-o", "%T"], timeout=60)
@@ -170,6 +182,15 @@ def array_states(self, array_id: str) -> dict[int, str]:
             raise SlurmError(f"squeue failed: {result.stderr.strip()}")
         return parse_array_squeue(result.stdout)
 
+    def pending_reasons(self, job_id: str) -> str:
+        """Return SLURM's reasons why the job, or the tasks of an array job, still wait, for the log."""
+        try:
+            result = self.run(["squeue", "-h", "-r", "-t", "PENDING", "-j", job_id, "-o", "%r"], timeout=60)
+        except (subprocess.TimeoutExpired, OSError):
+            return "unknown"
+        reasons = sorted({line.strip() for line in result.stdout.splitlines() if line.strip()})
+        return ", ".join(reasons) or "unknown"
+
     def wait(self, job_id: str, queue_limit_seconds: float) -> str:
         """Wait until the job leaves the queue. Return "ended" or "queue_limit".
 
@@ -194,14 +215,17 @@ def wait(self, job_id: str, queue_limit_seconds: float) -> str:
                 queued += now - last
             last = now
             if queued > queue_limit_seconds:
+                self.last_reasons = self.pending_reasons(job_id)
                 return "queue_limit"
             self.sleep(self.poll_seconds)
 
-    def wait_array(self, array_id: str, queue_limit_seconds: float) -> str:
+    def wait_array(self, array_id: str, queue_limit_seconds: float, max_parallel: int | None = None) -> str:
         """Wait until all tasks of an array job leave the queue. Return "ended" or "queue_limit".
 
-        Queue time counts while at least one task waits. After the queue limit, the waiting tasks are
-        canceled, also a task that a preemption put back in the queue, and the running tasks can end.
+        Queue time counts while at least one task waits and fewer than max_parallel tasks run. So a task
+        that waits only because max_parallel tasks of the array already run adds no queue time. After the
+        queue limit, the waiting tasks are canceled, also a task that a preemption put back in the queue,
+        and the running tasks can end.
         """
         queued = 0.0
         last = self.clock()
@@ -220,11 +244,13 @@ def wait_array(self, array_id: str, queue_limit_seconds: float) -> str:
             if states == {}:
                 return "queue_limit" if over_limit else "ended"
             waiting = bool(states) and any(state in QUEUE_WAIT_STATES for state in states.values())
-            if waiting and not over_limit:
+            running = sum(state not in QUEUE_WAIT_STATES for state in (states or {}).values())
+            if waiting and not over_limit and (max_parallel is None or running < max_parallel):
                 queued += now - last
             last = now
-            if queued > queue_limit_seconds:
+            if queued > queue_limit_seconds and not over_limit:
                 over_limit = True
+                self.last_reasons = self.pending_reasons(array_id)
             if over_limit and waiting:
                 self.cancel_waiting(array_id)
             self.sleep(self.poll_seconds)
diff --git a/tests/test_array.py b/tests/test_array.py
index f575410..f921aa8 100644
--- a/tests/test_array.py
+++ b/tests/test_array.py
@@ -2,7 +2,9 @@
 
 from __future__ import annotations
 
-from pruning_example.slurm import Slurm, parse_array_sacct, parse_array_squeue, task_indexes
+import pytest
+
+from pruning_example.slurm import Slurm, SlurmError, parse_array_sacct, parse_array_squeue, task_indexes
 
 TRES = "billing=8,cpu=8,gres/gpu=1,mem=32G,node=1"
 
@@ -53,6 +55,7 @@ def cancel_waiting(array_id, wait_seconds=60):
 
     slurm.cancel_waiting = cancel_waiting
     slurm.cancel = lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("must not cancel all tasks"))
+    slurm.pending_reasons = lambda array_id: "QOSMaxGRESPerUser"
     return slurm
 
 
@@ -75,6 +78,47 @@ def test_queue_limit_cancels_only_waiting_tasks_also_after_a_preemption():
     assert slurm.waiting_cancels == 2
 
 
+def test_a_task_that_waits_for_a_running_task_of_the_array_adds_no_queue_time():
+    # One GPU for each user: tasks 1 and 2 wait while task 0 runs, then the tasks run in turns.
+    turns = ([{0: "RUNNING", 1: "PENDING", 2: "PENDING"}] * 5 + [{1: "RUNNING", 2: "PENDING"}] * 5
+             + [{2: "RUNNING"}] * 5 + [{}])
+    slurm = fake_slurm(turns)
+    assert slurm.wait_array("500", queue_limit_seconds=100, max_parallel=1) == "ended"
+    assert slurm.waiting_cancels == 0
+    slurm = fake_slurm(turns)  # with room for 3 running tasks, the same waits count
+    assert slurm.wait_array("500", queue_limit_seconds=100, max_parallel=3) == "queue_limit"
+    assert slurm.waiting_cancels > 0
+
+
+def test_a_wait_with_no_running_task_counts_and_keeps_slurms_reason():
+    slurm = fake_slurm([{0: "PENDING", 1: "PENDING"}] * 3 + [{}])
+    assert slurm.wait_array("500", queue_limit_seconds=100, max_parallel=1) == "queue_limit"
+    assert slurm.last_reasons == "QOSMaxGRESPerUser"  # read before the waiting tasks are canceled
+
+
+def test_test_only_returns_slurms_error_and_raises_only_for_a_busy_controller():
+    answers = iter([
+        Result(0, err="sbatch: Job 1 to start at 2026-11-02T10:00:00 using 8 processors on nodes holygpu8a19505 "
+                      "in partition kempner_interactive"),
+        Result(1, err="sbatch: error: Batch job submission failed: Requested node configuration is not available"),
+        Result(1, err="sbatch: error: Batch job submission failed: Socket timed out on send/recv operation"),
+    ])
+    seen = []
+
+    def run(args, timeout=120, **kwargs):
+        seen.append(args)
+        return next(answers)
+
+    slurm = Slurm(sleep=lambda seconds: None)
+    slurm.run = run
+    options = ("--partition=kempner_interactive", "--constraint=a100-mig")
+    assert slurm.test_only("slurm/round.sbatch", ".", options) is None
+    assert seen[0] == ["sbatch", "--test-only", *options, "slurm/round.sbatch"]
+    assert "Requested node configuration is not available" in slurm.test_only("slurm/round.sbatch", ".")
+    with pytest.raises(SlurmError, match="Socket timed out"):
+        slurm.test_only("slurm/round.sbatch", ".")
+
+
 def test_submit_array_passes_the_range_and_the_throttle():
     seen = []
 
diff --git a/tests/test_loop.py b/tests/test_loop.py
index 79a7b2a..abe20bc 100644
--- a/tests/test_loop.py
+++ b/tests/test_loop.py
@@ -68,6 +68,13 @@ def __init__(self, state, results, gpu_minutes: float = 5.0):
         self.submitted = []
         self.canceled = []
         self.options = []
+        self.tested = []  # (script, options) of each sbatch --test-only
+        self.test_errors = {}  # script -> the error that sbatch --test-only returns
+        self.last_reasons = ""
+
+    def test_only(self, script, cwd, options=()):
+        self.tested.append((script, options))
+        return self.test_errors.get(script)
 
     def submit(self, script, args, cwd, options=()):
         self.options.append(options)
@@ -269,7 +276,9 @@ def test_queue_time_adds_up_over_requeues():
     times = iter(range(0, 1000, 60))
     slurm = Slurm(sleep=lambda seconds: None, clock=lambda: next(times))
     slurm.queue_state = lambda job_id: next(states)
+    slurm.pending_reasons = lambda job_id: "Priority"
     assert slurm.wait("1", queue_limit_seconds=150) == "queue_limit"
+    assert slurm.last_reasons == "Priority"  # read before the loop script cancels the job
 
 
 def test_select_round_rules():
@@ -447,13 +456,52 @@ def wait(self, job_id, queue_limit_seconds):
     assert not (state / finish.PENDING_FILE).exists()
 
 
-def test_round_constraint_in_the_experiment_file_reaches_each_round_job(started):
+def test_round_partition_and_constraint_in_the_experiment_file_reach_each_round_job(started):
     slurm = FakeSlurm(started / "state", lambda n: (0.91, SMALL))
     ctx = context(started, FakeAgent(lambda call: GOOD), slurm, plateau_rounds=1)
-    assert finish.round_job_options(ctx) == ()  # unset: the --constraint line of the job script applies
-    ctx.experiment["round"]["constraint"] = "a100-mig"
+    assert finish.round_job_options(ctx) == ()  # unset: the lines of the job script apply
+    ctx.experiment["round"].update(partition="kempner_interactive", constraint="a100-mig")
     assert loop.run_rounds(ctx) == "plateau"
-    assert slurm.options and all(options == ("--constraint=a100-mig",) for options in slurm.options)
+    expected = ("--partition=kempner_interactive", "--constraint=a100-mig")
+    assert slurm.options and all(options == expected for options in slurm.options)
+
+
+def test_the_loop_job_stops_before_any_agent_step_if_slurm_would_reject_the_round_jobs(started):
+    slurm = FakeSlurm(started / "state", lambda n: (0.91, SMALL))
+    slurm.test_errors["slurm/round.sbatch"] = ("sbatch: error: Batch job submission failed: "
+                                               "Requested node configuration is not available")
+    agent = FakeAgent(lambda call: GOOD)
+    ctx = context(started, agent, slurm)
+    ctx.experiment["round"]["partition"] = "kempner_interactive"  # the script asks for full GPUs, which it has not
+    with pytest.raises(SystemExit, match="Requested node configuration is not available"):
+        loop.run(ctx)
+    assert slurm.tested == [("slurm/round.sbatch", ("--partition=kempner_interactive",))]
+    assert agent.calls == 0 and not slurm.submitted
+
+
+def test_the_start_check_goes_on_if_the_slurm_controller_does_not_answer(started):
+    class BusySlurm(FakeSlurm):
+        def test_only(self, script, cwd, options=()):
+            raise SlurmError("sbatch --test-only failed: Socket timed out on send/recv operation")
+
+    messages = []
+    ctx = context(started, FakeAgent(lambda call: GOOD), BusySlurm(started / "state", lambda n: (0.91, SMALL)))
+    ctx.log = messages.append
+    loop.check_start(ctx)
+    assert any("could not check slurm/round.sbatch" in message for message in messages)
+
+
+def test_a_round_job_at_the_queue_limit_stops_the_loop_with_slurms_reason(started):
+    class FullSlurm(FakeSlurm):
+        def wait(self, job_id, queue_limit_seconds):
+            self.last_reasons = "QOSMaxGRESPerUser"  # another GPU job of the user holds the one GPU
+            return "queue_limit"
+
+    slurm = FullSlurm(started / "state", lambda n: None)
+    assert loop.run_rounds(context(started, FakeAgent(lambda call: GOOD), slurm)) == "queue_limit"
+    assert slurm.canceled == ["1001"]
+    row = records.read_history(started / "state")[-1]
+    assert row["status"] == "queue_limit" and "SLURM reason: QOSMaxGRESPerUser" in row["note"]
 
 
 def test_a_state_folder_with_final_steps_done_refuses_a_new_search(started, monkeypatch):
diff --git a/tests/test_multi_loop.py b/tests/test_multi_loop.py
index 46fc09e..be4eddb 100644
--- a/tests/test_multi_loop.py
+++ b/tests/test_multi_loop.py
@@ -110,6 +110,14 @@ def __init__(self, state, results, gpu_minutes: float = 2.0, outcome: str = "end
         self.arrays = []
         self.canceled = []
         self.options = []
+        self.waits = []  # the max_parallel of each wait_array
+        self.tested = []  # (script, options) of each sbatch --test-only
+        self.test_errors = {}  # script -> the error that sbatch --test-only returns
+        self.last_reasons = ""
+
+    def test_only(self, script, cwd, options=()):
+        self.tested.append((script, options))
+        return self.test_errors.get(script)
 
     def submit_array(self, script, tasks_file, count, max_parallel, cwd, options=()):
         self.options.append(options)
@@ -126,9 +134,12 @@ def submit_array(self, script, tasks_file, count, max_parallel, cwd, options=())
                                    int(seed), *outcome, job_id=f"{job}{index}")
         return job
 
-    def wait_array(self, job, queue_limit_seconds):
+    def wait_array(self, job, queue_limit_seconds, max_parallel=None):
+        self.waits.append(max_parallel)
         if self.stop_on_wave is not None and len(self.arrays) == self.stop_on_wave:
             raise Terminated()
+        if self.outcome == "queue_limit":
+            self.last_reasons = "QOSMaxGRESPerUser"
         return self.outcome
 
     def array_accounting(self, job, count):
@@ -372,10 +383,13 @@ def test_queue_limit_stops_after_the_wave(multi):
     slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL) if n == 2 else None,
                            outcome="queue_limit")
     ctx = context(multi, FakeRoles(), slurm)
+    messages = []
+    ctx.log = messages.append
     multi_loop.copy_round_one(ctx)
     assert multi_loop.run_waves(ctx) == "queue_limit"
     statuses = [row["status"] for row in records.read_history(ctx.paths.state)]
     assert statuses == ["ok", "ok", "queue_limit", "queue_limit"]
+    assert any("SLURM reason: QOSMaxGRESPerUser" in message for message in messages)
 
 
 # Stops and missing rows
@@ -470,7 +484,8 @@ def test_experiment_files_agree():
     assert multi["target"]["val_accuracy"] == single["target"]["val_accuracy"]
     assert multi["baseline"] == single["baseline"]
     for experiment in (single, multi):
-        experiment["round"].pop("constraint", None)  # the GPUs can differ; the other round settings cannot
+        for key in ("partition", "constraint"):  # the partition and the GPUs can differ; the other settings cannot
+            experiment["round"].pop(key, None)
     assert multi["round"] == single["round"]
     multi_loop.check_roles(multi)
     assert multi["roles"]["planner"]["model"] != multi["roles"]["critic"]["model"]
@@ -571,7 +586,7 @@ def test_finish_closes_the_wave_of_a_stopped_loop_job(multi):
     accepted = records.read_json(state_of(multi) / "waves" / "wave_01" / "accepted.json")
     # The loop job submits wave 1 and gets SIGKILL: no handler runs, so current_job stays.
     original_wait = slurm.wait_array
-    slurm.wait_array = lambda job, limit: (_ for _ in ()).throw(KeyboardInterrupt())
+    slurm.wait_array = lambda job, limit, max_parallel=None: (_ for _ in ()).throw(KeyboardInterrupt())
     with pytest.raises(KeyboardInterrupt):
         multi_loop.submit_wave(ctx, 1, accepted)
     slurm.wait_array = original_wait
@@ -667,15 +682,47 @@ def test_check_report_reads_the_baseline_metrics(tmp_path):
     assert check_report.unmatched_numbers(report, state, tmp_path / config.BASELINE_METRICS) == []
 
 
-def test_round_constraint_reaches_the_array_jobs(multi):
+def test_round_partition_and_constraint_reach_the_array_jobs(multi):
+    slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n))
+    planner = context(multi, FakeRoles(), slurm)
+    planner.plan_only = True
+    assert multi_loop.run(planner, "plan_only") == 0
+    ctx = context(multi, FakeRoles(), slurm, max_round_attempts=4)
+    ctx.experiment["round"].update(partition="kempner_interactive", constraint="a100-mig")
+    assert multi_loop.run(ctx, "loop") == 0
+    expected = ("--partition=kempner_interactive", "--constraint=a100-mig")
+    assert slurm.options == [expected] * 2  # the wave and the confirmation rounds
+    assert slurm.tested[-1] == ("slurm/round_array.sbatch", expected)  # the start check of the loop job
+
+
+def test_plan_only_stops_before_any_role_step_if_slurm_would_reject_the_array_jobs(multi):
+    slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: None)
+    slurm.test_errors["slurm/round_array.sbatch"] = ("sbatch: error: Batch job submission failed: "
+                                                     "Invalid partition name specified")
+    agent = FakeRoles()
+    ctx = context(multi, agent, slurm)
+    ctx.plan_only = True
+    ctx.experiment["round"]["partition"] = "kempner_interactiv"
+    with pytest.raises(SystemExit, match="SLURM would reject the round jobs of slurm/round_array.sbatch"):
+        multi_loop.run(ctx, "plan_only")
+    assert not agent.calls
+
+
+def test_one_gpu_for_each_user_runs_the_rounds_of_a_wave_in_turns(multi):
     slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n))
     planner = context(multi, FakeRoles(), slurm)
     planner.plan_only = True
     assert multi_loop.run(planner, "plan_only") == 0
     ctx = context(multi, FakeRoles(), slurm, max_round_attempts=4)
-    ctx.experiment["round"]["constraint"] = "a100-mig"
+    ctx.experiment["wave"]["max_parallel"] = 1
     assert multi_loop.run(ctx, "loop") == 0
-    assert slurm.options == [("--constraint=a100-mig",)] * 2  # the wave and the confirmation rounds
+    assert [max_parallel for _, _, max_parallel in slurm.arrays] == [1, 1]  # the throttle of the wave and the seeds
+    assert slurm.waits == [1, 1]  # so a task that waits for the task before it adds no queue time
+    # The loop job plans the time limit of each turn: 3 rounds take 3 turns of 10 minutes.
+    rest = ctx.budget["queue_limit_minutes"] + ctx.budget["final_reserve_minutes"]
+    assert multi_loop.array_minutes(ctx, 3) == 30 + rest
+    ctx.experiment["wave"]["max_parallel"] = 3
+    assert multi_loop.array_minutes(ctx, 3) == 10 + rest
 
 
 def test_a_state_folder_with_final_steps_done_refuses_a_new_search(multi):

From 14fda23ce6d1cbf71be001f6ed545b9908c960ef Mon Sep 17 00:00:00 2001
From: mmshad <31811010+mmshad@users.noreply.github.com>
Date: Mon, 2 Nov 2026 22:06:27 -0500
Subject: [PATCH 2/5] Document how to run the example on kempner_interactive

---
 AGENTS.md             |  4 ++--
 README.md             | 24 ++++++++++++++++++++++--
 experiment.yaml       |  1 +
 experiment_multi.yaml | 11 ++++++-----
 4 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/AGENTS.md b/AGENTS.md
index 419b990..f82e872 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -17,7 +17,7 @@ This project prunes ResNet-18 on CIFAR-10 on the Kempner AI cluster. Read this f
 ## Supervised round
 
 - Write the configuration only to `agent/next_config.json`. Use the keys and ranges in `configs/round.schema.json`.
-- Submit the round as a background command: `sbatch --wait slurm/round.sbatch agent/next_config.json`.
+- Submit the round as a background command: `sbatch --wait slurm/round.sbatch agent/next_config.json`. If `round.partition` or `round.constraint` is set in `experiment.yaml`, add it after `--wait`, for example `--partition=kempner_interactive --constraint=a100-mig`.
 - Run SLURM commands on their own, without pipes, redirects, `cd`, or `$(...)`. The sandbox needs this.
 - Query only your own jobs, with `squeue -j <id>` and `sacct -j <id>`. Do not run `squeue` without `-j`.
 - To check a finished round, run `/check-round N`.
@@ -33,5 +33,5 @@ This project prunes ResNet-18 on CIFAR-10 on the Kempner AI cluster. Read this f
 
 - Each role step has a role prompt from `prompts/multi/roles/`: the planner, the critic, the writer, or the auditor. Follow it.
 - The multi-agent loop uses `experiment_multi.yaml` and the state folder `state_multi/`. Read only the state folder that your prompt names, and `state/baseline/metrics.json`.
-- A wave has up to 3 rounds that run at the same time. The loop script stops after `target.plateau_rounds` complete rounds in a row, in round order, without a smaller passing model.
+- A wave has up to 3 rounds, planned together. They run at the same time, or in turns with `wave.max_parallel: 1`. The loop script stops after `target.plateau_rounds` complete rounds in a row, in round order, without a smaller passing model.
 - Return only the required JSON or text. Do not try to run commands or write files. The loop script runs the rounds.
diff --git a/README.md b/README.md
index b33e37c..02027d3 100644
--- a/README.md
+++ b/README.md
@@ -16,7 +16,7 @@ A second version of the loop uses four agent roles and runs 3 rounds at a time;
 
 ## Before you start
 
-- A Kempner AI cluster account, and a SLURM account that can use `kempner_requeue` and `shared`.
+- A Kempner AI cluster account, and a SLURM account that can use `kempner_requeue` and `shared`. To use `kempner_interactive` in place of `kempner_requeue`, see [Run on kempner_interactive](#run-on-kempner_interactive).
 - Claude Code from the native installer, because the loop job looks for `claude` in `~/.local/bin`. Sign in with the Claude access that HUIT provides through Harvard SSO, and leave `ANTHROPIC_API_KEY` unset.
 - [uv](https://docs.astral.sh/uv/).
 - The data is the shared copy of CIFAR-10 on the cluster. CIFAR-10 is public data (Level 1).
@@ -176,6 +176,26 @@ sbatch slurm/multi_loop.sbatch                               # 3. run the waves
 
 To audit the report of the single-agent loop, run `STATE_DIR=state sbatch slurm/multi_loop.sbatch --audit-only`. To compare designs, run each one in its own clone with only the baseline and round 1; see [Comparison of the designs](#comparison-of-the-designs).
 
+## Run on kempner_interactive
+
+You can run the GPU jobs of this example on `kempner_interactive` in batch mode, because each job is short and uses one GPU. The partition has only A100 MIG slices (3g.20gb), and each user can use one GPU at a time. The job scripts ask for `kempner_requeue`, and `slurm/baseline.sbatch` and `slurm/round.sbatch` ask for full GPUs, so set the partition and the constraint:
+
+1. In the `round` section of `experiment.yaml`, set `partition: kempner_interactive` and `constraint: a100-mig`. For the multi-agent loop, set the same two lines in `experiment_multi.yaml`.
+2. For the multi-agent loop, also set `wave.max_parallel: 1` and `budget.final_reserve_minutes: 120` in `experiment_multi.yaml`. The rounds of a wave and the two confirmation rounds then run in turns.
+3. Commit the changes.
+4. Train the baseline on a MIG slice:
+
+   ```bash
+   sbatch -p kempner_interactive --constraint=a100-mig slurm/baseline.sbatch
+   ```
+
+5. Run round 1 in the `test` partition, as in [Round 1](#round-1-supervised). The agent adds `--partition=kempner_interactive --constraint=a100-mig` to its `sbatch` command. Check the options before you approve the command.
+6. Start the loop job as usual. It runs on `shared`, and it submits the round jobs to `kempner_interactive`.
+
+- **GPU minutes.** A round uses about 2.6 GPU minutes on a MIG slice, so the default `gpu_minutes` of both loops are enough.
+- **One GPU.** While the loop job runs, do not keep another GPU job on `kempner_interactive`. The round jobs wait for it, and the loop stops at `queue_limit_minutes`. The loop log then shows the SLURM reason `QOSMaxGRESPerUser`.
+- **Start check.** At start, the loop job asks SLURM with `sbatch --test-only` if it would accept the round jobs. If not, the loop job stops before any agent step and shows the SLURM error.
+
 ## Test the loop with small limits
 
 Put test copies in `local/`, which Git ignores. `EXPERIMENT_FILE` and `STATE_DIR` point the jobs to them:
@@ -250,7 +270,7 @@ Claude Code 2.1.302 to 2.1.309 (the multi-agent loop: 2.1.307 to 2.1.309), uv 0.
 ## Adapt the example
 
 - **Task.** Change the data, the model, the metric, the round configuration schema, the accuracy threshold, and the limits.
-- **GPUs.** To run the round jobs on other GPUs, set `round.constraint` in the experiment file. For example, set `a100-mig` in `experiment.yaml` for MIG slices, or `a100|h100|h200|rtx6000pro` in `experiment_multi.yaml` for full GPUs. The loop scripts pass it to `sbatch`, and it replaces the `--constraint` line of the job scripts. A MIG slice is 3 to 4 times slower than a full GPU, so set `gpu_minutes` to match.
+- **GPUs.** To run the round jobs on other GPUs, set `round.constraint` in the experiment file. For example, set `a100-mig` in `experiment.yaml` for MIG slices, or `a100|h100|h200|rtx6000pro` in `experiment_multi.yaml` for full GPUs. The loop scripts pass it to `sbatch`, and it replaces the `--constraint` line of the job scripts. `round.partition` replaces the `--partition` line in the same way; see [Run on kempner_interactive](#run-on-kempner_interactive). A MIG slice is 3 to 4 times slower than a full GPU, so set `gpu_minutes` to match.
 - **Changes.** Commit each change to a tracked file. The jobs refuse to start while a tracked file has changes.
 - **Keep.** Keep the fixed evaluator, the read-only agent step, and the loop script that submits the jobs. Also keep the limits outside the model, the run records, the confirmation round, and the one-time test evaluation.
 
diff --git a/experiment.yaml b/experiment.yaml
index d5be90b..b9c84bf 100644
--- a/experiment.yaml
+++ b/experiment.yaml
@@ -17,6 +17,7 @@ round:
   batch_size: 256
   momentum: 0.9
   time_minutes: 10           # must agree with --time in slurm/round.sbatch
+  # partition: kempner_interactive   # optional: replaces the --partition line of the round job scripts
   # constraint: a100-mig     # optional: replaces the --constraint line of the round job scripts
 
 budget:
diff --git a/experiment_multi.yaml b/experiment_multi.yaml
index ea58ac2..7b9c016 100644
--- a/experiment_multi.yaml
+++ b/experiment_multi.yaml
@@ -1,7 +1,7 @@
 # Settings and limits of the multi-agent design.
 # Read this file before you start the multi-agent loop job: it sets the roles, the models, the effort levels,
-# and the budget. The sections target (val_accuracy), baseline, and round (except constraint) must agree with
-# experiment.yaml; a test checks this.
+# and the budget. The sections target (val_accuracy), baseline, and round (except partition and constraint) must
+# agree with experiment.yaml; a test checks this.
 
 target:
   val_accuracy: 0.905        # accuracy threshold: 0.900 plus a margin for noise in 5,000 validation images
@@ -20,11 +20,12 @@ round:
   batch_size: 256
   momentum: 0.9
   time_minutes: 10           # must agree with --time in slurm/round.sbatch and slurm/round_array.sbatch
+  # partition: kempner_interactive   # optional: replaces the --partition line of the round job scripts
   # constraint: a100|h100|h200|rtx6000pro   # optional: full GPUs in place of the MIG slices of slurm/round_array.sbatch
 
 wave:
   size: 3                    # configurations in each wave; fewer when a limit allows fewer
-  max_parallel: 3            # round jobs of one wave that can run at the same time
+  max_parallel: 3            # round jobs of one wave that can run at the same time; 1 on kempner_interactive
 
 # One model and one effort level for each role. Use full model names: an alias such as opus moves to newer models.
 # Effort levels: low, medium, high, xhigh, max. A checker uses a different model from the role that it checks.
@@ -62,7 +63,7 @@ budget:
   gpu_minutes: 150           # rounds after round 1, on A100 MIG slices; failed and requeued runs count
   max_round_attempts: 45     # round 1 included; all agents of a run share them
   max_failed_waves: 2        # waves in a row with no complete round
-  queue_limit_minutes: 45    # total queue time of one array job, requeues included
-  final_reserve_minutes: 110 # time kept for the final steps
+  queue_limit_minutes: 45    # total queue time of one array job, requeues included, without waits behind its own tasks
+  final_reserve_minutes: 110 # time kept for the final steps; 120 with max_parallel: 1, for 2 confirmation rounds in turns
 
 confirmation_seeds: [1, 2]   # the selected configuration runs again with each seed, as one array job

From 397c50effd8b19248ef972ab1f1de828def0069e Mon Sep 17 00:00:00 2001
From: mmshad <31811010+mmshad@users.noreply.github.com>
Date: Mon, 2 Nov 2026 22:30:18 -0500
Subject: [PATCH 3/5] Fix the review findings of the kempner_interactive change

---
 AGENTS.md                    |  2 +-
 README.md                    | 19 +++++++++------
 slurm/round_array.sbatch     |  6 ++---
 src/pruning_example/loop.py  |  4 +--
 src/pruning_example/slurm.py | 14 ++++++-----
 tests/test_array.py          | 47 +++++++++++++++++++++++++++++++++---
 tests/test_loop.py           | 10 ++++++++
 tests/test_multi_loop.py     | 16 ++++++++++++
 8 files changed, 94 insertions(+), 24 deletions(-)

diff --git a/AGENTS.md b/AGENTS.md
index f82e872..e6ff1c5 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -17,7 +17,7 @@ This project prunes ResNet-18 on CIFAR-10 on the Kempner AI cluster. Read this f
 ## Supervised round
 
 - Write the configuration only to `agent/next_config.json`. Use the keys and ranges in `configs/round.schema.json`.
-- Submit the round as a background command: `sbatch --wait slurm/round.sbatch agent/next_config.json`. If `round.partition` or `round.constraint` is set in `experiment.yaml`, add it after `--wait`, for example `--partition=kempner_interactive --constraint=a100-mig`.
+- Submit the round as a background command: `sbatch --wait slurm/round.sbatch agent/next_config.json`. If `round.partition` or `round.constraint` is set in `experiment.yaml`, add each one after `--wait`, for example `--partition=kempner_interactive --constraint=a100-mig`.
 - Run SLURM commands on their own, without pipes, redirects, `cd`, or `$(...)`. The sandbox needs this.
 - Query only your own jobs, with `squeue -j <id>` and `sacct -j <id>`. Do not run `squeue` without `-j`.
 - To check a finished round, run `/check-round N`.
diff --git a/README.md b/README.md
index 02027d3..55768aa 100644
--- a/README.md
+++ b/README.md
@@ -53,7 +53,7 @@ source cluster.env
 sbatch slurm/baseline.sbatch             # about 2 minutes on one GPU
 ```
 
-`state/baseline/metrics.json` then shows a validation accuracy of about 0.94.
+`state/baseline/metrics.json` then shows a validation accuracy of about 0.94. On `kempner_interactive`, use the baseline command in [Run on kempner_interactive](#run-on-kempner_interactive).
 
 ## Round 1 (supervised)
 
@@ -178,7 +178,9 @@ To audit the report of the single-agent loop, run `STATE_DIR=state sbatch slurm/
 
 ## Run on kempner_interactive
 
-You can run the GPU jobs of this example on `kempner_interactive` in batch mode, because each job is short and uses one GPU. The partition has only A100 MIG slices (3g.20gb), and each user can use one GPU at a time. The job scripts ask for `kempner_requeue`, and `slurm/baseline.sbatch` and `slurm/round.sbatch` ask for full GPUs, so set the partition and the constraint:
+You can run the GPU jobs of this example on `kempner_interactive` in batch mode, because each job is short and uses one GPU. The partition has only A100 MIG slices (3g.20gb), and each user can use one GPU at a time. The job scripts ask for `kempner_requeue`, and `slurm/baseline.sbatch` and `slurm/round.sbatch` ask for full GPUs, so set the partition and the constraint.
+
+Do the setup above, but use these steps in place of the baseline and round 1 commands:
 
 1. In the `round` section of `experiment.yaml`, set `partition: kempner_interactive` and `constraint: a100-mig`. For the multi-agent loop, set the same two lines in `experiment_multi.yaml`.
 2. For the multi-agent loop, also set `wave.max_parallel: 1` and `budget.final_reserve_minutes: 120` in `experiment_multi.yaml`. The rounds of a wave and the two confirmation rounds then run in turns.
@@ -186,15 +188,16 @@ You can run the GPU jobs of this example on `kempner_interactive` in batch mode,
 4. Train the baseline on a MIG slice:
 
    ```bash
-   sbatch -p kempner_interactive --constraint=a100-mig slurm/baseline.sbatch
+   source cluster.env
+   sbatch --partition=kempner_interactive --constraint=a100-mig slurm/baseline.sbatch
    ```
 
-5. Run round 1 in the `test` partition, as in [Round 1](#round-1-supervised). The agent adds `--partition=kempner_interactive --constraint=a100-mig` to its `sbatch` command. Check the options before you approve the command.
+5. Run round 1 in the `test` partition, as in [Round 1](#round-1-supervised). In the prompt, write the command as `sbatch --wait --partition=kempner_interactive --constraint=a100-mig slurm/round.sbatch agent/next_config.json`. Refuse the command if it does not have these two options.
 6. Start the loop job as usual. It runs on `shared`, and it submits the round jobs to `kempner_interactive`.
 
-- **GPU minutes.** A round uses about 2.6 GPU minutes on a MIG slice, so the default `gpu_minutes` of both loops are enough.
-- **One GPU.** While the loop job runs, do not keep another GPU job on `kempner_interactive`. The round jobs wait for it, and the loop stops at `queue_limit_minutes`. The loop log then shows the SLURM reason `QOSMaxGRESPerUser`.
-- **Start check.** At start, the loop job asks SLURM with `sbatch --test-only` if it would accept the round jobs. If not, the loop job stops before any agent step and shows the SLURM error.
+- **GPU minutes.** A round uses about 2.6 GPU minutes on a MIG slice. The single-agent loop job then uses about 40 GPU minutes, within its default 60 (`budget.gpu_minutes`). The budget of the multi-agent loop is for MIG slices already.
+- **One GPU.** While the loop job runs, do not keep another GPU job on `kempner_interactive`. A round job waits for it, and if the round job waits more than 45 minutes (`budget.queue_limit_minutes`), the loop stops. The loop log then shows the SLURM reason `QOSMaxGRESPerUser`.
+- **Start check.** Before a search starts, the loop job asks SLURM with `sbatch --test-only` if it would accept the round jobs. If not, the loop job stops before any agent step and shows the SLURM error.
 
 ## Test the loop with small limits
 
@@ -270,7 +273,7 @@ Claude Code 2.1.302 to 2.1.309 (the multi-agent loop: 2.1.307 to 2.1.309), uv 0.
 ## Adapt the example
 
 - **Task.** Change the data, the model, the metric, the round configuration schema, the accuracy threshold, and the limits.
-- **GPUs.** To run the round jobs on other GPUs, set `round.constraint` in the experiment file. For example, set `a100-mig` in `experiment.yaml` for MIG slices, or `a100|h100|h200|rtx6000pro` in `experiment_multi.yaml` for full GPUs. The loop scripts pass it to `sbatch`, and it replaces the `--constraint` line of the job scripts. `round.partition` replaces the `--partition` line in the same way; see [Run on kempner_interactive](#run-on-kempner_interactive). A MIG slice is 3 to 4 times slower than a full GPU, so set `gpu_minutes` to match.
+- **GPUs.** To run the round jobs on other GPUs, set `round.constraint` in the experiment file. For example, set `a100-mig` in `experiment.yaml` for MIG slices, or `a100|h100|h200|rtx6000pro` in `experiment_multi.yaml` for full GPUs. The loop scripts pass it to `sbatch`, and it replaces the `--constraint` line of the job scripts. `round.partition` replaces the `--partition` line in the same way; see [Run on kempner_interactive](#run-on-kempner_interactive). A round uses about 1.2 GPU minutes on a full GPU and about 2.6 on a MIG slice, so set `budget.gpu_minutes` to match.
 - **Changes.** Commit each change to a tracked file. The jobs refuse to start while a tracked file has changes.
 - **Keep.** Keep the fixed evaluator, the read-only agent step, and the loop script that submits the jobs. Also keep the limits outside the model, the run records, the confirmation round, and the one-time test evaluation.
 
diff --git a/slurm/round_array.sbatch b/slurm/round_array.sbatch
index 2b75d8f..19f5b85 100644
--- a/slurm/round_array.sbatch
+++ b/slurm/round_array.sbatch
@@ -20,9 +20,9 @@
 # kempner_requeue can preempt a task. The task saves a checkpoint after each epoch and continues after a requeue.
 # Keep --time equal to round.time_minutes in the experiment file.
 # The tasks use A100 MIG slices (3g.20gb). A round is small, and the slices are often free when the full GPUs
-# are busy. A round takes 1 to 3 minutes on a slice, 3 to 4 times longer than on a full GPU, so
-# experiment_multi.yaml has a larger GPU budget. The single-agent loop uses this script only for several
-# confirmation seeds; its rounds keep the full GPUs of slurm/round.sbatch.
+# are busy. A round takes 1 to 3 minutes on a slice, about twice as long as on a full GPU, so
+# experiment_multi.yaml has a larger GPU budget. The single-agent loop uses this script only for the
+# confirmation rounds when confirmation_seeds is not [1]; its rounds use slurm/round.sbatch.
 
 set -euo pipefail
 cd "$SLURM_SUBMIT_DIR"
diff --git a/src/pruning_example/loop.py b/src/pruning_example/loop.py
index 9a153a6..41cdf40 100644
--- a/src/pruning_example/loop.py
+++ b/src/pruning_example/loop.py
@@ -309,8 +309,8 @@ def check_start(ctx: Context) -> None:
         raise SystemExit("There is no baseline. Run 'sbatch slurm/baseline.sbatch' first.")
     if not (ctx.paths.state / "rounds" / "round_01" / "metrics.json").exists():
         raise SystemExit("There is no round 1. Run the supervised round first.")
-    several_seeds = len(finish.confirmation_seeds(ctx.experiment)) > 1  # their confirmation rounds run as an array job
-    scripts = ("slurm/round.sbatch",) + (("slurm/round_array.sbatch",) if several_seeds else ())
+    array = finish.confirmation_seeds(ctx.experiment) != [1]  # then the confirmation rounds run as an array job
+    scripts = ("slurm/round.sbatch",) + (("slurm/round_array.sbatch",) if array else ())
     check_time_lines(ctx, scripts)
     records.check_integrity(project)
     check_round_jobs(ctx, scripts)
diff --git a/src/pruning_example/slurm.py b/src/pruning_example/slurm.py
index e970ff6..42f898b 100644
--- a/src/pruning_example/slurm.py
+++ b/src/pruning_example/slurm.py
@@ -16,7 +16,8 @@
 }
 QUEUE_WAIT_STATES = {"PENDING", "REQUEUED", "REQUEUE_HOLD", "REQUEUE_FED", "CONFIGURING"}
 # sbatch errors of a busy SLURM controller. A new try can succeed.
-TRANSIENT_ERRORS = ("Socket timed out", "timed out", "Resource temporarily unavailable", "try again")
+TRANSIENT_ERRORS = ("Socket timed out", "timed out", "Resource temporarily unavailable", "try again",
+                    "Unable to contact slurm controller")
 
 
 class SlurmError(RuntimeError):
@@ -222,10 +223,11 @@ def wait(self, job_id: str, queue_limit_seconds: float) -> str:
     def wait_array(self, array_id: str, queue_limit_seconds: float, max_parallel: int | None = None) -> str:
         """Wait until all tasks of an array job leave the queue. Return "ended" or "queue_limit".
 
-        Queue time counts while at least one task waits and fewer than max_parallel tasks run. So a task
-        that waits only because max_parallel tasks of the array already run adds no queue time. After the
-        queue limit, the waiting tasks are canceled, also a task that a preemption put back in the queue,
-        and the running tasks can end.
+        Queue time counts while at least one task waits and fewer than max_parallel tasks are RUNNING. So a
+        task that waits only because max_parallel tasks of the array already run adds no queue time. A task
+        in another state, for example stuck in COMPLETING, does not stop the queue time. After the queue
+        limit, the waiting tasks are canceled, also a task that a preemption put back in the queue, and the
+        running tasks can end.
         """
         queued = 0.0
         last = self.clock()
@@ -244,7 +246,7 @@ def wait_array(self, array_id: str, queue_limit_seconds: float, max_parallel: in
             if states == {}:
                 return "queue_limit" if over_limit else "ended"
             waiting = bool(states) and any(state in QUEUE_WAIT_STATES for state in states.values())
-            running = sum(state not in QUEUE_WAIT_STATES for state in (states or {}).values())
+            running = sum(state == "RUNNING" for state in (states or {}).values())
             if waiting and not over_limit and (max_parallel is None or running < max_parallel):
                 queued += now - last
             last = now
diff --git a/tests/test_array.py b/tests/test_array.py
index f921aa8..23665d5 100644
--- a/tests/test_array.py
+++ b/tests/test_array.py
@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import subprocess
+
 import pytest
 
 from pruning_example.slurm import Slurm, SlurmError, parse_array_sacct, parse_array_squeue, task_indexes
@@ -73,9 +75,10 @@ def test_queue_limit_cancels_only_waiting_tasks_also_after_a_preemption():
         {0: "PENDING"},  # a preemption put task 0 back in the queue: cancel it too
         {},
     ]
-    slurm = fake_slurm(states)
-    assert slurm.wait_array("500", queue_limit_seconds=100) == "queue_limit"
-    assert slurm.waiting_cancels == 2
+    for max_parallel in (None, 3):  # 3 is the default of the waves
+        slurm = fake_slurm(states)
+        assert slurm.wait_array("500", queue_limit_seconds=100, max_parallel=max_parallel) == "queue_limit"
+        assert slurm.waiting_cancels == 2
 
 
 def test_a_task_that_waits_for_a_running_task_of_the_array_adds_no_queue_time():
@@ -92,8 +95,41 @@ def test_a_task_that_waits_for_a_running_task_of_the_array_adds_no_queue_time():
 
 def test_a_wait_with_no_running_task_counts_and_keeps_slurms_reason():
     slurm = fake_slurm([{0: "PENDING", 1: "PENDING"}] * 3 + [{}])
+
+    def reasons(array_id):
+        assert slurm.waiting_cancels == 0  # read before the waiting tasks are canceled
+        return "QOSMaxGRESPerUser"
+
+    slurm.pending_reasons = reasons
+    assert slurm.wait_array("500", queue_limit_seconds=100, max_parallel=1) == "queue_limit"
+    assert slurm.last_reasons == "QOSMaxGRESPerUser"
+    assert slurm.waiting_cancels > 0
+
+
+def test_a_task_stuck_in_completing_does_not_stop_the_queue_time():
+    slurm = fake_slurm([{0: "COMPLETING", 1: "PENDING"}] * 3 + [{}])
     assert slurm.wait_array("500", queue_limit_seconds=100, max_parallel=1) == "queue_limit"
-    assert slurm.last_reasons == "QOSMaxGRESPerUser"  # read before the waiting tasks are canceled
+
+
+def test_pending_reasons_gives_each_reason_once_or_unknown():
+    answers = iter([Result(0, out="QOSMaxGRESPerUser\nPriority\nQOSMaxGRESPerUser\n"), Result(0, out="")])
+    seen = []
+
+    def run(args, timeout=120, **kwargs):
+        seen.append(args)
+        return next(answers)
+
+    slurm = Slurm(sleep=lambda seconds: None)
+    slurm.run = run
+    assert slurm.pending_reasons("500") == "Priority, QOSMaxGRESPerUser"
+    assert seen[0] == ["squeue", "-h", "-r", "-t", "PENDING", "-j", "500", "-o", "%r"]
+    assert slurm.pending_reasons("500") == "unknown"  # no task waits any more
+
+    def slow(args, timeout=120, **kwargs):
+        raise subprocess.TimeoutExpired(args, timeout)
+
+    slurm.run = slow
+    assert slurm.pending_reasons("500") == "unknown"
 
 
 def test_test_only_returns_slurms_error_and_raises_only_for_a_busy_controller():
@@ -102,6 +138,7 @@ def test_test_only_returns_slurms_error_and_raises_only_for_a_busy_controller():
                       "in partition kempner_interactive"),
         Result(1, err="sbatch: error: Batch job submission failed: Requested node configuration is not available"),
         Result(1, err="sbatch: error: Batch job submission failed: Socket timed out on send/recv operation"),
+        Result(1, err="sbatch: error: Batch job submission failed: Unable to contact slurm controller (connect failure)"),
     ])
     seen = []
 
@@ -117,6 +154,8 @@ def run(args, timeout=120, **kwargs):
     assert "Requested node configuration is not available" in slurm.test_only("slurm/round.sbatch", ".")
     with pytest.raises(SlurmError, match="Socket timed out"):
         slurm.test_only("slurm/round.sbatch", ".")
+    with pytest.raises(SlurmError, match="Unable to contact slurm controller"):  # not a rejection of the job
+        slurm.test_only("slurm/round.sbatch", ".")
 
 
 def test_submit_array_passes_the_range_and_the_throttle():
diff --git a/tests/test_loop.py b/tests/test_loop.py
index abe20bc..55507ce 100644
--- a/tests/test_loop.py
+++ b/tests/test_loop.py
@@ -479,6 +479,16 @@ def test_the_loop_job_stops_before_any_agent_step_if_slurm_would_reject_the_roun
     assert agent.calls == 0 and not slurm.submitted
 
 
+def test_the_start_check_covers_each_job_script_that_the_loop_uses(started):
+    for seeds, scripts in (([1], ["slurm/round.sbatch"]),  # one confirmation round: one round job
+                           ([2], ["slurm/round.sbatch", "slurm/round_array.sbatch"])):  # any other seeds: an array job
+        slurm = FakeSlurm(started / "state", lambda n: (0.91, SMALL))
+        ctx = context(started, FakeAgent(lambda call: GOOD), slurm)
+        ctx.experiment["confirmation_seeds"] = seeds
+        loop.check_start(ctx)
+        assert [script for script, _ in slurm.tested] == scripts
+
+
 def test_the_start_check_goes_on_if_the_slurm_controller_does_not_answer(started):
     class BusySlurm(FakeSlurm):
         def test_only(self, script, cwd, options=()):
diff --git a/tests/test_multi_loop.py b/tests/test_multi_loop.py
index be4eddb..4d1e076 100644
--- a/tests/test_multi_loop.py
+++ b/tests/test_multi_loop.py
@@ -725,6 +725,22 @@ def test_one_gpu_for_each_user_runs_the_rounds_of_a_wave_in_turns(multi):
     assert multi_loop.array_minutes(ctx, 3) == 10 + rest
 
 
+@pytest.mark.parametrize("max_parallel, arrays", [(1, 0), (3, 1)])
+def test_a_wave_runs_only_if_the_loop_job_has_time_for_its_turns(multi, max_parallel, arrays):
+    slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n))
+    planner = context(multi, FakeRoles(), slurm)
+    planner.plan_only = True
+    assert multi_loop.run(planner, "plan_only") == 0
+    assert len(records.read_json(state_of(multi) / "waves" / "wave_01" / "accepted.json")) == 3
+    ctx = context(multi, FakeRoles(), slurm)
+    ctx.experiment["wave"]["max_parallel"] = max_parallel
+    # Time for one turn of 10 minutes, the queue limit, and the final reserve, but not for three turns.
+    ctx.clock = lambda: 0.0
+    ctx.job_end = 60 * (10 + ctx.budget["queue_limit_minutes"] + ctx.budget["final_reserve_minutes"] + 5)
+    assert multi_loop.run_waves(ctx) == "time_left"
+    assert len(slurm.arrays) == arrays  # in turns, wave 1 does not fit
+
+
 def test_a_state_folder_with_final_steps_done_refuses_a_new_search(multi):
     slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n))
     ctx, code = plan_then_run(multi, FakeRoles(), slurm, max_round_attempts=4)

From 5f02339f9aa3512920aca0ad6b55ef89674af711 Mon Sep 17 00:00:00 2001
From: mmshad <31811010+mmshad@users.noreply.github.com>
Date: Mon, 2 Nov 2026 22:36:21 -0500
Subject: [PATCH 4/5] Point to the kempner_interactive steps before the
 baseline command

---
 README.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 55768aa..83c9c04 100644
--- a/README.md
+++ b/README.md
@@ -46,14 +46,14 @@ exit
 
 `cluster.env` puts the uv cache and the Python interpreter in lab storage. The agent sandbox hides your home folder, so `.venv/bin/python` must not link to a Python in your home folder.
 
-Train the baseline:
+Train the baseline. On `kempner_interactive`, do not run this command: use the steps in [Run on kempner_interactive](#run-on-kempner_interactive) in place of the baseline and round 1 commands.
 
 ```bash
 source cluster.env
 sbatch slurm/baseline.sbatch             # about 2 minutes on one GPU
 ```
 
-`state/baseline/metrics.json` then shows a validation accuracy of about 0.94. On `kempner_interactive`, use the baseline command in [Run on kempner_interactive](#run-on-kempner_interactive).
+`state/baseline/metrics.json` then shows a validation accuracy of about 0.94.
 
 ## Round 1 (supervised)
 
@@ -196,7 +196,7 @@ Do the setup above, but use these steps in place of the baseline and round 1 com
 6. Start the loop job as usual. It runs on `shared`, and it submits the round jobs to `kempner_interactive`.
 
 - **GPU minutes.** A round uses about 2.6 GPU minutes on a MIG slice. The single-agent loop job then uses about 40 GPU minutes, within its default 60 (`budget.gpu_minutes`). The budget of the multi-agent loop is for MIG slices already.
-- **One GPU.** While the loop job runs, do not keep another GPU job on `kempner_interactive`. A round job waits for it, and if the round job waits more than 45 minutes (`budget.queue_limit_minutes`), the loop stops. The loop log then shows the SLURM reason `QOSMaxGRESPerUser`.
+- **One GPU.** While the loop job runs, do not keep another GPU job on `kempner_interactive`. Another GPU job makes the round jobs wait. If a round job waits more than 45 minutes (`budget.queue_limit_minutes`), the loop stops. The loop log then shows the SLURM reason `QOSMaxGRESPerUser`.
 - **Start check.** Before a search starts, the loop job asks SLURM with `sbatch --test-only` if it would accept the round jobs. If not, the loop job stops before any agent step and shows the SLURM error.
 
 ## Test the loop with small limits

From 6a384dfb68be480a2ce0672311b07724c71ab9bc Mon Sep 17 00:00:00 2001
From: mmshad <31811010+mmshad@users.noreply.github.com>
Date: Tue, 3 Nov 2026 00:22:03 -0500
Subject: [PATCH 5/5] Add the results of the kempner_interactive tests

---
 README.md                |  6 +++---
 tests/test_array.py      | 12 ++++++------
 tests/test_loop.py       |  3 +--
 tests/test_multi_loop.py |  4 ++--
 4 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index 83c9c04..1fd12f6 100644
--- a/README.md
+++ b/README.md
@@ -178,7 +178,7 @@ To audit the report of the single-agent loop, run `STATE_DIR=state sbatch slurm/
 
 ## Run on kempner_interactive
 
-You can run the GPU jobs of this example on `kempner_interactive` in batch mode, because each job is short and uses one GPU. The partition has only A100 MIG slices (3g.20gb), and each user can use one GPU at a time. The job scripts ask for `kempner_requeue`, and `slurm/baseline.sbatch` and `slurm/round.sbatch` ask for full GPUs, so set the partition and the constraint.
+You can run the GPU jobs of this example on `kempner_interactive` in batch mode, because each job is short and uses one GPU. The partition has only A100 MIG slices (3g.20gb), and each user can use one GPU at a time. The job scripts ask for `kempner_requeue`, and `slurm/baseline.sbatch` and `slurm/round.sbatch` ask for full GPUs. On `kempner_interactive`, SLURM rejects these two scripts with `Requested node configuration is not available`, so set the partition and the constraint.
 
 Do the setup above, but use these steps in place of the baseline and round 1 commands:
 
@@ -189,7 +189,7 @@ Do the setup above, but use these steps in place of the baseline and round 1 com
 
    ```bash
    source cluster.env
-   sbatch --partition=kempner_interactive --constraint=a100-mig slurm/baseline.sbatch
+   sbatch --partition=kempner_interactive --constraint=a100-mig slurm/baseline.sbatch   # about 5 minutes on a MIG slice
    ```
 
 5. Run round 1 in the `test` partition, as in [Round 1](#round-1-supervised). In the prompt, write the command as `sbatch --wait --partition=kempner_interactive --constraint=a100-mig slurm/round.sbatch agent/next_config.json`. Refuse the command if it does not have these two options.
@@ -268,7 +268,7 @@ Each design ran one time, from the same baseline and round 1, with the same sear
 
 ## Tested with
 
-Claude Code 2.1.302 to 2.1.309 (the multi-agent loop: 2.1.307 to 2.1.309), uv 0.8.19, Python 3.12.11, PyTorch 2.14.1 with CUDA 13.0, torchvision 0.29.1, and NVIDIA driver 610.57.04. GPUs: A100 (full GPUs and 3g.20gb MIG slices), H100, H200, and RTX PRO 6000.
+Claude Code 2.1.302 to 2.1.317 (the multi-agent loop: 2.1.307 to 2.1.317), uv 0.8.19, Python 3.12.11, PyTorch 2.14.1 with CUDA 13.0, torchvision 0.29.1, and NVIDIA driver 610.57.04. GPUs: A100 (full GPUs, and 3g.20gb MIG slices on `kempner_requeue` and `kempner_interactive`), H100, H200, and RTX PRO 6000.
 
 ## Adapt the example
 
diff --git a/tests/test_array.py b/tests/test_array.py
index 23665d5..5a17636 100644
--- a/tests/test_array.py
+++ b/tests/test_array.py
@@ -133,12 +133,12 @@ def slow(args, timeout=120, **kwargs):
 
 
 def test_test_only_returns_slurms_error_and_raises_only_for_a_busy_controller():
-    answers = iter([
-        Result(0, err="sbatch: Job 1 to start at 2026-11-02T10:00:00 using 8 processors on nodes holygpu8a19505 "
-                      "in partition kempner_interactive"),
-        Result(1, err="sbatch: error: Batch job submission failed: Requested node configuration is not available"),
-        Result(1, err="sbatch: error: Batch job submission failed: Socket timed out on send/recv operation"),
-        Result(1, err="sbatch: error: Batch job submission failed: Unable to contact slurm controller (connect failure)"),
+    answers = iter([  # the messages of sbatch --test-only on the cluster
+        Result(0, err="sbatch: Job 57206233 to start at 2026-11-03T01:02:51 using 8 processors on nodes "
+                      "holygpu8a19603 in partition kempner_interactive"),
+        Result(1, err="allocation failure: Requested node configuration is not available"),
+        Result(1, err="allocation failure: Socket timed out on send/recv operation"),
+        Result(1, err="allocation failure: Unable to contact slurm controller (connect failure)"),
     ])
     seen = []
 
diff --git a/tests/test_loop.py b/tests/test_loop.py
index 55507ce..71656e1 100644
--- a/tests/test_loop.py
+++ b/tests/test_loop.py
@@ -468,8 +468,7 @@ def test_round_partition_and_constraint_in_the_experiment_file_reach_each_round_
 
 def test_the_loop_job_stops_before_any_agent_step_if_slurm_would_reject_the_round_jobs(started):
     slurm = FakeSlurm(started / "state", lambda n: (0.91, SMALL))
-    slurm.test_errors["slurm/round.sbatch"] = ("sbatch: error: Batch job submission failed: "
-                                               "Requested node configuration is not available")
+    slurm.test_errors["slurm/round.sbatch"] = "allocation failure: Requested node configuration is not available"
     agent = FakeAgent(lambda call: GOOD)
     ctx = context(started, agent, slurm)
     ctx.experiment["round"]["partition"] = "kempner_interactive"  # the script asks for full GPUs, which it has not
diff --git a/tests/test_multi_loop.py b/tests/test_multi_loop.py
index 4d1e076..8ab15a4 100644
--- a/tests/test_multi_loop.py
+++ b/tests/test_multi_loop.py
@@ -697,8 +697,8 @@ def test_round_partition_and_constraint_reach_the_array_jobs(multi):
 
 def test_plan_only_stops_before_any_role_step_if_slurm_would_reject_the_array_jobs(multi):
     slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: None)
-    slurm.test_errors["slurm/round_array.sbatch"] = ("sbatch: error: Batch job submission failed: "
-                                                     "Invalid partition name specified")
+    slurm.test_errors["slurm/round_array.sbatch"] = ("sbatch: error: invalid partition specified: kempner_interactiv\n"
+                                                     "allocation failure: Invalid partition name specified")
     agent = FakeRoles()
     ctx = context(multi, agent, slurm)
     ctx.plan_only = True