在 Slurm 上部署#

使用 Ray 的 Slurm 可能有点不直观。

  • SLURM 要求将同一个程序的多个副本多次提交到同一个集群以进行集群编程。这对于基于 MPI 的工作负载特别适用。

  • 另一方面，Ray 期望采用一个带有单一入口点的头-工作节点架构。也就是说，你需要启动一个 Ray 头节点，多个 Ray 工作节点，并在头节点上运行你的 Ray 脚本。

警告

SLURM 支持仍在进行中。SLURM 用户应注意当前关于网络的限制。更多解释请参见 这里。

SLURM 支持由社区维护。维护者 GitHub 用户名：tupui。

本文档旨在阐明如何在SLURM上运行Ray。

使用 Ray 与 SLURM 的演练#

许多 SLURM 部署要求你通过 sbatch 与 slurm 交互，sbatch 在 SLURM 上执行一个批处理脚本。

要使用 sbatch 运行 Ray 作业，您需要在 sbatch 作业中使用多个 srun 命令（任务）启动 Ray 集群，然后执行使用 Ray 的 Python 脚本。每个任务将在单独的节点上运行，并启动/连接到 Ray 运行时。

以下演练将执行以下操作：

  1. 为 sbatch 脚本设置正确的头部。

  2. 加载适当的环境/模块。

  3. 获取可用计算节点的列表及其IP地址。

  4. 在一个节点（称为头节点）中启动一个头射线进程。

  5. 在 (n-1) 个工作节点上启动 Ray 进程，并通过提供头节点地址将它们连接到头节点。

  6. 在底层 ray 集群准备就绪后，提交用户指定的任务。

参见 slurm-basic.sh 以获取端到端示例。

sbatch 指令#

在你的 sbatch 脚本中，你会想要添加 指令以提供上下文 给 SLURM 以供你的作业使用。

#!/bin/bash
#SBATCH --job-name=my-workload

你需要告诉 SLURM 专门为 Ray 分配节点。然后，Ray 将找到并管理每个节点上的所有资源。

### Modify this according to your Ray workload.
#SBATCH --nodes=4
#SBATCH --exclusive

重要：为了确保每个 Ray 工作运行时将在单独的节点上运行，请设置 tasks-per-node。

#SBATCH --tasks-per-node=1

由于我们设置了 tasks-per-node = 1，这将用于确保每个 Ray 工作运行时将获得适当的资源。在这个例子中，我们请求每个节点至少 5 个 CPU 和 5 GB 的内存。

### Modify this according to your Ray workload.
#SBATCH --cpus-per-task=5
#SBATCH --mem-per-cpu=1GB
### Similarly, you can also specify the number of GPUs per node.
### Modify this according to your Ray workload. Sometimes this
### should be 'gres' instead.
#SBATCH --gpus-per-task=1

你也可以在你的 sbatch 指令中添加其他可选标志。

加载您的环境#

首先，您通常希望在脚本的开头加载模块或您自己的conda环境。

请注意，这是一个可选步骤，但它通常是启用正确依赖项集所必需的。

# Example: module load pytorch/v1.4.0-gpu
# Example: conda activate my-env

conda activate my-env

获取头部的IP地址#

接下来，我们需要为头节点获取一个主机名和节点IP地址。这样，当我们启动工作节点时，我们就能正确连接到正确的头节点。

# Getting the node names
nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
nodes_array=($nodes)

head_node=${nodes_array[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

# if we detect a space character in the head node IP, we'll
# convert it to an ipv4 address. This step is optional.
if [[ "$head_node_ip" == *" "* ]]; then
IFS=' ' read -ra ADDR <<<"$head_node_ip"
if [[ ${#ADDR[0]} -gt 16 ]]; then
  head_node_ip=${ADDR[1]}
else
  head_node_ip=${ADDR[0]}
fi
echo "IPV6 address detected. We split the IPV4 address as $head_node_ip"
fi

启动 Ray 头节点#

在检测到头节点主机名和头节点IP后，我们将希望创建一个Ray头节点运行时。我们将通过使用``srun``作为单任务/节点的后台任务来实现这一点（回想一下``tasks-per-node=1``）。

下面，您会看到我们明确指定了CPU数量（num-cpus）和GPU数量（num-gpus）给Ray，因为这将防止Ray使用超过分配的资源。我们还需要明确指示Ray头节点的``node-ip-address``：

port=6379
ip_head=$head_node_ip:$port
export ip_head
echo "IP Head: $ip_head"

echo "Starting HEAD at $head_node"
srun --nodes=1 --ntasks=1 -w "$head_node" \
    ray start --head --node-ip-address="$head_node_ip" --port=$port \
    --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_TASK}" --block &

通过将上述 srun 任务置于后台，我们可以继续启动 Ray 工作器运行时。

启动 Ray 工作节点#

下面，我们对每个工作节点做同样的事情。确保Ray头节点和Ray工作节点的进程不在同一个节点上启动。

# optional, though may be useful in certain versions of Ray < 1.0.
sleep 10

# number of nodes other than the head node
worker_num=$((SLURM_JOB_NUM_NODES - 1))

for ((i = 1; i <= worker_num; i++)); do
    node_i=${nodes_array[$i]}
    echo "Starting WORKER $i at $node_i"
    srun --nodes=1 --ntasks=1 -w "$node_i" \
        ray start --address "$ip_head" \
        --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_TASK}" --block &
    sleep 5
done

提交你的脚本#

最后，你可以调用你的 Python 脚本：

# ray/doc/source/cluster/doc_code/simple-trainer.py
python -u simple-trainer.py "$SLURM_CPUS_PER_TASK"

备注

-u 参数告诉 Python 以非缓冲方式打印到标准输出，这对于 slurm 处理重定向输出的方式非常重要。如果不包含此参数，您可能会遇到奇怪的打印行为，例如打印的语句在程序终止之前不会被 slurm 记录。

SLURM 网络注意事项#

在使用 SLURM 和 Ray 时，有两个重要的网络方面需要牢记：

  1. 端口绑定。

  2. IP 绑定。

SLURM 集群的一个常见用途是让多个用户在同一基础设施上运行并发作业。由于头节点与其工作节点通信的方式，这很容易与 Ray 发生冲突。

考虑到2个用户，如果他们同时使用Ray调度一个SLURM作业，他们都在创建一个头节点。在后端，Ray会为一些服务分配一些内部端口。问题在于，一旦第一个头节点被创建，它将绑定一些端口，并阻止其他头节点使用这些端口。为了防止任何冲突，用户必须手动指定不重叠的端口范围。以下端口需要调整。关于端口的解释，请参见 这里:

# used for all ports
--node-manager-port
--object-manager-port
--min-worker-port
--max-worker-port
# used for the head node
--port
--ray-client-server-port
--redis-shard-ports

例如，再次以2个用户为例，他们需要将上述指令调整为：

# user 1
# same as above
...
srun --nodes=1 --ntasks=1 -w "$head_node" \
    ray start --head --node-ip-address="$head_node_ip" \
        --port=6379 \
        --node-manager-port=6700 \
        --object-manager-port=6701 \
        --ray-client-server-port=10001 \
        --redis-shard-ports=6702 \
        --min-worker-port=10002 \
        --max-worker-port=19999 \
        --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_TASK}" --block &

# user 2
# same as above
...
srun --nodes=1 --ntasks=1 -w "$head_node" \
    ray start --head --node-ip-address="$head_node_ip" \
        --port=6380 \
        --node-manager-port=6800 \
        --object-manager-port=6801 \
        --ray-client-server-port=20001 \
        --redis-shard-ports=6802 \
        --min-worker-port=20002 \
        --max-worker-port=29999 \
        --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_TASK}" --block &

关于IP绑定，在某些集群架构中，网络接口不允许在节点之间使用外部IP。相反，存在内部网络接口（eth0、eth1`等）。目前，设置内部IP很困难（参见开放的 `问题）。

Python-接口 SLURM 脚本#

[由 @pengzhenghao 贡献] 下面，我们提供了一个辅助工具 (slurm-launch.py) 来自动生成 SLURM 脚本并启动。slurm-launch.py 使用一个底层模板 (slurm-template.sh) 并根据用户输入填充占位符。

你可以随意将这两个文件复制到你的集群中使用。也欢迎随时提交PR，为改进这个脚本做出贡献！

使用示例#

如果你想在slurm中利用多节点集群：

python slurm-launch.py --exp-name test --command "python your_file.py" --num-nodes 3

如果你想指定计算节点，只需在 sinfo 命令的输出格式中使用相同的节点名称：

python slurm-launch.py --exp-name test --command "python your_file.py" --num-nodes 3 --node NODE_NAMES

在调用 python slurm-launch.py 时，你可以使用其他选项：

  • --exp-name: 实验名称。将生成 {exp-name}_{date}-{time}.sh 和 {exp-name}_{date}-{time}.log。

  • --command: 你希望运行的命令。例如：rllib train XXX 或 python XXX.py。

  • --num-gpus: 每个计算节点中希望使用的GPU数量。默认值：0。

  • --node (-w): 你希望使用的特定节点，格式与 sinfo 的输出相同。如果未指定，节点将自动分配。

  • --num-nodes (-n): 你希望使用的节点数量。默认值：1。

  • --partition (-p): 你希望使用的分区。默认值：””，将使用用户的默认分区。

  • --load-env: 设置环境的命令。例如：module load cuda/10.1。默认值：””。

请注意，slurm-template.sh 与计算节点的 IPV4 和 IPV6 IP 地址兼容。

实现#

具体来说，(slurm-launch.py) 执行以下操作：

  1. 它会自动将您的需求（例如CPU数量、每个节点的GPU数量、节点数量等）写入一个名为 {exp-name}_{date}-{time}.sh 的sbatch脚本中。您启动自己作业的命令（--command）也会被写入sbatch脚本。

  2. 然后它将通过一个新进程将sbatch脚本提交给slurm管理器。

  3. 最后，Python 进程将自行终止，并留下一个名为 {exp-name}_{date}-{time}.log 的日志文件来记录您提交的命令的进度。与此同时，Ray 集群和您的作业正在 Slurm 集群中运行。

示例和模板#

以下是一些社区贡献的使用 SLURM 与 Ray 的模板：

  • Ray sbatch 提交脚本 用于 NERSC，一个美国国家实验室。

  • YASPI (yet another slurm python interface) 由 @albanie 开发。yaspi 的目标是提供一个提交 slurm 作业的接口，从而避免 sbatch 文件的繁琐。它通过配方来实现这一点——这些是用于生成 sbatch 脚本的模板和规则的集合。支持 Ray 的作业提交。

  • 便捷的 Python 接口 用于启动 Ray 集群并通过 @pengzhenghao 提交任务