Run the GPU jobs on kempner_interactive

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 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`.
@@ -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.

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).
@@ -46,7 +46,7 @@ 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
@@ -176,6 +176,29 @@ 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. 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:
+
+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
+   source cluster.env
+   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.
+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`. 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
 
 Put test copies in `local/`, which Git ignores. `EXPERIMENT_FILE` and `STATE_DIR` point the jobs to them:
@@ -245,12 +268,12 @@ 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
 
 - **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 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.
 

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:

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

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"

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):

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
@@ -291,9 +309,11 @@ 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
-    check_time_lines(ctx, ("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)
     notes = ctx.paths.state / "notes.md"
     if not notes.exists():
         first = records.read_json(ctx.paths.state / "rounds" / "round_01" / "config.json")