Esegui un carico di lavoro interattivo con Pathways

Anteprima

Questo prodotto o funzionalità è soggetto ai "Termini delle offerte pre-GA" nella sezione Termini di servizio generali dei Termini specifici dei servizi. I prodotti e le funzionalità pre-GA sono disponibili "così come sono" e potrebbero avere un supporto limitato. Per ulteriori informazioni, consulta le descrizioni delle fasi di lancio.

I workload interattivi di Pathways sono workload JAX remoti che vengono eseguiti all'interno di una VM che non fa parte del cluster GKE che ospita il cluster Pathways. A differenza dei workload batch, il completamento dell'operazione di workload interattivo non arresta i componenti del cluster Pathways, che rimangono disponibili per la connessione da parte di altri client JAX. Questo documento utilizza un blocco note Jupyter come esempio per dimostrare i workload interattivi.

Utilizzando l'interfaccia IFRT, gli utenti JAX inviano comandi a un cluster Pathways. Il codice JAX, eseguito da un terminale, un notebook o qualsiasi ambiente compatibile con Python, può interagire senza problemi con le risorse Pathways.

Prima di iniziare

Assicurati di avere:

• Creazione di un cluster GKE utilizzando XPK
• XPK installato
• Strumenti Kubernetes installati
• Installazione di gcloud CLI
• Abilitato l'API TPU
• Abilitato l'API Google Kubernetes Engine
• Assicurati che Pathways sia abilitato per il tuo progetto Google Cloud

Eseguire i percorsi in modalità interattiva

Puoi eseguire Pathways in modalità interattiva utilizzando xpk o kubectl.

XPK

1. Imposta le seguenti variabili di ambiente:

   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

   Sostituisci quanto segue:

   • WORKLOAD: impostalo su un nome univoco per identificare il workload
   • WORKLOAD_NODEPOOL_COUNT: il numero di node pool utilizzati da un workload Pathways
   • TPU_TYPE: il tipo di TPU specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per maggiori informazioni sui tipi di TPU supportati per ogni versione di TPU, consulta Versioni di TPU.
   • PROJECT: il tuo Google Cloud ID progetto
   • ZONE: la zona in cui prevedi di eseguire il tuo workload
   • CLUSTER: il nome del tuo cluster GKE

2. Crea i container Pathways sul cluster. Per eseguire un carico di lavoro headless, esegui questo comando:

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

A questo punto, il tuo workload JAX può connettersi al server proxy IFRT.

kubectl

Il seguente file YAML è identico al file YAML del workload batch, tranne per il fatto che non specifica il container main.

1. Sostituisci i segnaposto, copia il seguente codice YAML e incollalo in un file denominato 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
       

   Sostituisci quanto segue:
   • USERNAME : il tuo nome utente
   • MAX_RESTARTS : il numero massimo di riavvii del PathwaysJob
   • TPU_MACHINE_TYPE : il tipo di macchina TPU che vuoi utilizzare, ad esempio i valori supportati: "ct6e-standard-8t", "ct5p-hightpu-4t"
   • TOPOLOGY : la topologia TPU
   • WORKLOAD_NODEPOOL_COUNT : il numero di node pool utilizzati da un workload Pathways
   • BUCKET_NAME : un bucket Cloud Storage utilizzato per archiviare i file temporanei
   Per modificare il numero di node pool (repliche di pathways-worker) specificato da WORKLOAD_NODEPOOL_COUNT nel file YAML precedente, devi eliminare questo PathwaysJob e creare un nuovo PathwaysJob con il numero aggiornato di node pool. Devi anche riavviare i notebook connessi per stabilire la connessione con il nuovo cluster Pathways.
2. Applica il file pathways-headless-workload.yaml:

     kubectl apply -f pathways-headless-workload.yaml
     

3. Esegui kubectl get pods per verificare che tutti i container nel pod siano in esecuzione. L'output seguente è per una v5p 2x2x2 a 2 sezioni, dove USER è l'ID dell'utente che esegue il comando:

       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
       

Connessione al cluster Pathways in modalità interattiva

Puoi connetterti al cluster Pathways con o senza port forwarding. Utilizza una delle seguenti sezioni per connetterti al cluster Pathways.

Connettiti utilizzando il port forwarding

A questo punto puoi utilizzare l'inoltro delle porte (da qualsiasi host con accesso al piano di controllo del cluster) per accedere al server proxy:

Utilizza il comando appropriato per il tuo carico di lavoro:

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}

Dovresti visualizzare un output simile al seguente:

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}

Dovresti visualizzare un output simile al seguente:

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

Sullo stesso host, apri una nuova finestra del terminale. Imposta le variabili di ambiente JAX_PLATFORMS e JAX_BACKEND_TARGET ed esegui uno script Python che importa pathwaysutils e 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())'

Dovresti vedere un output simile al seguente:

[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

Connettiti dagli host nel VPC senza utilizzare l'inoltro delle porte

Se non vuoi utilizzare il port forwarding, puoi connetterti al cluster Pathways utilizzando Cloud DNS o un bilanciatore del carico interno.

Connettiti utilizzando Cloud DNS

L'attivazione di Cloud DNS nel cluster cambia il provider DNS da kube-dns a Cloud DNS. Se abilitata, viene creata una zona Cloud DNS privata in Virtual Private Cloud per i nomi Cloud DNS. Per maggiori informazioni, consulta Utilizzo di Cloud DNS per GKE.

Se abiliti Cloud DNS con l'ambito cluster, VPC additivo o VPC, i nomi Cloud DNS di Kubernetes sono risolvibili dalle VM non GKE all'interno di Virtual Private Cloud. I nomi hanno il formato <service_name>.<namespace>.svc.<custom_dns_domain>. Il pod principale di Pathways ha un servizio denominato <jobset_name>-pathways-head-0-0.<jobset_name>.<namespace>.svc.<custom_dns_domain>.

I seguenti comandi mostrano come connettersi al cluster Pathways utilizzando Cloud DNS:

1. Conferma che la voce Cloud DNS leader sia risolvibile da un host non GKE:

   XPK

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

   Dovresti visualizzare un output simile al seguente:

   <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

   Dovresti visualizzare un output simile al seguente:

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

2. Connettiti al cluster Pathways utilizzando il nome 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())'

   Dovresti visualizzare un output simile al seguente:

   [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

Connettiti utilizzando un bilanciatore del carico interno

Per un indirizzo IP privato nel tuo VPC che punta al deployment di percorsi, crea un servizio supportato da un bilanciatore del carico interno. Non è necessario che il cluster abbia Cloud DNS abilitato.

Per i cluster con molte VM, ti consigliamo di attivare il sottoinsieme ILB se stai creando bilanciatori del carico interni. Per maggiori informazioni, consulta Abilitare l'impostazione secondaria di GKE in un cluster esistente. Quando il sottoinsieme ILB non è abilitato, tutti i nodi del cluster faranno parte del gruppoistanza di backendd per tutti i bilanciatori del carico interni. Questa operazione non viene scalata oltre i 250 nodi. Con il subsetting ILB abilitato, GKE crea gruppi di endpoint di rete anziché gruppi di istanze e vengono inclusi solo i nodi che eseguono uno dei pod di servizio. L'abilitazione dell'impostazione secondaria del bilanciatore del carico interno ha una latenza di configurazione una tantum (~15 minuti). Il seguente comando mostra come abilitare il sottoinsieme ILB:

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

Una volta attivato il sottoinsieme ILB, puoi creare un servizio Kubernetes di tipo LoadBalancer utilizzando il seguente file YAML. In questo modo GKE creerà un bilanciatore del carico interno all'interno del VPC del cluster:

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

Aggiorna USER con il tuo ID utente Google Cloud e salva il file come pathways-headless-ilb.yaml.

Applica il manifest:

kubectl apply -f pathways-headless-ilb.yaml

Dopo la creazione del bilanciatore del carico (circa un minuto dopo), la colonna EXTERNAL-IP avrà un valore:

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

Puoi accedere al deployment dei percorsi senza port forwarding sugli host nello stesso VPC del cluster:

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

Nota: sostituisci l'indirizzo IP del backend JAX (10.0.0.22) nel comando precedente con il valore visualizzato da kubectl get services.

Dovresti visualizzare un output simile al seguente:

[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

Nota: i tuoi indirizzi IP saranno diversi da quelli mostrati nel comando precedente e negli esempi di output.

Blocchi note Jupyter

Puoi creare un notebook Jupyter utilizzando Vertex AI oppure puoi creare un notebook Jupyter self-hosted.

Crea un'istanza di Vertex AI Workbench

Dopo aver configurato e verificato il cluster Pathways, puoi accedere alle VM TPU GKE da un notebook Jupyter di Vertex AI. Le seguenti istruzioni di configurazione presuppongono che il cluster GKE Pathways si trovi nella stessa rete Virtual Private Cloud (che è la rete predefinita, a meno che tu non abbia configurato diversamente). Vai alla console Vertex AI Workbench.

Crea una nuova istanza di Workbench (dalla scheda Istanze) con il pulsante Crea nuova. Assicurati che la rete sia la stessa del cluster GKE. Puoi utilizzare la riga di comando per creare una nuova istanza di Workbench.

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

Nota: il flag --network è facoltativo. Se lo ometti, viene utilizzata la rete predefinita.

Una volta creata l'istanza, vai alla pagina e fai clic su Apri JupyterLab.

Crea un'istanza di blocco note Jupyter autogestita

Il seguente comando mostra come creare un'istanza di Jupyter Notebook self-hosted utilizzando 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"

Il seguente file YAML mostra come creare un'istanza di Jupyter Notebook autogestita utilizzando kubectl. Applica il seguente file YAML dopo aver creato un cluster Pathways headless. Per saperne di più, consulta Esegui percorsi in modalità interattiva con 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

Vai all'istanza del notebook

Connettiti al notebook dalla tua macchina locale utilizzando il port forwarding:

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

Nel browser locale, vai a http://localhost:8888?token=<var>your-token</var>. Sostituisci <your-token> con il token dei log del container del notebook Jupyter.

kubectl logs ${MAIN_POD}

che dovrebbe restituire:

...
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>

Connettività del notebook al cluster Pathways

1. Da JupyterLab, crea un nuovo notebook Python 3.
2. Connettiti al server proxy di Pathways

Nel blocco note, aggiungi una cella per installare pathwaysutils, imposta JAX_PLATFORMS su proxy e imposta JAX_BACKEND_TARGET su PROXY_ADDRESS.

Nota :il valore di PROXY_ADDRESS dipende dall'opzione di connessione scelta in Connetti al server proxy di Pathways.

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

Aggiungi una seconda cella come controllo di tipo "Hello World" e stampa i dispositivi nel cluster Pathways.

import pathwaysutils
import jax

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

Se tutto funziona correttamente, dovresti visualizzare un messaggio che indica che è stato rilevato il backend di Pathways on Cloud.

Il numero di dispositivi JAX elencati deve corrispondere al numero di chip TPU e al numero di sezioni specificate al momento della creazione del cluster Pathways.

Aggiungere il codice a un notebook

Aggiungi il tuo codice JAX ed eseguilo in modo interattivo sulle TPU nel cluster Pathways. Il seguente codice mostra come eseguire calcoli su due sezioni da un singolo blocco note.

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)

Eliminare il cluster interattivo di Pathways

XPK

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

kubectl

kubectl delete -f pathways-headless-workload.yaml

Passaggi successivi

• Crea un cluster GKE con Pathways
• Inferenza multihost con Pathways
• Carichi di lavoro batch con percorsi
• Modalità interattiva di Pathways
• Portare i carichi di lavoro JAX su Pathways
• Formazione resiliente con Pathways
• Percorsi di risoluzione dei problemi