Configura KubeRay con TPU Trillium

Autopilot Standard

Questo tutorial mostra come configurare KubeRay con TPU Trillium su Google Kubernetes Engine (GKE). Scopri come configurare le TPU single-host e multi-host, incluse le variabili di ambiente e le specifiche del pod necessarie per TPU Trillium.

Questo tutorial è destinato agli amministratori e agli operatori della piattaforma e agli specialisti di dati e AI che vogliono imparare a configurare l'inizializzazione di TPU Trillium con KubeRay per i node pool single-host e multi-host. Questo tutorial mostra come eseguire uno script con Jax che verifica l'inizializzazione della TPU. Questo tutorial non esegue il deployment di un modello.

Prima di configurare KubeRay in GKE, assicurati di conoscere le definizioni e la terminologia di Ray in GKE.

Panoramica

Questo tutorial mostra come eseguire uno script Python con Jax che verifica che l'inizializzazione della TPU Trillium con KubeRay sia andata a buon fine. Jax è una libreria per il calcolo numerico ad alte prestazioni che supporta i carichi di lavoro di machine learning. KubeRay è un operatore Kubernetes che fornisce un modo unificato per eseguire il deployment, gestire e monitorare le applicazioni Ray su Kubernetes.

Le TPU Trillium (v6e) richiedono variabili di ambiente e specifiche Pod specifiche diverse dalle generazioni di TPU precedenti. Questo tutorial fornisce le configurazioni necessarie per eseguire il deployment di un workload con KubeRay sulle TPU Trillium.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti operazioni:

• Attiva l'API Google Kubernetes Engine.
Attiva l'API Google Kubernetes Engine
• Se vuoi utilizzare Google Cloud CLI per questa attività, installala e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, scarica l'ultima versione eseguendo gcloud components update.
  Nota:per le installazioni esistenti di gcloud CLI, assicurati di impostare la proprietà compute/region. Se utilizzi principalmente cluster zonali, imposta invece compute/zone. Se imposti una località predefinita, puoi evitare errori in gcloud CLI come il seguente: One of [--zone, --region] must be supplied: Please specify location. Potresti dover specificare la posizione in determinati comandi se la posizione del cluster è diversa da quella predefinita che hai impostato.

• Assicurati di aver installato la CLI Ray (versione 2.37.0).

Attiva Cloud Shell

Cloud Shell include gli strumenti a riga di comando gcloud, helm e kubectl utilizzati in questo tutorial.

1. Vai alla consoleGoogle Cloud .

2. Nella parte superiore della finestra della console Google Cloud , fai clic sul pulsante Attiva Cloud Shell .

   All'interno di un nuovo frame nella console Google Cloud si apre una sessione di Cloud Shell e viene visualizzato un prompt della riga di comando.

   

Crea un cluster GKE e un pool di nodi

Puoi configurare KubeRay sulle TPU in un cluster GKE Autopilot o Standard. Ti consigliamo di utilizzare un cluster Autopilot per un'esperienza Kubernetes completamente gestita. Per scegliere la modalità operativa GKE più adatta ai tuoi workload, consulta la sezione Informazioni sulle modalità operative di GKE.

Autopilot

1. In Cloud Shell, esegui questo comando:

   gcloud container clusters create-auto CLUSTER_NAME \
       --enable-ray-operator \
       --release-channel=rapid \
       --location=LOCATION
   

   Sostituisci quanto segue:

   • CLUSTER_NAME: il nome del nuovo cluster.
   • LOCATION: la regione in cui è disponibile la capacità TPU Trillium. Per maggiori informazioni, consulta Disponibilità di TPU in GKE.

   GKE crea un cluster Autopilot con il componente aggiuntivo dell'operatore Ray abilitato. Il componente aggiuntivo installa automaticamente l'webhook Ray TPU nel piano di controllo del cluster.

2. Per comunicare con il cluster, configura kubectl :

   gcloud container clusters get-credentials CLUSTER_NAME --location=LOCATION
   

Standard

1. In Cloud Shell, crea un cluster Standard che attiva il componente aggiuntivo dell'operatore Ray eseguendo questo comando :

   gcloud container clusters create CLUSTER_NAME \
     --location LOCATION \
     --addons=RayOperator \
     --cluster-version=1.33 \
     --machine-type=n1-standard-16
   

   Sostituisci quanto segue:

   • CLUSTER_NAME: il nome del nuovo cluster.
   • LOCATION: la regione in cui è disponibile la capacità TPU Trillium. Per maggiori informazioni, consulta la sezione Disponibilità delle TPU in GKE.

   La creazione del cluster potrebbe richiedere diversi minuti.

2. Per comunicare con il cluster, configura kubectl :

   gcloud container clusters get-credentials CLUSTER_NAME --location=LOCATION
   

3. Puoi creare un pool di nodi di slice TPU single-host o multi-host:

Host singolo

In Cloud Shell, esegui questo comando:

gcloud container node-pools create v6e-4 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=1 \
    --threads-per-core=1 \
    --tpu-topology=2x2

Multi-host

In Cloud Shell, esegui questo comando:

gcloud container node-pools create v6e-16 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=4 \
    --threads-per-core=1 \
    --tpu-topology=4x4

Esegui una risorsa personalizzata RayJob

Definendo un manifest RayJob, indichi a KubeRay di eseguire le seguenti operazioni:

• Crea un RayCluster:la specifica RayJob include un rayClusterSpec che definisce la configurazione del cluster Ray (gruppi di head e worker) che vuoi.
• Esegui un job specifico:il campo entrypoint all'interno di RayJob specifica il comando o lo script da eseguire all'interno del cluster Ray creato. In questo tutorial, entrypoint è uno script Python (tpu_list_devices.py) progettato per verificare l'inizializzazione della TPU Trillium.

Per creare una risorsa personalizzata RayJob, completa i seguenti passaggi:

Host singolo

1. Crea il seguente manifest ray-job.tpu-v6e-singlehost.yaml:

   apiVersion: ray.io/v1
   kind: RayJob
   metadata:
     name: v6e-4-job
   spec:
     entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
     runtimeEnvYAML: |
       working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
       pip:
         - jax[tpu]==0.4.33
         - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
     rayClusterSpec:
       rayVersion: '2.43.0'
       headGroupSpec:
         rayStartParams: {}
         template:
           spec:
             containers:
             -   name: ray-head
                 image: rayproject/ray:2.43.0-py310
                 ports:
                   - containerPort: 6379
                     name: gcs-server
                   - containerPort: 8265
                     name: dashboard
                   - containerPort: 10001
                     name: client
                 resources:
                   limits:
                     cpu: "8"
                     memory: 40G
                   requests:
                     cpu: "8"
                     memory: 40G
       workerGroupSpecs:
       -   replicas: 1
           minReplicas: 1
           maxReplicas: 1
           numOfHosts: 1
           groupName: tpu-group
           rayStartParams: {}
           template:
             spec:
               containers:
               -   name: ray-worker
                   image: rayproject/ray:2.43.0-py310
                   resources:
                     limits:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
                     requests:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
               nodeSelector:
                 cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                 cloud.google.com/gke-tpu-topology: 2x2

2. Applica il manifest:

   kubectl apply -f ray-job.tpu-v6e-singlehost.yaml
   

3. Verifica che RayJob sia stato creato e sia in esecuzione:

   kubectl get rayjobs v6e-4-job
   

   L'output è simile al seguente:

   NAME      JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME       START TIME  END TIME   AGE
   v6e-4-job PENDING      Running             v6e-4-job-raycluster   2024-10-15T23:15:22Z  20s
   

4. Stampa l'output di RayJob.

   kubectl logs -l=job-name=v6e-4-job
   

   L'output è simile al seguente:

   2024-10-15 16:15:40,222 INFO cli.py:300 -- ray job stop v6e-4-job-hzq5q
   2024-10-15 16:15:40,246 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
   2024-10-15 16:15:40,112 INFO job_manager.py:528 -- Runtime env is setting up.
   2024-10-15 16:15:50,181 INFO worker.py:1461 -- Using address 10.84.1.25:6379 set in the environment variable RAY_ADDRESS
   2024-10-15 16:15:50,181 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.25:6379...
   2024-10-15 16:15:50,186 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.25:8265
   ['TPU cores:4']
   2024-10-15 16:16:12,349 SUCC cli.py:63 -- -------------------------------------
   2024-10-15 16:16:12,349 SUCC cli.py:64 -- Job 'v6e-4-job-hzq5q' succeeded
   2024-10-15 16:16:12,349 SUCC cli.py:65 -- -------------------------------------
   

Multi-host

1. Crea il seguente manifest ray-job.tpu-v6e-multihost.yaml:

   apiVersion: ray.io/v1
   kind: RayJob
   metadata:
     name: v6e-16-job
   spec:
     entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
     runtimeEnvYAML: |
       working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
       pip:
         - jax[tpu]==0.4.33
         - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
     rayClusterSpec:
       rayVersion: '2.43.0'
       headGroupSpec:
         rayStartParams: {}
         template:
           spec:
             containers:
             -   name: ray-head
                 image: rayproject/ray:2.43.0-py310
                 ports:
                   - containerPort: 6379
                     name: gcs-server
                   - containerPort: 8265
                     name: dashboard
                   - containerPort: 10001
                     name: client
                 resources:
                   limits:
                     cpu: "8"
                     memory: 40G
                   requests:
                     cpu: "8"
                     memory: 40G
       workerGroupSpecs:
         - replicas: 1
           minReplicas: 1
           maxReplicas: 1
           numOfHosts: 4
           groupName: tpu-group
           rayStartParams: {}
           template:
             spec:
               containers:
                 - name: ray-worker
                   image: rayproject/ray:2.43.0-py310
                   resources:
                     limits:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
                     requests:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
                   env:
                   - name: NODE_IP
                     valueFrom:
                       fieldRef:
                         fieldPath: status.hostIP
                   - name: VBAR_CONTROL_SERVICE_URL
                     value: $(NODE_IP):8353
                   - name: JAX_PLATFORMS
                     value: tpu,cpu
                   - name: ENABLE_PJRT_COMPATIBILITY
                     value: "true"
                   ports:
                   - containerPort: 8081
                     name: mxla
               nodeSelector:
                 cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                 cloud.google.com/gke-tpu-topology: 4x4

2. Applica il manifest:

   kubectl apply -f ray-job.tpu-v6e-multihost.yaml
   

3. Verifica che RayJob v6e-16 sia stato creato e sia in esecuzione:

   kubectl get rayjobs v6e-16-job
   

   L'output è simile al seguente:

   NAME         JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME              START TIME             END TIME   AGE
   v6e-16-job                Running             v6e-16-job-raycluster-qr6vk   2024-10-16T19:28:19Z              66s
   

4. Stampa l'output di v6e-16 RayJob:

   kubectl logs -l=job-name=v6e-16-job
   

   L'output è simile al seguente:

   2024-10-16 12:21:33,986 INFO cli.py:300 -- ray job stop v6e-16-job-z44s7
   2024-10-16 12:21:34,011 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
   2024-10-16 12:21:33,826 INFO job_manager.py:528 -- Runtime env is setting up.
   2024-10-16 12:21:46,327 INFO worker.py:1461 -- Using address 10.84.1.61:6379 set in the environment variable RAY_ADDRESS
   2024-10-16 12:21:46,327 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.61:6379...
   2024-10-16 12:21:46,333 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.61:8265
   ['TPU cores:16', 'TPU cores:16', 'TPU cores:16', 'TPU cores:16']
   2024-10-16 12:22:12,156 SUCC cli.py:63 -- ---------------------------------
   2024-10-16 12:22:12,156 SUCC cli.py:64 -- Job 'v6e-16-job-z44s7' succeeded
   2024-10-16 12:22:12,156 SUCC cli.py:65 -- ---------------------------------
   

Visualizza RayJob nella dashboard Ray

Verifica che GKE abbia creato il servizio RayCluster e connettiti all'istanza RayCluster.

Host singolo

1. Recupera il nome di RayCluster generato per RayJob:

   export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-4-job -o jsonpath='{.status.rayClusterName}')
   

2. Recupera il nome del servizio head di RayCluster:

   export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')
   

3. Connettiti alla dashboard Ray inoltrando la porta del servizio head:

   kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
   

4. Apri un browser web e inserisci il seguente URL:

   http://localhost:8265/#/jobs
   

5. Visualizza lo stato di RayJob e i log pertinenti.

Multi-host

1. Recupera il nome di RayCluster generato per RayJob:

   export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-16-job -o jsonpath='{.status.rayClusterName}')
   

2. Recupera il nome del servizio head di RayCluster:

   export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')
   

3. Connettiti alla dashboard Ray inoltrando la porta del servizio head:

   kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
   

4. Apri un browser web e inserisci il seguente URL:

   http://localhost:8265/#/jobs
   

5. Visualizza lo stato di RayJob e i log pertinenti.

Ray imposta una risorsa TPU-{accelerator}-Head per identificare il nodo worker Ray che corrisponde al valore TPU_WORKER_ID=0. Nel gruppo TPU multihost, il nodo Ray con TPU_WORKER_ID=0 ha TPU-v6e-16-head: 1.0 impostato nelle risorse. Questa variabile di ambiente TPU_WORKER_ID è impostata da un webhook di mutazione GKE per KubeRay.

Operazione riuscita: hai eseguito correttamente un RayJob su un pool di nodi di slice TPU che esegue TPU Trillium. Il completamento di RayJob indica che Jax è stato inizializzato correttamente e che i nodi TPU distribuiti sono abilitati correttamente.

Esegui la pulizia

Dopo aver completato il tutorial, per evitare addebiti indesiderati sul tuo account, elimina RayJob:

Host singolo

kubectl delete rayjobs v6e-4-job

Multi-host

kubectl delete rayjobs v6e-16-job

Passaggi successivi

• Scopri di più su Ray su Kubernetes.
• Scopri come gestire vLLM su GKE con le TPU.
• Scopri come gestire SDXL su GKE con le TPU.
• Scopri di più sulle TPU in GKE.