diff --git a/AGENTS.md b/AGENTS.md index 419b990..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`. +- 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 ` and `sacct -j `. 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..1fd12f6 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). @@ -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. 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 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/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..41cdf40 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 @@ -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") 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..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): @@ -130,6 +131,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 +153,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 +183,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 +216,18 @@ 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 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() @@ -220,11 +246,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 == "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 - 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..5a17636 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2,7 +2,11 @@ from __future__ import annotations -from pruning_example.slurm import Slurm, parse_array_sacct, parse_array_squeue, task_indexes +import subprocess + +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 +57,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 @@ -70,9 +75,87 @@ 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(): + # 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 + [{}]) + + 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" + + +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(): + 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 = [] + + 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", ".") + 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 79a7b2a..71656e1 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,61 @@ 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"] = "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 + 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_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=()): + 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..8ab15a4 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,63 @@ 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"]["constraint"] = "a100-mig" + ctx.experiment["round"].update(partition="kempner_interactive", constraint="a100-mig") assert multi_loop.run(ctx, "loop") == 0 - assert slurm.options == [("--constraint=a100-mig",)] * 2 # the wave and the confirmation rounds + 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: invalid partition specified: kempner_interactiv\n" + "allocation failure: 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["wave"]["max_parallel"] = 1 + assert multi_loop.run(ctx, "loop") == 0 + 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 + + +@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):