Menjalankan workload interaktif dengan Pathways

Workload interaktif Pathways adalah workload JAX jarak jauh yang berjalan di luar komponen pengontrol Pathways cluster, misalnya, di VM. Tidak seperti workload batch, saat workload interaktif selesai, workload tersebut tidak akan menonaktifkan komponen cluster Pathways, yang tetap tersedia untuk koneksi dari klien JAX lainnya. Dokumen ini menggunakan notebook Jupyter sebagai contoh untuk mendemonstrasikan workload interaktif. Anda dapat menjalankan workload interaktif dengan Pathways di cluster Google Kubernetes Engine (GKE) atau di cluster Cluster Director dengan Slurm.

Halaman ini berisi petunjuk untuk kedua platform. Di setiap bagian, pilih tab GKE atau Cluster Director.

Dengan menggunakan antarmuka IFRT, pengguna JAX mengirim perintah ke cluster Pathways. Kode JAX, baik yang dieksekusi dari terminal, notebook, atau lingkungan yang kompatibel dengan Python, dapat berinteraksi dengan resource Pathways.

Sebelum memulai

• Pilih tab untuk melihat bagaimana Anda berencana menggunakan contoh di halaman ini:

  gcloud

  Instal Google Cloud CLI, lalu login ke gcloud CLI dengan identitas gabungan Anda. Setelah login, inisialisasi Google Cloud CLI dengan menjalankan perintah berikut:

  gcloud init

  REST

  Untuk menggunakan contoh REST API di halaman ini dalam lingkungan pengembangan lokal, gunakan kredensial yang Anda berikan ke gcloud CLI.

  Instal Google Cloud CLI, lalu login ke gcloud CLI dengan identitas gabungan Anda.

  Untuk mengetahui informasi selengkapnya, lihat Melakukan autentikasi untuk menggunakan REST dalam dokumentasi autentikasi Google Cloud .

Sebelum memulai

GKE

Selesaikan prasyarat berikut:

• Membuat cluster GKE dengan Pathways.
• Instal Cluster Toolkit.
• Instal alat Kubernetes.
• Menginstal Google Cloud CLI.
• Aktifkan Kubernetes Engine API.
• Pastikan Pathways diaktifkan untuk project Google Cloud Anda.

Cluster Director

Selesaikan prasyarat berikut:

• Buat cluster Cluster Director dan siapkan lingkungan Pathways.
• Instal gcloud CLI.
• Pastikan Anda memiliki akses SSH ke node login cluster.
• Instal Python 3 dengan modul venv di host yang menjalankan klien JAX. Di Debian dan Ubuntu, instal paket python3-venv.

Menjalankan Pathways dalam mode interaktif

GKE

1. Tetapkan variabel lingkungan berikut:

   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

   Ganti kode berikut:

   • WORKLOAD: nama unik untuk mengidentifikasi workload Anda (22 karakter atau kurang karena batas label DNS).
   • WORKLOAD_NODEPOOL_COUNT: jumlah node pool yang digunakan oleh workload Pathways.
   • COMPUTE_TYPE: jenis mesin TPU—misalnya, ct6e-standard-4t. Untuk mengetahui informasi selengkapnya tentang jenis TPU yang didukung untuk setiap versi TPU, lihat Versi TPU.
   • TOPOLOGY: topologi penempatan TPU—misalnya, 2x4.
   • PROJECT_ID:Google Cloud project ID Anda.
   • REGION: region tempat bucket dan cluster Anda berada.
   • ZONE: zona tempat Anda berencana menjalankan workload.
   • CLUSTER: nama cluster GKE Anda.
   • BUCKET_NAME: bucket Cloud Storage untuk artefak Pathways.

2. Jika Anda belum membuat bucket Cloud Storage untuk artefak Pathways, buatlah:

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

3. Buat container Pathways di cluster. Untuk menjalankan workload headless, gunakan perintah gcluster job submit (untuk mengetahui informasi selengkapnya, lihat Panduan Pengiriman Tugas 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
   

   Perintah gcluster job submit mencari node pool CPU bernama cpu-np atau pathways-np untuk menjalankan Pod head Pathways. Jika node pool CPU Anda memiliki nama yang berbeda (seperti cpu-pathways-np dari panduan pembuatan cluster), tambahkan flag --pathways-head-np=NODE_POOL_NAME ke perintah gcluster job submit di halaman ini. Jika tidak, perintah akan gagal dengan error failed to auto-detect a suitable CPU node pool.

Pada tahap ini, workload JAX Anda dapat terhubung ke server proxy IFRT.

Cluster Director

1. Hubungkan ke node login cluster:

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

   Ganti kode berikut:

   • CLUSTER_ID: nama cluster Cluster Director Anda.
   • PROJECT_ID:Google Cloud project ID Anda.
   • ZONE: zona tempat cluster Anda berada.

2. Tetapkan variabel lingkungan berikut:

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

   Ganti kode berikut:

   • TOPOLOGY: topologi TPU setiap slice—misalnya, 2x4. Topologi ini harus cocok dengan topologi set node TPU Anda.
   • NUM_SLICES: jumlah slice TPU. Tetapkan MIN_SLICES dan MAX_SLICES ke nilai yang sama agar cluster Pathways hanya siap saat semua slice berjalan.
   • BUCKET_NAME: bucket Cloud Storage yang ada untuk artefak Pathways.

   Jika partisi TPU Anda bukan partisi default cluster, tetapkan juga PARTITION ke namanya.

3. Mulai cluster Pathways dalam mode headless:

   cd ~/pathways
   sbatch start_pathways.sh
   

4. Tunggu hingga log layanan, pathways-service-JOB_ID.out (dengan JOB_ID adalah ID pekerjaan Slurm), menampilkan output yang mirip dengan berikut ini:

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

   RM_IP adalah alamat IP internal server proxy. File ~/pathways/cluster.env berisi alamat yang sama.

Pada tahap ini, workload JAX Anda dapat terhubung ke server proxy IFRT.

Menghubungkan ke cluster Pathways dalam mode interaktif

Anda dapat terhubung ke cluster Pathways dengan atau tanpa penerusan port. Gunakan salah satu bagian berikut untuk terhubung ke cluster Pathways.

Di Cluster Director, klien JAX harus menggunakan versi JAX yang sama dengan server Pathways, 0.11.1.

Terhubung menggunakan penerusan port

GKE

Dari host mana pun yang memiliki akses ke bidang kontrol cluster Anda, teruskan port lokal 29000 ke Pod head Pathways untuk mengakses server proxy:

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

Outputnya mirip dengan hal berikut ini:

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

Di host yang sama, buka jendela terminal baru. Dengan menggunakan Python 3.11 atau yang lebih baru, buat lingkungan virtual, instal paket pathwaysutils dan jax, tetapkan variabel lingkungan JAX_PLATFORMS dan JAX_BACKEND_TARGET, lalu jalankan skrip Python yang mengimpor pathwaysutils dan jax:

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())'

Outputnya mirip dengan hal berikut ini:

[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

Dari host mana pun yang dapat melakukan SSH ke node login cluster, seperti workstation Anda atau VM di luar subnet cluster, teruskan port proxy melalui node login:

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

Catatan: Jika Anda menjalankan gcloud CLI di VM Compute Engine dengan akun layanan defaultnya, VM memerlukan cakupan akses cloud-platform, atau gcloud compute ssh akan gagal dengan insufficient authentication scopes. Atau, jalankan gcloud auth login di VM untuk menggunakan akun Anda sendiri.

Di host yang sama, buka jendela terminal baru. Tetapkan variabel lingkungan JAX_PLATFORMS dan JAX_BACKEND_TARGET, lalu jalankan skrip Python yang mengimpor pathwaysutils dan jax. Host tidak memiliki rute ke RM_IP, sehingga klien terhubung ke ujung lokal tunnel, 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())'

Outputnya mirip dengan hal berikut ini:

[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)

Peringatan IFRT proxy server disconnected akan muncul saat skrip keluar dan menutup koneksinya. Hal ini wajar.

Catatan: Setiap panggilan JAX melewati tunnel SSH. Gunakan penerusan port untuk memverifikasi konektivitas. Untuk notebook, jalankan klien JAX di subnet cluster.

Menghubungkan dari host di VPC tanpa menggunakan penerusan port

GKE

Jika tidak ingin menggunakan penerusan port, Anda dapat terhubung ke cluster Pathways menggunakan Cloud DNS atau load balancer internal.

Terhubung menggunakan Cloud DNS

Mengaktifkan Cloud DNS di cluster Anda akan mengalihkan penyedia DNS dari kube-dns ke Cloud DNS. Jika diaktifkan, GKE akan membuat zona Cloud DNS pribadi di jaringan Virtual Private Cloud (VPC) Anda untuk nama DNS. Untuk mengetahui informasi selengkapnya, lihat Menggunakan Cloud DNS untuk GKE.

Jika Anda mengaktifkan Cloud DNS dengan cakupan VPC atau cakupan VPC tambahan, nama DNS Kubernetes dapat di-resolve dari VM non-GKE di dalam jaringan VPC Anda. Dengan cakupan cluster, hanya node dalam cluster yang dapat me-resolve nama ini. Nama memiliki format SERVICE_NAME.NAMESPACE.svc.CUSTOM_DNS_DOMAIN. Pod head Pathways memiliki layanan bernama WORKLOAD-pathways-head-0-0.WORKLOAD.NAMESPACE.svc.CUSTOM_DNS_DOMAIN.

Perintah berikut menunjukkan cara terhubung ke cluster Pathways menggunakan Cloud DNS:

1. Pastikan entri DNS leader dapat di-resolve dari host non-GKE:

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

   Ganti kode berikut:

   • WORKLOAD: nama workload Pathways Anda.
   • CUSTOM_DNS_DOMAIN: domain DNS kustom yang dikonfigurasi untuk cluster Anda.

   Outputnya mirip dengan hal berikut ini:

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

2. Hubungkan ke cluster Pathways menggunakan nama DNS:

   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())'

   Outputnya mirip dengan hal berikut ini:

   [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
   

Menghubungkan menggunakan load balancer internal

Untuk alamat IP pribadi di jaringan VPC Anda yang mengarah ke deployment Pathways, buat Service yang didukung oleh load balancer internal (ILB). Hal ini tidak mengharuskan cluster Anda mengaktifkan Cloud DNS.

Untuk cluster dengan banyak VM, aktifkan subsetelan ILB jika Anda membuat load balancer internal. Untuk mengetahui informasi selengkapnya, lihat Mengaktifkan subsetelan GKE di cluster yang ada. Jika subsetting ILB tidak diaktifkan, semua node dalam cluster akan menjadi bagian dari grup instance backend untuk semua load balancer internal, yang tidak dapat diskalakan melebihi 250 node. Dengan mengaktifkan subsetelan ILB, GKE membuat grup endpoint jaringan, bukan grup instance, dan hanya menyertakan node yang menjalankan salah satu Pod aktif Service. Mengaktifkan subsetelan ILB memiliki latensi penyiapan satu kali sekitar 15 menit. Perintah berikut mengaktifkan subsetting ILB:

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

Ganti CLUSTER_LOCATION dengan region atau zona tempat cluster Anda berada.

Setelah mengaktifkan subsetting ILB, buat Service Kubernetes berjenis LoadBalancer menggunakan manifes YAML berikut. Hal ini menyebabkan GKE membuat load balancer internal di dalam jaringan VPC cluster Anda:

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

Ganti WORKLOAD dengan nama workload Pathways Anda, dan simpan file sebagai pathways-headless-ilb.yaml.

Terapkan manifes:

kubectl apply -f pathways-headless-ilb.yaml

Setelah GKE membuat load balancer (sekitar satu menit kemudian), pastikan kolom EXTERNAL-IP menampilkan alamat IP:

kubectl get services

Outputnya mirip dengan hal berikut ini:

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

Anda dapat mengakses deployment Pathways tanpa penerusan port dari host di jaringan VPC yang sama dengan cluster Anda:

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

Ganti EXTERNAL_IP dengan nilai EXTERNAL-IP yang ditampilkan oleh kubectl get services—misalnya, 10.0.0.22.

Outputnya mirip dengan hal berikut ini:

[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

Catatan: Alamat IP Anda berbeda dengan alamat yang ditampilkan dalam contoh perintah dan output sebelumnya.

Cluster Director

Host di subnet cluster, seperti VM Compute Engine, dapat terhubung ke server proxy secara langsung. Firewall Cluster Director mengizinkan traffic dari subnet cluster, tetapi tidak dari subnet lain di jaringan VPC yang sama:

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

Ganti RM_IP dengan alamat proxy IFRT dari log layanan.

Catatan: Server proxy IFRT tidak mengautentikasi klien. Host mana pun yang dapat menjangkau RM_IP:29000 dapat menggunakan cluster Pathways.

Notebook Jupyter

GKE

Anda dapat membuat notebook Jupyter menggunakan Vertex AI, atau membuat notebook Jupyter yang dihosting sendiri.

Membuat instance Vertex AI Workbench

Setelah menyiapkan dan memverifikasi cluster Pathways, Anda dapat mengakses VM TPU GKE dari notebook Jupyter Vertex AI Workbench. Petunjuk penyiapan berikut mengasumsikan bahwa cluster GKE Pathways Anda berada di jaringan VPC yang sama (yang merupakan jaringan default, kecuali jika Anda telah mengonfigurasinya secara berbeda).

1. Buka halaman Vertex AI Workbench di konsol Google Cloud .

2. Di tab Instances, klik Create new, dan pastikan jaringan cocok dengan jaringan cluster GKE Anda.

   Atau, gunakan command line untuk membuat instance Vertex AI Workbench baru:

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

   Ganti kode berikut:

   • INSTANCE_NAME: nama instance Vertex AI Workbench Anda.
   • PROJECT_ID:Google Cloud project ID Anda.
   • ZONE: zona tempat Anda ingin membuat instance.
   • NETWORK: nama jaringan VPC cluster GKE Anda. Flag --network bersifat opsional; jika Anda menghilangkannya, perintah akan menggunakan jaringan default.

3. Setelah Vertex AI membuat instance, buka instance dan klik Open JupyterLab.

Membuat instance notebook Jupyter yang dihosting sendiri

Perintah berikut menunjukkan cara membuat instance notebook Jupyter yang dihosting sendiri menggunakan perintah gcluster job submit:

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"

Manifes YAML berikut menunjukkan cara membuat instance notebook Jupyter yang dihosting sendiri menggunakan kubectl. Terapkan manifes YAML berikut setelah Anda membuat cluster Pathways headless. Untuk mengetahui informasi selengkapnya, lihat Menjalankan Pathways dalam mode interaktif:

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

Kolom hostNetwork: true memungkinkan klien JAX di notebook menyelesaikan handshake Application Layer Transport Security (ALTS) dengan server proxy Pathways. Tanpa kolom ini, koneksi akan gagal pada node pool yang menggunakan Workload Identity Federation for GKE. Kolom dnsPolicy: ClusterFirstWithHostNet memungkinkan notebook me-resolve nama DNS dalam cluster, seperti nama Pod head Pathways. Untuk mengetahui informasi selengkapnya, lihat Memecahkan Masalah Pathways di Cloud.

Cluster Director

Anda dapat menjalankan notebook Jupyter di VM Compute Engine yang ada di subnet cluster, atau membuat notebook Jupyter yang dihosting sendiri di node login cluster. Dalam kedua kasus tersebut, notebook terhubung ke server proxy tanpa penerusan port.

Menggunakan VM yang ada di subnet cluster

Di VM, instal JupyterLab di lingkungan virtual dan mulai:

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

Membuat instance notebook Jupyter yang dihosting sendiri

Di node login cluster, instal JupyterLab dan mulai. Buat lingkungan virtual di direktori utama Anda dengan Python cluster. Python sistem node login terlalu lama untuk rilis JAX saat ini, sehingga pip tidak dapat menemukan jax==0.11.1 di sana:

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

Menghubungkan ke instance notebook

GKE

Hubungkan ke notebook dari komputer lokal Anda menggunakan penerusan port:

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

Di browser lokal Anda, buka http://localhost:8888?token=TOKEN. Ganti TOKEN dengan token dari log container notebook Jupyter:

kubectl logs ${MAIN_POD}

Outputnya mirip dengan hal berikut ini:

...
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 adalah nama host node GKE yang menjalankan Pod notebook (misalnya, gke-my-cluster-cpu-np-1a2b3c4d-ef56).

Cluster Director

Untuk notebook di VM atau di node login, hubungkan dari mesin lokal Anda dengan menggunakan penerusan port ke host tersebut. Misalnya, untuk node login:

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

Di browser lokal Anda, buka http://localhost:8888?token=TOKEN. Ganti TOKEN dengan token yang dicetak JupyterLab saat dimulai.

Konektivitas notebook ke cluster Pathways

1. Dari dalam JupyterLab, buat notebook Python 3 baru.
2. Hubungkan ke server proxy Pathways.

Di notebook, tambahkan sel untuk menginstal pathwaysutils, tetapkan JAX_PLATFORMS ke proxy, dan tetapkan JAX_BACKEND_TARGET ke grpc://PROXY_ADDRESS:29000:

Catatan: Nilai untuk PROXY_ADDRESS bergantung pada opsi koneksi yang Anda pilih di Menghubungkan ke cluster Pathways dalam mode interaktif. Di Cluster Director, instal jax==0.11.1 di sel yang sama. Pada notebook yang dihosting sendiri yang Anda mulai setelah source ~/pathways/cluster.env, variabelnya sudah ditetapkan.

!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

Ganti PROXY_ADDRESS dengan host atau alamat IP server proxy Pathways Anda dari Menghubungkan ke cluster Pathways dalam mode interaktif (misalnya, 127.0.0.1, nama host Cloud DNS, atau alamat IP load balancer internal). Di Cluster Director, gunakan 127.0.0.1 jika Anda menggunakan penerusan port, atau alamat proxy IFRT dari log layanan. Jika Anda men-deploy notebook yang dihosting sendiri menggunakan kubectl, ganti PROXY_ADDRESS dengan nama DNS dalam cluster Pod head Pathways, sehingga nilai JAX_BACKEND_TARGET adalah grpc://WORKLOAD-pathways-head-0-0.WORKLOAD:29000.

Tambahkan sel kedua untuk memverifikasi konektivitas dan mencetak perangkat di cluster Pathways:

import pathwaysutils
import jax

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

Jika koneksi berhasil, output akan menyertakan pesan yang menunjukkan bahwa JAX mendeteksi backend Pathways-on-Cloud.

Jumlah perangkat JAX yang tercantum sesuai dengan jumlah chip TPU dan jumlah slice yang Anda tentukan saat membuat cluster Pathways.

Menambahkan kode ke notebook

Tambahkan kode JAX Anda sendiri dan jalankan secara interaktif di TPU dalam cluster Pathways. Kode berikut menunjukkan cara melakukan komputasi di dua slice dari satu notebook. Untuk menjalankan contoh lengkap, buat workload Pathways Anda dengan setidaknya dua slice dengan menetapkan WORKLOAD_NODEPOOL_COUNT ke 2 atau lebih. Di Cluster Director, tetapkan MIN_SLICES dan MAX_SLICES ke 2 atau lebih. Dengan satu slice, slice1_devices kosong, dan bagian terakhir contoh gagal dengan 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)

Menghapus workload interaktif Pathways

GKE

Cluster Toolkit

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

kubectl

kubectl delete job WORKLOAD

Cluster Director

Di node login cluster, batalkan tugas layanan Pathways:

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

Perintah ini juga membatalkan tugas worker Pathways. VM TPU adalah node statis, sehingga terus berjalan dan tersedia untuk tugas Slurm lainnya.

Jika Anda mengalami masalah, lihat Memecahkan masalah Pathways di Cloud.

Langkah berikutnya

• Membuat cluster GKE dengan Pathways
• Membuat cluster Cluster Director dengan Pathways
• Menjalankan workload batch dengan Pathways
• Mem-porting workload JAX ke Pathways
• Pelatihan yang tangguh dengan Pathways
• Memecahkan masalah Pathways di Cloud