Pathways でインタラクティブなワークロードを実行する

Pathways インタラクティブ ワークロードは、クラスタの Pathways コントローラ コンポーネントの外部（VM など）で実行されるリモート JAX ワークロードです。バッチ ワークロードとは異なり、インタラクティブ ワークロードが終了しても、Pathways クラスタ コンポーネントはシャットダウンされず、他の JAX クライアントからの接続に使用できます。このドキュメントでは、Jupyter ノートブックを例として使用して、インタラクティブなワークロードを示します。Pathways を使用したインタラクティブなワークロードは、Google Kubernetes Engine（GKE）クラスタまたは Slurm を使用した Cluster Director クラスタで実行できます。

このページでは、両方のプラットフォームの手順について説明します。各セクションで、[GKE] タブまたは [Cluster Director] タブを選択します。

JAX ユーザーは、IFRT インターフェースを使用して Pathways クラスタにコマンドを送信します。JAX コードは、ターミナル、ノートブック、Python 互換環境のいずれから実行されても、Pathways リソースとやり取りできます。

始める前に

• このページのサンプルをどのように使うかに応じて、タブを選択してください。

  gcloud

  Google Cloud CLI をインストールし、 フェデレーション ID を使用して gcloud CLI にログインします。ログイン後、次のコマンドを実行して Google Cloud CLI を初期化します。

  gcloud init

  REST

  このページの REST API サンプルをローカル開発環境で使用するには、gcloud CLI に指定した認証情報を使用します。

  Google Cloud CLI をインストールし、 フェデレーション ID を使用して gcloud CLI にログインします。

  詳細については、 Google Cloud 認証ドキュメントの REST を使用して認証するをご覧ください。

始める前に

GKE

次の前提条件を満たしていることを確認してください。

• Pathways を使用して GKE クラスタを作成する。
• Cluster Toolkit をインストールします。
• Kubernetes ツールをインストールする。
• Google Cloud CLI をインストールします。
• Kubernetes Engine API を有効にします。
• Google Cloud プロジェクトで Pathways が有効になっていることを確認します。

Cluster Director

次の前提条件を満たしていることを確認してください。

• Cluster Director クラスタを作成し、Pathways 環境を設定する。
• gcloud CLI をインストールします。
• クラスタのログインノードに SSH アクセスできることを確認します。
• JAX クライアントを実行するホストに、venv モジュールを含む Python 3 をインストールします。Debian と Ubuntu では、python3-venv パッケージをインストールします。

インタラクティブ モードで Pathways を実行する

GKE

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

   export WORKLOAD=WORKLOAD
   export WORKLOAD_NODEPOOL_COUNT=WORKLOAD_NODEPOOL_COUNT
   export COMPUTE_TYPE="COMPUTE_TYPE"
   export TOPOLOGY="TOPOLOGY"
   export PROJECT_ID=PROJECT_ID
   export REGION=REGION
   export ZONE=ZONE
   export CLUSTER=CLUSTER
   export BUCKET_NAME=BUCKET_NAME

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

   • WORKLOAD: ワークロードを識別する一意の名前（DNS ラベルの制限により 22 文字以下）。
   • WORKLOAD_NODEPOOL_COUNT: Pathways ワークロードで使用されるノードプールの数。
   • COMPUTE_TYPE: TPU マシンタイプ（例: ct6e-standard-4t）。TPU の各バージョンでサポートされている TPU タイプの詳細については、TPU のバージョンをご覧ください。
   • TOPOLOGY: TPU の配置トポロジ（例: 2x4）。
   • PROJECT_ID: 実際のGoogle Cloud プロジェクト ID。
   • REGION: バケットとクラスタが配置されているリージョン。
   • ZONE: ワークロードを実行するゾーン。
   • CLUSTER: GKE クラスタの名前。
   • BUCKET_NAME: Pathways アーティファクトの Cloud Storage バケット。

2. Pathways アーティファクト用の Cloud Storage バケットをまだ作成していない場合は、作成します。

   gcloud storage buckets create gs://${BUCKET_NAME} \
       --project=${PROJECT_ID} \
       --location=${REGION}
   

3. クラスタに Pathways コンテナを作成します。ヘッドレス ワークロードを実行するには、gcluster job submit コマンドを使用します（詳細については、Cluster Toolkit ジョブ送信ガイドをご覧ください）。

   gcluster job submit \
       --pathways \
       --pathways-gcs-location="gs://${BUCKET_NAME}/pathways-artifacts" \
       --name=${WORKLOAD} \
       --cluster=${CLUSTER} \
       --project=${PROJECT_ID} \
       --location=${ZONE} \
       --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
       --compute-type=${COMPUTE_TYPE} \
       --topology=${TOPOLOGY} \
       --pathways-headless
   

   gcluster job submit コマンドは、Pathways ヘッド Pod を実行するために、cpu-np または pathways-np という名前の CPU ノードプールを探します。CPU ノードプールの名前が異なる場合（クラスタ作成ガイドの cpu-pathways-np など）、このページの gcluster job submit コマンドに --pathways-head-np=NODE_POOL_NAME フラグを追加します。そうしないと、コマンドはエラー failed to auto-detect a suitable CPU node pool で失敗します。

この時点で、JAX ワークロードは IFRT プロキシ サーバーに接続できます。

Cluster Director

1. クラスタのログインノードに接続します。

   gcloud compute ssh CLUSTER_ID-login-001 \
       --project=PROJECT_ID \
       --zone=ZONE \
       --tunnel-through-iap

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

   • CLUSTER_ID: Cluster Director クラスタの名前。
   • PROJECT_ID: 実際のGoogle Cloud プロジェクト ID。
   • ZONE: クラスタが配置されているゾーン。

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

   export SLICE_SHAPE=TOPOLOGY
   export MAX_SLICES=NUM_SLICES
   export MIN_SLICES=NUM_SLICES
   export GCS_BUCKET=gs://BUCKET_NAME

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

   • TOPOLOGY: 各スライスの TPU トポロジ（例: 2x4）。TPU ノードセットのトポロジと一致している必要があります。
   • NUM_SLICES: TPU スライスの数。すべてのスライスが実行されている場合にのみ Pathways クラスタが準備完了になるように、MIN_SLICES と MAX_SLICES を同じ値に設定します。
   • BUCKET_NAME: Pathways アーティファクト用の既存の Cloud Storage バケット。

   TPU パーティションがクラスタのデフォルト パーティションでない場合は、PARTITION をその名前に設定します。

3. ヘッドレス モードで Pathways クラスタを起動します。

   cd ~/pathways
   sbatch start_pathways.sh
   

4. サービスログ pathways-service-JOB_ID.out（JOB_ID は Slurm ジョブ ID）に次のような出力が表示されるまで待ちます。

   Pathways Runtime Service is READY!
    Discovery file    : /home/USER/pathways/cluster.env
    IFRT Proxy Target : grpc://RM_IP:29000
   

   RM_IP は、プロキシ サーバーの内部 IP アドレスです。ファイル ~/pathways/cluster.env には同じアドレスが含まれています。

この時点で、JAX ワークロードは IFRT プロキシ サーバーに接続できます。

インタラクティブ モードで Pathways クラスタに接続する

Pathways クラスタには、ポート転送の有無にかかわらず接続できます。次のいずれかのセクションを使用して、Pathways クラスタに接続します。

クラスタ ディレクターでは、JAX クライアントは Pathways サーバー（0.11.1）と同じ JAX バージョンを使用する必要があります。

ポート転送を使用して接続する

GKE

クラスタのコントロール プレーンにアクセスできるホストから、ローカルポート 29000 を Pathways ヘッド Pod に転送して、プロキシ サーバーにアクセスします。

PROXY_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
PROXY_PORT=29000
kubectl port-forward ${PROXY_POD} ${PROXY_PORT}:${PROXY_PORT}

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

Forwarding from 127.0.0.1:29000 -> 29000
Forwarding from [::1]:29000 -> 29000

同じホストで、新しいターミナル ウィンドウを開きます。Python 3.11 以降を使用して、仮想環境を作成し、pathwaysutils パッケージと jax パッケージをインストールします。JAX_PLATFORMS 環境変数と JAX_BACKEND_TARGET 環境変数を設定し、pathwaysutils と jax をインポートする Python スクリプトを実行します。

python3 -m venv .venv
source .venv/bin/activate
pip install pathwaysutils jax

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

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

[device(144,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(145,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(146,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(147,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(148,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(149,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(150,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(151,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(162,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(163,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(164,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(165,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(166,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(167,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(168,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(169,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
Waiting up to 5 seconds.
Sent all pending logs.
2024-11-13 21:38:51.267523: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled

Cluster Director

クラスタのログインノードに SSH 接続できるホスト（ワークステーションやクラスタのサブネット外の VM など）から、ログインノードを介してプロキシポートを転送します。

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 29000:RM_IP:29000

注: デフォルトのサービス アカウントを使用して Compute Engine VM で gcloud CLI を実行する場合、VM には cloud-platform アクセス スコープが必要です。そうでない場合、gcloud compute ssh は insufficient authentication scopes で失敗します。または、VM で gcloud auth login を実行して、独自のアカウントを使用します。

同じホストで、新しいターミナル ウィンドウを開きます。JAX_PLATFORMS 環境変数と JAX_BACKEND_TARGET 環境変数を設定し、pathwaysutils と jax をインポートする Python スクリプトを実行します。ホストには RM_IP へのルートがないため、クライアントはトンネルのローカルエンド 127.0.0.1:29000 に接続します。

python3 -m venv .venv
source .venv/bin/activate
pip install pathwaysutils jax==0.11.1

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

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

[device(288,TPU_DEVICE,coords=[0,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 device(289,TPU_DEVICE,coords=[1,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 ...
 device(313,TPU_DEVICE,coords=[1,3,0,0],logical_task=1,slice=1,default_mem=device,mem_spaces=3)]
W1029 22:52:22.897483    2361 grpc_client.cc:91] IFRT proxy server disconnected: CANCELLED: Stream removed (CANCELLED)

IFRT proxy server disconnected 警告は、スクリプトが終了して接続を閉じるときに表示されます。これは想定された動作です。

注: 各 JAX 呼び出しは SSH トンネルを通過します。ポート転送を使用して接続を確認します。ノートブックの場合は、クラスタのサブネットで JAX クライアントを実行します。

ポート転送を使用せずに VPC 内のホストから接続する

GKE

ポート転送を使用しない場合は、Cloud DNS または内部ロードバランサを使用して Pathways クラスタに接続できます。

Cloud DNS を使用して接続する

クラスタで Cloud DNS を有効にすると、DNS プロバイダが kube-dns から Cloud DNS に切り替わります。有効にすると、GKE は DNS 名の Virtual Private Cloud（VPC）ネットワークに限定公開 Cloud DNS ゾーンを作成します。詳細については、GKE 向け Cloud DNS の使用をご覧ください。

VPC スコープまたは追加の VPC スコープで Cloud DNS を有効にすると、VPC ネットワーク内の GKE 以外の VM から Kubernetes DNS 名を解決できます。クラスタ スコープの場合、これらの名前を解決できるのはクラスタ内のノードのみです。名前の形式は SERVICE_NAME.NAMESPACE.svc.CUSTOM_DNS_DOMAIN です。Pathways ヘッド Pod には、WORKLOAD-pathways-head-0-0.WORKLOAD.NAMESPACE.svc.CUSTOM_DNS_DOMAIN という名前のサービスがあります。

次のコマンドは、Cloud DNS を使用して Pathways クラスタに接続する方法を示しています。

1. リーダー DNS エントリが GKE 以外のホストから解決可能であることを確認します。

   host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN

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

   • WORKLOAD: Pathways ワークロードの名前。
   • CUSTOM_DNS_DOMAIN: クラスタに構成されているカスタム DNS ドメイン。

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

   WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN has address 10.0.2.75

2. DNS 名を使用して Pathways クラスタに接続します。

   JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

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

   [device(216,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(217,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(218,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(219,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(220,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(221,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(222,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(223,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(234,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(235,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(236,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(237,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(238,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
   device(239,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
   device(240,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
   device(241,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
   Waiting up to 5 seconds.
   Sent all pending logs.
   2024-11-14 00:02:49.882044: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled
   

内部ロードバランサを使用して接続する

Pathways デプロイを指す VPC ネットワーク内のプライベート IP アドレスに、内部ロードバランサ（ILB）をバックエンドとする Service を作成します。この場合、クラスタで Cloud DNS を有効にする必要はありません。

VM が多数あるクラスタで内部ロードバランサを作成する場合は、ILB のサブセット化を有効にします。詳細については、既存のクラスタで GKE のサブセット化を有効にするをご覧ください。ILB のサブセット化が有効になっていない場合、クラスタ内のすべてのノードがすべての内部ロードバランサのバックエンド インスタンス グループの一部になります。この場合、250 ノードを超えてスケーリングすることはできません。ILB のサブセット化を有効にすると、GKE はインスタンス グループではなくネットワーク エンドポイント グループを作成し、Service のサービスを提供する Pod のいずれかを実行するノードのみを含めます。ILB のサブセット化を有効にすると、1 回限りの設定レイテンシが約 15 分発生します。次のコマンドで ILB のサブセット化を有効にします。

gcloud container clusters update ${CLUSTER} \
    --project=${PROJECT_ID} \
    --location=CLUSTER_LOCATION \
    --enable-l4-ilb-subsetting

CLUSTER_LOCATION は、クラスタが配置されているリージョンまたはゾーンに置き換えます。

ILB のサブセット化を有効にしたら、次の YAML マニフェストを使用して LoadBalancer タイプの Kubernetes Service を作成します。これにより、GKE はクラスタの VPC ネットワーク内に内部ロードバランサを作成します。

apiVersion: v1
kind: Service
metadata:
  name: WORKLOAD-ilb
  annotations:
    networking.gke.io/load-balancer-type: "Internal"
    networking.gke.io/internal-load-balancer-allow-global-access: "true"
spec:
  type: LoadBalancer
  externalTrafficPolicy: Local
  selector:
    jobset.sigs.k8s.io/jobset-name: WORKLOAD
    jobset.sigs.k8s.io/replicatedjob-name: pathways-head
  ports:
  - name: tcp-port
    protocol: TCP
    port: 29000
    targetPort: 29000

WORKLOAD は、Pathways ワークロードの名前に置き換え、ファイルを pathways-headless-ilb.yaml として保存します。

次のようにマニフェストを適用します。

kubectl apply -f pathways-headless-ilb.yaml

GKE がロードバランサを作成したら（約 1 分後）、EXTERNAL-IP 列に IP アドレスが表示されていることを確認します。

kubectl get services

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

NAME             TYPE           CLUSTER-IP      EXTERNAL-IP   PORT(S)           AGE
WORKLOAD       ClusterIP      None            <none>        <none>            30m
WORKLOAD-ilb   LoadBalancer   34.118.232.46   10.0.0.22     29000:31246/TCP   2m41s

クラスタと同じ VPC ネットワーク内のホストからポート転送を行わずに、Pathways Deployment にアクセスできます。

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://EXTERNAL_IP:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

EXTERNAL_IP は、kubectl get services で表示される EXTERNAL-IP 値（例: 10.0.0.22）に置き換えます。

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

[device(288,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(289,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(290,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(291,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(292,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(293,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(294,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(295,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(306,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(307,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(308,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(309,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(310,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(311,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(312,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(313,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
Waiting up to 5 seconds.
Sent all pending logs.
2024-11-14 00:30:07.296917: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled

注: 実際の IP アドレスは、上記のコマンドと出力例に示されているアドレスとは異なります。

Cluster Director

Compute Engine VM などのクラスタのサブネット内のホストは、プロキシ サーバーに直接接続できます。クラスタ ディレクターのファイアウォールは、クラスタのサブネットからのトラフィックを許可しますが、同じ VPC ネットワークの他のサブネットからのトラフィックは許可しません。

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://RM_IP:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

RM_IP は、サービスログの IFRT プロキシアドレスに置き換えます。

注: IFRT プロキシ サーバーはクライアントを認証しません。RM_IP:29000 にアクセスできるホストは、Pathways クラスタを使用できます。

Jupyter ノートブック

GKE

Jupyter ノートブックは、Vertex AI を使用して作成することも、セルフホスト型の Jupyter ノートブックを作成することもできます。

Vertex AI Workbench インスタンスを作成する

Pathways クラスタを設定して検証したら、Vertex AI Workbench Jupyter ノートブックから GKE TPU VM にアクセスできます。次の設定手順では、GKE Pathways クラスタが同じ VPC ネットワーク（特に構成していない場合はデフォルト ネットワーク）に存在することを前提としています。

1. Google Cloud コンソールで [Vertex AI Workbench] ページに移動します。

2. [インスタンス] タブで [新規作成] をクリックし、ネットワークが GKE クラスタのネットワークと一致していることを確認します。

   または、コマンドラインを使用して新しい Vertex AI Workbench インスタンスを作成します。

   gcloud workbench instances create INSTANCE_NAME \
       --project=PROJECT_ID \
       --machine-type=e2-standard-4 \
       --data-disk-size=100 \
       --location=ZONE \
       --network=NETWORK

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

   • INSTANCE_NAME: Vertex AI Workbench インスタンスの名前。
   • PROJECT_ID: 実際のGoogle Cloud プロジェクト ID。
   • ZONE: インスタンスを作成するゾーン。
   • NETWORK: GKE クラスタの VPC ネットワークの名前。--network フラグは省略可能です。省略すると、コマンドは default ネットワークを使用します。

3. Vertex AI がインスタンスを作成したら、インスタンスに移動して [JupyterLab を開く] をクリックします。

セルフホスト型 Jupyter ノートブック インスタンスを作成する

次のコマンドは、gcluster job submit コマンドを使用してセルフホスト型 Jupyter ノートブック インスタンスを作成する方法を示しています。

gcluster job submit \
    --pathways \
    --name=${WORKLOAD} \
    --cluster=${CLUSTER} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
    --compute-type=${COMPUTE_TYPE} \
    --topology=${TOPOLOGY} \
    --pathways-gcs-location="gs://${BUCKET_NAME}/pathways-artifacts" \
    --image=jupyter/base-notebook \
    --command="start-notebook.sh"

次の YAML マニフェストは、kubectl を使用してセルフホスト型の Jupyter ノートブック インスタンスを作成する方法を示しています。ヘッドレス Pathways クラスタを作成したら、次の YAML マニフェストを適用します。詳細については、インタラクティブ モードで Pathways を実行するをご覧ください。

apiVersion: batch/v1
kind: Job
metadata:
  name: WORKLOAD
spec:
  template:
    spec:
      hostNetwork: true
      dnsPolicy: ClusterFirstWithHostNet
      restartPolicy: OnFailure
      containers:
      - name: jupyter-notebook
        image: jupyter/base-notebook  # Use the appropriate Jupyter image
        ports:
        - containerPort: 8888

hostNetwork: true フィールドを使用すると、ノートブックの JAX クライアントは、Pathways プロキシ サーバーとの Application Layer Transport Security（ALTS）ハンドシェイクを完了できます。このフィールドがないと、Workload Identity Federation for GKE を使用するノードプールで接続が失敗します。dnsPolicy: ClusterFirstWithHostNet フィールドを使用すると、ノートブックは Pathways ヘッド Pod の名前などのクラスタ内 DNS 名を解決できます。詳細については、Cloud での Pathways のトラブルシューティングをご覧ください。

Cluster Director

クラスタのサブネット内の既存の Compute Engine VM で Jupyter ノートブックを実行することも、クラスタのログインノードにセルフホスト型の Jupyter ノートブックを作成することもできます。どちらの場合も、ノートブックはポート転送なしでプロキシ サーバーに接続します。

クラスタのサブネット内の既存の VM を使用する

VM で、仮想環境に JupyterLab をインストールして起動します。

python3 -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
jupyter lab --no-browser --port=8888

セルフホスト型 Jupyter ノートブック インスタンスを作成する

クラスタのログインノードに JupyterLab をインストールして起動します。クラスタの Python を使用して、ホーム ディレクトリに仮想環境を作成します。ログインノードのシステム Python は、現在の JAX リリースに対して古すぎるため、pip は jax==0.11.1 を見つけることができません。

cd ~
/home/pathways/venv/bin/python -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
source ~/pathways/cluster.env
jupyter lab --no-browser --port=8888

ノートブック インスタンスに接続する

GKE

ポート転送を使用して、ローカルマシンからノートブックに接続します。

Cluster Toolkit

MAIN_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
kubectl port-forward pod/${MAIN_POD} 8888:8888

kubectl

MAIN_POD=$(kubectl get pods \
    --selector=batch.kubernetes.io/job-name=WORKLOAD \
    --output=jsonpath='{.items[0].metadata.name}')
kubectl port-forward pod/${MAIN_POD} 8888:8888

ローカル ブラウザで http://localhost:8888?token=TOKEN にアクセスします。TOKEN は、Jupyter ノートブック コンテナのログのトークンに置き換えます。

kubectl logs ${MAIN_POD}

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

...
Or copy and paste one of these URLs:
  http://NODE_NAME:8888/lab?token=TOKEN
  http://127.0.0.1:8888/lab?token=TOKEN

NODE_NAME は、ノートブック Pod を実行する GKE ノードのホスト名です（例: gke-my-cluster-cpu-np-1a2b3c4d-ef56）。

Cluster Director

VM またはログインノードのノートブックの場合は、そのホストへのポート転送を使用してローカルマシンから接続します。たとえば、ログインノードの場合:

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 8888:localhost:8888

ローカル ブラウザで http://localhost:8888?token=TOKEN にアクセスします。TOKEN は、JupyterLab の起動時に出力されたトークンに置き換えます。

Pathways クラスタへのノートブックの接続

1. JupyterLab 内から、新しい Python 3 ノートブックを作成します。
2. Pathways プロキシ サーバーに接続します。

ノートブックにセルを追加して pathwaysutils をインストールし、JAX_PLATFORMS を proxy に設定し、JAX_BACKEND_TARGET を grpc://PROXY_ADDRESS:29000 に設定します。

注: PROXY_ADDRESS の値は、インタラクティブ モードで Pathways クラスタに接続するで選択した接続オプションによって異なります。Cluster Director で、同じセルに jax==0.11.1 をインストールします。source ~/pathways/cluster.env の後に起動したセルフホスト ノートブックでは、変数はすでに設定されています。

!pip install pathwaysutils
# Set JAX_PLATFORMS and JAX_BACKEND_TARGET if you're not using `gcluster job submit` to deploy the notebook
%env JAX_PLATFORMS=proxy
  # Replace your proxy address in the following line:
%env JAX_BACKEND_TARGET=grpc://PROXY_ADDRESS:29000

PROXY_ADDRESS は、インタラクティブ モードで Pathways クラスタに接続するの Pathways プロキシ サーバーのホストまたは IP アドレス（127.0.0.1、Cloud DNS ホスト名、内部ロードバランサの IP アドレスなど）に置き換えます。Cluster Director で、ポート転送を使用する場合は 127.0.0.1 を使用します。または、サービスログの IFRT プロキシ アドレスを使用します。kubectl を使用してセルフホスト型ノートブックをデプロイした場合は、PROXY_ADDRESS を Pathways ヘッド Pod のクラスタ内 DNS 名に置き換えます。これにより、JAX_BACKEND_TARGET 値は grpc://WORKLOAD-pathways-head-0-0.WORKLOAD:29000 になります。

2 つ目のセルを追加して接続を確認し、Pathways クラスタ内のデバイスを出力します。

import pathwaysutils
import jax

pathwaysutils.initialize()
print(jax.devices())

接続が成功すると、JAX が Pathways-on-Cloud バックエンドを検出したことを示すメッセージが出力に含まれます。

リストされている JAX デバイスの数は、Pathways クラスタの作成時に指定した TPU チップの数とスライスの数と一致します。

ノートブックにコードを追加する

独自の JAX コードを追加し、Pathways クラスタの TPU でインタラクティブに実行します。次のコードは、1 つのノートブックから 2 つのスライスにわたって計算を実行する方法を示しています。実証済みの例を実行するには、WORKLOAD_NODEPOOL_COUNT を 2 以上に設定して、2 つ以上のスライスを含む Pathways ワークロードを作成します。Cluster Director で、MIN_SLICES と MAX_SLICES を 2 以上に設定します。1 つのスライスでは、slice1_devices は空になり、例の最後の部分は ValueError で失敗します。

import jax
import jax.numpy as jnp
from jax import lax
import numpy as np

# You can use JAX APIs as usual across any of the devices.
x = jax.device_put(np.pi / 2., jax.devices()[-1])
jax.jit(jnp.sin)(x)

# pmap can run across all devices on all slices
num_tpus = jax.device_count()
f = jax.pmap(lambda x: lax.psum(1, 'i'), 'i')
x = jnp.arange(num_tpus)
y = f(x)
print(y)

# You can also target devices from a specific slice
slice0_devices = [d for d in jax.devices() if d.slice_index == 0]
f = jax.pmap(lambda x: lax.psum(1, 'i'), 'i', devices=slice0_devices)
x = jnp.arange(len(slice0_devices))
y = f(x)
print(y)
print(y.global_shards)

# You can send data produced on one slice to another slice
slice1_devices = [d for d in jax.devices() if d.slice_index == 1]
g = jax.pmap(lambda x: x + lax.axis_index('i'), 'i', devices=slice1_devices)
z = g(y)
print(z)
print(z.global_shards)

Pathways インタラクティブ ワークロードを削除する

GKE

Cluster Toolkit

gcluster job cancel WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=ZONE

kubectl

kubectl delete job WORKLOAD

Cluster Director

クラスタのログインノードで、Pathways サービスジョブをキャンセルします。

scancel -u ${USER} -n pathways-service

このコマンドは、Pathways ワーカーのジョブもキャンセルします。TPU VM は静的ノードであるため、実行を継続し、他の Slurm ジョブで使用できるようになります。

問題が発生した場合は、Cloud での Pathways のトラブルシューティングをご覧ください。

次のステップ

• Pathways を使用して GKE クラスタを作成する
• Pathways を使用して Cluster Director クラスタを作成する
• Pathways を使用してバッチ ワークロードを実行する
• JAX ワークロードを Pathways に移植する
• Pathways を使用した復元力のあるトレーニング
• Cloud での Pathways のトラブルシューティング