Menjalankan workload interaktif dengan Pathways

Pratinjau

Produk atau fitur ini tunduk pada "Persyaratan Penawaran Pra-GA" di bagian Persyaratan Layanan Umum dari Persyaratan Khusus Layanan. Produk dan fitur pra-GA tersedia "sebagaimana adanya" dan mungkin memiliki dukungan terbatas. Untuk mengetahui informasi selengkapnya, lihat deskripsi tahap peluncuran.

Workload interaktif Pathways adalah workload JAX jarak jauh yang berjalan dalam VM yang bukan bagian dari cluster GKE yang menghosting cluster Pathways. Tidak seperti workload batch, penyelesaian operasi workload interaktif tidak akan mematikan komponen cluster Pathways, yang tetap tersedia untuk koneksi oleh klien JAX lainnya. Dokumen ini menggunakan notebook Jupyter sebagai contoh untuk mendemonstrasikan workload interaktif.

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

Sebelum memulai

Pastikan Anda memiliki:

• Membuat cluster GKE menggunakan XPK
• XPK yang diinstal
• Alat Kubernetes yang terinstal
• Menginstal gcloud CLI
• Mengaktifkan TPU API
• Mengaktifkan Google Kubernetes Engine API
• Pastikan Pathways diaktifkan untuk project Google Cloud Anda

Menjalankan Pathways dalam mode interaktif

Anda dapat menjalankan Pathways dalam mode interaktif menggunakan xpk atau kubectl.

XPK

1. Tetapkan variabel lingkungan berikut:

   export WORKLOAD=WORKLOAD
   export WORKLOAD_NODEPOOL_COUNT=WORKLOAD_NODEPOOL_COUNT
   export TPU_TYPE=TPU_TYPE
   export PROJECT_ID=PROJECT
   export ZONE=ZONE \
   export CLUSTER=CLUSTER

   Ganti kode berikut:

   • WORKLOAD: tetapkan ini ke nama unik untuk mengidentifikasi workload Anda
   • WORKLOAD_NODEPOOL_COUNT: jumlah node pool yang digunakan oleh beban kerja Pathways
   • TPU_TYPE: jenis TPU menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis TPU yang didukung untuk setiap versi TPU, lihat Versi TPU
   • PROJECT: Project ID Google Cloud Anda
   • ZONE: zona tempat Anda berencana menjalankan workload
   • CLUSTER: nama cluster GKE Anda

2. Buat container Pathways di cluster. Untuk menjalankan beban kerja tanpa tampilan, jalankan perintah berikut:

   xpk workload create-pathways \
   --headless \
   --workload=${WORKLOAD} \
   --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
   --tpu-type=${TPU_TYPE} \
   --project=${PROJECT} \
   --zone=${ZONE} \
   --cluster=${CLUSTER}

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

kubectl

YAML berikut sama dengan YAML workload batch, kecuali tidak menentukan container main.

1. Ganti placeholder, salin YAML berikut, lalu tempel ke dalam file bernama pathways-headless-workload.yaml.

   apiVersion: pathways-job.pathways.domain/v1
   kind: PathwaysJob
   metadata:
     name: pathways-USERNAME
   spec:
     maxRestarts: MAX_RESTARTS
     workers:
       - type: TPU_MACHINE_TYPE
         topology: TOPOLOGY
         numSlices: WORKLOAD_NODEPOOL_COUNT
     pathwaysDir: gs://BUCKET_NAME
     controller:
       deploymentMode: default
       

   Ganti kode berikut:
   • USERNAME : nama pengguna Anda
   • MAX_RESTARTS : jumlah maksimum PathwaysJob dapat dimulai ulang
   • TPU_MACHINE_TYPE : jenis mesin TPU yang ingin Anda gunakan, contoh nilai yang didukung: "ct6e-standard-8t", "ct5p-hightpu-4t"
   • TOPOLOGY : topologi TPU
   • WORKLOAD_NODEPOOL_COUNT : jumlah node pool yang digunakan oleh beban kerja Pathways
   • BUCKET_NAME : bucket Cloud Storage yang digunakan untuk menyimpan file sementara
   Untuk mengubah jumlah node pool, (replika pathways-worker) yang ditentukan oleh WORKLOAD_NODEPOOL_COUNT di YAML sebelumnya, Anda harus menghapus PathwaysJob ini dan membuat PathwaysJob baru dengan jumlah node pool yang diperbarui. Anda juga perlu memulai ulang notebook yang terhubung untuk membuat koneksi dengan cluster Pathways baru.
2. Terapkan file pathways-headless-workload.yaml:

     kubectl apply -f pathways-headless-workload.yaml
     

3. Jalankan kubectl get pods untuk memeriksa apakah semua container di Pod sedang berjalan. Output berikut adalah untuk v5p 2x2x2 2 slice, dengan USER adalah ID pengguna yang menjalankan perintah:

       NAME                                         READY   STATUS    RESTARTS   AGE
       pathways-USER-pathways-head-0-0-n848j      2/2     Running   0          49s
       pathways-USER-pathways-workers-0-0-jxt2z   1/1     Running   0          71s
       pathways-USER-pathways-workers-0-1-cxmhc   1/1     Running   0          70s
       pathways-USER-pathways-workers-1-0-5kmz9   1/1     Running   0          71s
       pathways-USER-pathways-workers-1-1-vg5n4   1/1     Running   0          71s
       

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.

Menghubungkan menggunakan penerusan port

Pada tahap ini, Anda dapat menggunakan penerusan port (dari host mana pun yang memiliki akses ke bidang kontrol cluster Anda) untuk mengakses server proxy:

Gunakan perintah yang sesuai untuk beban kerja Anda:

XPK

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

Anda akan melihat output yang mirip dengan:

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

kubectl

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

Anda akan melihat output yang mirip dengan:

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

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:

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

Anda akan melihat output seperti berikut:

[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

Menghubungkan dari host di VPC tanpa menggunakan penerusan port

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

Menghubungkan menggunakan Cloud DNS

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

Jika Anda mengaktifkan Cloud DNS dengan cakupan cluster, VPC aditif, atau VPC, nama Cloud DNS Kubernetes dapat di-resolve dari VM non-GKE di dalam Virtual Private Cloud Anda. Nama memiliki format <service_name>.<namespace>.svc.<custom_dns_domain>. Pod head Pathways memiliki layanan bernama <jobset_name>-pathways-head-0-0.<jobset_name>.<namespace>.svc.<custom_dns_domain>.

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

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

   XPK

   host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.USERNAME-test

   Anda akan melihat output yang mirip dengan:

   <WORKLOAD>-pathways-head-0-0.<WORKLOAD>.default.svc.<user>-test has address 10.0.2.75

   kubectl

   host pathways-USERNAME-pathways-head-0-0.pathways-USERNAME.default.svc.USERNAME-test

   Anda akan melihat output yang mirip dengan:

   pathways-<user>-pathways-head-0-0.pathways-<user>.default.svc.<user>-test has address 10.0.2.75

2. Hubungkan ke cluster Pathways menggunakan nama Cloud DNS:

   XPK

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

   kubectl

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

   Anda akan melihat output yang mirip dengan:

   [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 VPC Anda yang mengarah ke deployment jalur pembelajaran Anda, buat layanan yang didukung oleh load balancer internal. Anda tidak perlu mengaktifkan Cloud DNS di cluster.

Untuk cluster dengan banyak VM, sebaiknya aktifkan subsetelan ILB jika Anda membuat load balancer internal. Untuk mengetahui informasi selengkapnya, lihat Mengaktifkan subsetelan GKE di cluster yang ada. Jika subset ILB tidak diaktifkan, semua node dalam cluster akan menjadi bagian dari grup instance backend untuk semua load balancer internal. Hal ini tidak dapat diskalakan lebih dari 250 node. Dengan subset ILB yang diaktifkan, GKE membuat grup endpoint jaringan, bukan grup instance, dan hanya node yang menjalankan salah satu Pod penayangan layanan yang disertakan. Pengaktifan subset ILB memiliki latensi penyiapan satu kali (~15 menit). Perintah berikut menunjukkan cara mengaktifkan subsetting ILB:

gcloud container clusters update ${CLUSTER} \
  --project=${PROJECT} \
  [--zone=${ZONE} | --region=${REGION}] \
  --enable-l4-ilb-subsetting

Setelah subset ILB diaktifkan, Anda dapat membuat layanan Kubernetes jenis LoadBalancer menggunakan YAML berikut. Tindakan ini akan menyebabkan GKE membuat load balancer internal di dalam VPC cluster Anda:

apiVersion: v1
kind: Service
metadata:
  name: pathways-USERNAME-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: pathways-USER
    jobset.sigs.k8s.io/replicatedjob-name: pathways-head
  ports:
  - name: tcp-port
    protocol: TCP
    port: 29000
    targetPort: 29000

Perbarui USER dengan ID pengguna Google Cloud Anda, lalu simpan file sebagai pathways-headless-ilb.yaml.

Terapkan manifes:

kubectl apply -f pathways-headless-ilb.yaml

Setelah load balancer dibuat (sekitar satu menit kemudian), kolom EXTERNAL-IP akan memiliki nilai:

kubectl get services
NAME                  TYPE           CLUSTER-IP      EXTERNAL-IP   PORT(S)        AGE
pathways-$USER       ClusterIP      None            <none>        <none>         30m
pathways-$USER-ilb   LoadBalancer   34.118.232.46   10.0.0.22     80:31246/TCP   2m41s

Anda dapat mengakses deployment jalur tanpa penerusan port di host dalam VPC yang sama dengan cluster Anda:

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

Catatan: Ganti alamat IP backend JAX (10.0.0.22) dalam perintah sebelumnya dengan nilai yang ditampilkan oleh kubectl get services.

Anda akan melihat output yang mirip dengan:

[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 akan berbeda dengan yang ditampilkan dalam perintah sebelumnya dan contoh output.

Notebook Jupyter

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. Petunjuk penyiapan berikut mengasumsikan bahwa cluster GKE Pathways Anda berada di jaringan Virtual Private Cloud yang sama (yang merupakan jaringan default kecuali jika Anda telah mengonfigurasi sebaliknya). Buka konsol Vertex AI Workbench.

Buat instance Workbench baru (dari tab Instances) dengan tombol Create new. Pastikan jaringan sama dengan jaringan cluster GKE Anda. Anda dapat menggunakan command line untuk membuat instance Workbench baru.

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

Catatan: flag --network bersifat opsional. Jika Anda menghapusnya, jaringan default akan digunakan.

Setelah instance dibuat, buka instance tersebut, lalu klik Open Jupyterlab.

Membuat instance notebook Jupyter yang dihosting sendiri

Perintah berikut menunjukkan cara membuat instance notebook Jupyter yang dihosting sendiri menggunakan XPK:

xpk workload create-pathways \
--workload=${WORKLOAD} \
--num-slices=${WORKLOAD_NODEPOOL_COUNT} \
--tpu-type=${TPU_TYPE} \
--project=${PROJECT} \
--zone=${ZONE} \
--cluster=${CLUSTER} \
--docker-image=jupyter/base-notebook \
--command "start-notebook.sh"

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

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

Buka instance notebook

Hubungkan ke notebook dari mesin lokal Anda menggunakan penerusan port:

XPK

  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 | grep jupyter-notebook-USERNAME | awk '{print $1}')
  kubectl port-forward pod/${MAIN_POD} 8888:8888

Buka http://localhost:8888?token=<var>your-token</var> di browser lokal Anda. Ganti <your-token> dengan token dari log penampung notebook Jupyter.

kubectl logs ${MAIN_POD}

Yang akan menghasilkan output:

...
Or copy and paste one of these URLs:
  http://jupyter-notebook-<user>-bbbdh:8888/lab?token=<token>
  http://127.0.0.1:8888/lab?token=<token>

Konektivitas notebook ke cluster Pathways

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

Di notebook, tambahkan sel untuk menginstal pathwaysutils, tetapkan JAX_PLATFORMS ke proxy, dan tetapkan JAX_BACKEND_TARGET ke PROXY_ADDRESS.

Catatan: Nilai untuk PROXY_ADDRESS bergantung pada opsi koneksi yang Anda pilih di Menghubungkan ke server proxy Pathways.

!pip install pathwaysutils
%env JAX_PLATFORMS=proxy
# Replace your proxy address below:
%env JAX_BACKEND_TARGET=PROXY_ADDRESS

Tambahkan sel kedua sebagai pemeriksaan jenis "hello world" dan cetak perangkat di cluster Pathways.

import pathwaysutils
import jax

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

Jika semuanya berfungsi dengan baik, Anda akan melihat pesan yang menunjukkan bahwa backend Pathways-on-Cloud terdeteksi.

Jumlah perangkat JAX yang tercantum harus cocok 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.

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.
jax.jit(jnp.sin, device=jax.devices()[-1])(np.pi / 2.)

# 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 cluster interaktif Pathways Anda

XPK

xpk workload delete --workload=WORKLOAD --cluster=CLUSTER --project=PROJECT --zone=ZONE

kubectl

kubectl delete -f pathways-headless-workload.yaml

Langkah berikutnya

• Membuat Cluster GKE dengan Pathways
• Inferensi multihost dengan Pathways
• Batch workload dengan Pathways
• Mode interaktif jalur
• Memindahkan beban kerja JAX ke Pathways
• Pelatihan yang tangguh dengan Pathways
• Pemecahan Masalah Jalur