Skip to content

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

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

Add configurable Dask frame execution with SLURM-backed HPC support#354
harryswift01 wants to merge 19 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
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

1 participant