Pathways を使用して GKE クラスタを作成する

このページでは、プラットフォーム管理者と機械学習(ML)エンジニア向けに、Pathways ベースのワークロード用に構成された Google Kubernetes Engine(GKE)クラスタを作成する方法について説明します。Cluster Toolkit を使用して、事前構成されたクラスタを自動的にプロビジョニングすることも、Google Cloud CLI を使用してクラスタを手動で作成することもできます。

始める前に

先に進む前に、次の前提条件を満たしてください。

ローカル環境を設定する

Google Cloud 認証情報を使用して gcloud CLI にログインします。

gcloud auth login

Cluster Toolkit を使用する場合は、Terraform が使用するアプリケーションのデフォルト認証情報(ADC)も生成します。

gcloud auth application-default login

GKE クラスタを作成する

次の例では、TPU v6e 2x4 ノードプールを含むクラスタを作成します。クラスタを作成するには、Cluster Toolkit または gcloud CLI を使用します。

Cluster Toolkit

  1. GKE TPU v6e クラスタをデプロイするの手順に沿って、GKE TPU クラスタをデプロイします。すべての TPU デプロイ オプションの概要については、Cloud TPU のデプロイの概要をご覧ください。

    デプロイ ブループリントを構成する際は、次の操作を行います。

    • enable_pathways_for_tpus を true に設定します(Cluster Toolkit TPU ブループリントではデフォルトで有効になっています)。この設定により、Pathways コントローラに必要な cpu-np CPU ノードプール(n4-standard-64)と Kueue 構成が自動的にプロビジョニングされます。Cluster Toolkit は、kueue_configuration_path 変数でカスタム Kueue 構成ファイルを指定しない場合にのみ、デフォルトの Pathways Kueue 構成を適用します。
    • num_slices は、ワークロードに必要な TPU スライス ノードプールの数に設定します。
    • tpu_topology を、選択した TPU スライス トポロジ(2x4 や 4x4 など)に設定します。
    • 使用オプションを選択します。デフォルトでは、TPU v6e デプロイ ファイルは特定の TPU 予約を使用するため、reservation_affinity 変数に予約名を指定します。別の使用オプションを使用するには、GKE の使用オプションをご覧ください。
  2. クラスタノードの STATUS が Ready であることを確認します。

    kubectl get nodes
    

    クラスタを作成したら、--pathways フラグと --pathways-gcs-location フラグを指定して gcluster job submit コマンドを使用し、Pathways ワークロードを送信できます。詳細については、Pathways でワークロードを実行するをご覧ください。

gcloud

GKE クラスタを作成し、Pathways ワークロードを実行するように設定する手順は次のとおりです。

  1. 次の環境変数を設定します。

    CLUSTER=GKE_CLUSTER_NAME
    PROJECT=PROJECT_ID
    ZONE=ZONE
    REGION=REGION
    CLUSTER_VERSION=GKE_CLUSTER_VERSION
    PW_CPU_MACHINE_TYPE="n2-standard-64"
    NETWORK=NETWORK
    SUBNETWORK=SUB_NETWORK
    CLUSTER_NODEPOOL_COUNT=3
    TPU_MACHINE_TYPE="ct6e-standard-4t"
    WORKERS_PER_SLICE=2
    TOPOLOGY="2x4"
    NUM_CPU_NODES=1

    次のように置き換えます。

    • GKE_CLUSTER_NAME: GKE クラスタの名前。
    • PROJECT_ID: 実際の Google Cloud プロジェクト ID。
    • ZONE: クラスタを作成するゾーン。TPU v6e が利用可能なゾーンを選択します。詳細については、Cloud TPU のリージョンとゾーンをご覧ください。
    • REGION: 省略可。リージョン クラスタを使用する場合にリソースを作成するリージョン。
    • GKE_CLUSTER_VERSION: 省略可: GKE クラスタのバージョン。1.32.2-gke.1475000 以降を使用し、標準サポートのバージョンを選択します。詳細については、GKE リリース スケジュールをご覧ください。バージョンを指定しない場合は、次の gcloud container clusters create コマンドから --cluster-version フラグを削除します。
    • NETWORK: 省略可。Virtual Private Cloud(VPC)ネットワークの名前。クラスタを作成する前に、ネットワークが存在している必要があります。ネットワークを指定しない場合は、次の gcloud container clusters create コマンドから --network フラグを削除します。
    • SUB_NETWORK: 省略可。サブネットワークの名前。サブネットワークは、クラスタを作成する前に存在している必要があります。サブネットワークを指定しない場合は、次の gcloud container clusters create コマンドから --subnetwork フラグを削除します。

    この例では、残りの変数にデフォルト値を設定しています。これらの値はワークロードに合わせて変更できます。

    • PW_CPU_MACHINE_TYPE: クラスタのデフォルト ノードプールと、Pathways ヘッド(Pathways リソース マネージャーとプロキシ)を実行する CPU ノードプールのマシンタイプ。
    • NUM_CPU_NODES: CPU ノードプール内のノード数。
    • CLUSTER_NODEPOOL_COUNT: TPU スライスの数。各スライスは個別のノードプールです。
    • TPU_MACHINE_TYPE と TOPOLOGY: 各スライスの TPU マシンタイプと TPU トポロジ。
    • WORKERS_PER_SLICE: 各スライスのノード数。これは、トポロジ内の TPU チップ数をノードあたりのチップ数で割った値です。たとえば、2x4 トポロジには 8 個のチップがあり、ct6e-standard-4t ノードには 4 個のチップがあるため、各スライスには 2 個のノードがあります。
  2. gcloud container clusters create コマンドを使用して GKE クラスタを作成します。

    gcloud container clusters create ${CLUSTER} \
        --project=${PROJECT} \
        --zone=${ZONE} \
        --cluster-version=${CLUSTER_VERSION} \
        --scopes=storage-full,gke-default,cloud-platform \
        --machine-type=${PW_CPU_MACHINE_TYPE} \
        --network=${NETWORK} \
        --subnetwork=${SUBNETWORK}
    
  3. gcloud container node-pools create コマンドを使用して、TPU スライス ノードプールを作成します。

    for i in $(seq 1 ${CLUSTER_NODEPOOL_COUNT}); do
        gcloud container node-pools create "tpu-np-${i}" \
            --project=${PROJECT} \
            --zone=${ZONE} \
            --cluster=${CLUSTER} \
            --machine-type=${TPU_MACHINE_TYPE} \
            --num-nodes=${WORKERS_PER_SLICE} \
            --placement-type=COMPACT \
            --tpu-topology=${TOPOLOGY} \
            --scopes=storage-full,gke-default,cloud-platform \
            --workload-metadata=GCE_METADATA
    done
    

    これらのコマンドは、特定の予約を対象としません。TPU 予約を使用するには、--reservation-affinity=specific フラグと --reservation=RESERVATION_NAME フラグを追加します。ここで、RESERVATION_NAME は TPU 予約の名前です。Spot VM を使用するには、--spot フラグを追加します。

  4. Pathways リソース マネージャーとプロキシを実行する Pathways ヘッド用の CPU ノードプールを作成します。gcluster job submit --pathways コマンドで検出できるように、ノードプールに cpu-np という名前を付けます。

    gcloud container node-pools create cpu-np \
        --project=${PROJECT} \
        --zone=${ZONE} \
        --cluster=${CLUSTER} \
        --machine-type=${PW_CPU_MACHINE_TYPE} \
        --num-nodes=${NUM_CPU_NODES} \
        --scopes=storage-full,gke-default,cloud-platform \
        --workload-metadata=GCE_METADATA
    
  5. クラスタの認証情報を取得し、kubectl のデフォルトの Namespace を設定します。

    gcloud container clusters get-credentials ${CLUSTER} \
        --zone=${ZONE} \
        --project=${PROJECT} \
        && kubectl config set-context --current --namespace=default
    
  6. GKE クラスタで Pathways アーキテクチャを使用するには、JobSet API をインストールします。

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.12.0/manifests.yaml
    
  7. ノードの STATUS が Ready で、jobset-controller-manager Pod の STATUS が Running であることを確認します。

    kubectl get nodes
    kubectl get pods --namespace=jobset-system
    

    出力は次のようになります。

    NAME                                      STATUS   ROLES    AGE   VERSION
    gke-cluster-cpu-np-12345678-abcd          Ready    <none>   5m    v1.32.2-gke.1475000
    gke-cluster-tpu-np-1-12345678-efgh        Ready    <none>   5m    v1.32.2-gke.1475000
    ...
    NAME                                         READY   STATUS    RESTARTS   AGE
    jobset-controller-manager-6f7b8c9d4-xyz12    1/1     Running   0          2m
    

クリーンアップ

このページで使用したリソースについて、 Google Cloud アカウントに課金されないようにするには、完了時にクラスタを削除します。

  • Cluster Toolkit: 次のコマンドを実行して、デプロイを破棄します。

    ./gcluster destroy DEPLOYMENT_FOLDER_NAME

    DEPLOYMENT_FOLDER_NAME は、デプロイ フォルダの名前に置き換えます。

  • gcloud CLI: 次のコマンドを実行して、GKE クラスタとそのノードプールを削除します。

    gcloud container clusters delete ${CLUSTER} \
        --project=${PROJECT} \
        --zone=${ZONE}
    

次のステップ