submitit

You can run a single job that spans multiple tasks across one or several nodes using tasks_per_node and nodes in update_parameters().

• Execution: The same function is executed in every task.
• Accessing Task Info: Use submitit.JobEnvironment() inside your function to access cluster-agnostic metadata like num_tasks, local_rank, node, and global_rank. This is typically used to chunk inputs so each task processes a subset of data.
• Accessing Individual Tasks: You can use job.task(rank) to access a specific task as if it were a standalone Job object (e.g., to get its specific stdout() or result()).

import submitit
from math import ceil

def my_func(inputs):
    job_env = submitit.JobEnvironment()
    # Use job_env to chunk inputs
    num_items_per_task = int(ceil(len(inputs) / job_env.num_tasks))
    r = job_env.local_rank
    task_chunk = inputs[r * num_items_per_task: (r + 1) * num_items_per_task]
    return process(task_chunk)

executor = submitit.AutoExecutor(folder="log_test")
# 3 tasks per node * 2 nodes = 6 total tasks
executor.update_parameters(tasks_per_node=3, nodes=2, timeout_min=1, slurm_partition="dev")
job = executor.submit(my_func, my_large_input_list)

# Access specific task results
print(job.task(2).result())
# Access concatenated stdout of all tasks
print(job.stdout())

When you submit a function and its arguments, submitit pickles the function and arguments. A batch file is then executed on the cluster, which loads the pickled object, computes the function with the provided arguments, and pickles the output into a new file. Once this file is available, your Job instance can recover it.

Important Requirements:

• The computation in the cluster uses the current conda environment. You must ensure all necessary dependencies (including submitit) are installed in that environment.
• If a computation fails and the error is catchable, the traceback is available via the Job instance. If it is not catchable, you must inspect the log files.

Generated Files: For each job, the following files are typically created:

• <job_id>_submitted.pkl: The task file.
• <job_id>_result.pkl: The output file.
• batchfile_<uuid>.sh: The batch file (where <uuid> is generated by submitit).
• <job_id>_<task_id>_log.out: The stdout log file.
• <job_id>_<task_id>_log.err: The stderr log file.

Note: <job_id> refers to the ID from the scheduler (e.g., Slurm).

Checkpointing allows submitit to requeue a job (after preemption or timeout) by editing the submitted task according to the current state of the computation.

To enable checkpointing, you must submit a callable (an instance of a class with a __call__ method) rather than a standard function. This is because a standard function's state cannot be accessed during requeueing.

When a job is requeued, submitit checks if the callable has a __submitit_checkpoint__ or checkpoint method. If found, it calls this method with the same arguments used in the original __call__ method. The checkpoint method is responsible for preparing the new submission and must return a submitit.helpers.DelayedSubmission object (which behaves like executor.submit) or None if the job should not be requeued.

Cluster Requirement: For preemptions to be recognized, your Slurm cluster must be configured with SlurmctldParameters=preempt_send_user_signal.

To submit many jobs efficiently, use executor.map_array. This is preferred over individual submit calls because it submits all jobs in a single Slurm call, avoiding scheduler flooding and allowing you to control parallelism.

• Parallelism: Use executor.update_parameters(slurm_array_parallelism=N) to cap how many jobs run concurrently.
• Job IDs: Array jobs use the format <array job id>_<array task id> (e.g., 17390420_15).
• Warning: map_array creates one pickle per job. If your function uses large objects (like PyTorch models), serialize them once and pass the file path instead of the object itself.

Batch Context Manager: You can convert a standard for loop of submit calls into a job array by wrapping the loop in with executor.batch():. The jobs are only actually submitted when exiting the context.

a = [1, 2, 3, 4]
b = [10, 20, 30, 40]
executor = submitit.AutoExecutor(folder=log_folder)

# Limit parallelism to 2 jobs at once
executor.update_parameters(slurm_array_parallelism=2)

# Submit as an array
jobs = executor.map_array(add, a, b)

# Alternatively, using the batch context manager:
jobs = []
with executor.batch():
    for arg in whatever:
        job = executor.submit(myfunc, arg)
        jobs.append(job)
# Jobs are submitted here, upon exiting the context

To ensure reliable execution and clear error reporting, follow these best practices for function definition:

1. Define functions in modules: Always prefer submitting functions defined within a module rather than locally defined functions. This ensures tracebacks are explicit and easier to debug.
2. Avoid sys.path.append for imports: Modules added to the path via sys.path.append before submission often cannot be correctly pickled. If you must use custom paths, use a 'lazy import' pattern by performing the sys.path.append and the import inside the function being submitted.
3. Beware of module changes: Since submitted functions are references to module functions, if the module file is modified between the time of submission and the time the job actually starts, the computation may behave unexpectedly or fail.
4. Avoid non-picklable arguments: Do not use non-picklable objects (such as threading.Lock) as default arguments in the functions you intend to submit.

The submitit.AutoExecutor class is the primary interface for submitting Python functions to a cluster (like Slurm) or running them locally. It follows the concurrent.futures.Executor API.

Key features:

• Log Management: Specify a folder where job information, logs, and results are stored. Use %j in the path to automatically include the job ID at runtime.
• Parameter Updates: Use update_parameters() to configure job settings. Cluster-specific options must be prefixed with the cluster name (e.g., slurm_partition for Slurm).
• Result Retrieval: job.result() waits for completion and returns the output. If the job fails, it raises a FailedJobError containing the traceback.

import submitit

def add(a, b):
    return a + b

log_folder = "log_test/%j"
executor = submitit.AutoExecutor(folder=log_folder)

# Configure parameters (e.g., timeout and Slurm partition)
executor.update_parameters(timeout_min=4, slurm_partition="dev")

# Submit the job
job = executor.submit(add, 5, 7)
print(job.job_id)

# Get the result
output = job.result()
assert output == 12