Workflow Module

Daisy Chain Submitter

Daisy-chain job submission for HPC SLURM scheduler.

This module provides utilities for breaking long MD simulations into smaller dependent jobs that are automatically chained together using SLURM job dependencies.

class polyzymd.workflow.daisy_chain.SegmentInfo(index, duration_ns, samples, is_initial, cumulative_time_ns)[source]

Bases: object

Information about a single simulation segment.

index

Segment index (0-based for initial, 1+ for continuations)

Type:: int

duration_ns

Duration of this segment in nanoseconds

Type:: float

samples

Number of trajectory frames to save

Type:: int

is_initial

Whether this is the initial (build + equilibration + first prod) segment

Type:: bool

cumulative_time_ns

Total simulated time up to and including this segment

Type:: float

index: int

duration_ns: float

samples: int

is_initial: bool

cumulative_time_ns: float

__init__(index, duration_ns, samples, is_initial, cumulative_time_ns)

class polyzymd.workflow.daisy_chain.DaisyChainConfig(slurm_config, total_production_time_ns, total_segments=10, total_samples=2500, equilibration_time_ns=0.5, replicates=<factory>, dry_run=False, output_script_dir=PosixPath('daisy_chain_scripts'), config_path='config.yaml')[source]

Bases: object

Configuration for daisy-chain submission.

slurm_config

SLURM job configuration

Type:: polyzymd.workflow.slurm.SlurmConfig

total_production_time_ns

Total production time in nanoseconds

Type:: float

total_segments

Number of segments to split production into

Type:: int

total_samples

Total trajectory frames across all segments

Type:: int

equilibration_time_ns

Equilibration time (only for initial segment)

Type:: float

replicates

List of replicate numbers to run

Type:: List[int]

dry_run

If True, create scripts but don’t submit

Type:: bool

output_script_dir

Directory for generated job scripts

Type:: pathlib.Path

config_path

Path to the YAML configuration file

Type:: str

slurm_config: SlurmConfig

total_production_time_ns: float

total_segments: int = 10

total_samples: int = 2500

equilibration_time_ns: float = 0.5

replicates: List[int]

dry_run: bool = False

output_script_dir: Path = PosixPath('daisy_chain_scripts')

config_path: str = 'config.yaml'

property segment_duration_ns: float: Get the duration of each segment in nanoseconds.

property samples_per_segment: int: Get the number of frames per segment.

get_segments()[source]

Generate segment information for all segments.

Returns:: List of SegmentInfo objects for each segment.
Return type:: List[SegmentInfo]

classmethod from_simulation_config(sim_config, slurm_config, replicates='1', dry_run=False, output_script_dir='daisy_chain_scripts', config_path='config.yaml')[source]

Create DaisyChainConfig from a SimulationConfig.

Parameters:

sim_config (SimulationConfig) – Simulation configuration
slurm_config (SlurmConfig) – SLURM configuration
replicates (str | List[int]) – Replicate range string (e.g., “1-5”) or list of ints
dry_run (bool) – If True, don’t submit jobs
output_script_dir (str | Path) – Directory for job scripts
config_path (str) – Path to the YAML configuration file

Returns:

Configured DaisyChainConfig

Return type:

DaisyChainConfig

__init__(slurm_config, total_production_time_ns, total_segments=10, total_samples=2500, equilibration_time_ns=0.5, replicates=<factory>, dry_run=False, output_script_dir=PosixPath('daisy_chain_scripts'), config_path='config.yaml')

class polyzymd.workflow.daisy_chain.SubmissionResult(job_id, script_path, segment_index, replicate, is_dry_run=False)[source]

Bases: object

Result of job submission.

job_id

SLURM job ID (or dummy ID for dry run)

Type:: str

script_path

Path to the generated script

Type:: pathlib.Path

segment_index

Segment index for this job

Type:: int

replicate

Replicate number

Type:: int

is_dry_run

Whether this was a dry run

Type:: bool

job_id: str

script_path: Path

segment_index: int

replicate: int

is_dry_run: bool = False

__init__(job_id, script_path, segment_index, replicate, is_dry_run=False)

class polyzymd.workflow.daisy_chain.DaisyChainSubmitter(sim_config, dc_config, conda_env='polymerist-env', openff_logs=False, skip_build=False)[source]

Bases: object

Handles daisy-chain job submission for MD simulations.

This class generates SLURM job scripts and submits them with proper dependencies so that continuation jobs run after their prerequisites.

Example

>>> sim_config = SimulationConfig.from_yaml("config.yaml")
>>> slurm_config = SlurmConfig.from_preset("aa100", email="user@example.com")
>>> dc_config = DaisyChainConfig.from_simulation_config(
...     sim_config, slurm_config, replicates="1-3"
... )
>>> submitter = DaisyChainSubmitter(sim_config, dc_config)
>>> results = submitter.submit_all()

__init__(sim_config, dc_config, conda_env='polymerist-env', openff_logs=False, skip_build=False)[source]

Initialize the DaisyChainSubmitter.

Parameters:

sim_config (SimulationConfig) – Simulation configuration
dc_config (DaisyChainConfig) – Daisy-chain configuration
conda_env (str) – Conda environment name
openff_logs (bool) – Enable verbose OpenFF logs in generated scripts
skip_build (bool) – Skip system building in generated scripts (use pre-built system)

property sim_config: SimulationConfig: Get the simulation configuration.

property dc_config: DaisyChainConfig: Get the daisy-chain configuration.

property job_chains: Dict[int, List[SubmissionResult]]: Get the job chains for all replicates.

generate_initial_script(replicate)[source]

Generate the initial job script content.

Parameters:: replicate (int) – Replicate number
Returns:: Script content string
Return type:: str

generate_continuation_script(segment_index, replicate)[source]

Generate a continuation job script content.

Parameters:

segment_index (int) – Segment index (1 or higher)
replicate (int) – Replicate number

Returns:

Script content string

Return type:

str

submit_replicate_chain(replicate)[source]

Submit all jobs for a single replicate.

Parameters:: replicate (int) – Replicate number
Returns:: List of SubmissionResults for all segments
Return type:: List[SubmissionResult]

submit_all()[source]

Submit jobs for all replicates.

Returns:: Dictionary mapping replicate numbers to their job chains
Return type:: Dict[int, List[SubmissionResult]]

polyzymd.workflow.daisy_chain.submit_daisy_chain(config_path, slurm_preset='aa100', replicates='1', email='', dry_run=False, conda_env='polymerist-env', output_dir=None, scratch_dir=None, projects_dir=None, time_limit=None, memory=None, openff_logs=False, skip_build=False)[source]

Convenience function to submit daisy-chain jobs from a YAML config.

Parameters:

config_path (str | Path) – Path to simulation YAML config
slurm_preset (str) – SLURM preset name (aa100, al40, blanca-shirts, testing)
replicates (str) – Replicate range string (e.g., “1-5”, “1,3,5”)
email (str) – Email for job notifications
dry_run (bool) – If True, don’t submit jobs
conda_env (str) – Conda environment name
output_dir (str | Path | None) – Directory for job scripts (default: from config or “job_scripts”)
scratch_dir (str | Path | None) – Override scratch directory for simulation output
projects_dir (str | Path | None) – Override projects directory for scripts/logs
time_limit (str | None) – Override SLURM time limit (format: HH:MM:SS or M:SS)
memory (str | None) – Override SLURM memory allocation (e.g., “4G”, “8G”)
openff_logs (bool) – Enable verbose OpenFF logs in generated scripts
skip_build (bool) – Skip system building in generated scripts (use pre-built system)

Returns:

Dictionary mapping replicate numbers to submission results

Return type:

Dict[int, List[SubmissionResult]]

Example

>>> results = submit_daisy_chain(
...     config_path="simulation.yaml",
...     slurm_preset="aa100",
...     replicates="1-5",
...     email="user@example.com",
...     dry_run=True,
... )

polyzymd.workflow.daisy_chain.main()[source]

Main entry point for daisy-chain submission CLI.

Returns:: Exit code (0 for success, 1 for failure).
Return type:: int

SLURM Configuration

SLURM job script generation for HPC cluster submission.

This module provides templates and utilities for generating SLURM batch scripts for MD simulations.

class polyzymd.workflow.slurm.SlurmConfig(partition='aa100', qos='normal', account='ucb625_asc1', time_limit='23:59:59', email='', nodes=1, ntasks=1, memory='3G', gpus=1, exclude=None)[source]

Bases: object

Configuration for SLURM job submission.

partition

SLURM partition(s) to use.

Type:: str

qos

Quality of service.

Type:: str

account

Account for resource allocation.

Type:: str

time_limit

Wall time limit (HH:MM:SS).

Type:: str

email

Email for notifications.

Type:: str

nodes

Number of nodes.

Type:: int

ntasks

Number of tasks.

Type:: int

memory

Memory allocation (e.g., “3G”).

Type:: str

gpus

Number of GPUs.

Type:: int

exclude

Nodes to exclude.

Type:: str | None

partition: str = 'aa100'

qos: str = 'normal'

account: str = 'ucb625_asc1'

time_limit: str = '23:59:59'

email: str = ''

nodes: int = 1

ntasks: int = 1

memory: str = '3G'

gpus: int = 1

exclude: str | None = None

classmethod from_preset(preset, email='')[source]

Create a SlurmConfig from a preset.

Parameters:

preset (Literal['aa100', 'al40', 'blanca-shirts', 'testing']) – Preset name.
email (str) – Email for notifications.

Returns:

SlurmConfig with preset values.

Return type:

SlurmConfig

__init__(partition='aa100', qos='normal', account='ucb625_asc1', time_limit='23:59:59', email='', nodes=1, ntasks=1, memory='3G', gpus=1, exclude=None)

class polyzymd.workflow.slurm.JobContext(job_name, output_file, scratch_dir, projects_dir='.', segment_index=0, replicate_num=1, extra_vars=<factory>)[source]

Bases: object

Context for job script template rendering.

job_name

SLURM job name.

Type:: str

output_file

Output file pattern (for SLURM logs).

Type:: str

scratch_dir

Directory for simulation output (trajectories, checkpoints).

Type:: str

projects_dir

Directory for scripts and logs.

Type:: str

segment_index

Current segment index.

Type:: int

replicate_num

Replicate number.

Type:: int

extra_vars

Additional template variables.

Type:: Dict

job_name: str

output_file: str

scratch_dir: str

projects_dir: str = '.'

segment_index: int = 0

replicate_num: int = 1

extra_vars: Dict

property working_dir: str: Alias for scratch_dir for backwards compatibility.

__init__(job_name, output_file, scratch_dir, projects_dir='.', segment_index=0, replicate_num=1, extra_vars=<factory>)

class polyzymd.workflow.slurm.SlurmScriptGenerator(config, conda_env='polymerist-env', openff_logs=False, skip_build=False)[source]

Bases: object

Generator for SLURM batch scripts.

Supports separate directories for: - projects_dir: Where scripts live and jobs are submitted from - scratch_dir: Where simulation output goes (trajectories, checkpoints)

Example

>>> config = SlurmConfig.from_preset("aa100", email="user@example.com")
>>> generator = SlurmScriptGenerator(config)
>>> script = generator.generate_initial_job(
...     context=JobContext(
...         job_name="my_sim",
...         output_file="logs/output.log",
...         scratch_dir="/scratch/user/sim_output",
...         projects_dir="/projects/user/polyzymd",
...     ),
...     python_script="run_simulation.py",
...     python_args={"temperature": 300},
... )

INITIAL_JOB_TEMPLATE = '#!/bin/bash\n#SBATCH --partition={partition}\n#SBATCH --job-name=i_{job_name}\n#SBATCH --output={output_file}\n#SBATCH --qos={qos}\n#SBATCH --nodes={nodes}\n#SBATCH --ntasks={ntasks}\n#SBATCH --mem={memory}\n#SBATCH --time={time_limit}\n#SBATCH --gres=gpu:{gpus}\n#SBATCH --mail-type=FAIL\n#SBATCH --mail-user={email}\n#SBATCH --account={account}\n{exclude_line}\n\n# =============================================================================\n# PolyzyMD Initial Simulation Job\n# Segment: {segment_index}\n# =============================================================================\n\n# Load conda environment (ignore module warnings on some HPC systems)\nmodule purge 2>/dev/null || true\nmodule load miniforge 2>/dev/null || true\n\n# Initialize conda/mamba for non-interactive shell\neval "$(conda shell.bash hook)"\nmamba activate {conda_env}\n\n# Enable strict error handling after environment setup\nset -e\n\n# Projects directory (scripts, configs, logs)\nPROJECTS_DIR="{projects_dir}"\n\n# Scratch directory (simulation output)\nSCRATCH_DIR="{scratch_dir}"\n\n# Ensure scratch directory exists\nmkdir -p "$SCRATCH_DIR"\n\n# Change to projects directory where config and scripts live\ncd "$PROJECTS_DIR"\n\necho "Starting initial simulation segment {segment_index}"\necho "Projects dir: $PROJECTS_DIR"\necho "Scratch dir: $SCRATCH_DIR"\necho "Config: {config_path}"\necho "Replicate: {replicate}"\necho "Timestamp: $(date)"\n\n# Run the initial simulation using polyzymd CLI\n# This builds the system, runs equilibration, and runs the first production segment\npolyzymd{openff_logs_flag} run -c "{config_path}" \\\n --replicate {replicate} \\\n --scratch-dir "$SCRATCH_DIR" \\\n --segment-time {segment_time} \\\n --segment-frames {segment_frames}{skip_build_flag}\n\necho "Segment {segment_index} completed successfully at $(date)"\n'

CONTINUATION_JOB_TEMPLATE = '#!/bin/bash\n#SBATCH --partition={partition}\n#SBATCH --job-name=c_{job_name}\n#SBATCH --output={output_file}\n#SBATCH --qos={qos}\n#SBATCH --nodes={nodes}\n#SBATCH --ntasks={ntasks}\n#SBATCH --mem={memory}\n#SBATCH --time={time_limit}\n#SBATCH --gres=gpu:{gpus}\n#SBATCH --mail-type=FAIL\n#SBATCH --mail-user={email}\n#SBATCH --account={account}\n{exclude_line}\n\n# =============================================================================\n# PolyzyMD Continuation Job\n# Segment: {segment_index}\n# =============================================================================\n\n# Load conda environment (ignore module warnings on some HPC systems)\nmodule purge 2>/dev/null || true\nmodule load miniforge 2>/dev/null || true\n\n# Initialize conda/mamba for non-interactive shell\neval "$(conda shell.bash hook)"\nmamba activate {conda_env}\n\n# Enable strict error handling after environment setup\nset -e\n\n# Projects directory (scripts, configs, logs)\nPROJECTS_DIR="{projects_dir}"\n\n# Scratch directory (simulation output - where previous segment data lives)\nSCRATCH_DIR="{scratch_dir}"\n\n# Change to projects directory\ncd "$PROJECTS_DIR"\n\necho "Starting continuation segment {segment_index}"\necho "Projects dir: $PROJECTS_DIR"\necho "Scratch dir: $SCRATCH_DIR"\necho "Timestamp: $(date)"\n\n# Continue simulation from previous segment using polyzymd CLI\n# Reads checkpoint from previous segment in SCRATCH_DIR\n# Writes new trajectory and checkpoint to SCRATCH_DIR\npolyzymd{openff_logs_flag} continue \\\n -w "$SCRATCH_DIR" \\\n -s {segment_index} \\\n -t {segment_time} \\\n -n {num_samples}\n\necho "Segment {segment_index} completed successfully at $(date)"\n'

__init__(config, conda_env='polymerist-env', openff_logs=False, skip_build=False)[source]

Initialize the generator.

Parameters:

config (SlurmConfig) – SLURM configuration.
conda_env (str) – Conda environment name.
openff_logs (bool) – Enable verbose OpenFF logs in generated scripts.
skip_build (bool) – Skip system building in generated scripts (use pre-built system).

property config: SlurmConfig: Get the SLURM configuration.

generate_initial_job(context, config_path, replicate, segment_time, segment_frames)[source]

Generate an initial simulation job script.

Parameters:

context (JobContext) – Job context information.
config_path (str) – Path to the YAML configuration file.
replicate (int) – Replicate number.
segment_time (float) – Duration of first segment in nanoseconds.
segment_frames (int) – Number of frames to save in first segment.

Returns:

SLURM batch script content.

Return type:

str

generate_continuation_job(context, segment_time, num_samples)[source]

Generate a continuation job script.

Parameters:

context (JobContext) – Job context information.
segment_time (float) – Duration of this segment in nanoseconds.
num_samples (int) – Number of frames to save.

Returns:

SLURM batch script content.

Return type:

str

save_script(script_content, output_path, make_executable=True)[source]

Save a script to a file.

Parameters:

script_content (str) – Script content.
output_path (str | Path) – Output file path.
make_executable (bool) – Whether to make the script executable.

Returns:

Path to the saved script.

Return type:

Path

polyzymd.workflow.slurm.parse_replicate_range(replicate_range)[source]

Parse a SLURM array range into a list of replicate numbers.

Parameters:: replicate_range (str) – SLURM array format (e.g., “1-5”, “1,3,5”, “1-10:2”).
Returns:: List of replicate numbers.
Return type:: List[int]

Example

>>> parse_replicate_range("1-5")
[1, 2, 3, 4, 5]
>>> parse_replicate_range("1,3,5")
[1, 3, 5]
>>> parse_replicate_range("1-10:2")
[1, 3, 5, 7, 9]

polyzymd.workflow.slurm.validate_replicate_range(replicate_range)[source]

Validate that a replicate range is in proper SLURM array format.

Parameters:: replicate_range (str) – Range string to validate.
Returns:: True if valid.
Raises:: ValueError – If the format is invalid.
Return type:: bool