Skip to content

Adding persistent workers for 3dvar_atmos suite - #777

Open
ftgoktas wants to merge 21 commits into
developfrom
feature/slurm-optimization
Open

Adding persistent workers for 3dvar_atmos suite#777
ftgoktas wants to merge 21 commits into
developfrom
feature/slurm-optimization

Conversation

@ftgoktas

@ftgoktas ftgoktas commented May 5, 2026

Copy link
Copy Markdown
Member

When running a workflow, every heavy task (RunJediVariationalExecutable, BuildJedi, EvaObservations) normally submits its own sbatch job to SLURM. Each job waits in the queue before it can run, even if the previous task just finished 2 seconds ago.

The changes adds a persistent_workers flag to slurm.yaml. When persistent_workers is set to true, Swell reserves compute nodes at the start of the workflow using salloc --no-shell and holds them until all cycles finish. Every heavy task then runs as an srun --jobid step inside that allocation without individual queue waits. The nodes are released via scancel once all cycles finish.

How to test:

Before creating an experiment, make your src/swell/deployment/platforms/nccs_discover_cascade/slurm.yaml has:

  • persistent_workers: true

Some things I've tried and encountered errors during testing (fixed now):

  1. MPI processes couldn't talk across nodes: Calling srun --jobid from a login node (outside a batch job) requires --mpi=pmi2 to set up MPI across nodes. Without it, each node starts its own independent group of MPI processes instead of one single communicator, and the run fails.
  2. Wrong --ntasks-per-node: The global ntasks-per-node: 64 from slurm.yaml was being passed to the srun step. The correct value is np // nodes (96 / 3 = 32). Passing 64 would request 192 total tasks and cause the run to fail.
  3. Single-task steps killed by CPU cgroup: srun --ntasks=1 allocates exactly 1 CPU. EvaObservations creates 40 Python workers and BuildJedi runs parallel compile threads. SLURM's cgroup killed both. Fixed by adding --cpus-per-task from each task's srun_directives.
  4. EvaObservations blocked by RunJedi's --exclusive lock: RunJediVariationalExecutable uses --exclusive which locks all allocated nodes. EvaObservations with --exclusive would queue behind it indefinitely. Fixed by running EvaObservations as srun --ntasks=1 without --exclusive it only needs CPUs on one node and doesn't need to lock all three.

#697

@ftgoktas ftgoktas self-assigned this May 5, 2026
@ftgoktas ftgoktas added the core development design related issues and improvements label May 5, 2026
@ftgoktas
ftgoktas requested review from a user and mranst May 5, 2026 18:57
@mer-a-o

mer-a-o commented Jun 1, 2026

Copy link
Copy Markdown
Contributor

@ftgoktas, here's the branch I mentioned today: #800
Each background tar ball includes 6 background files. I need to run ConvertState task for each of them to change their resolutions. These tasks are independent of each other. Right now each takes one node which is not efficient.

image image

@ftgoktas

Copy link
Copy Markdown
Member Author

Ran 100 trials of each approach on Discover. Both sets submitted simultaneously so they were under same cluster conditions.

Mean queue wait Max queue wait
Non-persistent (RunJediVariationalExecutable, 3 nodes, 1hr limit) ~25 min (1508s) ~40 min (2411s)
Persistent (salloc, 3 nodes, 40min limit) ~4.6 min (278s) ~10 min (604s)

So it's 5.4x improvement in average queue wait. The max on non-persistent (~40 min) matches what was being observed on real 3dvar_atmos runs before this approach.

@ftgoktas
ftgoktas requested a review from Dooruk August 21, 2026 17:12
@ftgoktas
ftgoktas marked this pull request as ready for review August 21, 2026 17:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core development design related issues and improvements

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants