使用 GPU

本文档提供了有关在 KubeRay 中使用 GPU 的提示。

要在 Kubernetes 上使用 GPU，需要配置您的 Kubernetes 环境，并向您的 Ray 集群配置添加额外的值。

要了解在不同云上如何使用 GPU，请参阅 GKE、EKS 和 AKS 的说明。

快速入门：服务基于 GPU 的 StableDiffusion 模型

您可以在文档的示例部分找到几个 GPU 工作负载示例。其中，StableDiffusion 示例是一个不错的起点。

基于 GPU 机器学习的依赖项

Ray Docker Hub 提供了包含 Ray 和特定机器学习库的基于 CUDA 的容器镜像。例如，镜像 rayproject/ray-ml:2.6.3-gpu 非常适合使用 Ray 2.6.3 运行基于 GPU 的 ML 工作负载。Ray ML 镜像包含了这些文档中使用的 Ray Libraries 所需的依赖项（例如 TensorFlow 和 PyTorch）。要添加自定义依赖项，请使用以下一种或两种方法：

• 使用官方 Ray docker 镜像之一作为基础构建 docker 镜像。

• 使用 Ray 运行时环境。

配置 Ray Pod 以使用 GPU

使用 Nvidia GPU 需要在您的 RayCluster 的 headGroupSpec 和/或 workerGroupSpecs 的容器字段中指定 nvidia.com/gpu 资源 limits 和 requests。

这是一个最多包含 5 个 GPU worker 的 RayCluster workerGroup 配置片段。

groupName: gpu-group
replicas: 0
minReplicas: 0
maxReplicas: 5
...
template:
    spec:
     ...
     containers:
      - name: ray-node
        image: rayproject/ray-ml:2.6.3-gpu
        ...
        resources:
         nvidia.com/gpu: 1 # Optional, included just for documentation.
         cpu: 3
         memory: 50Gi
        limits:
         nvidia.com/gpu: 1 # Required to use GPU.
         cpu: 3
         memory: 50Gi
         ...

该组中的每个 Ray Pod 都可以调度到 AWS p2.xlarge 实例（1 个 GPU，4 个 vCPU，61Gi RAM）上。

提示

GPU 实例价格昂贵 – 考虑为您的 GPU Ray worker 设置自动扩缩容，如上面 minReplicas:0 和 maxReplicas:5 设置所示。要启用自动扩缩容，请记住在 RayCluster 的 spec 中设置 enableInTreeAutoscaling:True。最后，请确保您配置了 GPU Kubernetes 节点组或池以进行自动扩缩容。有关自动扩缩容节点池的详细信息，请参阅您的云提供商文档。

GPU 多租户

如果 Pod 的资源配置中不包含 nvidia.com/gpu，用户通常会期望该 Pod 不感知任何 GPU 设备，即使它被调度到 GPU 节点上。然而，当未指定 nvidia.com/gpu 时，NVIDIA_VISIBLE_DEVICES 的默认值变为 all，使得 Pod 感知节点上的所有 GPU 设备。此行为并非 KubeRay 特有，而是 NVIDIA 的已知问题。一个 workaround 是在不需要 GPU 设备的 Pod 中将 NVIDIA_VISIBLE_DEVICES 环境变量设置为 void。

一些有用的链接

• NVIDIA/k8s-device-plugin#61

• NVIDIA/k8s-device-plugin#87

• [NVIDIA] 防止在 Kubernetes 中对 GPU 的非特权访问

• ray-project/ray#29753

GPU 与 Ray

本节讨论在 Kubernetes 上运行的 Ray 应用程序的 GPU 使用。有关 Ray 中 GPU 使用的通用指南，另请参阅加速器支持。

KubeRay operator 会将容器的 GPU 资源限制通知给 Ray 调度器和 Ray 自动扩缩容器。特别是，Ray 容器的 ray start 入口点会自动配置适当的 --num-gpus 选项。

GPU 工作负载调度

部署具有 GPU 访问权限的 Ray Pod 后，它将能够执行带有 GPU 请求注解的任务和 Actor。例如，装饰器 @ray.remote(num_gpus=1) 会为需要 1 个 GPU 的任务或 Actor 添加注解。

GPU 自动扩缩容

Ray 自动扩缩容器知道每个 Ray worker 组的 GPU 容量。假设我们有一个按照上述配置片段配置的 RayCluster

• 有一个由 Ray Pod 组成的 worker 组，每个 Pod 具有 1 个 GPU 容量单位。

• Ray 集群当前没有来自该组的任何 worker。

• 该组的 maxReplicas 至少为 2。

那么下面的 Ray 程序将触发扩展 2 个 GPU worker。

import ray

ray.init()

@ray.remote(num_gpus=1)
class GPUActor:
    def say_hello(self):
        print("I live in a pod with GPU access.")

# Request actor placement.
gpu_actors = [GPUActor.remote() for _ in range(2)]
# The following command will block until two Ray pods with GPU access are scaled
# up and the actors are placed.
ray.get([actor.say_hello.remote() for actor in gpu_actors])

程序退出后，Actor 会被垃圾回收。GPU worker Pod 会在空闲超时（默认为 60 秒）后缩减。如果 GPU worker Pod 运行在 Kubernetes 节点的自动扩缩容池上，Kubernetes 节点也会随之缩减。

请求 GPU

您也可以直接向自动扩缩容器发出请求以扩展 GPU 资源。

import ray

ray.init()
ray.autoscaler.sdk.request_resources(bundles=[{"GPU": 1}] * 2)

节点扩展后，它们会一直存在，直到请求被明确覆盖。下面的程序将移除资源请求。

import ray

ray.init()
ray.autoscaler.sdk.request_resources(bundles=[])

然后 GPU worker 可以缩减。

覆盖 Ray GPU 容量 (高级)

对于特殊用例，可以覆盖向 Ray 通知的 Ray Pod GPU 容量。为此，可以在 head 或 worker 组的 rayStartParams 的 num-gpus 键设置一个值。例如：

rayStartParams:
    # Note that all rayStartParam values must be supplied as strings.
    num-gpus: "2"

Ray 调度器和自动扩缩容器随后会为组中的每个 Ray Pod 计算 2 个 GPU 容量单位，即使容器限制并未表明存在 GPU。

GPU Pod 调度 (高级)

GPU 污点和容忍度

注意

托管式 Kubernetes 服务通常会为您处理与 GPU 相关的污点和容忍度。如果您正在使用托管式 Kubernetes 服务，则可能无需担心本节内容。

Kubernetes 的 Nvidia gpu 插件会给 GPU 节点应用污点；这些污点可以阻止非 GPU Pod 被调度到 GPU 节点上。GKE、EKS 和 AKS 等托管式 Kubernetes 服务会自动为请求 GPU 资源的 Pod 应用匹配的容忍度。容忍度是通过 Kubernetes 的 ExtendedResourceToleration 准入控制器实现的。如果您的 Kubernetes 集群未启用此准入控制器，您可能需要手动将 GPU 容忍度添加到您的每个 GPU Pod 配置中。例如：

apiVersion: v1
kind: Pod
metadata:
 generateName: example-cluster-ray-worker
 spec:
 ...
 tolerations:
 - effect: NoSchedule
   key: nvidia.com/gpu
   operator: Exists
 ...
 containers:
 - name: ray-node
   image: rayproject/ray:nightly-gpu
   ...

节点选择器和节点标签

为了确保 Ray Pod 绑定到满足特定条件（例如存在 GPU 硬件）的 Kubernetes 节点，您可能希望使用 workerGroup 的 Pod 模板 spec 的 nodeSelector 字段。有关 Pod 到节点的分配的更多信息，请参阅 Kubernetes 文档。

进一步参考和讨论

在此处了解 Kubernetes 设备插件 here，在此处了解 Kubernetes GPU 插件 here，在此处了解 Nvidia 的 Kubernetes GPU 插件 here。