Parameter space discovery via slurm#

This module allows submitting a slurm job array where each individual job runs a simulation with a different set of parameters.

For a given problem ExampleProblem with design type ExampleDesign, the following code will submit a slurm job array with 3 jobs:

static_args = slurm.Args(simulate_args={"config": {"sim_arg": 10}}, problem_args={"some_arg": True})
parameter_space = [
     slurm.Args(problem_args={"problem_arg": 1}, design_args={"design_arg": -1}),
     slurm.Args(problem_args={"problem_arg": 2}, design_args={"design_arg": -2}),
     slurm.Args(problem_args={"problem_arg": 3}, design_args={"design_arg": -3}),
]
slurm.submit(
    problem=ExampleProblem,
    static_args=static_args,
    parameter_space=parameter_space),
)

For example the first job would execute:

problem = ExampleProblem(some_arg=True, problem_arg=1)
design = ExampleDesign(design_arg=-1)
problem.simulate(design, config={"sim_arg": 10})

and so on.

The main steps needed to submit a job is to declare the arguments as a Args object:

class engibench.utils.slurm.Args(problem_args: dict[str, ~typing.Any] = <factory>, simulate_args: dict[str, ~typing.Any] = <factory>, optimize_args: dict[str, ~typing.Any] = <factory>, design_args: dict[str, ~typing.Any] = <factory>)[source]#

Collection of arguments passed to Problem(), Problem.simulate() and DesignType().

design_args: dict[str, Any]#

Keyword arguments to be passed to DesignType() or the design_factory argument of submit().

optimize_args: dict[str, Any]#

Keyword arguments to be passed to engibench.core.Problem.optimize().

problem_args: dict[str, Any]#

Keyword arguments to be passed to engibench.core.Problem().

simulate_args: dict[str, Any]#

Keyword arguments to be passed to engibench.core.Problem.simulate().

passing it to submit():

slurm.submit(problem: type[Problem], parameter_space: list[Args], design_factory: Callable[[...], DesignType] | None = None, config: SlurmConfig | None = None) None#

Submit a job array for a parameter discovery to slurm.

  • job_type - The type of the job to be submitted: ‘simulate’, ‘optimize’, or ‘render’.

  • problem - The problem type for which the simulation should be run.

  • parameter_space - One Args instance per simulation run to be submitted.

  • design_factory - If not None, pass Args.design_args to design_factory instead of DesignType().

  • design_factory - Custom arguments passed to sbatch.

To tweak the arguments passed to sbatch, the config argument can be passed to submit():

class engibench.utils.slurm.SlurmConfig(sbatch_executable: str = 'sbatch', log_dir: str | None = None, name: str | None = None, account: str | None = None, runtime: str | None = None, constraint: str | None = None, mem_per_cpu: str | None = None, mem: str | None = None, nodes: int | None = None, ntasks: int | None = None, cpus_per_task: int | None = None, extra_args: Sequence[str] = ())[source]#

Collection of slurm parameters passed to sbatch.

account: str | None = None#

Slurm account to use

constraint: str | None = None#

Optional constraint

extra_args: Sequence[str] = ()#

Extra arguments passed to sbatch.

log_dir: str | None = None#

Path of the log directory

mem: str | None = None#

E.g. “4G”.

mem_per_cpu: str | None = None#

E.g. “4G”.

name: str | None = None#

Optional name for the jobs

runtime: str | None = None#

Optional runtime in the format hh:mm:ss.

sbatch_executable: str = 'sbatch'#

Path to the sbatch executable if not in PATH