Skip to content

Add configurable Dask frame execution with SLURM-backed HPC support#354

Open
harryswift01 wants to merge 22 commits into
mainfrom
306-dask-parallel-implementation
Open

Add configurable Dask frame execution with SLURM-backed HPC support#354
harryswift01 wants to merge 22 commits into
mainfrom
306-dask-parallel-implementation

Conversation

@harryswift01
Copy link
Copy Markdown
Member

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.

@harryswift01 harryswift01 added this to the 2.3.0 milestone May 29, 2026
@harryswift01 harryswift01 self-assigned this May 29, 2026
@harryswift01 harryswift01 added the feature request New feature or request label May 29, 2026
@harryswift01 harryswift01 linked an issue May 29, 2026 that may be closed by this pull request
@harryswift01 harryswift01 requested a review from jimboid June 2, 2026 08:08
Copy link
Copy Markdown
Member

@jimboid jimboid left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature request New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature]: Dask-based Parallelisation for CodeEntropy

2 participants