使用 PyTorch、Ray 和 Google Kubernetes Engine (GKE) 训练模型

您可以在稳健、可伸缩的环境中高效地训练和部署机器学习 (ML) 模型。通过在 Google Kubernetes Engine (GKE) 上组合使用Ray、PyTorch 和 Ray Operator 插件，可以实现此目的。GKE 提供了要求严苛的 AI/机器学习工作负载所需的控制和性能。这些工作负载包括训练大型基础模型和大规模提供推理请求。在本教程中，您将使用 GKE 的自动扩缩和 Ray 的分布式计算来简化机器学习工作流。

Ray 简介

Ray 是一个适用于 AI/机器学习应用的开源可伸缩计算框架。Ray Train 是 Ray 中的一个组件，专为分布式模型训练和微调而设计。您可以使用 Ray Train API 在多台机器上规模化训练，并与 PyTorch 等机器学习库集成。

您可以使用 RayCluster 或 RayJob 资源部署 Ray 训练作业。在生产环境中部署 Ray 作业时，您应使用 RayJob 资源，原因如下：

  • RayJob 资源创建一个临时 Ray 集群，该集群可以在作业完成后自动删除。
  • RayJob 资源支持重试政策，以弹性执行作业。
  • 您可以使用熟悉的 Kubernetes API 模式管理 Ray 作业。

目标

本指南适用于对通过 Kubernetes 容器编排功能来使用 Ray 应用模型感兴趣的生成式 AI 客户、GKE 的新用户或现有用户、机器学习工程师、MLOps (DevOps) 工程师或平台管理员。

  • 创建 GKE 集群。
  • 使用 RayCluster 自定义资源创建 Ray 集群。
  • 使用 Ray 作业训练模型。
  • 使用 RayJob 自定义资源部署 Ray 作业。

费用

在本文档中，您将使用的以下收费组件： Google Cloud

您可使用 价格计算器 根据您的预计使用情况来估算费用。

新 Google Cloud 用户可能有资格申请免费试用。

完成本文档中描述的任务后，您可以通过删除所创建的资源来避免继续计费。如需了解详情，请参阅清理。

准备工作

Cloud Shell 中预安装了本教程所需的软件，包括 kubectl 和 gcloud CLI。如果您不使用 Cloud Shell，则必须安装 gcloud CLI。

  1. 登录您的 Google Cloud 账号。如果您是新手 Google Cloud， 请创建一个账号来评估我们的产品在 实际场景中的表现。新客户还可获享 $300 赠金，用于 运行、测试和部署工作负载。

  2. 安装 Google Cloud CLI。

  3. 如果您使用的是外部身份提供方 (IdP)，则必须先使用联合身份登录 gcloud CLI。

  4. 如需初始化 gcloud CLI，请运行以下命令： 

    gcloud init

  5. 创建或选择 Google Cloud 项目。

    选择或创建项目所需角色

    • 选择项目：选择项目不需要特定的 IAM 角色，您可以选择已获授予角色的任何项目。
    • 创建项目：如需创建项目，您需要拥有 Project Creator 角色 (roles/resourcemanager.projectCreator)，该角色包含 resourcemanager.projects.create 权限。了解如何授予 角色。

    • 创建 Google Cloud 项目：

      gcloud projects create PROJECT_ID

      将 PROJECT_ID 替换为您要创建的 Google Cloud 项目名称。

    • 选择您创建的 Google Cloud 项目：

      gcloud config set project PROJECT_ID

      将 PROJECT_ID 替换为您的 Google Cloud 项目名称。

  6. 验证是否已为您的 Google Cloud 项目启用结算功能。

  7. 启用 GKE API：

    启用 API 所需的角色

    如需启用 API，您需要拥有 Service Usage Admin IAM 角色 (roles/serviceusage.serviceUsageAdmin)，该角色包含 serviceusage.services.enable 权限。了解如何授予 角色。 

    gcloud services enable container.googleapis.com

  8. 安装 Google Cloud CLI。

  9. 如果您使用的是外部身份提供方 (IdP)，则必须先使用联合身份登录 gcloud CLI。

  10. 如需初始化 gcloud CLI，请运行以下命令： 

    gcloud init

  11. 创建或选择 Google Cloud 项目。

    选择或创建项目所需角色

    • 选择项目：选择项目不需要特定的 IAM 角色，您可以选择已获授予角色的任何项目。
    • 创建项目：如需创建项目，您需要拥有 Project Creator 角色 (roles/resourcemanager.projectCreator)，该角色包含 resourcemanager.projects.create 权限。了解如何授予 角色。

    • 创建 Google Cloud 项目：

      gcloud projects create PROJECT_ID

      将 PROJECT_ID 替换为您要创建的 Google Cloud 项目名称。

    • 选择您创建的 Google Cloud 项目：

      gcloud config set project PROJECT_ID

      将 PROJECT_ID 替换为您的 Google Cloud 项目名称。

  12. 验证是否已为您的 Google Cloud 项目启用结算功能。

  13. 启用 GKE API：

    启用 API 所需的角色

    如需启用 API，您需要拥有 Service Usage Admin IAM 角色 (roles/serviceusage.serviceUsageAdmin)，该角色包含 serviceusage.services.enable 权限。了解如何授予 角色。 

    gcloud services enable container.googleapis.com
    14.

  14. 为您的用户账号授予角色。对以下每个 IAM 角色运行以下命令一次：roles/container.clusterAdmin, roles/container.admin

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

    替换以下内容：

    • PROJECT_ID：您的项目 ID。
    • USER_IDENTIFIER：您的用户账号的标识符。 例如，myemail@example.com。
    • ROLE：您授予用户账号的 IAM 角色。

准备环境

如需准备环境，请按照以下步骤操作：

  1. 点击Google Cloud 控制台中的 Cloud Shell 激活图标 激活 Cloud Shell，从 Google Cloud 控制台启动 Cloud Shell 会话。此操作会在 Google Cloud 控制台的底部窗格中启动会话。

  2. 设置环境变量：

    export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=ray-cluster
export COMPUTE_REGION=us-central1
export COMPUTE_ZONE=us-central1-c
export CLUSTER_VERSION=CLUSTER_VERSION
export TUTORIAL_HOME=`pwd`

    替换以下内容：

    • PROJECT_ID：您的 Google Cloud项目 ID。
    • CLUSTER_VERSION：要使用的 GKE 版本。必须为 1.30.1 或更高版本。

  3. 克隆 GitHub 代码库：

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples

  4. 切换到工作目录：

    cd kubernetes-engine-samples/ai-ml/gke-ray/raytrain/pytorch-mnist

  5. 创建 Python 虚拟环境：

    python -m venv myenv && \
source myenv/bin/activate

  6. 安装 Ray。

创建 GKE 集群

创建 Autopilot 或 Standard GKE 集群：

Autopilot

创建 Autopilot 集群：

gcloud container clusters create-auto ${CLUSTER_NAME}  \
    --enable-ray-operator \
    --cluster-version=${CLUSTER_VERSION} \
    --location=${COMPUTE_REGION}

标准

创建 Standard 集群：

gcloud container clusters create ${CLUSTER_NAME} \
    --addons=RayOperator \
    --cluster-version=${CLUSTER_VERSION}  \
    --machine-type=e2-standard-8 \
    --location=${COMPUTE_ZONE} \
    --num-nodes=4

部署 RayCluster 资源

将 RayCluster 资源部署到集群：

  1. 请查看以下清单：

    apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: pytorch-mnist-cluster
spec:
  rayVersion: '2.37.0'
  headGroupSpec:
    rayStartParams:
      dashboard-host: '0.0.0.0'
    template:
      metadata:
      spec:
        containers:
        - name: ray-head
          image: rayproject/ray:2.37.0
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "2"
              ephemeral-storage: "9Gi"
              memory: "4Gi"
            requests:
              cpu: "2"
              ephemeral-storage: "9Gi"
              memory: "4Gi"
  workerGroupSpecs:
  - replicas: 4
    minReplicas: 1
    maxReplicas: 5
    groupName: worker-group
    rayStartParams: {}
    template:
      spec:
        containers:
        - name: ray-worker
          image: rayproject/ray:2.37.0
          resources:
            limits:
              cpu: "4"
              ephemeral-storage: "9Gi"
              memory: "8Gi"
            requests:
              cpu: "4"
              ephemeral-storage: "9Gi"
              memory: "8Gi"

    此清单描述了一个 RayCluster 自定义资源。

  2. 将清单应用于 GKE 集群：

    kubectl apply -f ray-cluster.yaml

  3. 验证 RayCluster 资源是否已准备就绪：

    kubectl get raycluster

    输出类似于以下内容：

    NAME                    DESIRED WORKERS   AVAILABLE WORKERS   CPUS   MEMORY   GPUS   STATUS   AGE
pytorch-mnist-cluster   2                 2                   6      20Gi     0      ready    63s

    在此输出中，STATUS 列中的 ready 表示 RayCluster 资源已准备就绪。

连接到 RayCluster 资源

连接到 RayCluster 资源以提交 Ray 作业。

  1. 验证 GKE 是否已创建 RayCluster 服务：

    kubectl get svc pytorch-mnist-cluster-head-svc

    输出类似于以下内容：

    NAME                             TYPE        CLUSTER-IP       EXTERNAL-IP   PORT(S)                                AGE
pytorch-mnist-cluster-head-svc   ClusterIP   34.118.238.247   <none>        10001/TCP,8265/TCP,6379/TCP,8080/TCP   109s

  2. 建立到 Ray 头节点的端口转发会话：

    kubectl port-forward svc/pytorch-mnist-cluster-head-svc 8265:8265 2>&1 >/dev/null &

  3. 验证 Ray 客户端是否可以使用 localhost 连接到 Ray 集群：

    ray list nodes --address http://localhost:8265

    输出类似于以下内容：

    Stats:
------------------------------
Total: 3

Table:
------------------------------
    NODE_ID                                                   NODE_IP     IS_HEAD_NODE    STATE    NODE_NAME    RESOURCES_TOTAL                 LABELS
0  1d07447d7d124db641052a3443ed882f913510dbe866719ac36667d2  10.28.1.21  False           ALIVE    10.28.1.21   CPU: 2.0                        ray.io/node_id: 1d07447d7d124db641052a3443ed882f913510dbe866719ac36667d2
# Several lines of output omitted

训练模型

使用 Fashion MNIST 数据集训练 PyTorch 模型：

  1. 提交 Ray 作业并等待作业完成：

    ray job submit --submission-id pytorch-mnist-job --working-dir . --runtime-env-json='{"pip": ["torch", "torchvision"], "excludes": ["myenv"]}' --address http://localhost:8265 -- python train.py

    输出类似于以下内容：

    Job submission server address: http://localhost:8265

--------------------------------------------
Job 'pytorch-mnist-job' submitted successfully
--------------------------------------------

Next steps
  Query the logs of the job:
    ray job logs pytorch-mnist-job
  Query the status of the job:
    ray job status pytorch-mnist-job
  Request the job to be stopped:
    ray job stop pytorch-mnist-job

Handling connection for 8265
Tailing logs until the job exits (disable with --no-wait):
...
...

  2. 验证作业状态：

    ray job status pytorch-mnist

    输出类似于以下内容：

    Job submission server address: http://localhost:8265
Status for job 'pytorch-mnist-job': RUNNING
Status message: Job is currently running.

    等待 Status for job 变为 COMPLETE。此过程可能需要 15 分钟或更长时间。

  3. 查看 Ray 作业日志：

    ray job logs pytorch-mnist

    输出类似于以下内容：

    Training started with configuration:
╭─────────────────────────────────────────────────╮
│ Training config                                  │
├──────────────────────────────────────────────────┤
│ train_loop_config/batch_size_per_worker       8  │
│ train_loop_config/epochs                     10  │
│ train_loop_config/lr                      0.001  │
╰─────────────────────────────────────────────────╯

# Several lines omitted

Training finished iteration 10 at 2024-06-19 08:29:36. Total running time: 9min 18s
╭───────────────────────────────╮
│ Training result                │
├────────────────────────────────┤
│ checkpoint_dir_name            │
│ time_this_iter_s      25.7394  │
│ time_total_s          351.233  │
│ training_iteration         10  │
│ accuracy               0.8656  │
│ loss                  0.37827  │
╰───────────────────────────────╯

# Several lines omitted
-------------------------------
Job 'pytorch-mnist' succeeded
-------------------------------

部署 RayJob

RayJob 自定义资源会在执行单个 Ray 作业期间管理 RayCluster 资源的生命周期。

  1. 请查看以下清单：

    apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: pytorch-mnist-job
spec:
  shutdownAfterJobFinishes: true
  entrypoint: python ai-ml/gke-ray/raytrain/pytorch-mnist/train.py
  runtimeEnvYAML: |
    pip:
      - torch
      - torchvision
    working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/main.zip"
    env_vars:
      NUM_WORKERS: "4"
      CPUS_PER_WORKER: "2"
  rayClusterSpec:
    rayVersion: '2.37.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
            - name: ray-head
              image: rayproject/ray:2.37.0
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "2"
                  ephemeral-storage: "9Gi"
                  memory: "4Gi"
                requests:
                  cpu: "2"
                  ephemeral-storage: "9Gi"
                  memory: "4Gi"
    workerGroupSpecs:
      - replicas: 4
        minReplicas: 1
        maxReplicas: 5
        groupName: small-group
        rayStartParams: {}
        template:
          spec:
            containers:
              - name: ray-worker
                image: rayproject/ray:2.37.0
                resources:
                  limits:
                    cpu: "4"
                    ephemeral-storage: "9Gi"
                    memory: "8Gi"
                  requests:
                    cpu: "4"
                    ephemeral-storage: "9Gi"
                    memory: "8Gi"

    此清单描述了一个 RayJob 自定义资源。

  2. 将清单应用于 GKE 集群：

    kubectl apply -f ray-job.yaml

  3. 验证 RayJob 资源是否正在运行：

    kubectl get rayjob

    输出类似于以下内容：

    NAME                JOB STATUS   DEPLOYMENT STATUS   START TIME             END TIME   AGE
pytorch-mnist-job   RUNNING      Running             2024-06-19T15:43:32Z              2m29s

    在此输出中，DEPLOYMENT STATUS 列表示 RayJob 资源处于 Running 状态。

  4. 查看 RayJob 资源状态：

    kubectl logs -f -l job-name=pytorch-mnist-job

    输出类似于以下内容：

    Training started with configuration:
╭─────────────────────────────────────────────────╮
│ Training config                                  │
├──────────────────────────────────────────────────┤
│ train_loop_config/batch_size_per_worker       8  │
│ train_loop_config/epochs                     10  │
│ train_loop_config/lr                      0.001  │
╰─────────────────────────────────────────────────╯

# Several lines omitted

Training finished iteration 10 at 2024-06-19 08:29:36. Total running time: 9min 18s
╭───────────────────────────────╮
│ Training result                │
├────────────────────────────────┤
│ checkpoint_dir_name            │
│ time_this_iter_s      25.7394  │
│ time_total_s          351.233  │
│ training_iteration         10  │
│ accuracy               0.8656  │
│ loss                  0.37827  │
╰───────────────────────────────╯

# Several lines omitted
-------------------------------
Job 'pytorch-mnist' succeeded
-------------------------------

  5. 验证 Ray 作业是否已完成：

    kubectl get rayjob

    输出类似于以下内容：

    NAME                JOB STATUS   DEPLOYMENT STATUS   START TIME             END TIME               AGE
pytorch-mnist-job   SUCCEEDED    Complete            2024-06-19T15:43:32Z   2024-06-19T15:51:12Z   9m6s

    在此输出中，DEPLOYMENT STATUS 列表示 RayJob 资源处于 Complete 状态。

观察 Ray 工作负载

如需查看 RayJob 的详细信息，您可以在 Google Cloud 控制台中前往 Kubernetes Engine > AI/机器学习 > 作业部分。

在 Google Cloud 控制台中查看 RayJob

清理

删除项目

    删除项目： Google Cloud 

    gcloud projects delete PROJECT_ID

删除各个资源

如果您使用的是现有项目，并且不想将其删除，则可逐个删除不再需要的资源。 如需删除集群，请输入以下命令：

gcloud container clusters delete ${CLUSTER_NAME}

后续步骤

  • 探索有关 Google Cloud 的参考架构、图表和最佳做法。查看我们的 Cloud 架构中心。