Esegui un carico di lavoro interattivo con Pathways

I workload interattivi di Pathways sono workload JAX remoti eseguiti all'esterno dei componenti del controller Pathways del cluster, ad esempio su una VM. A differenza dei workload batch, quando un workload interattivo termina, non arresta i componenti del cluster Pathways, che rimangono disponibili per le connessioni da altri client JAX. Questo documento utilizza un blocco note Jupyter come esempio per illustrare i workload interattivi. Puoi eseguire workload interattivi con Pathways su un cluster Google Kubernetes Engine (GKE) o su un cluster Cluster Director con Slurm.

Questa pagina contiene le istruzioni per entrambe le piattaforme. In ogni sezione, seleziona la scheda GKE o Cluster Director.

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 con le risorse Pathways.

Prima di iniziare

• Seleziona la scheda relativa al modo in cui prevedi di utilizzare gli esempi in questa pagina:

  gcloud

  Installa Google Cloud CLI, quindi accedi a gcloud CLI con la tua identità federata. Dopo aver eseguito l'accesso, inizializza Google Cloud CLI eseguendo il comando seguente:

  gcloud init

  REST

  Per utilizzare gli esempi di API REST in questa pagina in un ambiente di sviluppo locale, utilizzi le credenziali che fornisci a gcloud CLI.

  Installa Google Cloud CLI, quindi accedi a gcloud CLI con la tua identità federata.

  Per saperne di più, consulta Autenticati per usare REST nella documentazione sull'autenticazione di Google Cloud .

Prima di iniziare

GKE

Completa i seguenti prerequisiti:

• Crea un cluster GKE con Pathways.
• Installa Cluster Toolkit.
• Installa gli strumenti Kubernetes.
• Installa Google Cloud CLI.
• Abilita l'API Kubernetes Engine.
• Assicurati che Pathways sia abilitato per il tuo progetto Google Cloud .

Cluster Director

Completa i seguenti prerequisiti:

• Crea un cluster Cluster Director e configura l'ambiente Pathways.
• Installa gcloud CLI.
• Verifica di avere accesso SSH al nodo di accesso del cluster.
• Installa Python 3 con il modulo venv sull'host che esegue il client JAX. Su Debian e Ubuntu, installa il pacchetto python3-venv.

Esegui Pathways in modalità interattiva

GKE

1. 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: 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 progettoGoogle 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 Pathways.

2. Se non hai ancora creato un bucket Cloud Storage per gli artefatti Pathways, creane uno:

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

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

   Il comando gcluster job submit cerca un pool di nodi CPU denominato cpu-np o pathways-np per eseguire il pod head di Pathways. Se il tuo pool di nodi CPU ha un nome diverso (ad esempio cpu-pathways-np dalla guida alla creazione del cluster), aggiungi il flag --pathways-head-np=NODE_POOL_NAME ai comandi gcluster job submit in questa pagina. In caso contrario, il comando non va a buon fine e viene visualizzato l'errore failed to auto-detect a suitable CPU node pool.

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

Cluster Director

1. Connettiti al nodo di accesso del cluster:

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

   Sostituisci quanto segue:

   • CLUSTER_ID: il nome del tuo cluster Cluster Director.
   • PROJECT_ID: il tuo ID progettoGoogle Cloud .
   • ZONE: la zona in cui si trova il cluster.

2. Imposta le seguenti variabili di ambiente:

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

   Sostituisci quanto segue:

   • TOPOLOGY: la topologia TPU di ogni sezione, ad esempio 2x4. Deve corrispondere alla topologia del set di nodi TPU.
   • NUM_SLICES: il numero di sezioni TPU. Imposta MIN_SLICES e MAX_SLICES sullo stesso valore in modo che il cluster Pathways sia pronto solo quando tutte le sezioni sono in esecuzione.
   • BUCKET_NAME: un bucket Cloud Storage esistente per gli artefatti Pathways.

   Se la partizione TPU non è la partizione predefinita del cluster, imposta anche PARTITION sul suo nome.

3. Avvia il cluster Pathways in modalità headless:

   cd ~/pathways
   sbatch start_pathways.sh
   

4. Attendi finché il log del servizio, pathways-service-JOB_ID.out (dove JOB_ID è l'ID job Slurm), non mostra un output simile al seguente:

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

   RM_IP è l'indirizzo IP interno del server proxy. Il file ~/pathways/cluster.env contiene lo stesso indirizzo.

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

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

In Cluster Director, il client JAX deve utilizzare la stessa versione JAX del server Pathways, 0.11.1.

Connettiti utilizzando il port forwarding

GKE

Da qualsiasi host con accesso al control plane del cluster, inoltra la porta locale 29000 al pod head di Pathways per accedere al 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}

L'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. Utilizzando Python 3.11 o versioni successive, crea un ambiente virtuale, installa i pacchetti pathwaysutils e jax, 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())'

L'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

Cluster Director

Da qualsiasi host che può connettersi tramite SSH al nodo di accesso del cluster, ad esempio la tua workstation o una VM al di fuori della subnet del cluster, inoltra la porta proxy tramite il nodo di accesso:

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

Nota: se esegui gcloud CLI su una VM di Compute Engine con il suo account di servizio predefinito, la VM deve avere l'ambito di accesso cloud-platform, altrimenti gcloud compute ssh non va a buon fine con insufficient authentication scopes. In alternativa, esegui gcloud auth login sulla VM per utilizzare il tuo account.

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. L'host non ha un percorso per RM_IP, quindi il client si connette all'estremità locale del 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())'

L'output è simile al seguente:

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

L'avviso IFRT proxy server disconnected viene visualizzato quando lo script esce e chiude la connessione. È un comportamento previsto.

Nota: ogni chiamata JAX passa attraverso il tunnel SSH. Utilizza il port forwarding per verificare la connettività. Per i notebook, esegui il client JAX nella subnet del cluster.

Connettiti dagli host nel VPC senza utilizzare il port forwarding

GKE

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 da kube-dns a Cloud DNS. Se abilitato, GKE crea una zona Cloud DNS privata nella tua rete Virtual Private Cloud (VPC) per i nomi DNS. Per saperne di più, consulta Utilizzo di Cloud DNS per GKE.

Se abiliti Cloud DNS con ambito VPC o ambito VPC aggiuntivo, i nomi DNS di Kubernetes sono risolvibili dalle VM non GKE all'interno della tua rete VPC. Con l'ambito del cluster, solo i nodi nel cluster possono risolvere questi nomi. I nomi hanno il formato SERVICE_NAME.NAMESPACE.svc.CUSTOM_DNS_DOMAIN. Il pod head di Pathways ha un servizio denominato WORKLOAD-pathways-head-0-0.WORKLOAD.NAMESPACE.svc.CUSTOM_DNS_DOMAIN.

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

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

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

   Sostituisci quanto segue:

   • WORKLOAD: il nome del tuo workload Pathways.
   • CUSTOM_DNS_DOMAIN: il dominio DNS personalizzato configurato per il tuo cluster.

   L'output è simile al seguente:

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

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

   L'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 nella tua rete VPC che punta al tuo deployment Pathways, crea un servizio supportato da un bilanciatore del carico interno (ILB). Non è necessario che nel cluster sia abilitato Cloud DNS.

Per i cluster con molte VM, attiva la suddivisione in sottoinsiemi ILB se crei bilanciatori del carico interni. Per saperne di più, consulta Abilita il subsetting GKE in un cluster esistente. Se la suddivisione in sottoinsiemi ILB non è abilitata, tutti i nodi del cluster fanno parte del gruppo di istanza di backend per tutti i bilanciatori del carico interni, che 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 include solo i nodi che eseguono uno dei pod di distribuzione del servizio. L'attivazione della suddivisione in sottoinsiemi ILB ha una latenza di configurazione una tantum di circa 15 minuti. Il seguente comando attiva la suddivisione in sottoinsiemi ILB:

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

Sostituisci CLUSTER_LOCATION con la regione o la zona in cui si trova il cluster.

Dopo aver attivato la creazione di sottoinsiemi ILB, crea un servizio Kubernetes di tipo LoadBalancer utilizzando il seguente manifest YAML. In questo modo, GKE crea un bilanciatore del carico interno nella rete VPC del cluster:

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

Sostituisci WORKLOAD con il nome del tuo workload Pathways e salva il file come pathways-headless-ilb.yaml.

Applica il manifest:

kubectl apply -f pathways-headless-ilb.yaml

Dopo che GKE ha creato il bilanciatore del carico (circa un minuto dopo), verifica che la colonna EXTERNAL-IP mostri un indirizzo IP:

kubectl get services

L'output è simile al seguente:

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

Puoi accedere al deployment di Pathways senza port forwarding dagli host nella stessa rete VPC del cluster:

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

Sostituisci EXTERNAL_IP con il valore EXTERNAL-IP visualizzato da kubectl get services, ad esempio 10.0.0.22.

L'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 sono diversi da quelli mostrati negli esempi di comandi e output precedenti.

Cluster Director

Gli host nella subnet del cluster, ad esempio una VM di Compute Engine, possono connettersi direttamente al server proxy. Il firewall di Cluster Director consente il traffico dalla subnet del cluster, ma non da altre subnet della stessa rete VPC:

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

Sostituisci RM_IP con l'indirizzo proxy IFRT del log del servizio.

Nota: il server proxy IFRT non autentica i client. Qualsiasi host che può raggiungere RM_IP:29000 può utilizzare il cluster Pathways.

Blocchi note Jupyter

GKE

Puoi creare un notebook Jupyter utilizzando Vertex AI o un notebook 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 notebook Jupyter di Vertex AI Workbench. Le seguenti istruzioni di configurazione presuppongono che il cluster GKE Pathways si trovi nella stessa rete VPC (che è la rete predefinita, a meno che tu non abbia configurato diversamente).

1. Vai alla pagina Vertex AI Workbench nella console Google Cloud .

2. Nella scheda Istanze, fai clic su Crea nuova e assicurati che la rete corrisponda a quella del tuo cluster GKE.

   In alternativa, utilizza la riga di comando per creare una nuova istanza di Vertex AI Workbench:

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

   Sostituisci quanto segue:

   • INSTANCE_NAME: il nome della tua istanza Vertex AI Workbench.
   • PROJECT_ID: il tuo ID progettoGoogle Cloud .
   • ZONE: la zona in cui vuoi creare l'istanza.
   • NETWORK: il nome della rete VPC del tuo cluster GKE. Il flag --network è facoltativo. Se lo ometti, il comando utilizza la rete default.

3. Dopo che Vertex AI ha creato l'istanza, vai all'istanza 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 manifest YAML mostra come creare un'istanza di blocco note Jupyter in hosting autonomo utilizzando kubectl. Applica il seguente manifest YAML dopo aver creato un cluster Pathways headless. Per saperne di più, consulta Esegui Pathways in modalità interattiva:

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

Il campo hostNetwork: true consente al client JAX nel notebook di completare l'handshake ALTS (Application Layer Transport Security) con il server proxy Pathways. Senza questo campo, la connessione non riesce sui node pool che utilizzano Workload Identity Federation for GKE. Il campo dnsPolicy: ClusterFirstWithHostNet consente al notebook di risolvere i nomi DNS nel cluster, ad esempio il nome del pod head di Pathways. Per saperne di più, consulta Risolvi i problemi relativi a Pathways on Cloud.

Cluster Director

Puoi eseguire un blocco note Jupyter su una VM di Compute Engine esistente nella subnet del cluster oppure puoi creare un blocco note Jupyter in hosting autonomo sul nodo di accesso del cluster. In entrambi i casi, il notebook si connette al server proxy senza port forwarding.

Utilizza una VM esistente nella subnet del cluster

Nella VM, installa JupyterLab in un ambiente virtuale e avvialo:

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

Crea un'istanza di blocco note Jupyter self-hosted

Sul nodo di accesso del cluster, installa JupyterLab e avvialo. Crea l'ambiente virtuale nella tua home directory con Python del cluster. Il Python di sistema del nodo di accesso è troppo vecchio per le release JAX attuali, quindi pip non riesce a trovare jax==0.11.1:

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

Connettiti all'istanza del blocco note

GKE

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 \
    --selector=batch.kubernetes.io/job-name=WORKLOAD \
    --output=jsonpath='{.items[0].metadata.name}')
kubectl port-forward pod/${MAIN_POD} 8888:8888

Nel browser locale, vai a http://localhost:8888?token=TOKEN. Sostituisci TOKEN con il token dei log del container del blocco note Jupyter:

kubectl logs ${MAIN_POD}

L'output è simile al seguente:

...
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 è il nome host del nodo GKE che esegue il pod del blocco note (ad esempio, gke-my-cluster-cpu-np-1a2b3c4d-ef56).

Cluster Director

Per un notebook su una VM o sul nodo di accesso, connettiti dalla tua macchina locale utilizzando il port forwarding a quell'host. Ad esempio, per il nodo di accesso:

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

Nel browser locale, vai a http://localhost:8888?token=TOKEN. Sostituisci TOKEN con il token stampato da JupyterLab all'avvio.

Connettività del notebook al cluster Pathways

1. In JupyterLab, crea un nuovo blocco note Python 3.
2. Connettiti al server proxy Pathways.

Nel notebook, aggiungi una cella per installare pathwaysutils, imposta JAX_PLATFORMS su proxy e imposta JAX_BACKEND_TARGET su grpc://PROXY_ADDRESS:29000:

Nota: il valore di PROXY_ADDRESS dipende dall'opzione di connessione scelta in Connettiti al cluster Pathways in modalità interattiva. Su Cluster Director, installa jax==0.11.1 nella stessa cella. In un notebook in hosting autonomo avviato dopo il giorno source ~/pathways/cluster.env, le variabili sono già impostate.

!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

Sostituisci PROXY_ADDRESS con l'host o l'indirizzo IP del tuo server proxy Pathways da Connettiti al cluster Pathways in modalità interattiva (ad esempio, 127.0.0.1, il nome host di Cloud DNS o l'indirizzo IP del bilanciatore del carico interno). In Cluster Director, utilizza 127.0.0.1 se utilizzi il port forwarding o l'indirizzo proxy IFRT dal log del servizio. Se hai eseguito il deployment di un notebook in hosting autonomo utilizzando kubectl, sostituisci PROXY_ADDRESS con il nome DNS nel cluster del pod head di Pathways, in modo che il valore JAX_BACKEND_TARGET sia grpc://WORKLOAD-pathways-head-0-0.WORKLOAD:29000.

Aggiungi una seconda cella per verificare la connettività e stampare i dispositivi nel cluster Pathways:

import pathwaysutils
import jax

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

Se la connessione va a buon fine, l'output include un messaggio che indica che JAX ha rilevato il backend Pathways-on-Cloud.

Il numero di dispositivi JAX elencati corrisponde 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. Per eseguire l'esempio completo, crea il tuo workload Pathways con almeno due sezioni impostando WORKLOAD_NODEPOOL_COUNT su 2 o un valore superiore. In Cluster Director, imposta MIN_SLICES e MAX_SLICES su 2 o un valore superiore. Con una sezione, slice1_devices è vuoto e l'ultima parte dell'esempio non va a buon fine con un 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)

Elimina il workload interattivo di Pathways

GKE

Cluster Toolkit

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

kubectl

kubectl delete job WORKLOAD

Cluster Director

Sul nodo di accesso del cluster, annulla il job del servizio Pathways:

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

Questo comando annulla anche i job dei worker Pathways. Le VM TPU sono nodi statici, quindi continuano a essere eseguite e diventano disponibili per altri job Slurm.

Se riscontri problemi, consulta Risolvi i problemi relativi a Pathways on Cloud.

Passaggi successivi

• Crea un cluster GKE con Pathways
• Crea un cluster Cluster Director con Pathways
• Esegui un workload batch con Pathways
• Porting dei workload JAX su Pathways
• Addestramento resiliente con Pathways
• Risolvi i problemi relativi a Pathways on Cloud