Investigate Parsl as Python-native workflow manager for build/simulate orchestration

[maxscheurer]

Summary

Compute-heavy tasks in MDFactory (building systems, running simulations) are currently handed off to Nextflow (workflows/simulate.nf, workflows/build.nf) and require manual orchestration via shell scripts (workflows/run_full_pipeline.sh). Nextflow will be removed completely. There is no smooth Python-native integration of compute workflow management. For analysis execution, we use submitit to wrap lightweight Python functions for trajectory analyses (mdfactory/analysis/submit.py), but this doesn't extend to the heavier build/simulate workloads.

This issue tracks a proof-of-concept investigation of Parsl as a Python-native workflow manager for MDFactory, particularly for handling GROMACS MD runs and system builds through its bash_app, SLURM provider, retry functionality, and checkpoint/restart support.

Current Behavior

• Simulation orchestration is done via Nextflow processes (minimization → nvt → npt → production) defined in workflows/simulate.nf with SLURM config in workflows/simulate.config.
• Build orchestration uses workflows/build.nf to run mdfactory prepare-build then mdfactory build per hash.
• Analysis dispatch uses submitit (mdfactory/analysis/submit.py) — one SLURM job per simulation path.
• There is no Python-level MD run orchestration; the user must manually invoke Nextflow and adapt SLURM configuration.
• No unified status tracking, retry handling, or restart capability across build/simulate/analysis.

Proposed Behavior

A Parsl-based execution layer inside MDFactory that:

1. Replaces Nextflow entirely for both build and simulation orchestration.
2. Runs system builds (mdfactory build) on GPU nodes via SLURM — builds are compute-heavy (OpenMM compression) and belong on HPC, not login nodes.
3. Chains GROMACS simulation stages (EM → NVT → NPT → production) via @bash_app + futures.
4. Submits work to SLURM via HighThroughputExecutor + SlurmProvider (pilot-job model).
5. Supports custom workflow composition on the fly — e.g., packing multiple simulations onto a single node/job, or composing arbitrary build→simulate→analysis pipelines from Python.
6. Closely integrates with and reuses existing MDFactory building blocks (build.py, run_schedules/, Pydantic models) rather than reimplementing logic.
7. Handles embarrassingly parallel campaigns (hundreds of independent builds/simulations).
8. Provides retry, checkpoint/restart, and monitoring without external tooling.

Note: SLURM autodetection (detecting whether we're on a SLURM cluster) is out of scope — handled separately in another PR.

Requirements Matrix

Capability	Parsl Mechanism	PoC Scope
Python-native orchestration	@bash_app/@python_app decorators, futures-based DAG	✅ Core
CLI integration	Config objects generated from MDFactory models	✅ Core
System-building workflow support	@python_app or @bash_app invoking mdfactory build on GPU nodes	✅ Core
Simulation workflow support	@bash_app chaining EM→NVT→NPT→production	✅ Core
SLURM submission	SlurmProvider + HighThroughputExecutor	✅ Core
Embarrassingly parallel execution	Loop of app calls → auto-parallel futures	✅ Core
External command execution	@bash_app returns shell command string (with module loads, cd, etc.)	✅ Core
Dynamic workflow generation	Plain Python generating app calls from Pydantic models	✅ Core
Pydantic-compatible planning	MDFactory models → Parsl config/task generation	✅ Core
Reusable command/task templates	Parameterized @bash_app functions reusing MDFactory building blocks	✅ Core
Custom workflow composition	Composing tasks flexibly (multi-sim per node, arbitrary pipelines)	✅ Core
Monitoring / job status tracking	MonitoringHub → SQLite DB	✅ Core
User-facing status command	Query monitoring.db or MDFactory state	🔲 Stretch
Retries	Config(retries=N, retry_handler=...)	✅ Core
Restart / resume support	checkpoint_mode='task_exit' + get_all_checkpoints()	✅ Core
GROMACS checkpoint restart	@bash_app generating gmx mdrun -cpi ... -append	🔲 Stretch
Output-aware completion checks	Pre-check output files before submitting task	🔲 Stretch
MDFactory-owned persistent state	Sync Parsl monitoring.db to MDFactory registry	🔲 Future
Dry-run / inspectable plan	Generate tasks without parsl.load()	✅ Core
Multi-stage workflow composition	Future-chaining: step2(step1(x))	✅ Core
Many-simulation campaign execution	Loop over hashes/replicas	✅ Core
Minimal boilerplate	Single Config object, no external files	✅ Core
Low conceptual overhead	Hidden behind MDFactory CLI/API	✅ Core
Backend abstraction	LocalProvider ↔ SlurmProvider swap	✅ Core
Future extensibility	Multiple executors, @join_app for dynamic sub-workflows	🔲 Future

Acceptance Criteria

• A working @bash_app that runs gmx grompp + gmx mdrun for a single simulation stage, including proper HPC context (module loads, working directory, environment setup)
• System build (mdfactory build) running on GPU nodes via SLURM through Parsl
• Demonstration of chaining multiple stages (at minimum EM → production) via futures
• SLURM submission via SlurmProvider on the HPC cluster
• Embarrassingly parallel execution of multiple independent simulations (≥10 concurrent)
• Custom workflow composition demonstrated (e.g., multiple simulations packed onto a single node/job)
• Deferred (see also Investigate Parsl as Python-native workflow manager for build/simulate orchestration #16 (comment))
• Retry behavior validated (task fails → automatic resubmission)
• Checkpoint/restart demonstrated (re-run skips completed tasks)
• Integration with existing MDFactory building blocks (reuses build.py, Pydantic models, run schedules — not reimplementing)
• Dry-run mode that prints resolved commands without submitting
• Brief write-up of findings: what works, what doesn't, gaps vs. requirements

Technical Notes

Relevant existing code to reuse/integrate with:

• mdfactory/build.py — build pipeline (build_mixedbox, build_bilayer, build_lnp) — should be callable from Parsl tasks
• mdfactory/workflows.py — run_build_from_file() / run_build_from_dict() dispatch
• mdfactory/run_schedules/gromacs/ — MDP templates for EM/NVT/NPT/production
• mdfactory/models/input.py — BuildInput Pydantic model with SHA-1 hash
• mdfactory/analysis/submit.py — existing submitit pattern (SlurmConfig dataclass, path resolution)
• mdfactory/cli.py — cyclopts CLI (integration point for workflow commands)
• workflows/simulate.nf — current Nextflow workflow (reference for what to replace)
• workflows/simulate.config — SLURM config (maps to SlurmProvider params)

Parsl architecture for this use case:

• HighThroughputExecutor with SlurmProvider + SrunLauncher for GPU nodes
• worker_init for environment setup: module load gromacs-gpu, conda/pixi activation, etc.
• scheduler_options='#SBATCH --gres=gpu:1' for GPU allocation
• Pilot-job model: few SLURM jobs hosting many tasks (reduces scheduler pressure vs. submitit's 1-job-per-task)
• MonitoringHub for SQLite-backed status tracking

HPC-aware bash_app patterns to validate:

@bash_app
def grompp(work_dir, mdp, gro, top, tpr_out, stdout=parsl.AUTO_LOGNAME):
    # Real HPC apps need: module loads, cd to workdir, proper env
    return f'''
    module load gromacs-gpu/2024
    cd {work_dir}
    gmx grompp -f {mdp} -c {gro} -p {top} -o {tpr_out} -maxwarn 1
    '''

@bash_app
def mdrun(work_dir, tpr, deffnm, ntomp=8, stdout=parsl.AUTO_LOGNAME):
    return f'''
    module load gromacs-gpu/2024
    cd {work_dir}
    gmx mdrun -s {tpr} -deffnm {deffnm} -ntomp {ntomp} -nb gpu -pme gpu
    '''

@python_app
def build_system(build_input_dict):
    """Run mdfactory build on a GPU node — reuses existing building blocks."""
    from mdfactory.workflows import run_build_from_dict
    return run_build_from_dict(build_input_dict)

# Chain: grompp → mdrun, implicit dependency via futures
tpr = grompp('/scratch/sim001', 'em.mdp', 'system.gro', 'topology.top', 'em.tpr')
em_result = mdrun('/scratch/sim001', tpr, 'min')

# Custom workflow: multiple sims on one node
@bash_app
def multi_sim_node(work_dirs, deffnms, stdout=parsl.AUTO_LOGNAME):
    """Pack multiple short simulations onto a single GPU node."""
    commands = []
    for wd, name in zip(work_dirs, deffnms):
        commands.append(f'cd {wd} && gmx mdrun -deffnm {name} -ntomp 4 &')
    commands.append('wait')
    return '\n'.join(['module load gromacs-gpu/2024'] + commands)

Risks to investigate:

• Working directory management: Parsl workers may not cd to the expected simulation directory — must be explicit in bash_app commands
• Environment setup: module load and pixi/conda activation must work inside worker processes
• GPU allocation: ensuring correct GPU assignment when multiple tasks share a node
• File staging: whether parsl.File objects are needed or if absolute paths suffice on shared filesystems
• Custom node packing: whether HighThroughputExecutor supports multiple tasks per node cleanly, or if WorkQueueExecutor / TaskVineExecutor is better
• Overhead of monitoring for large campaigns (1000+ tasks)
• Serialization of MDFactory objects (Pydantic models, OpenFF molecules) when using @python_app

[maxscheurer]

Issue Assessment

Summary

This issue proposes replacing the entire Nextflow-based orchestration layer (build.nf, simulate.nf, run_full_pipeline.sh) with Parsl, a Python-native workflow manager. The scope covers build orchestration on GPU nodes, GROMACS simulation chaining, SLURM submission via pilot-job model, embarrassingly parallel campaigns, retry/checkpoint/restart, monitoring, dry-run mode, and custom workflow composition — essentially a full rewrite of the compute dispatch layer.

Current State

The orchestration is split across three independent mechanisms:

Layer	Tool	Files
Build dispatch	Nextflow DSL2	workflows/build.nf
Simulation chain	Nextflow DSL2	workflows/simulate.nf, simulate.config
Analysis dispatch	submitit	mdfactory/analysis/submit.py

The Python-side build machinery (build.py, workflows.py, models/input.py) is clean and well-structured with clear entry points (run_build_from_dict, run_build_from_file). The Nextflow layer is relatively thin — build.nf parallelizes mdfactory prepare-build → mdfactory build, and simulate.nf chains 4 GROMACS processes (EM → NVT → NPT → production). No Parsl or any other workflow engine exists in the codebase today.

Gaps

1. No unified orchestration — three different tools for three stages, no Python-level control over the full build→simulate→analyze pipeline.
2. No retry/checkpoint at the Python level — Nextflow has errorStrategy='retry' with maxRetries=1, but it's opaque to MDFactory.
3. Missing LNP run schedule — run_schedules.yaml declares system_type: lnp but gromacs/lnp/ directory doesn't exist; builds silently skip copying MDP files.
4. No dry-run — can't preview what would be submitted without actually submitting.
5. Manual SLURM config — users must edit simulate.config (Nextflow) and SlurmConfig (submitit) separately.

Ambiguities

1. Analysis layer fate — Should Parsl also replace submitit in analysis/submit.py, or does that remain separate? The issue mentions "build/simulate" scope but the requirements matrix implies "unified".
2. CLI design — Where does the user invoke Parsl workflows? New CLI subcommands? Python API only?
3. Config source of truth — run_schedules.yaml already has HPC settings (GPUs, cores, walltime) per system type. Should Parsl config read from this, or is it a separate config?
4. Nextflow removal timing — Remove Nextflow files in the final PR, or keep as fallback during transition?
5. Serialization — @python_app for builds requires serializing BuildInput and all transitive dependencies (OpenFF molecules, MDAnalysis Universes). Whether this works cleanly with Parsl's serialization is a significant unknown.

Scope & Proposed Decomposition

This is a large issue — 11 acceptance criteria, 20+ requirements matrix items, touching orchestration, CLI, config, SLURM, monitoring, and testing. It should be split into 4 incremental PRs:

PR 1 — Parsl infrastructure + build orchestration (replace build.nf)

• Add parsl as optional dependency (mdfactory[parsl])
• Create mdfactory/orchestration/ module
• Pydantic models for HPC/executor config → Parsl Config objects
• Backend abstraction: LocalProvider ↔ SlurmProvider swap
• @python_app or @bash_app for mdfactory build on GPU nodes
• Parallel builds from CSV/YAML reusing run_build_from_dict
• Integration with BuildInput model and RunScheduleManager
• Dry-run mode (generate tasks without parsl.load())
• Goal: working end-to-end slice testable on the cluster immediately

PR 2 — Simulation orchestration (replace simulate.nf)

• @bash_app wrappers for gmx grompp + gmx mdrun with HPC context
• Future-chaining for EM → NVT → NPT → production
• Embarrassingly parallel multi-simulation execution
• SLURM submission via SlurmProvider

PR 3 — Retry, checkpoint/restart, monitoring

• Config(retries=N, retry_handler=...)
• checkpoint_mode='task_exit' + get_all_checkpoints()
• MonitoringHub → SQLite for status tracking

PR 4 — Custom composition, CLI, Nextflow removal

• Multi-sim-per-node packing, arbitrary pipeline composition
• Full CLI integration
• Remove workflows/build.nf, simulate.nf, simulate.config, run_full_pipeline.sh
• Write-up of findings

Stretch (separate issues): GROMACS checkpoint restart (-cpi -append), output-aware completion checks, user-facing status command, sync monitoring.db to MDFactory registry.

Risks

1. Serialization — @python_app requires pickling everything. BuildInput is Pydantic (fine), but build functions pull in OpenMM, MDAnalysis, OpenFF — these may not serialize cleanly. Fallback: @bash_app invoking mdfactory build CLI.
2. Working directory management — Parsl workers don't cd automatically. Every @bash_app must explicitly handle paths — fragile and a common bug source.
3. GPU allocation — HighThroughputExecutor GPU assignment works for single-task-per-GPU, but multi-task-per-node packing needs careful config. May need WorkQueueExecutor for complex scenarios.
4. Environment setup — module load + pixi/conda activation inside Parsl workers is notoriously tricky. Must go in worker_init, and failures there are hard to debug.
5. Monitoring overhead — MonitoringHub adds a DB write per task state change. For 1000+ task campaigns this may bottleneck — needs benchmarking.
6. Test coverage — Hard to unit-test SLURM submission without a cluster. Need mock/local provider tests for CI + integration tests on HPC.

---

Automated assessment by mach6

[gregorweiss]

Custom workflow composition — separate issue

The "custom workflow composition" acceptance criterion (multiple simulations packed onto a single node/job via explicit &/wait in a bash_app) is being deferred to a separate issue rather than addressed in PR #25. Two reasons:

A) The current implementation achieves multi-task-per-node execution via Parsl-native max_workers_per_node + available_accelerators, making explicit bash packing a workaround rather than a necessity for basic full-node utilisation.

B) Explicit node packing becomes genuinely required in the GPU MPS scenario (Issue #10), where multiple simulations share a single GPU via NVIDIA MPS. The bash-level &/wait pattern sketched in this issue fits naturally as part of the MPS implementation, not the base orchestration layer.

The remaining acceptance criteria are met or closed:

• Dry-run with resolved commands — now prints actual gmx grompp/gmx mdrun command lines (not just the plan summary)
• Findings write-up — added as docs/content/docs/user-guide/parsl-findings.mdx
• All other criteria verified as delivered across PR feat: add Parsl-based parallel build orchestration #17 (build) and PR feat: Parsl-based simulation orchestration (Issue #16) #25 (simulate)

[gregorweiss]

Parsl Investigation: Findings

Summary of the Parsl proof-of-concept for build and simulation orchestration, replacing the legacy Nextflow layer. Delivered across PR #17 (build) and PR #25 (simulate).

What works

Build orchestration (PR #17)

• @python_app wrapping run_build_from_dict() — heavy imports inside workers, only plain dicts cross serialisation boundary
• Embarrassingly parallel CSV campaigns via HighThroughputExecutor + SlurmProvider
• Rich progress TUI with success/fail counts and SLURM block status
• Ctrl+C → scancel of SLURM jobs + parsl.clear()
• min_blocks=0 — idle allocations released, no wasted GPU-hours

Simulation orchestration (PR #25)

• @bash_app for grompp + mdrun with pure script builders (_build_grompp_script, _build_mdrun_script) — all conditional logic resolved at generation time, no bash branching
• EM → NVT → NPT → Production chained via inputs=[prev_future]; independent simulations in parallel
• Checkpoint/restart: three modes (auto/skip/force), disk-state detection, gmx mdrun -cpi <cpt> -append
• Adaptive rescue retry: physics failures retried with conservative MDP parameters up to --max-rescue tiers
• Per-stage overrides for cpus_per_node, gres, gmx_binary (affect mdrun command line, not SLURM allocation)
• Dry-run prints resolved gmx grompp/gmx mdrun command lines per stage
• Partial campaigns: missing-build systems skipped with status="skipped" in results

Infrastructure

• parsl_session() context manager — DFK guard, guaranteed parsl.clear(), best-effort scancel
• EnvironmentConfig — structured model for compute-node env (modules, pixi/conda/venv, extra init); global persistence to ~/.config/mdfactory/environment.yaml
• SLURM autodiscovery via SlurmExecutorConfig.from_cluster() (3-tier precedence)
• Interactive TUI wizard (mdfactory config slurm, --slurm tui)
• Parsl as optional dependency (pip install 'mdfactory[parsl]')

What doesn't work / limitations

Per-stage SLURM allocations
Parsl builds one SLURM allocation per HighThroughputExecutor. Per-stage walltime, mem, nodes, and partition overrides are not honoured at the SLURM level — they would require separate executors per stage. The stage_overrides field is restricted to cpus_per_node, gres, and gmx_binary (mdrun command line only). Invalid override keys raise ValueError at config time.

Note: a gres override (e.g. gres: null for EM) only changes the mdrun flags (-nb cpu instead of -nb gpu) — the stage still runs on the same GPU partition with a GPU allocated but unused. This is a pragmatic trade-off since EM is typically fast.

Custom node packing (multi-sim per bash_app)
Not implemented. Parsl-native max_workers_per_node + available_accelerators achieve multi-task-per-node execution. Explicit bash-level packing (&/wait) deferred to #10 (GPU MPS) where it fits naturally.

Monitoring
MonitoringHub supported via run_dir config field, but no mdfactory status command. Status visible through showdb and the Rich progress TUI.

Gaps vs. requirements

Requirement	Status	Notes
Python-native orchestration	✅	@bash_app/@python_app + futures DAG
CLI integration	✅	mdfactory build/simulate --slurm
System builds on GPU nodes	✅	PR #17
Simulation stage chaining	✅	PR #25
SLURM submission	✅	SlurmProvider + HighThroughputExecutor
Embarrassingly parallel execution	✅	Validated on cluster
Custom workflow composition	⚠️	Via max_workers_per_node, not explicit bash packing. Deferred to #10
Retry behaviour	✅	Config(retries=N) + adaptive rescue retry
Checkpoint/restart	✅	Disk-state detection + -cpi -append
Integration with building blocks	✅	Reuses run_build_from_dict, Pydantic models, MDP templates
Dry-run mode	✅	Plan summary + resolved gmx commands
Nextflow removal	✅	PR #29

Architecture decisions

1. Pure script builders — fully testable without a Parsl DFK; generated scripts have no conditional logic
2. StageSpec registry — frozen dataclass per stage as single source of truth; new stages = one entry
3. Sequential DFK sessions — parsl.clear() on exit for clean subsequent sessions
4. Fail-fast prerequisite validation — input files checked on login node before Parsl session

Test coverage

• PR feat: add Parsl-based parallel build orchestration #17: 68 tests (build orchestration, config, TUI, CLI)
• PR feat: Parsl-based simulation orchestration (Issue #16) #25: 227 tests (stages at 100% coverage, simulate, apps, CLI, rescue retry, config)
• Integration: rescue retry validated with real GROMACS execution (3-water fixture)