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

Question

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

现在，问题是我需要通过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 脚本以正确分配资源。 settings srun --ntasks=1不起作用，因为那时脚本只会“看到”其中一个节点。

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

Answer 1

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

However, my initial confusion and my assumption that it must be done as stated in the original question came from TensorFlow reporting out of memory errors whenever I tried to use the MultiWorkerMirrored strategy, whereas I knew that without this the model fitted perfectly within the available memory .

我在代码中的某个地方调用了tf.config.get_visible_devices("GPU") 。 现在为了让 TensorFlow 获得 GPU，它将分配它们，默认情况下通过填充完整的 GPU memory 来分配它们。 但是，由于所有脚本都是并行运行的，因此每个脚本都会尝试自己执行此操作（因为这是在策略的 scope 之外完成的），从而导致 memory (OOM) 错误。

删除这段代码后，一切运行良好。

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

如果您摆脱了 memory 错误； 确保没有任何代码分配 scope 之外的所有 GPU。 这可能是 TensorFlow 某处的一些启动，但如果这存在于所有脚本中，则会导致 OOM 错误。 一个方便的测试方法是使用tf.config.experimental.set_memory_growth ，以允许 memory 增长而不是完整的 memory 分配。

在我的代码中，我使用了 tf.distribute.cluster_resolver.SlurmClusterResolver 的get_task_info() tf.distribute.cluster_resolver.SlurmClusterResolver ，并且仅在任务号为 0 时运行分配 memory 的函数，即主要工作人员。

（以上功能和注释基于TensorFlow 2.2.0和Python 3.7.7）