python - 在 Slurm 集群上运行 TensorFlow？

Question

我可以访问一个计算集群，特别是一个有两个 12 核 CPU 的节点，它运行 Slurm Workload Manager .

我想运行 TensorFlow在该系统上，但不幸的是，我无法找到有关如何执行此操作或什至是否可能的任何信息。我对此很陌生，但据我了解，我必须通过创建 Slurm 作业来运行 TensorFlow，并且不能通过 ssh 直接执行 python/tensorflow。

有人有关于这个主题的想法、教程或任何类型的资源吗？

Answer 1

相对简单。

在您请求每个主机一个进程的简化假设下，slurm 将为您提供环境变量中所需的所有信息，特别是 SLURM_PROCID、SLURM_NPROCS 和 SLURM_NODELIST。

例如，您可以按如下方式初始化您的任务索引、任务数和节点列表:

from hostlist import expand_hostlist
task_index  = int( os.environ['SLURM_PROCID'] )
n_tasks     = int( os.environ['SLURM_NPROCS'] )
tf_hostlist = [ ("%s:22222" % host) for host in
                expand_hostlist( os.environ['SLURM_NODELIST']) ]  

请注意，slurm 会以压缩格式为您提供主机列表(例如，“myhost[11-99]”)，您需要对其进行扩展。我通过模块主机列表来做到这一点 Kent Engström，请点击此处 https://pypi.python.org/pypi/python-hostlist

到那时，您可以直接使用现有信息创建您的 TensorFlow 集群规范和服务器，例如:

cluster = tf.train.ClusterSpec( {"your_taskname" : tf_hostlist } )
server  = tf.train.Server( cluster.as_cluster_def(),
                           job_name   = "your_taskname",
                           task_index = task_index )

你已经准备好了!您现在可以使用通常的语法在分配的特定主机上执行 TensorFlow 节点放置:

for idx in range(n_tasks):
   with tf.device("/job:your_taskname/task:%d" % idx ):
       ...

上面报告的代码的一个缺陷是你所有的作业都会指示 Tensorflow 安装服务器监听固定端口 22222。如果多个这样的作业恰好被安排到同一个节点，第二个将无法监听 22222。

更好的解决方案是让 slurm 为每个作业预留端口。您需要让您的 slurm 管理员参与并要求他配置 slurm，以便它允许您使用 --resv-ports 选项请求端口。实际上，这需要要求他们在他们的 slurm.conf 中添加如下一行:

MpiParams=ports=15000-19999

在你打扰你的 slurm 管理员之前，检查已经配置了哪些选项，例如:

scontrol show config | grep MpiParams

如果您的站点已经在使用旧版本的 OpenMPI，则可能已经存在这样的选项。

然后，修改我的第一段代码如下:

from hostlist import expand_hostlist
task_index  = int( os.environ['SLURM_PROCID'] )
n_tasks     = int( os.environ['SLURM_NPROCS'] )
port        = int( os.environ['SLURM_STEP_RESV_PORTS'].split('-')[0] )
tf_hostlist = [ ("%s:%s" % (host,port)) for host in
                expand_hostlist( os.environ['SLURM_NODELIST']) ]  

祝你好运!