Submitit: 轻量级的Python作业提交工具

submitit

Submitit:轻量级的Python作业提交工具

Submitit是一个轻量级的Python工具,用于向Slurm集群提交Python函数进行计算。它为研究人员和开发人员提供了一种简单而强大的方式,可以轻松地在高性能计算环境中执行Python代码。

什么是Submitit?

Submitit是Facebook AI Research (FAIR)开发的开源项目,旨在简化在Slurm集群上提交和管理计算作业的过程。它的核心功能包括:

将Python函数包装并提交到Slurm集群执行
提供访问作业结果、日志等信息的接口
支持在Slurm和本地环境之间无缝切换执行

Slurm(Simple Linux Utility for Resource Management)是一个开源的、容错的、高度可扩展的集群管理和作业调度系统,广泛应用于大型和小型Linux集群。Submitit充分利用了Slurm的功能,同时提供了更加友好的Python接口。

Submitit的主要特性

简单易用:只需几行代码即可提交Python函数到Slurm集群执行。
灵活性:支持提交任何Python函数,包括lambda函数和脚本中定义的函数。
错误处理:如果作业失败,会抛出带有堆栈跟踪的错误,便于调试。
重新排队:对于Slurm,支持重新排队被抢占的作业。
兼容性:可以方便地在Submitit执行器和concurrent.futures执行器之间切换。
检查点:支持对有状态的可调用对象进行检查点,在被抢占或超时时可以从当前状态重新排队(高级功能)。
多节点/多任务支持:易于访问多节点/多任务作业的本地/全局排名。
可移植性:同一份代码可以在不同的集群上工作,这得益于其插件系统。

安装Submitit

Submitit支持Python 3.8及以上版本。安装过程非常简单,可以通过pip或conda进行:

# 使用pip安装
pip install submitit

# 使用conda安装
conda install -c conda-forge submitit

对于希望使用最新开发版本的用户,可以直接从GitHub安装:

pip install git+https://github.com/facebookincubator/submitit@main#egg=submitit

使用示例

下面是一个简单的示例,展示了如何使用Submitit提交一个加法函数到Slurm集群执行:

import submitit

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

# 创建执行器(日志会保存在指定文件夹)
executor = submitit.AutoExecutor(folder="log_test")

# 设置超时时间(分钟)和运行分区
executor.update_parameters(timeout_min=1, slurm_partition="dev")

# 提交作业
job = executor.submit(add, 5, 7)  # 将计算add(5, 7)

print(job.job_id)  # 打印作业ID

output = job.result()  # 等待完成并返回结果
assert output == 12  # 5 + 7 = 12, 加法在集群中完成

这个例子展示了Submitit的核心功能:

定义一个简单的Python函数
创建Submitit执行器
设置作业参数
提交作业
获取结果

Submitit还提供了更多高级功能,如读取日志文件(job.stdout() 和 job.stderr())、处理命令行程序等。

Submitit vs dask.distributed

虽然Submitit和dask.distributed都提供了分布式计算的能力,但它们的工作方式有所不同:

Submitit直接在集群上创建独立的作业,而dask.distributed将作业分配给一组工作进程。
Submitit提供了更底层的接口,用户可以直接控制每个作业,包括独立的stdout和stderr,以及在抢占和超时情况下的检查点功能。
使用Submitit时应避免提交多个小任务,因为这可能会创建许多独立的作业并可能使集群过载。而通过dask.distributed可以轻松处理这种情况。

结语

Submitit为Python开发者提供了一种简单而强大的方式来利用Slurm集群的计算能力。它特别适合那些需要在高性能计算环境中运行Python代码的研究人员和数据科学家。通过简化作业提交和管理过程,Submitit让用户能够专注于他们的核心工作,而不是被复杂的集群操作所困扰。

随着数据规模的增长和计算需求的增加,像Submitit这样的工具将在科学计算和机器学习领域扮演越来越重要的角色。无论您是在进行大规模数据分析、训练复杂的机器学习模型,还是运行密集的科学模拟,Submitit都能为您提供所需的灵活性和效率。

如果您正在寻找一种简单而有效的方法来管理和执行大规模Python计算任务,Submitit无疑是一个值得考虑的选择。它不仅简化了与Slurm集群的交互,还提供了灵活性和可扩展性,以满足各种计算需求。无论您是经验丰富的HPC用户还是刚刚开始探索分布式计算的新手,Submitit都能为您的工作流程带来显著的改进。