submitit

import submitit

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

# executor is the submission interface
executor = submitit.AutoExecutor(folder="log_test")
# set timeout in min, and partition for running the job
executor.update_parameters(
    timeout_min=1,
    slurm_partition="dev",
)
job = executor.submit(add, 5, 7)  # will compute add(5, 7)
print(job.job_id)  # ID of your job

output = job.result()  # waits for completion and returns output
assert output == 12  # 5 + 7 = 12... your addition was computed in the cluster

import bz2
import cloudpickle

job_name = "my_cool_job"

def expensive_calc(a):
    return a + 1

# executor is the submission interface (logs are dumped in the folder)
executor = submitit.AutoExecutor(folder="log_test")
# set timeout in min, and partition for running the job
executor.update_parameters(
    timeout_min=1,
    slurm_partition="dev", #misc slurm args
    tasks_per_node=4,  # number of cores
    mem=8,  # memory in GB I think
)

#submit 10 jobs:
jobs = [executor.submit(expensive_calc, num) for num in range(10)]

# but actually maybe we want to come back to those jobs later, so let’s save them to disk
with bz2.open(job_name + ".job.pkl.bz2", "wb") as f:
    cloudpickle.dump(jobs, f)

# We can quit this python session and do something else now
# Resume after quitting
with bz2.open(job_name + ".job.pkl.bz2", "rb") as f:
    jobs = cloudpickle.load(f)

# wait for all jobs to finish
[job.wait() for job in jobs]

res_list = []
# # Alternatively append these results to a previous run
# with bz2.open(job_name + ".pkl.bz2", "rb") as f:
#     res_list.extend(pickle.load(f))

# Examine job outputs for failures etc
fail_ids = [job.job_id for job in jobs if job.state not in ('DONE', 'COMPLETED')]
res_list.extend([job.result() for job in jobs if job.job_id not in fail_ids])
failures = [job for job in jobs if job.job_id in fail_ids]

if failures:
    print("failures")
    print("===")
    for job in failures:
        print(job.state, job.stderr())

# Save results to disk
with bz2.open(job_name + ".result.pkl.bz2", "wb") as f:
    cloudpickle.dump(res_list, f)