Google utilizza la tecnologia AI per tradurre i contenuti nella tua lingua preferita. Le traduzioni generate dall'AI potrebbero contenere errori.

Esegui un carico di lavoro interattivo con Pathways

I workload interattivi di Pathways sono workload JAX remoti 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 del 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 illustrare 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:

Creato un cluster GKE con Pathways
Cluster Toolkit installato
Strumenti Kubernetes installati
gcloud CLI installata
Abilitato l'API Google Kubernetes Engine
Assicurati che Pathways sia abilitato per il tuo progetto Google Cloud

Esegui Pathways in modalità interattiva

Imposta le seguenti variabili di ambiente:
```
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
```
Sostituisci quanto segue:
- WORKLOAD: imposta un nome univoco per identificare il workload (massimo 22 caratteri a causa dei limiti delle etichette DNS)
- WORKLOAD_NODEPOOL_COUNT: il numero di node pool utilizzati da un workload Pathways
- COMPUTE_TYPE: il tipo di macchina TPU (ad esempio ct6e-standard-4t). Per ulteriori informazioni sui tipi di TPU supportati per ogni versione di TPU, consulta Versioni di TPU.
- TOPOLOGY: la topologia di posizionamento della TPU (ad esempio, 2x4)
- PROJECT_ID: il tuo ID progetto Google Cloud
- REGION: la regione in cui si trovano il bucket e il cluster
- ZONE: la zona in cui prevedi di eseguire il workload
- CLUSTER: il nome del tuo cluster GKE
- BUCKET_NAME: il bucket Cloud Storage per gli artefatti di Pathways
Se non hai ancora creato un bucket Cloud Storage per gli artefatti di Pathways, creane uno:
```
gcloud storage buckets create gs://${BUCKET_NAME} --location=${REGION}
```

Crea i container Pathways nel cluster. Per eseguire un workload headless, utilizza il comando gcluster job submit (per ulteriori informazioni, consulta la guida all'invio dei job di 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

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

Connessione al cluster Pathways in modalità interattiva

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

Connettiti utilizzando il port forwarding

A questo punto puoi utilizzare il port forwarding (da qualsiasi host con accesso al control plane del cluster) per accedere al server proxy:

Utilizza il comando appropriato per il tuo workload:

Cluster Toolkit

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 il port forwarding

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'abilitazione di Cloud DNS nel cluster cambia il provider DNS Cloud da kube-dns a Cloud DNS. Se abilitata, viene creata una zona Cloud DNS privata nel tuo Virtual Private Cloud per i nomi Cloud DNS. Per saperne di più, 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 del tuo Virtual Private Cloud. I nomi hanno il formato <service_name>.<namespace>.svc.<custom_dns_domain>. Il pod head di Pathways ha un servizio denominato <jobset_name>-pathways-head-0-0.<jobset_name>.<namespace>.svc.<custom_dns_domain>.

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

Verifica che la voce Cloud DNS leader sia risolvibile da un host non GKE:

Cluster Toolkit

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

Connettiti al cluster Pathways utilizzando il nome Cloud DNS:

Cluster Toolkit

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 VPC che punta al deployment di Pathways, crea un servizio supportato da un bilanciatore del carico interno. Non è necessario che Cloud DNS sia abilitato nel cluster.

Per i cluster con molte VM, ti consigliamo di attivare la suddivisione in sottoinsiemi ILB se stai creando bilanciatori del carico interni. Per saperne di più, consulta Abilita il subsetting GKE in un cluster esistente. Quando la suddivisione in sottoinsiemi ILB non è abilitata, tutti i nodi del cluster faranno parte del gruppo di istanza di backend per tutti i bilanciatori del carico interni. Non scala oltre i 250 nodi. Con la suddivisione in sottoinsiemi ILB abilitata, GKE crea gruppi di endpoint di rete anziché gruppi di istanze e vengono inclusi solo i nodi che eseguono uno dei pod di distribuzione del servizio. L'abilitazione della suddivisione in sottoinsiemi ILB prevede una latenza di configurazione una tantum (~15 minuti). Il seguente comando mostra come abilitare la suddivisione in sottoinsiemi ILB:

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

Una volta abilitato il subsetting ILB, puoi creare un servizio Kubernetes di tipo LoadBalancer utilizzando il seguente yaml. In questo modo, GKE creerà un bilanciatore del carico interno nel 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 di Pathways senza port forwarding sugli host nello stesso VPC del tuo 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())'

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

Blocchi note Jupyter

Puoi creare un blocco note Jupyter utilizzando Vertex AI o un blocco note Jupyter in hosting autonomo.

Crea un'istanza di Vertex AI Workbench

Dopo aver configurato e verificato il cluster Pathways, puoi accedere alle VM TPU GKE da un blocco note 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 di Vertex AI Workbench.

Crea una nuova istanza Workbench (dalla scheda Istanze) con il pulsante Crea nuova. Assicurati che la rete sia la stessa del tuo 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]

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

Crea un'istanza di blocco note Jupyter self-hosted

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

Il seguente YAML mostra come creare un'istanza di blocco note Jupyter self-hosted utilizzando kubectl. Applica il seguente YAML dopo aver creato un cluster Pathways headless. Per saperne di più, consulta Esegui Pathways 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 blocco note

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

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 | 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 blocco note Jupyter.

kubectl logs ${MAIN_POD}

L'output dovrebbe essere:

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

In JupyterLab, crea un nuovo notebook Python 3
Connettiti al server proxy Pathways

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

!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 Pathways-on-Cloud.

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

Aggiungi il codice a un notebook

Aggiungi il tuo codice JAX ed eseguilo in modo interattivo sulle TPU nel cluster Pathways. Il codice seguente mostra come eseguire calcoli su due sezioni da un singolo 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)

Elimina il workload interattivo di Pathways

Cluster Toolkit

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

kubectl

kubectl delete -f pathways-headless-workload.yaml

Esegui un carico di lavoro interattivo con Pathways Mantieni tutto organizzato con le raccolte Salva e classifica i contenuti in base alle tue preferenze.

Prima di iniziare

Esegui Pathways in modalità interattiva

Connessione al cluster Pathways in modalità interattiva

Connettiti utilizzando il port forwarding

Cluster Toolkit

kubectl

Connettiti dagli host nel VPC senza utilizzare il port forwarding

Connettiti utilizzando Cloud DNS

Cluster Toolkit

kubectl

Cluster Toolkit

kubectl

Connettiti utilizzando un bilanciatore del carico interno

Blocchi note Jupyter

Crea un'istanza di Vertex AI Workbench

Crea un'istanza di blocco note Jupyter self-hosted

Vai all'istanza del blocco note

Cluster Toolkit

kubectl

Connettività del notebook al cluster Pathways

Aggiungi il codice a un notebook

Elimina il workload interattivo di Pathways

Cluster Toolkit

kubectl

Passaggi successivi

Esegui un carico di lavoro interattivo con Pathways