OMSH Cluster Submission Scripts¶

The OMSH cluster helpers live in notebooks/experiments/scripts/:

omsh.sh
env.sh
env_hpc.sh
pipeline.sh
pipeline_all_envs.sh
pipeline_all_envs_hpc.sh
pipeline_env.sh
pipeline_env_hpc.sh
vis.sh

All launchers source support.bash, which owns the supported environment names and scheduler-target validation policy.

Script filenames intentionally omit the redundant omsh_ prefix. Internal shell functions and environment variables retain their omsh_* and OMSH_* namespaces to avoid collisions when launchers source shell helpers. User-facing README commands use scripts/omsh/ as the alias for this source directory; maintainer examples in this document use the repository-relative notebooks/experiments/scripts/ path.

They submit notebook paths with the $HOME/Projects convention used by the Topaz runner. All Slurm and PBS launchers default to omsh/<notebook-path>, matching the deployed checkouts at $HOME/Projects/omsh on both gpucluster and CX3. Every launcher accepts a PROJECT_PREFIX override for a deliberately nested deployment. Each helper appends venv=omsh unless the caller already passes an explicit venv=... override.

Curated Batch Launcher¶

omsh.sh [key=value ...] submits the curated OMSH experiment set:

training: gathering, markov_stag_hunt, pursuit
a30: bertrand, chicken
a40: empty after the August 2026 benchmark retirement

Override partition assignment with comma-separated notebook keys:

TRAINING_ENVS=gathering,pursuit A30_ENVS=bertrand A40_ENVS= bash notebooks/experiments/scripts/omsh.sh num_runs=10

Per-Environment GPU Mode¶

env.sh <notebook-key-or-path> <partition> [key=value ...] submits GPU jobs for one notebook.

The primary suite has four independently named algorithms: IPPO, IPPO-Lagrangian, plain CPO, and OmSh-Monotone. The default is one job per primary algorithm on every partition. OmSh-Credible remains available as an explicit diagnostic; when selected alongside OmSh-Monotone, the two modes must remain in different jobs because one notebook execution has one opponent-mode setting.

Use algorithms=<comma-separated names> to select a subset or opt into a diagnostic, for example algorithms=omsh_credible. The launcher-owned bundles=<positive integer> can combine compatible selected algorithms; it rejects a grouping that would place both OmSh modes in one job. Neither control is forwarded to Marimo.

Every standard experiment notebook defaults all five run_* switches to False. Each job therefore passes only its selected algorithms as True; omitted algorithms remain disabled. Launch commands place n_agents first when present, followed by the selected run_* switches and then other notebook/runner arguments.

Known notebook keys are gathering, markov_stag_hunt, pursuit, bertrand, and chicken. A direct notebook path can be passed instead. Retired benchmark history and the rl-engine handoff are documented in retired-benchmarks-2026-08.md.

The internal OMSH_DEPENDENCY_JOB_ID=<job-id> environment control adds --dependency=afterok:<job-id> to every algorithm bundle. The pipeline launcher sets this after it submits OM; ordinary callers normally should not set it.

Full WM -> OM -> Algorithm Pipeline¶

pipeline_env.sh <notebook-key-or-path> <partition> [key=value ...] is the normal launcher when all artifacts need to be rebuilt. It submits one dependency-ordered pipeline for an environment:

train_wm_mo.py -> train_om_mo.py -> four experiments_mo.py jobs
      WM              OM          IPPO, Lag, CPO, Monotone

The Slurm edges use afterok, so OM does not start when WM fails and no algorithms start when OM fails. The WM and OM submission commands request parsable Slurm output; the launcher fails locally if it cannot recover a numeric job id, rather than submitting an unsafe dependency-free continuation. This has been checked against the DoC submit helper, which passes --parsable and --dependency=... through to sbatch.

This all-algorithm barrier is intentionally stricter than the minimum artifact dependency. IPPO, IPPO-Lag, and CPO cells stop or run without loading learned shielding artifacts. Selected OmSh modes load wm/env_transition_graph.pkl and om/iop_stack.pt. Waiting for OM before every bundle nevertheless gives the full-pipeline command one simple success boundary and ensures that every resulting environment export belongs to the same completed WM/OM rebuild.

Examples:

bash notebooks/experiments/scripts/pipeline_env.sh pursuit training n_agents=3 num_runs=3
bash notebooks/experiments/scripts/pipeline_env.sh chicken a30 n_steps=25000 num_runs=1
DRY_RUN=1 bash notebooks/experiments/scripts/pipeline_env.sh gathering a16

Arguments are routed by notebook contract rather than blindly copied to all three notebook types:

n_agents=... goes to WM, OM, and algorithm notebooks.
n_steps=... goes only to WM.
venv=... and html=... are runner controls and go to every job.
Other key=value arguments, including num_runs and algorithm hyperparameters, go only to experiments_mo.py.
bundles=... is validated before WM submission, then consumed by env.sh as an algorithm-grouping control.
The launcher rejects run_*, run_algorithms, and just_vis, because it owns the complete algorithm suite and training mode.

The path argument may be a known key, an environment notebook directory, or any one of the environment's train_wm_mo.py, train_om_mo.py, and experiments_mo.py files. The sibling paths are derived automatically.

pipeline.sh [key=value ...] applies the same pipeline to the complete curated environment set using the same default partition assignment as omsh.sh. TRAINING_ENVS, A30_ENVS, and A40_ENVS can restrict or override the environment lists. It is intentionally a thin wrapper around omsh.sh, selecting pipeline_env.sh through the internal OMSH_ENV_LAUNCHER hook so the two curated launchers cannot drift in environment or partition policy. Because each environment is submitted as its own Slurm dependency graph, preparation is serial within that environment, then the four primary algorithm jobs become runnable together; the cluster's per-user partition limit governs runnable bundles across multiple environments.

pipeline_all_envs.sh <partition> [key=value ...] applies pipeline_env.sh to the retained experiment matrix: four fixed-default environments once and Pursuit with n_agents=2,3. This creates six independent dependency graphs and 36 Slurm jobs: two preparation jobs and four isolated algorithm jobs per graph. The wrapper owns n_agents and validates caller arguments before any submissions.

The wrapper deliberately takes one explicit partition and honors JOB_SCRIPT, so it can target either DoC Slurm cluster. gpucluster uses the normal GPU runner. cpucluster must use the CPU runner and deliberately extend the partition policy, for example:

# gpucluster
bash notebooks/experiments/scripts/pipeline_all_envs.sh a100 num_runs=5

# cpucluster
bash notebooks/experiments/scripts/pipeline_all_envs.sh \
  amd48 bundles=4 num_runs=5

The CPU host provides its runner configuration, so the user-facing command does not require inline environment overrides. amd48 is part of the default Slurm allow-list. OMSH_PIPELINE_ENV_LAUNCHER is an internal launcher override intended for tests.

Normal Slurm pipelines accept stage-specific memory requests through OMSH_SLURM_WM_MEMORY, OMSH_SLURM_OM_MEMORY, and OMSH_SLURM_ALGORITHM_MEMORY. OMSH_SLURM_OMSH_MEMORY can override the last value only for an algorithm bundle containing OmSh. Each non-empty value is passed to sbatch --mem for that stage. This is important for the larger tractable exact graphs: the runner's default 4 GiB-per-CPU allocation is not enough for Gathering or Markov Stag Hunt. For example:

OMSH_SLURM_WM_MEMORY=64G \
OMSH_SLURM_OM_MEMORY=64G \
OMSH_SLURM_ALGORITHM_MEMORY=32G \
OMSH_SLURM_OMSH_MEMORY=64G \
  bash notebooks/experiments/scripts/pipeline_env.sh markov_stag_hunt amd48

The corrected final-paper launcher reserves 16 GiB for Gathering OmSh and 64 GiB for MSH OmSh. The August 2026 campaign showed that both shield setups exceeded an 8 GiB allocation before policy training, while their three baseline jobs did not need the larger requests. OMSH_SLURM_OMSH_MEMORY can raise the Gathering reservation, while MSH remains pinned to 64 GiB. Keeping these overrides condition-specific avoids wasting the same memory on every algorithm.

These overrides do not make three-agent Pursuit viable with the current materialized stochastic-edge layout. Its default build is preflight-rejected; use compact or procedural graph storage before relaunching that condition.

HPC CPU Mode¶

pipeline_env_hpc.sh <notebook-key-or-path> [queue] [key=value ...] is the full Imperial HPC launcher. It uses native PBS Pro qsub calls because the deployed convenience submit.sh wrapper does not expose dependency attributes. The submitted graph is:

GPU WM --afterok--> GPU OM --afterok--> four CPU algorithm jobs

The launcher parses full PBS job IDs such as 12345.pbs, attaches dependencies with -W depend=afterok:<job-id>, and fails before submitting a downstream stage if a job ID cannot be recovered. It invokes the deployed Python runner at HPC_JOB_SCRIPT=$HOME/py.sh by passing REL_PATH and PY_ARGS through qsub -v. Its default PROJECT_PREFIX=omsh produces paths beneath the deployed $HOME/Projects/omsh checkout; nested deployments can override the prefix explicitly.

Default resources and overrides are:

Stage	Defaults	Overrides
WM	16 CPUs, 64 GB, 1 GPU, 71:59:59	`OMSH_HPC_WM_CPUS`, `OMSH_HPC_WM_MEMORY`, `OMSH_HPC_WM_GPUS`, `OMSH_HPC_WM_WALLTIME`
OM	16 CPUs, 64 GB, 1 GPU, 71:59:59	`OMSH_HPC_OM_CPUS`, `OMSH_HPC_OM_MEMORY`, `OMSH_HPC_OM_GPUS`, `OMSH_HPC_OM_WALLTIME`
Each algorithm	8 CPUs, 32 GB, 0 GPUs, 71:59:59	`OMSH_HPC_ALGORITHM_CPUS`, `OMSH_HPC_ALGORITHM_MEMORY`, `OMSH_HPC_ALGORITHM_WALLTIME`

Set OMSH_HPC_GPU_TYPE=L40S or A100 only when a concrete GPU type is required. The optional positional queue is passed with qsub -q; leaving it absent uses Imperial's routing policy.

pipeline_all_envs_hpc.sh [queue] [key=value ...] calls pipeline_env_hpc.sh for the complete experiment matrix. It submits the five retained environments, with Pursuit at n_agents=2 and 3. This is six dependency graphs and 36 PBS jobs. The wrapper owns n_agents and rejects a caller override before submitting anything; other notebook arguments are forwarded to every graph. OMSH_HPC_ENV_LAUNCHER is an internal launcher override intended for tests.

Both HPC launchers default SUBMIT_DELAY_SECONDS to two seconds. pipeline_env_hpc.sh pauses after each of its six submissions, including the final algorithm submission, so consecutive graphs launched by pipeline_all_envs_hpc.sh retain the same two-second spacing.

env_hpc.sh <notebook-key-or-path> [queue] [key=value ...] submits CPU-only jobs. Each algorithm gets its own job with ngpus=0 walltime=71:59:59 and exactly one owned run_* flag set to true. The default command uses the Imperial wrapper shape, submit <notebook> ...; an optional PBS routing queue becomes queue=<name>. Queue names receive syntax validation, but the script intentionally does not maintain a static queue allow-list because Imperial routing availability changes and normally should be left unset.

This helper remains the algorithm-only path for exports whose WM/OM artifacts already exist.

Visualization Mode¶

vis.sh [key=value ...] submits CPU visualization reruns for the curated notebooks. It defaults to PARTITION=amd48, JOB_SCRIPT=py_cpu.sh, and passes just_vis=True.

Common Controls¶

All helpers support:

DRY_RUN=1 to print commands without submitting.
SUBMIT_DELAY_SECONDS=<seconds> to control the sleep between submissions.
SUBMIT_CMD=<command> to override topaz job submit.
JOB_SCRIPT=<script> to override the job script.
PROJECT_PREFIX=<prefix> to override the default omsh project path.
OMSH_SLURM_WM_MEMORY, OMSH_SLURM_OM_MEMORY, and OMSH_SLURM_ALGORITHM_MEMORY to add stage-specific sbatch --mem requests to normal Slurm pipelines, with OMSH_SLURM_OMSH_MEMORY as the shield-bundle override.

Pipeline launchers additionally rely on numeric Slurm job ids returned by SUBMIT_CMD; use a command compatible with sbatch --parsable when overriding it.

The deployed Slurm runners update one shared checkout before executing a notebook. Multiple dependency siblings can become runnable together, so their checkout update and environment sync must be protected by a shared filesystem lock. Without that lock, concurrent git pull processes can fail while updating origin/main, or observe the worktree changing between fetch and merge. The cpucluster runner uses one lock per project around git pull --ff-only and uv sync; notebook execution happens after releasing the lock.

The CPU runner must also select the JAX CPU backend before starting Python (JAX_PLATFORMS=cpu and an empty CUDA_VISIBLE_DEVICES). The shared OMSH environment contains the JAX CUDA plugin, and merely probing that plugin on a GPU-less CPU node can terminate the native kernel rather than falling back cleanly. The deployed Marimo batch helper treats its Python kernel ... died console notification as a failed run; Marimo does not send a completed-run notification after this kind of native crash, so failing to recognize it would leave the Slurm allocation alive until its time limit.

Validation Policy¶

Bare environment arguments must match one of the nine entries in OMSH_ENVIRONMENTS. A value containing / or ending in .py is treated as an intentional direct notebook path. This prevents a typo such as gathring from silently becoming a submitted relative path while preserving advanced path-based use.

Training partitions default to the allow-list training,a16,a30,a40,a100,amd48; visualization CPU partitions default to amd48. Override OMSH_GPU_PARTITIONS or OMSH_CPU_PARTITIONS with a comma- or space-separated deployment list when the live Slurm configuration changes. Curated TRAINING_ENVS, A30_ENVS, and A40_ENVS overrides accept known environment names only.