GPU で PyTorch、Ray、Google Kubernetes Engine(GKE)を使用してモデルをトレーニングする

このガイドでは、RayPyTorchRay Operator アドオンを使用して Google Kubernetes Engine(GKE)でモデルをトレーニングする方法について説明します。

Ray について

Ray は、AI / ML アプリケーション向けのオープンソースのスケーラブルなコンピューティング フレームワークです。Ray Train は、分散モデルのトレーニングと微調整用に設計された Ray のコンポーネントです。Ray Train API を使用すると、複数のマシンにわたってトレーニングをスケーリングし、PyTorch などの機械学習ライブラリと統合できます。

Ray トレーニング ジョブをデプロイするには、RayCluster リソースまたは RayJob リソースを使用します。次の理由から、本番環境で Ray ジョブをデプロイするときは、RayJob リソースを使用する必要があります。

  • RayJob リソースは、ジョブの完了時に自動的に削除できるエフェメラル Ray クラスタを作成します。
  • RayJob リソースは、復元力のあるジョブ実行の再試行ポリシーをサポートしています。
  • Ray ジョブは、使い慣れた Kubernetes API パターンを使用して管理できます。

環境を準備する

環境を準備する手順は次のとおりです。

  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 クラスタを作成する

GKE の Autopilot または Standard GKE クラスタを作成します。

Autopilot

Autopilot クラスタを作成します。

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

Standard

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 Service を作成したことを確認します。

    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. Job のステータスを確認します。

    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 jobCOMPLETE になるまで待ちます。この処理には 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 であることを示しています。