nntool.slurm.wrap

Functions

slurm_fn(submit_fn)

A decorator to wrap a function to be run on slurm.

slurm_function(submit_fn)

A decorator to annoate a function to be run in slurm.

slurm_launcher(ArgsType[, parser, ...])

A slurm launcher decorator for distributed or non-distributed job (controlled by use_distributed_env in slurm field).
nntool.slurm.wrap.slurm_fn(submit_fn)[source]

A decorator to wrap a function to be run on slurm. The function decorated by this decorator should be launched on the way below. The decorated function submit_fn is non-blocking now. To block and get the return value, you can call job.result().

Parameters:

submit_fn (Callable) – the function to be run on slurm

Returns:

the function to be run on slurm

Return type:

SlurmFunction

Example

>>> @slurm_fn
... def run_on_slurm(a, b):
...     return a + b
>>> slurm_config = SlurmConfig(
...     mode="slurm",
...     partition="PARTITION",
...     job_name="EXAMPLE",
...     tasks_per_node=1,
...     cpus_per_task=8,
...     mem="1GB",
... )
>>> job = run_on_slurm[slurm_config](1, b=2)
>>> result = job.result()  # block and get the result
nntool.slurm.wrap.slurm_launcher(ArgsType, parser='tyro', slurm_key='slurm', slurm_params_kwargs={}, slurm_submit_kwargs={}, slurm_task_kwargs={}, *extra_args, **extra_kwargs)[source]

A slurm launcher decorator for distributed or non-distributed job (controlled by use_distributed_env in slurm field). This decorator should be used as the program entry. The decorated function is non-blocking in the mode of slurm, while other modes cause blocking.

Parameters:

  • ArgsType (Type[Any]) – the experiment arguments type, which should be a dataclass (it mush have a slurm field defined by slurm_key)

  • slurm_key (str) – the key of the slurm field in the ArgsType, defaults to “slurm”

  • parser (str | Callable) – the parser for the arguments, defaults to “tyro”

  • slurm_config – SlurmConfig, the slurm configuration dataclass

  • slurm_params_kwargs (dict) – extra slurm arguments for the slurm configuration, defaults to {}

  • slurm_submit_kwargs (dict) – extra slurm arguments for srun or sbatch, defaults to {}

  • slurm_task_kwargs (dict) – extra arguments for the setting of distributed task, defaults to {}

  • extra_args – extra arguments for the parser

  • extra_kwargs – extra keyword arguments for the parser

Returns:

decorator function with main entry

Return type:

Callable[[Callable[[…], Any]], SlurmFunction]

Exported Distributed Enviroment Variables:

  1. NNTOOL_SLURM_HAS_BEEN_SET_UP is a special environment variable to indicate that the slurm has been set up.

  2. After the set up, the distributed job will be launched and the following variables are exported: num_processes: int, num_machines: int, machine_rank: int, main_process_ip: str, main_process_port: int.

nntool.slurm.wrap.slurm_function(submit_fn)[source]

A decorator to annoate a function to be run in slurm. The function decorated by this decorator should be launched in the way below.

Deprecated:

This function is deprecated and will be removed in future versions. Please use slurm_fn instead.

Example

>>> @slurm_function
... def run_on_slurm(a, b):
...     return a + b
>>> slurm_config = SlurmConfig(
...     mode="slurm",
...     partition="PARTITION",
...     job_name="EXAMPLE",
...     tasks_per_node=1,
...     cpus_per_task=8,
...     mem="1GB",
... )
>>> job = run_on_slurm(slurm_config)(1, b=2)
>>> result = job.result()  # block and get the result