Add configurable Dask frame execution with SLURM-backed HPC support

[harryswift01]

Summary

This PR adds configurable parallel frame-by-frame execution to CodeEntropy using Dask, with support for both local Dask workers and SLURM-backed HPC execution. It also adds runtime submission support for SLURM master jobs, validates the new parallel/HPC configuration options, and updates the documentation to explain the new arguments and parallel execution modes.

Changes

Parallel frame execution:

• Added Dask-based frame-level execution in LevelDAG.
• Added a parent-side reduction model where Dask workers compute frame-local covariance outputs and the main process reduces them into shared accumulators.
• Avoided broadcasting mutable trajectory/shared state to workers to prevent frame interference.
• Excluded parent-only objects and covariance accumulator aliases from worker payloads.
• Added sequential fallback behaviour when parallel execution is disabled, no Dask client is available, or only one frame is selected.

Local Dask and SLURM/HPC configuration:

• Added new configuration arguments for local Dask execution, including parallel_frames, use_dask, dask_workers, and dask_threads_per_worker.
• Added SLURM/HPC configuration options, including hpc, submit, hpc_queue, hpc_nodes, hpc_cores, hpc_processes, hpc_memory, hpc_walltime, hpc_account, hpc_qos, hpc_constraint, conda_path, conda_exec, and conda_env.
• Added HPCDaskManager for configuring SLURM-backed Dask clusters and generating/submitting master SLURM scripts.
• Added runtime handling for submit=True, allowing CodeEntropy to submit a master SLURM job and exit before starting local analysis.
• Added validation for local Dask, HPC, submit, and conda-related options.

Tests and documentation:

• Added unit tests for parallel frame execution, worker payload filtering, Dask submission, failure handling, and completed-frame validation.
• Added tests for local Dask and HPC workflow configuration paths.
• Added tests for SLURM helper behaviour, submit script generation, and submit error handling.
• Merged overlapping LevelDAG tests into a single test file to reduce duplication.
• Updated the getting started documentation with the new configuration arguments.
• Added a new documentation section explaining local Dask execution, SLURM/HPC Dask execution, and submit mode.

Impact

• Enables CodeEntropy to process frame-local covariance calculations in parallel.
• Provides a scalable path for larger trajectories using local Dask or SLURM-backed Dask workers.
• Keeps covariance reduction deterministic by performing reductions in the parent process.
• Improves HPC usability by allowing users to generate and submit SLURM master jobs directly from CodeEntropy.
• Adds validation to catch invalid parallel/HPC configuration before execution.
• Expands test coverage for the new parallel, HPC, and submit code paths.
• Updates user-facing documentation so users can configure and run sequential, local parallel, HPC, or submit-mode workflows.

[jimboid]

This PR looks very complete. The code additions to core functionality look like they are in line with the parallel version I wrote in the water-entropy code. It looks like the code has been sufficiently adapted to this code base with complete and thorough treatment of parameter handling. Documentation looks complete and detailed, the tests look detailed and complete like on the WE implementation. I don't see any surprises with the HPC set up logic, however we should not leave this to chance. In such a PR I would like to see the following data/screenshots or have someone verify the following points (since unit tests only go so far when there is an element of infrastructure involved):

1. Performance characterisation difference for serial implementation and parallel implementation for a local run. Total execution time is sufficient for say 20 - 50 frames.
2. Testing that the submission script on a known hpc platform is generated correctly for any automatic submission mode
3. Testing that the auto submission and conda detection machinery is working correctly and that the dask scheduler is submitted correctly.
4. That the performance data for a single node is inline and likely much stronger than those in point 1.

[harryswift01]

This PR looks very complete. The code additions to core functionality look like they are in line with the parallel version I wrote in the water-entropy code. It looks like the code has been sufficiently adapted to this code base with complete and thorough treatment of parameter handling. Documentation looks complete and detailed, the tests look detailed and complete like on the WE implementation. I don't see any surprises with the HPC set up logic, however we should not leave this to chance. In such a PR I would like to see the following data/screenshots or have someone verify the following points (since unit tests only go so far when there is an element of infrastructure involved):

1. Performance characterisation difference for serial implementation and parallel implementation for a local run. Total execution time is sufficient for say 20 - 50 frames.

2. Testing that the submission script on a known hpc platform is generated correctly for any automatic submission mode

3. Testing that the auto submission and conda detection machinery is working correctly and that the dask scheduler is submitted correctly.

4. That the performance data for a single node is inline and likely much stronger than those in point 1.

HPC validation and benchmark evidence

Thanks for the feedback. I have now run additional validation on SCARF to check the points you raised around performance, generated SLURM scripts, conda detection, and Dask worker submission.

Validation performed

Local serial vs local Dask performance

I ran 20-frame validation benchmarks for both DNA and benzene using serial execution and local Dask frame parallelism.

For the DNA 20-frame case:

Mode	Workers	Runtime / s
Serial	0	45
Local Dask	1	28
Local Dask	2	17
Local Dask	4	12
Local Dask	8	10
Local Dask	16	9

This gives a best local Dask speedup of around 5.0x over serial for this small validation case.

For the benzene 20-frame case:

Mode	Workers	Runtime / s
Serial	0	355
Local Dask	1	377
Local Dask	2	293
Local Dask	4	252
Local Dask	8	249
Local Dask	16	244

This case shows the expected Dask overhead at low worker counts, but still improves to around 1.45x over serial at 16 local workers.

Generated SLURM worker script

I also checked the generated Dask worker submission script on SCARF. For a small smoke test, the generated dask-cluster-submit.sh contained:

#SBATCH -J dask-worker
#SBATCH -p scarf
#SBATCH -n 1
#SBATCH --cpus-per-task=1
#SBATCH -t 00:20:00
unset SLURM_MEM_PER_CPU
unset SLURM_MEM_PER_GPU
unset SLURM_MEM_PER_NODE
unset SLURM_CPU_BIND
eval "$(/home/vol07/scarf1307/miniconda3/bin/conda shell.bash hook)"
conda activate CodeEntropy312
export SLURM_CPU_FREQ_REQ=2250000
srun python -m distributed.cli.dask_worker tcp://... --nthreads 1 --memory-limit 3.73GiB --nanny --death-timeout 60 --local-directory $PWD --interface eth0

This verifies that the worker script is generated correctly on SCARF, uses the expected partition, launches one core per worker, activates the expected conda environment, and includes the SLURM environment cleanup needed for nested srun calls.

Conda auto-detection

I tested the conda auto-detection by omitting conda_env, conda_exec, and conda_path from the run configuration.

The generated worker script correctly detected:

conda_path = /home/vol07/scarf1307/miniconda3/bin/conda
conda_env  = CodeEntropy312
conda_exec = conda

and wrote:

eval "$(/home/vol07/scarf1307/miniconda3/bin/conda shell.bash hook)"
conda activate CodeEntropy312

into the generated SLURM worker script.

These runtime checks are also backed up by unit tests around conda auto-detection, generated master submit scripts, generated Dask worker scripts, optional SLURM directives, SLURM cleanup, and sbatch failure handling.

I also tested the submit: true master submission path. This generated a CodeEntropy-master-submit.sh script with the detected conda environment and the recursion guard used to prevent repeated self-submission:

#SBATCH --job-name=codeentropy-master
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=2
#SBATCH --time=00:20:00
#SBATCH --partition=scarf
#SBATCH --output=CodeEntropy-master-%j.out
#SBATCH --error=CodeEntropy-master-%j.err

eval "$(/home/vol07/scarf1307/miniconda3/bin/conda shell.bash hook)"
conda activate CodeEntropy312

export CODEENTROPY_SUBMITTED_JOB=1

srun CodeEntropy

Main scaling benchmark

I then ran a larger benzene 128-frame benchmark as the main scaling case.

Mode	Workers	Runtime / s	Speedup vs serial
Serial	0	1644	1.00x
Local Dask	1	1788	0.92x
Local Dask	2	1154	1.42x
Local Dask	4	836	1.97x
Local Dask	8	696	2.36x
Local Dask	16	656	2.51x
Local Dask	32	595	2.76x
Local Dask	64	652	2.52x
HPC Dask	4	1234	1.33x
HPC Dask	8	726	2.26x
HPC Dask	16	603	2.73x
HPC Dask	32	547	3.01x
HPC Dask	64	542	3.03x

This larger case gives a clearer view of the scaling behaviour than the 20-frame validation runs.

The best local run was 32 workers at 595 s, which is around 2.76x faster than serial. The best HPC run was 64 workers at 542 s, which is around 3.03x faster than serial.

The local 64-worker run was slightly slower than the 32-worker run, which is likely due to the useful parallelism limit being reached for this 128-frame case. The HPC 32- and 64-worker results are very similar, which suggests the useful scaling limit for this benchmark is around 32 workers.

[jimboid]

Scaling behaviour is not as immediately strong as for WE, but this suggests that as we discussed there are other parts of the code that might need to be optimised. Please go ahead and merge