nemo_gym.orchestration.executors.slurm
nemo_gym.orchestration.executors.slurm
Module Contents
Classes
Functions
Data
_READER_UV_LOCK_TIMEOUT_DEFAULT_S
API
Bases: RuntimeError
A job script would be rendered again by a launcher of another Gym revision than the one the job runs.
The launcher’s job.sh templates may need modules the pinned revision does not have, and the capability checks run later inspect only the newly rendered script, so the mismatch is refused before anything changes.
How a heavy reader runs in a Slurm step: the caller’s srun arguments (each quoted), how long the step may take to start, how many reader steps the user may have running or pending at once, how long a reader waits for one of them to end when they are at that cap, and whether a reader that cannot run in its step may run on the login host instead.
Bases: KeyboardInterrupt
A signal stopped recovery while the login-host Store reader ran; the reader was stopped by its token.
A KeyboardInterrupt, so the reader’s ordinary error handling does not turn it into evidence and the open
connection unwinds and closes its master.
Bases: BaseExecutor
Slurm executor for Pyxis-enabled clusters (https://github.com/NVIDIA/pyxis).
Every service and the driver are launched via srun --container-image so they
run inside the container specified in their config. Health checks run as plain
bash inside the sbatch script (no container needed — they just poll HTTP).
The locked steps that record extend_attempts more attempts for every rollout, restored unless queued.
The Store reads the record (artifacts/rollouts_attempt_extension) as an addition to the run’s saved
max_rollout_attempts; a rollout that ended with a model failure is never run again.
The locked steps that record extend_retries more infra retries, restored unless the job is queued.
The chain marker a fresh restart writes: why, the recorded walltime and the new one.
Ask the scheduler about the run’s judge-only jobs that have no receipt (CollectedEvidence.judge_jobs,
state queued), set each one’s state and scheduler_state from its record, and name the newest live
one as judge-only job <id> (STATE); None when none is live, or when the query failed or knows none.
The locked steps of an operator continuation: fill the counters it names, replace the unknown marker and journal the continuation with its epoch before sbatch, so a queued continuation always has its epoch.
Each file is kept aside first and put back unless the job is queued. A readable counter is never named, except
the retry extension an explicit --extend-retries raises when the spent count is readable.
Queue the judge-only job of a run that is not CPU-only, on the judge execution profile derived from its
saved config (judge_profile), and journal it: the served-judge profile when the run serves services beside
its policy (its judge on a node pool of its own), else the CPU profile.
The job’s command is the same gym eval reverify --judge-failed-only --settle-into-source as a CPU-only
run’s, with the run’s saved settings, the policy’s --model-type so the run’s config composes as it did
(the policy server it names is never called) and a receipt the job writes when it ends. Under the resubmit
lock, no job of the run may wait or run in its job dir, no allocation of the chain may be live, and the saved
sources, resolved config and chain accounting must be as the evidence saw them; then the script is staged
beside the run’s and queued with --no-requeue, and its id, partition, walltime, rows and receipt path are
appended to the job dir’s judge-job journal. The job is not a segment of the policy chain: nothing is added
to .gym_job_chain and no retry or walltime budget is spent or needed. When the last
JUDGE_JOB_FAILURE_BOUND judge-only jobs ended without settling a saved answer, nothing is queued and the
stop names them and the way on: acknowledge_failed_judge_jobs with a nonempty acknowledge_reason appends
one line (time, user, reason, those jobs’ ids) to the acknowledgement journal once the new job is accepted by the
scheduler, so exactly one more job queues; the acknowledgement covers only the jobs it names.
The locked shell steps that render one benchmark’s job.sh and saved config again, with a repin, a walltime
raise or both, and the journal entries to append once its job is queued. A fresh restart raises the walltime
above the recorded one without keep_budget.
The locked steps that journal the approval to run masked results again, taken back unless the job is queued.
The line goes to the rerun-masked journal beside the rollouts file, where the Store reads it when the job starts: its time, user, reason and each approved identity with the attempt it covers and the attempts counted for it.
The resolved config persist() saved for this run, so a resubmit repeats what submit did.
Run gym eval export mlflow on the run’s saved artifacts from the login host, and report its result.
The export runs where the remote Store reader runs, with the run’s saved pinned Gym install (and its MLflow
extra) in a temporary checkout: it only reads the saved artifacts and talks to the tracking server, so it needs
no allocation. It takes the benchmark’s saved mlflow_* and upload_rollouts run settings, and its
credentials as the driver did: the submission’s forwarded environment and the driver’s env. Gym’s export takes
the run lock, prepares the health metrics and writes the receipt bound to the artifact digests, and skips a run
that already has one. The resubmit lock is held and no job of the chain may be live meanwhile.
The full Gym commit the run’s saved config pins its driver to; None when there is none or it is unreadable.
The Gym commit this launcher runs; the same value the judge-only job and a repin are checked against.
The decision of recovery_evidence without a score key.
Collect the job dir’s identities, failures, quality summary, named score and chain budget over SSH.
reader_ref runs the Store reader at that full commit instead of the saved one: a repinned run resumes there.
max_walltime is a walltime raise the resubmit will render: the plan checks the chain cap it renders.
server_spinup_timeout_seconds (with spinup_timeout_reason) is a spin-up limit raise the resubmit
will render: the decision names it.
rerun_masked_without_artifact plans the masked results that kept a saved answer but have no gradeable
artifact and no grade as retries; the resubmit journals them (see resubmit).
operator_budget lets a run whose chain accounting is unknown continue (see decide_from_collected); with
it, a failed scheduler query is a stop that names it instead of an error.
Queue benchmark’s job.sh again from its job dir and record the new job id.
With keep_budget (--recover) the job keeps the spent retry count, walltime and per-rollout attempts;
a job.sh rendered before that option would reset them, so it is refused. accounting_snapshot is the
collected evidence’s accounting_snapshot; it is required with keep_budget and rechecked under the
resubmit lock before anything in the job dir changes, so a decision made on stale accounting is refused.
accounting_reconciled (with keep_budget) admits the requeue marker that the collected evidence
reconciled from terminal scheduler records; it is replaced by a note just before the job is queued.
repin (with keep_budget and repin_reason) continues the run on that full Gym commit, which must be
this launcher’s: under the lock, after the live-job and snapshot checks, job.sh and resolved-config.yaml are
rendered again with the new ref (and the MLflow pin tags that named the old one), the old files are kept
beside them, and both are restored if the job is not queued. Rows, attempts and budget files stay as they are;
the queued job’s record carries the new commit and the run dir’s repin journal records the move.
extend_retries (with keep_budget) adds that many infra retries to the recorded extension that the
job’s prologue adds to the configured cap; it is written just before sbatch, restored if the job is not
queued, and journaled once queued. The spent retry count is never reset.
max_walltime (with keep_budget and walltime_reason) raises the Slurm walltime the same way a repin
moves the commit, in the same render: compute.<cluster>.walltime becomes it, the chain cap becomes at least
it, and a value below the recorded walltime is refused. The run dir’s walltime journal records the raise.
segment_walltime (from walltime_plan) is what each Slurm segment requests instead, when the partition’s
MaxTime is below the raise: the raise stays the chain cap.
rerun_masked (with keep_budget) is the operator’s approval to run again the masked results that kept a
saved answer but have no gradeable artifact and no grade: its reason and identities (identity, attempt
id and attempts counted) are appended as one line to the run’s rerun-masked journal in artifacts/, which the
Store reads. The line is written just before sbatch and taken back if the job is not queued.
A job the scheduler refuses raises SchedulerRefusedError with sbatch’s text and the restored job dir.
unsafe_resume_reason (with keep_budget) renders the job with allow_unsafe_resume the same way, so
its resume proceeds past a refused run identity check; the run dir’s unsafe-resume journal records the reason
and the user once the job is queued.
server_spinup_timeout_seconds (with spinup_timeout_reason) renders the job with that server
spin-up limit the same way; it only raises the recorded limit (0 removes it), and the job dir’s
spin-up timeout journal records both limits, the reason and the user once the job is queued.
operator_continuation (with keep_budget) is a decision’s OPERATOR-CONTINUATION record for a run whose
accounting is unknown: under the lock it fills the counters the record names, replaces the unknown-accounting
marker with a note and keeps the readable counters; each step is undone unless the job is queued, and the job
dir’s operator journal records the queued continuation.
fresh_restart (with max_walltime, walltime_reason and accounting_snapshot, without
keep_budget) queues a run that has no task inventory again as a fresh run: a fresh retry and walltime
budget, a walltime longer than the recorded one, the snapshot rechecked under the lock. Its walltime journal
entry, the receipt, carries plan FRESH-RESTART and why; .gym_chain_incomplete names the restart
and both walltimes until the queued job continues the chain.
Refuses while any of the caller’s jobs runs in that job dir, so two jobs never write one run, and
refuses a run with rollouts whose gym eval run job.sh predates resubmit: it would start the run
fresh and delete them. A command benchmark’s job.sh runs its own command, so it is not refused.
Queue a judge-only job that grades the run’s selected saved answers again, never the policy.
A CPU-only submission (no service, no policy model, CPU node pools only) runs it on its own compute with the
chain’s remaining walltime allowance, as before. Any other submission, one that served its policy on GPUs,
runs it on a judge execution profile derived from its saved config (_queue_judge_job): the run’s own judge
and verifier servers on one node of a CPU partition (judge_partition, a saved CPU pool or the launcher’s
reader placement, with judge_qos or a QOS derived for it), or, for a run that serves its judge on a node
pool of its own, on that saved pool with those services started and the run’s QOS, for the run’s segment
walltime bounded by the partition’s MaxTime, at the run’s pinned Gym commit or repin.
acknowledge_failed_judge_jobs (with acknowledge_reason) lets one more judge-only job queue when the last
ones ended without settling a saved answer (see _queue_judge_job). Either way the job settles through
gym eval reverify --judge-failed-only --settle-into-source. This never starts a policy service or turns an
unavailable judge into a model run.
Check a --max-walltime raise against the partition’s MaxTime, read with sinfo on the cluster, before
any evidence is read (see _plan_walltime). A failed query keeps the whole raise per segment and says so.
Shell lines that fail unless the scheduler confirms that no job of chain waits or runs.
Slurm drops a finished job from the controller after MinJobAge, and squeue -j then exits 1 with “Invalid job
id specified” once it knows none of the ids. That is no answer about the jobs, so sacct, which keeps finished
jobs, decides: a job whose latest incarnation (highest Restarts; sacct -D lists every requeue) is in one of the
LIVE_JOB_STATES is live, and one it does not know is not. Any other squeue failure, or a failing sacct, leaves
the scheduler unknown.
A nonnegative number of seconds from the environment; the default for an unset or invalid value.
The used walltime of a chain after an operator continuation: its baseline plus every allocation since.
The continuation filled or kept the counter as the baseline, so the allocations before it are not accounted again. None when the records since then are incomplete or the counter does not match them.
Stop this user’s login-host processes whose environment or command line carries token: TERM, then KILL
for any still running after 10 s. It prints __GYM_READER_KILL:term=<n> kill=<n>.
Benchmark name to (job_id, error), exactly one of which is set.
The partition a job of compute queues on and its MaxTime, read with sinfo over conn.
The partitions are the node pools’ ones, else extra_args.partition (a list a,b allowed), else the
cluster’s default one. Every one must be listed, and the lowest limit applies: partition names it,
max_time is sinfo’s text and max_time_s its seconds (None when unlimited); error says why they are
unknown.
How the raise max_walltime fits the partition limit: each segment requests the smaller of the raise
and the partition’s MaxTime, and the raise is the chain cap the segments continue under. Refused when no segment
fits: a run that does not chain segments needs the whole raise in one job, and a chained one at least its
recorded segment walltime.
The chain cap after a raise to walltime: at least the raise, and the raise itself when there was none.
The login-host steps: size the bound from the rollout files, run the reader, report its exit code and stderr.
The reader’s own failure is reported on stdout (__GYM_READER_* lines) so a timeout, a MemoryError or a
failed bootstrap reaches the decision with its exit code and the end of its stderr; only SSH itself fails the
command. With token, the reader runs with it as GYM_READER_TOKEN in its environment. With step, it
runs in a Slurm step when it can (_reader_step_lines) and reports where it ran.
With watchdog, for a connection that holds the shell’s stdin open while it runs: the reader runs in its own
session (setsid), and a watcher in its own session waits for EOF on the shell’s stdin, which comes when the
client process or its SSH connection ends, then kills the reader’s process group (TERM, KILL after 10 s). It is
an EOF loop, not a hangup: without a terminal nothing sends SIGHUP. The steps run as one function that exits the
shell, so the watcher never reads a line of the script itself. In a step, the watched process is srun, and the
watcher also cancels the step by name. On the login host, a progress loop in its own session prints
gym-reader progress: pass=<n> bytes=<b> elapsed_s=<s> on stderr every progress_interval_s (60 s by
default) while the reader runs; a reader in a step has none. Stdout and its __GYM_READER_* lines are
unchanged.
_reader_command with the reader in a Slurm step: the payload is saved once, sent on the step’s stdin, and
read again from the file by a login-host reader that runs instead. With watchdog, srun and that login-host
reader each run under the EOF watcher, as the reader of _reader_command(..., watchdog=True) does.
The settings a remote reader takes from its caller: the reader’s command starts from a cleared environment on another host, so without this a cold install there waits uv’s own default for the lock and gives up early.
The reader failure named from its report with its exit code and stderr, and whether it reached its deadline (the only failure retried with twice the bound); (None, False) when it exited 0.
Bash that runs script in a Slurm step, else leaves it to the login host.
It sets _gym_host (“login” or “cpu:<job id>”) and _gym_fallback (why the reader stays on the login host,
empty when it ran in the step); after a step that ran, _gym_rc is its exit code and $_gym_step_out holds
its output. Inside the step the input on stdin becomes FD 3, because srun forwards stdin but no other descriptor,
and script reads the reader bound as $1. The step is named gym-reader-<token>, and its time limit is the
reader bound plus 60 s, so Slurm ends an orphan. A step that has not started after the wait is cancelled by name,
and the reader runs on the login host. With watchdog (FD 4 is the shell’s stdin), srun runs under the EOF
watcher (_watch_lines), which also cancels the step by name. Without step.login_fallback, a reader that
would fall back runs nowhere: _gym_host is none, _gym_rc 1, and _gym_fallback says why.
When the user’s reader steps are at step.max_readers, the reader queues: it counts them again every
_READER_QUEUE_POLL_S seconds for up to step.queue_wait_s, saying on stderr how many it is queued behind
(a progress line), and runs in its step once one ends. Only a cap still reached after that wait keeps it out.
A step srun refused outright (_SRUN_REFUSAL: a submit or QOS limit, or no node free to start it) was never
queued, so the reader submits it again after _READER_SRUN_RETRY_S seconds, doubling up to
_READER_SRUN_RETRY_MAX_S, with a progress line each time, for as long as step.queue_wait_s allows. Any
other srun failure, or a refusal past that wait, falls through to the cancellation below as before.
The benchmark’s server spin-up limit in seconds as the run starts its servers; 0 or less is no limit.
Use the saved pinned install over the existing transport; never import Gym into plain host Python.
The reader with the caller’s metadata helper defined after its imports, sent as one script.
The explicit reader bound from NEMO_GYM_REMOTE_READER_TIMEOUT_S, clamped to [300, 14400], else None.
Without it the login host sizes the bound from the run’s rollout files (_reader_command).
Name the reader failure with its cause, not only the exception type: the end of the message, where it is.
config with the driver’s Gym install at ref, and the benchmark’s MLflow pin tags that named the old one.
Run the reader. A signal on the client (raised as KeyboardInterrupt) leaves the remote reader running,
since ssh bash -s has no terminal to hang up, so it is stopped by its token before the error unwinds.
A connection that holds_input keeps the remote shell’s stdin open while the reader runs, and the command’s
watcher (_reader_command(..., watchdog=True)) stops the reader when that stdin ends.
One sbatch that reports its own benchmark, exit status and output.
rc is captured immediately: $? after the tr pipeline would be tr’s
status, which is zero however badly sbatch failed. tr flattens sbatch’s
multi-line error messages so the whole result stays on one marker line.
A partition time limit as sinfo prints it (%l) in seconds; None for UNLIMITED or infinite.
Takes Slurm’s duration spellings (MM, MM:SS, HH:MM:SS, D-HH, D-HH:MM, D-HH:MM:SS) and
raises ValueError for anything else, such as n/a.
config with the benchmark’s server spin-up limit seconds; every other field as saved.
Why seconds is not a raise of the recorded limit (0 is no limit), or None when it is.
The sentence a decision carries for a spin-up timeout raise.
Forward explicitly selected values over SSH stdin, never into artifacts.
What run() exports ahead of sbatch: the forwarded values, then the telemetry token when it is on.
config with the benchmark’s job resuming past a refused run identity check; every other field as saved.
Fail before anything is staged or copied.
_sbatch_command raises this same check, but only once submission is
already underway — after staging, connecting, mount validation and the
rsync. Checking here first means a bad name fails before any of that, and
fails on --dry-run too, instead of the dry run printing a clean script
listing for a benchmark that could never actually be submitted.
config with a chain cap of at least walltime and each Slurm segment requesting segment (by default
walltime itself); every other field as saved. A chain without a cap takes walltime as its cap: a raise
always bounds the total.
Bash that waits for the background reader $_gym_reader, started in its own session, and sets _gym_rc.
A watcher in its own session reads the shell’s stdin through FD 4 and at EOF kills the reader’s process group:
TERM, then KILL after 10 s. With cancel, a quoted Slurm job name, it cancels that job before the TERM: the step
runs on another host, and srun may exit at the TERM, which ends the wait and stops the watcher. progress
starts a loop that sets $_gym_progress; it is stopped after the wait.
The collector in a reader step, else on the login host; its output keeps only the collector’s own lines after
the placement report is taken off (reader_placement, without_placement). Its exit code is the
collector’s.
Where a reader ran (“login” or “cpu:<job id>”) and why it fell back to the login host, from its report; both None for a reader that ran without a reader step.
The reader step from NEMO_GYM_REMOTE_READER_SRUN (extra srun arguments, for example -p cpu -c 4),
NEMO_GYM_REMOTE_READER_SRUN_WAIT_S (default 120), NEMO_GYM_REMOTE_READER_MAX (default 8),
NEMO_GYM_REMOTE_READER_QUEUE_WAIT_S (default 600; 0 does not wait) and
NEMO_GYM_REMOTE_READER_LOGIN_FALLBACK (allow; anything else forbids it); None when the arguments are unset
or empty.
The token this process’s login-host Store readers carry in their environment as GYM_READER_TOKEN: the
caller’s own value when it sets one, else one generated uuid for the process. Recovery reports it as
reader_token, so whoever stops a recovery can find its reader.
The output without the reader placement report.