SLURM srun 不并行运行 python 脚本，但可以访问并行资源答案

【问题标题】：SLURM srun run a python script not in parallel, but have access to the parallel resourcesSLURM srun 不并行运行 python 脚本，但可以访问并行资源
【发布时间】：2020-09-05 20:03:00
【问题描述】：

我正在尝试运行一个 python 脚本，其中部分代码将根据一些 SLURM 环境变量进行并行化。我认为确切的代码并不重要，但作为参考，我想use this 来训练我的网络。

现在，问题是我需要通过srun 运行我的脚本，但是这会产生我不想要的脚本的多个并行实例。

最基本的例子是这样的：

#!/bin/sh
#SBATCH -N 2
#SBATCH --ntasks=2

srun python myscript.py

现在我将有 2 个节点和 2 个任务，这意味着当我运行 python myscript.py 时，将有 2 个 myscript.py 实例并行运行。

然而，这不是我想要的。我希望只有一个myscript.py 实例正在运行，但它应该可以访问srun 设置的环境变量，并将其留给python 脚本以正确分配资源。设置srun --ntasks=1 不起作用，此后脚本将仅“看到”其中一个节点。

是否可以使用srun 运行脚本的单个实例，同时它仍然可以“访问” SLURM 环境变量？我查看了诸如 --exclusive 和 --preserve-env 之类的选项，但在这种情况下它们似乎对我没有帮助。

【问题讨论】：

从SlurmClusterResolver 中的代码来看，预计会有多个并行运行的脚本副本。例如：self._rank = int(os.environ['SLURM_PROCID']) 或 self._rank = int(os.environ['OMPI_COMM_WORLD_RANK'])。除非使用mpiexec 启动进程，否则您在环境中没有 Open MPI 世界排名。
见here。您应该分配多个节点，然后启动脚本的多个副本。然后，分发策略将负责在每个副本中做正确的事情。
这也是我最初的期望。然而，实际发生的是脚本的每次运行都会尝试为自己分配所有 4 个资源，并将其分配给它们。所以我最终得到了 4 个进程，每个进程都试图分配给 4 个 so 任务。显然这不起作用，因为他们试图为自己分配相同的资源。也许我应该看看 SLURM_PROCID 和 OMPI_COMM_WORLD_RANK 并确保它们设置正确。
也许您应该提取脚本的相关部分并将它们放入您的问题中？
你是对的，应该有多个进程。我的代码中的其他内容导致 TensorFlow 过度急切地分配导致错误的资源。我在下面给出了详细的答案以供将来使用。感谢您的帮助！

标签： python tensorflow mpi slurm

【解决方案1】：

事实证明，Hristo Iliev 在 cmets 中是正确的，要正确使用 SlurmClusterResolver，需要并行运行多个作业。这可能有点令人困惑，因为所有内容都会被打印多次，因为所有内容都是并行运行的，但这是正常的。

但是，我最初的困惑和我认为必须按照原始问题中所述完成的假设来自 TensorFlow 在我尝试使用 MultiWorkerMirrored 策略时报告内存不足错误，而我知道没有这个模型完全适合可用内存。

我在代码中的某个地方调用了tf.config.get_visible_devices("GPU")。现在，为了让 TensorFlow 获得 GPU，它会分配它们，并且默认情况下会填满整个 GPU 内存。但是，由于所有脚本都是并行运行的，因此每个脚本都会尝试自己执行此操作（因为此操作超出了策略范围），从而导致内存不足 (OOM) 错误。

删除这段代码后，一切正常。

对将来可能偶然发现此帖子的人的建议： - 脚本应该并行运行，你会看到多次相同的输出 - 确保一切都在strategy.scope() 下完成，即模型编译、数据生成设置（使用 tf.data） - 特别注意保存模型；只有“主要”工作人员应该将模型保存到真正的保存文件中，其他人应该写入临时文件see here

如果出现内存不足错误；确保没有某些代码将所有 GPU 分配到范围之外。这可能是 TensorFlow 某处的一些启动，但如果这存在于所有脚本中，则会导致 OOM 错误。一个方便的测试方法是使用tf.config.experimental.set_memory_growth，以允许内存增长而不是完全分配内存。

在我的代码中，我使用了tf.distribute.cluster_resolver.SlurmClusterResolver的get_task_info()函数，并且只运行在任务号为0时分配内存的函数，主worker。

（以上函数和cmets基于TensorFlow 2.2.0和Python 3.7.7）

【讨论】：