pyantz.infrastructure.runner.slurm package

Submodules

pyantz.infrastructure.runner.slurm.slurm_runner module

Submit jobs to a slurm HPC grid.

Each job is saved as JSON in the shared working directory and submitted with sbatch, with slurm dependencies matching the pipeline’s. On the compute node, this module is run as a script (see ENTRY_MODULE_NAME) to load and run the job; jobs it submits are sbatched from there.

CONFIG_FILE_ENVIRON_VAR_NAME: str = 'PYANTZ_SLURM_CONFIG_FILE'

Environment variable holding the path of a file (the saved slurm runner config), not a directory. It is only passed on to jobs through the environment (--export=ALL), so prefer passing the path explicitly where possible.

ENTRY_BASH_FILE_PATH = PosixPath('/home/runner/work/pyantz/pyantz/src/pyantz/infrastructure/runner/slurm/run.sh')

Script which sbatch runs; it runs the command it is given.

ENTRY_MODULE_NAME = 'pyantz.infrastructure.runner.slurm.slurm_runner'

Module run (python -m) on the compute node to run a job.

SECONDS_PER_MINUTE = 60

Used to convert the job time limit into sbatch’s minutes.

_add_dependency(new_job_id: str, job_deps_to_update: dict[str, list[str]])

Add new job id as a dep to the job deps listed.

Parameters:

• new_job_id (str) – slurm job id of the newly submitted job.

• job_deps_to_update (dict[str, list[str]]) – slurm job ids of the jobs to make wait for it, mapped to the ids they already wait for.

Return type:

None

Returns:

None.

Raises:

subprocess.CalledProcessError – if scontrol fails.

_add_dependency_to_sbatch(args: list[str], dependencies: list[str], dependency_type: str = 'afterok')

Add the listed dependencies to the job submission.

Any dependency the user gave (--dependency=X, --dependency X, -d X or -dX) is merged with the new ones into a single --dependency=X,afterok:.. (or the job’s dependency type) argument added at the end. The other arguments are left as they were.

Parameters:

• args (list[str]) – the sbatch arguments.

• dependencies (list[str]) – slurm job ids the job must wait for.

• dependency_type (str) – how to wait for them, e.g., afterok.

Returns:

the arguments with the combined --dependency.

Return type:

list[str]

Raises:

ValueError – if a dependency flag has no value, or the user’s dependency uses ? (OR), which can’t be combined.

Example

from pyantz.infrastructure.runner.slurm.slurm_runner import (
    _add_dependency_to_sbatch,
)

args = ["--partition", "short", "-d", "afterany:7"]
print(_add_dependency_to_sbatch(args, ["10", "11"]))

['--partition', 'short', '--dependency=afterany:7,afterok:10,afterok:11']

_create_job_submitter(config: SlurmRunnerConfig, slurm_config_path: Path)

Create a method to accept a job config and submit it.

Must be called within a slurm job. With inherit_dependencies, the jobs waiting for the current job are made to wait for each submitted job too.

Parameters:

• config (SlurmRunnerConfig) – the slurm runner’s configuration.

• slurm_config_path (Path) – the saved slurm runner configuration.

Returns:

function which sbatches a job.

Return type:

SubmissionFnType

Raises:

• RuntimeError – if SLURM_JOB_ID isn’t set (not within a slurm job).

• subprocess.CalledProcessError – if squeue fails.

_create_sbatch(config: SlurmRunnerConfig, slurm_config_path: Path)

Sbatch the provided job config and return its new job id.

Parameters:

• config (SlurmRunnerConfig) – the slurm runner’s configuration.

• slurm_config_path (Path) – the saved slurm runner configuration.

Returns:

function which sbatches a job, waiting for the given slurm job ids, and returns its slurm job id.

Return type:

Callable[[JobConfig, list[str] | None], str]

Raises:

None. –

_get_dependent_jobs(curr_job_id: str)

Get slurm job ids of jobs which depend on the job in which this is run.

Parameters:

curr_job_id (str) – the slurm job id of the current job.

Returns:

pending jobs which wait on curr_job_id, mapped to

the ids of all the jobs they wait on (including curr_job_id)

Return type:

dict[str, list[str]]

Raises:

subprocess.CalledProcessError – if squeue fails.

_get_slurm_env_variables()

Get slurm environment variables.

Parameters:

None.

Returns:

the job’s id, name, nodes, number of tasks and cpus per task, each None if not set (e.g., outside of slurm).

Return type:

dict[str, str | None]

Raises:

None. –

Example

import os

from pyantz.infrastructure.runner.slurm.slurm_runner import (
    _get_slurm_env_variables,
)

os.environ["SLURM_JOB_ID"] = "1234"
try:
    print(_get_slurm_env_variables()["job_id"])
finally:
    del os.environ["SLURM_JOB_ID"]

1234

_has_time_limit(slurm_args: list[str])

Check whether the sbatch arguments already set a time limit.

Parameters:

slurm_args (list[str]) – the sbatch arguments.

Returns:

True if they include --time or -t in any form.

Return type:

bool

Raises:

None. –

Example

from pyantz.infrastructure.runner.slurm.slurm_runner import _has_time_limit

print(_has_time_limit(["-t10"]), _has_time_limit(["--test-only"]))

True False

_initial_submit(config: InitialConfig, runner_config: SlurmRunnerConfig, slurm_config_path: Path)

Submit our initial set of jobs.

Done differently because rather than depending on current, the initial jobs have their own dependency graph that we must resolve up front.

Parameters:

• config (InitialConfig) – the validated configuration; its variables are given to the jobs.

• runner_config (SlurmRunnerConfig) – the slurm runner’s configuration.

• slurm_config_path (Path) – the saved slurm runner configuration.

Return type:

None

Returns:

None.

Raises:

• subprocess.CalledProcessError – if sbatch fails.

• OSError – if a job can’t be saved in the working directory.

_load_job_config()

Load our job config from shared storage using sys args.

Reads -i/--slurm_config (path of the saved slurm runner configuration) and -c/--job_config (path of the job, relative to the working directory) from the command line.

Parameters:

None.

Returns:

the job, the slurm runner configuration and the path of the saved configuration.

Return type:

tuple[JobWithContext, SlurmRunnerConfig, Path]

Raises:

• ValueError – if either file doesn’t exist.

• SystemExit – if the command line arguments are invalid.

• pydantic.ValidationError – if a file’s contents are invalid.

_run_on_node()

Run the job on this node.

Parameters:

None.

Returns:

True if the job succeeded

Return type:

bool

Raises:

• ValueError – if the job or configuration file doesn’t exist.

• RuntimeError – if not run within a slurm job.

_save_slurm_config(runner_config: SlurmRunnerConfig)

Save the slurm configuration in a global location and store in environment.

The configuration is saved in the working directory under a unique name, and its path put in CONFIG_FILE_ENVIRON_VAR_NAME.

Parameters:

runner_config (SlurmRunnerConfig) – the slurm runner’s configuration.

Returns:

where the slurm configuration was saved

Return type:

Path

Raises:

OSError – if the file can’t be written.

Example

import os
from pathlib import Path

from pyantz.infrastructure.config.runners import SlurmRunnerConfig
from pyantz.infrastructure.runner.slurm.slurm_runner import (
    _save_slurm_config,
)

path = _save_slurm_config(SlurmRunnerConfig(working_dir="."))
print(path.name.startswith("global_slurm_config_"))
print(os.environ["PYANTZ_SLURM_CONFIG_FILE"] == os.fspath(path))

True
True

start(config: InitialConfig)

Start the slurm pipeline up.

The jobs are submitted in dependency order and this returns; they run (and submit their jobs) on the grid afterwards.

Parameters:

config (InitialConfig) – the validated configuration, with a slurm_sbatch submitter.

Return type:

None

Returns:

None.

Raises:

• TypeError – if the submitter of the configuration is not a slurm runner

• NotADirectoryError – if the slurm working directory does not exist

• subprocess.CalledProcessError – if sbatch fails.

Module contents

Run jobs on the slurm grid.