如何在 Google Colab 上运行 mpi4py 代码？

Question

我有一段使用 mpi4py 库分发的任务编写的代码。本质上，代码只是分配循环并在多个核心中运行它们，无法矢量化，它看起来像这样：

from mpi4py import MPI
comm = MPI.COMM_WORLD
numtasks = comm.Get_size()
taskid = comm.Get_rank()
MASTER = 0
TAG1 = 1
TAG2 = 2
TAG3 = 3
TAG4 = 4

通常，当我必须在自己的本地计算机上运行时，我只需在终端中输入类似的内容

mpiexec -np 4 python my-program-name.py

如果我要在所有 4 个核心上运行。我的问题是，我如何在 Google VM（例如 Google Colab）上做同样的事情？我正在寻找一个外行指令，用于基本上在谷歌虚拟机上使用 mpi4py 进行分布式编程。

我在 stackoverflow 上能找到的唯一东西是这个 使用 Python (mpi4py) 在 Google Cloud Engine 上进行分布式编程，但它现在确实很有帮助。我该如何做所有这些事情，例如“通过 Google 容器引擎设置集群节点”和“从我的集群 ssh 连接到一台虚拟机并运行代码”？

Answer 1

在任何 Jupyter 代码窗口中，您只需键入“！”接下来是命令，如下所示：

! pip install mpi4py

它将被解释为像在终端中一样运行。所以你只需添加“！”在命令之前您需要运行：

! mpiexec -np 4 python my-program-name.py

这样您就可以做任何您喜欢的事情：复制/创建文件、克隆存储库...

另一种选择是使用 python 子进程库。 作为子进程运行此命令，如下所示：

import subprocess
subprocess.call(['mpiexec', '-np', '4', 'python', 'my-program-name.py'])

Answer 2

使用 sudo 运行命令 例如

sudo mpirun --allow-run-as-root --oversubscribe -np <num> <executable>