使用 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 库所需的依赖项（例如 TensorFlow 和 PyTorch）。要添加自定义依赖项，请使用以下一种或两种方法：

• 使用官方 Ray Docker 镜像 作为基础来构建 Docker 镜像。

• 使用 Ray Runtime 环境。

配置 Ray Pod 以使用 GPU

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

这是最多 5 个 GPU 工作程序的 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 工作程序设置自动缩放，如上面的 `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 的一个已知问题。一种变通方法是将 `NVIDIA_VISIBLE_DEVICES` 环境变量设置为 `void`，以用于不需要 GPU 设备的 Pod。

一些有用的链接

• NVIDIA/k8s-device-plugin#61

• NVIDIA/k8s-device-plugin#87

• [NVIDIA] Prevent unprivileged access to GPUs in Kubernetes

• 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 工作程序组的 GPU 容量。假设我们有一个 RayCluster 配置如上面的配置片段所示

• 有一个 Ray Pod 工作程序组，每个 Pod 具有 1 个 GPU 容量单元。

• Ray 集群目前没有该组中的任何工作程序。

• 该组的 `maxReplicas` 至少为 2。

然后，以下 Ray 程序将触发 2 个 GPU 工作程序的向上扩展。

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 工作程序 Pod 将在空闲超时（默认为 60 秒）后向下缩放。如果 GPU 工作程序 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 工作程序可以向下缩放。

覆盖 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 Taints 和 Tolerations

注意

托管的 Kubernetes 服务通常会为您处理与 GPU 相关的 Taints 和 Tolerations。如果您使用的是托管的 Kubernetes 服务，您可能不需要担心本节。

Kubernetes 的 NVIDIA GPU 插件 会将 Taints 应用于 GPU 节点；这些 Taints 会阻止非 GPU Pod 调度到 GPU 节点上。GKE、EKS 和 AKS 等托管 Kubernetes 服务会自动将匹配的 Tolerations 应用于请求 GPU 资源的 Pod。Tolerations 通过 Kubernetes 的 ExtendedResourceToleration Admission Controller 应用。如果您的 Kubernetes 集群未启用此 Admission Controller，您可能需要手动将 GPU Toleration 添加到每个 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 设备插件的信息，在此处阅读有关 Kubernetes GPU 插件的信息，以及在此处阅读有关 NVIDIA 的 Kubernetes GPU 插件的信息。