使用 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 工作。

目標

本指南適用於生成式 AI 客戶、GKE 新手或現有使用者、機器學習工程師、MLOps (DevOps) 工程師，或是有意使用 Kubernetes 容器自動化調度管理功能，透過 Ray 提供模型的平台管理員。

  • 建立 GKE 叢集。
  • 使用 RayCluster 自訂資源建立 Ray 叢集。
  • 使用 Ray 工作訓練模型。
  • 使用 RayJob 自訂資源部署 Ray 工作。

費用

在本文件中，您會使用下列 Google Cloud的計費元件：

如要根據預測用量估算費用，請使用 Pricing Calculator。

初次使用 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 角色，只要您在專案中獲派角色，即可選取該專案。
    • 建立專案：如要建立專案，您需要「專案建立者」角色 (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，您需要具備服務使用情形管理員 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 角色，只要您在專案中獲派角色，即可選取該專案。
    • 建立專案：如要建立專案，您需要「專案建立者」角色 (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，您需要具備服務使用情形管理員 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：使用者帳戶的 ID。 myemail@example.com。
    • ROLE：授予使用者帳戶的 IAM 角色。

準備環境

如要準備環境，請按照下列步驟操作：

  1. 在 Google Cloud 控制台中，按一下Cloud Shell 啟用圖示Google Cloud 控制台中的「啟用 Cloud Shell」，啟動 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}

標準

建立標準叢集：

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/ML」>「Jobs」部分。

在 Google Cloud 控制台中查看 RayJobs

清除所用資源

刪除專案

    刪除 Google Cloud 專案：

    gcloud projects delete PROJECT_ID

刪除個別資源

如果您使用現有專案，但不想刪除專案，可以刪除個別資源。如要刪除叢集，請輸入：

gcloud container clusters delete ${CLUSTER_NAME}

後續步驟