Workflow deadlocks when walker count exceeds available GPU slots

[chrisdjscott]

Summary

run_westpa_workflow launches one SimulationAgent per walker by default. When the resampler grows the walker population beyond the number of available GPU slots, excess agents queue indefinitely waiting for a slot while the orchestrator waits for results from all of them — a circular block. The workflow silently hangs with no error or exception until the Slurm wall time is reached.

Observed behaviour

Slurm job 5540615 (NTL9 example, 72 walkers, 4 GPU slots): 72 Parsl tasks were dispatched simultaneously, only 4 simulation results ever returned, then the job was SIGTERM'd after hitting wall time:

Parsl task 0 try 0 launched on executor htex with executor id 1
...
Parsl task 71 try 0 launched on executor htex with executor id 72
...
received sim 1 iter 11. batch: 1/72
received sim 0 iter 11. batch: 2/72
received sim 4 iter 11. batch: 3/72
received sim 3 iter 11. batch: 4/72
*** JOB 5540615 ON g11 CANCELLED AT 2026-04-15T20:36:35 DUE to SIGNAL Terminated ***

Root cause

run_westpa_workflow hardcodes num_agents = len(initial_sims) — one agent per initial walker. When walkers > GPU slots (e.g. 72 walkers, 4 GPUs), 68 of the 72 agents sit in the Parsl executor queue waiting for a free slot. WestpaAgent.receive_simulation_data never sees 72 results, so _batch_ready never fires, and the workflow is stuck.

The existing code even had a TODO comment acknowledging this limitation:

# TODO: Generalize this so we don't have to assume one agent per sim.
# This is the case where we reuse the same hardware for multiple sims.
num_agents = len(initial_sims)

Fix

Add a num_sim_agents: int | None parameter to run_westpa_workflow (and ExperimentSettings for YAML-driven runs) so the agent pool can be sized to match available GPU slots. Simulations are already distributed round-robin across agents, so each agent handles multiple simulations sequentially.

Also harden dispatch_round_robin with:

• A semaphore to limit concurrent SSE sends (avoids connection-pool exhaustion when many sims are dispatched to few agents)
• Per-attempt exponential-backoff retries for transient exchange errors

Files to change: deepdrivewe/workflows/westpa.py, examples/openmm_ntl9_hk/config.yaml, examples/openmm_ntl9_hk/main.py, examples/openmm_ntl9_hk/workflow.py

Test plan

Run the openmm_ntl9_hk example with num_sim_agents set to the number of GPU slots (e.g., 4) and a walker count that exceeds the slot count. Confirm the workflow runs to completion rather than stalling silently.