在 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 作业中启动一个 Ray 集群，包含多个 srun 命令（任务），然后执行使用 Ray 的 Python 脚本。每个任务将运行在一个单独的节点上，并启动/连接到 Ray 运行时。

以下演练将执行以下操作

  1. 为 sbatch 脚本设置适当的头部。

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

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

  4. 在其中一个节点（称为头节点）中启动一个 Ray 主进程。

  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 worker 运行时都在单独的节点上运行，请设置 tasks-per-node。

#SBATCH --tasks-per-node=1

由于我们已将 tasks-per-node = 1，这将用于保证每个 Ray worker 运行时获得适当的资源。在此示例中，我们请求每个节点至少分配 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）。

您将在下面看到，我们明确指定了 Ray 使用的 CPU 数量 (num-cpus) 和 GPU 数量 (num-gpus)，这将防止 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 worker 运行时。

启动 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 以无缓冲方式打印到 stdout，这对于 slurm 处理输出重定向非常重要。如果未包含此参数，您可能会遇到奇怪的打印行为，例如直到程序终止后，slurm 才记录打印的语句。

SLURM 网络注意事项#

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

  1. 端口绑定。

  2. IP 绑定。

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

考虑两个用户，如果他们同时使用 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

例如，再次以两个用户为例，他们必须按照上面看到的说明进行调整

# 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 结合使用的模板

  • 在 NERSC（美国国家实验室）使用的Ray sbatch 提交脚本。

  • @albanie 的 YASPI（又一个 slurm python 接口）。yaspi 的目标是提供一个用于提交 slurm 作业的接口，从而避免使用 sbatch 文件的麻烦。它通过 recipes（即用于生成 sbatch 脚本的模板和规则集合）来实现这一点。支持 Ray 的作业提交。

  • @pengzhenghao 提供的用于启动 Ray 集群和提交任务的便捷 python 接口