Slurm with srun¶
Wire DaskPool into a Slurm allocation
via dask-jobqueue's SLURMCluster. Each
Dask worker becomes one Slurm task, the sweep dispatches across them, and
the per-run subprocess isolation that gmat-sweep already does on a
laptop carries over node-by-node unchanged.
Prerequisites¶
- A Slurm cluster you can submit to (
sbatch,srun,squeueon PATH). - A shared filesystem (NFS, Lustre, GPFS, …) mounted at the same path on the login node and every compute node. The driver writes the manifest there and every worker reads/writes the per-run output dir from there.
- A working GMAT install reachable on every compute node — either the
same shared filesystem or a node-local install at the same path. GMAT
discovery is done by
gmat-run; see its install guide for the search order.gmat-sweepdoes not install GMAT. - Python with
gmat-sweep[dask]anddask-jobqueueinstalled in an env reachable from the worker nodes (a sharedvenv,condaenv, ormodule loadline is the usual choice).dask-jobqueueis not agmat-sweepdependency — install it yourself.
Worked example¶
driver.py¶
from dask.distributed import Client
from dask_jobqueue import SLURMCluster
from gmat_sweep import sweep
from gmat_sweep.backends import DaskPool
cluster = SLURMCluster(
cores=1,
processes=1,
memory="4 GB",
walltime="01:00:00",
job_extra_directives=["--ntasks=1", "--cpus-per-task=1"],
local_directory="/scratch/$USER/dask-worker-space",
log_directory="/scratch/$USER/dask-logs",
)
cluster.scale(jobs=8)
client = Client(cluster)
with DaskPool(client=client) as pool:
df = sweep(
"/shared/missions/mission.script",
grid={"Sat.SMA": [7000.0, 7100.0, 7200.0, 7300.0]},
backend=pool,
out="/shared/sweeps/sma-scan",
)
print(df.head())
SLURMCluster submits one Slurm job per Dask worker; cluster.scale(jobs=8)
asks for eight. The worker pool is elastic — Dask drops idle workers back
to Slurm and re-requests them as new specs arrive — so the same driver
also works for shorter or much longer sweeps without rewiring.
Submitting the driver¶
The driver itself is a single Python process. Submit it as a one-task allocation that lives long enough to manage the worker pool:
#!/bin/bash
#SBATCH --job-name=gmat-sweep-driver
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --mem=2G
#SBATCH --time=02:00:00
#SBATCH --output=/shared/sweeps/sma-scan/driver.log
srun python -u driver.py
sbatch driver.sbatch queues it; squeue -u $USER shows the driver job
plus the eight worker jobs as Dask scales them up.
Caveats¶
Shared filesystem is non-negotiable¶
Every worker reads the script, writes per-run Parquet outputs, and
appends manifest entries through the same out= directory. That
directory must be visible at the same path from the driver and from
every compute node. NFS, Lustre, BeeGFS, and GPFS all work; node-local
SSD scratch does not unless you stage the result back at the end.
If out= resolves to a different location on different nodes, runs will
appear to "succeed" but their Parquet files end up scattered across
worker-local disks and the aggregated DataFrame will be empty or partial.
GMAT must be reachable from every worker¶
gmat-run discovers a local GMAT install at import time. Each Dask
worker is a fresh Python process on a compute node, so each worker has
to find GMAT independently. Two patterns work:
- Shared-filesystem install — a single GMAT tree under
/shared/gmat/visible to every node. SetGMAT_BIN_DIR(or whatever discovery hookgmat-runexposes) in the worker environment. - Node-local install at the same canonical path on every node.
Both are operationally fine; pick whichever your cluster admins prefer.
The discovery itself is gmat-run's problem, not gmat-sweep's — but
a misconfigured worker shows up as every run failing with the same
import error, so it's worth checking in a srun python -c "import
gmat_run" before launching the sweep.
Subprocess isolation still applies inside each worker¶
Each Dask worker is one Python process. The sweep's per-run subprocess
hop runs inside that worker — _worker_entrypoint spawns a child
Python that bootstraps GMAT, runs the script, and exits. No special
Slurm-side configuration is needed; the per-task fresh GMAT context
that protects sweeps on a laptop protects them on Slurm in exactly the
same way.
The reuse_gmat_context=True default still amortises bootstrap across
the runs assigned to a single worker, so worker-level reuse remains the
fast path on Slurm too.
When this isn't enough¶
The recipe above covers the common Slurm-via-dask-jobqueue setup. For
exotic schedulers, MPI-style launches, or Slurm features that don't fit
the SLURMCluster model, write a custom Pool against the
Pool ABC — its only requirement is that
each task runs through the subprocess hop.