Esegui il deployment dei carichi di lavoro TPU su GKE Autopilot

Questa pagina descrive come eseguire il deployment di carichi di lavoro che utilizzano acceleratori Cloud TPU (TPU) nei cluster Autopilot di Google Kubernetes Engine (GKE). Dovresti avere già familiarità con i seguenti concetti:

  1. Introduzione a Cloud TPU
  2. Architettura di sistema Cloud TPU
  3. Informazioni sulle TPU in GKE

Come funzionano le TPU in Autopilot

Per utilizzare le TPU nei carichi di lavoro Autopilot, richiedi una versione TPU e una topologia supportata per quella versione TPU nel manifest del tuo carico di lavoro. Successivamente, utilizzerai i campi resources.requests e resources.limits di Kubernetes per specificare il numero di chip TPU da utilizzare nel carico di lavoro. Quando esegui il deployment del carico di lavoro, GKE esegue il provisioning dei nodi che hanno la configurazione TPU richiesta e pianifica i pod sui nodi. GKE posiziona ciascun carico di lavoro sul proprio nodo in modo che ogni pod possa accedere a tutte le risorse del nodo con il rischio di interruzione ridotto al minimo.

Le TPU in Autopilot sono compatibili con le seguenti funzionalità:

  1. Pod spot
  2. Prenotazioni di capacità specifiche
  3. Pod con tempo di esecuzione esteso

Pianifica la configurazione TPU

Prima di richiedere le TPU, decidi la configurazione che preferisci in base ai requisiti di CPU e memoria del carico di lavoro. Devi decidere quanto segue:

  • Versione TPU: la versione specifica di Cloud TPU, ad esempio v5e.
  • Topologia per la versione TPU selezionata: disposizione e numero di TPU.

La versione di TPU e la topologia che selezioni determinano se GKE esegue il provisioning dei nodi come sezioni con singolo host o sezioni con più host. Nelle sezioni con singolo host, ogni nodo è indipendente dagli altri nodi TPU. Nelle sezioni multi-host, GKE crea un gruppo di nodi con VM TPU interconnesse. Le sezioni con più host sono atomiche, il che significa che GKE scala l'intero gruppo interconnesso di nodi come una singola unità.

Per informazioni sulle versioni di TPU disponibili, le topologie corrispondenti, la capacità di CPU e memoria e il tipo di sezione risultante, consulta Scegliere una configurazione TPU Autopilot.

Prezzi

Per informazioni sui prezzi, vedi Prezzi di Autopilot.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti attività:

  • Abilita l'API Google Kubernetes Engine.
    • Abilita l'API Google Kubernetes Engine
  • Se vuoi utilizzare Google Cloud CLI per questa attività, installa e initialize gcloud CLI. Se hai già installato gcloud CLI, scarica la versione più recente eseguendo gcloud components update.
  • Assicurati di avere un cluster Autopilot che esegue GKE versione 1.29.2-gke.1521000+ o successiva.
  • Per utilizzare le TPU prenotate, assicurati di avere una prenotazione di capacità specifica esistente. Per le istruzioni, consulta Utilizzo di risorse di zona riservate.

Assicurati di disporre della quota TPU

Per creare nodi TPU, devi disporre di una quota TPU, a meno che non utilizzi una prenotazione di capacità esistente. Se utilizzi TPU riservate, salta questa sezione.

La creazione di nodi TPU in GKE richiede la quota per l'API Compute Engine (compute.googleapis.com), non la quota API Cloud TPU (tpu.googleapis.com). Il nome della quota è diverso nei pod Autopilot normali e nei pod Spot.

Per verificare il limite e l'utilizzo attuale della tua quota API Compute Engine per le TPU, segui questi passaggi:

  1. Vai alla pagina Quote nella console Google Cloud:

    Vai a Quote

  2. Nella casella Filtro, segui questi passaggi:

    1. Seleziona la proprietà Servizio, inserisci API Compute Engine e premi Invio.

    2. Seleziona la proprietà Tipo e scegli Quota.

    3. Seleziona la proprietà Nome e inserisci un nome quota in base al tipo di TPU desiderato, come segue:

      • TPU v5p (anteprima): chip TPU v5p
      • TPU v5e (tpu-v5-lite-podslice): chip PodSlice TPU v5 Lite
      • TPU v5e (tpu-v5-lite-device): chip di dispositivo TPU v5 Lite
      • TPU v4 (tpu-v4-podslice): chip PodSlice TPU v4

      Per i pod Spot, seleziona la quota "Prerilasciabile" corrispondente.

    4. Seleziona la proprietà Dimensioni (ad esempio località) e inserisci region: seguito dal nome della regione in cui prevedi di creare TPU in GKE. Ad esempio, inserisci region:us-west4 se prevedi di creare nodi TPU nella zona us-west4-a. La quota TPU è a livello di regione, pertanto tutte le zone all'interno della stessa regione consumano la stessa quota TPU.

Se nessuna quota corrisponde al filtro inserito, al progetto non è stata concessa nessuna delle quote specificate per la regione desiderata e devi richiedere un aumento della quota TPU.

Prepara l'applicazione TPU

I carichi di lavoro TPU hanno i seguenti requisiti di preparazione.

  1. Framework come JAX, PyTorch e TensorFlow accedono alle VM TPU utilizzando la libreria condivisa libtpu. libtpu include il compilatore XLA, il software di runtime TPU e il driver TPU. Ogni release di PyTorch e JAX richiede una determinata versione di libtpu.so. Per utilizzare le TPU in GKE, assicurati di utilizzare le seguenti versioni:
    Tipo di TPU Versione di libtpu.so
    TPU v5e
    tpu-v5-lite-podslice
    tpu-v5-lite-device
    TPU v5p
    tpu-v5p-slice
    • Versione jax[tpu] consigliata: 0.4.19 o successiva.
    • Versione torchxla[tpuvm] consigliata: è consigliabile utilizzare una build della versione notturna il 23 ottobre 2023.
    TPU v4
    tpu-v4-podslice
  2. Imposta le seguenti variabili di ambiente per il container che richiede le risorse TPU:
    • TPU_WORKER_ID: un numero intero univoco per ogni pod. Questo ID indica un ID worker univoco nella sezione TPU. I valori supportati per questo campo sono compresi tra zero e il numero di pod meno uno.
    • TPU_WORKER_HOSTNAMES: un elenco separato da virgole di nomi host o indirizzi IP delle VM TPU che devono comunicare tra loro all'interno della sezione. Deve essere presente un nome host o un indirizzo IP per ogni VM TPU nella sezione. L'elenco di indirizzi IP o nomi host è ordinato e non è indicizzato da TPU_WORKER_ID.

      • GKE inserisce automaticamente queste variabili di ambiente utilizzando un webhook mutante quando viene creato un job con le proprietà completionMode: Indexed, subdomain, parallelism > 1 e richiedendo le proprietà google.com/tpu. GKE aggiunge un servizio headless in modo che i record DNS vengano aggiunti per i pod che supportano il servizio.

Dopo aver completato la preparazione del carico di lavoro, puoi eseguire un job che utilizza TPU.

Richiedi TPU in un carico di lavoro

Questa sezione mostra come creare un job che richiede TPU in Autopilot. In qualsiasi carico di lavoro che richiede TPU, devi specificare quanto segue:

  • Selettori di nodi per la versione e la topologia TPU
  • Il numero di chip TPU per un container nel tuo carico di lavoro

Per un elenco delle versioni TPU supportate, delle topologie e del numero corrispondente di chip e nodi TPU in una sezione, consulta Scegliere una configurazione TPU Autopilot.

Considerazioni sulle richieste TPU nei carichi di lavoro

Solo un container in un pod può utilizzare le TPU. Il numero di chip richiesti da un container deve essere uguale al numero di chip associati a un nodo nella sezione. Ad esempio, se richiedi TPU v5e (tpu-v5-lite-podslice) con una topologia 2x4, puoi richiedere quanto segue:

  • 4 chip, che crea due nodi multi-host con 4 chip ciascuno
  • 8 chip, che crea un nodo host singolo con 8 chip

Come best practice per massimizzare l'efficienza in termini di costi, utilizza sempre tutti i chip nella sezione richiesta. Se richiedi una sezione multi-host di due nodi con 4 chip ciascuno, dovresti eseguire il deployment di un carico di lavoro che viene eseguito su entrambi i nodi e che utilizza tutti gli 8 chip della sezione.

Crea un carico di lavoro che richiede TPU

I passaggi seguenti creano un job che richiede TPU. Se hai carichi di lavoro che vengono eseguiti su sezioni TPU multi-host, devi anche creare un servizio headless che selezioni il carico di lavoro per nome. Questo servizio headless consente ai pod su nodi diversi nella sezione multi-host di comunicare tra loro aggiornando la configurazione DNS di Kubernetes in modo che punti ai pod nel carico di lavoro.

  1. Salva il seguente manifest come tpu-autopilot.yaml:

    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS

    Sostituisci quanto segue:

    • TPU_TYPE: il tipo di TPU da utilizzare, ad esempio tpu-v4-podslice. Deve essere un valore supportato da GKE.
    • TOPOLOGY: la disposizione dei chip TPU nella sezione, ad esempio 2x2x4. Deve essere una topologia supportata per il tipo di TPU selezionato.
    • NUMBER_OF_CHIPS: il numero di chip TPU che il container deve utilizzare. Deve essere lo stesso valore per limits e requests.

  2. Esegui il deployment del job:

    kubectl create -f tpu-autopilot.yaml

Quando crei questo job, GKE esegue automaticamente le seguenti operazioni:

  1. Esegue il provisioning dei nodi per l'esecuzione dei pod. A seconda del tipo di TPU, della topologia e delle richieste di risorse specificate, questi nodi sono sezioni con singolo host o più host.
  2. Aggiunge incompatibilità ai pod e le tolleranze ai nodi per impedire l'esecuzione di altri carichi di lavoro sugli stessi nodi dei carichi di lavoro TPU.

Esempio: visualizza i chip TPU totali in una sezione con più host

Il seguente carico di lavoro restituisce il numero di chip TPU in tutti i nodi in una sezione TPU multi-host. Per creare una sezione multi-host, il carico di lavoro ha i seguenti parametri:

  • Versione TPU: TPU v4
  • Topologia: 2x2x4

Questa selezione di versioni e topologia determina una sezione con più host.

  1. Salva il seguente manifest come available-chips-multihost.yaml: 
    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
  2. Esegui il deployment del manifest: 
    kubectl create -f available-chips-multihost.yaml

    GKE esegue una sezione TPU v4 con quattro VM TPU (sezione TPU multi-host). La sezione ha 16 chip interconnessi.

  3. Verifica che il job abbia creato quattro pod: 
    kubectl get pods

    L'output è simile al seguente:

    NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
  4. Recupera i log di uno dei pod: 
    kubectl logs POD_NAME

    Sostituisci POD_NAME con il nome di uno dei pod creati. Ad esempio, tpu-job-podslice-0-5cd8r.

    L'output è simile al seguente: 

    TPU cores: 16

Esempio: visualizza i chip TPU in un singolo nodo

Il seguente carico di lavoro è un pod statico che mostra il numero di chip TPU collegati a un nodo specifico. Per creare un nodo host singolo, il carico di lavoro ha i seguenti parametri:

  • Versione TPU: TPU v5e
  • Topologia: 2 x 4

Questa selezione di versioni e topologia determina una sezione con host singolo.

  1. Salva il seguente manifest come available-chips-singlehost.yaml: 
    apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8
  2. Esegui il deployment del manifest: 
    kubectl create -f available-chips-singlehost.yaml

    GKE esegue il provisioning dei nodi con otto sezioni TPU con host singolo che utilizzano TPU v5e. Ogni VM TPU ha otto chip (sezione TPU con singolo host).

  3. Recupera i log del pod: 
    kubectl logs tpu-job-jax-v5

    L'output è simile al seguente: 

    Total TPU chips: 8

Osservabilità e metriche

Dashboard

Nella pagina Cluster Kubernetes della console Google Cloud, la scheda Osservabilità mostra le metriche di osservabilità delle TPU. Per saperne di più, consulta Metriche di osservabilità di GKE.

La dashboard TPU viene compilata solo se hai abilitato le metriche di sistema nel tuo cluster GKE.

Metriche di runtime

In GKE versione 1.27.4-gke.900 o successive, i carichi di lavoro TPU che utilizzano JAX versione 0.4.14 o successive e specificano le metriche di utilizzo delle TPU di containerPort: 8431 esportazione come metriche di sistema GKE. In Cloud Monitoring sono disponibili le seguenti metriche per monitorare le prestazioni di runtime del carico di lavoro TPU:

  • Ciclo di servizio: percentuale di tempo nell'ultimo periodo di campionamento (60 secondi) durante il quale i TensorCore hanno eseguito attivamente l'elaborazione su un chip TPU. Una percentuale maggiore significa un migliore utilizzo delle TPU.
  • Memoria utilizzata: quantità di memoria dell'acceleratore allocata in byte. Campionamento eseguito ogni 60 secondi.
  • Memoria totale: memoria totale dell'acceleratore in byte. Campionamento eseguito ogni 60 secondi.

Queste metriche si trovano nello schema del nodo Kubernetes (k8s_node) e del container Kubernetes (k8s_container).

Container Kubernetes:

  • kubernetes.io/container/accelerator/duty_cycle
  • kubernetes.io/container/accelerator/memory_used
  • kubernetes.io/container/accelerator/memory_total

Nodo Kubernetes:

  • kubernetes.io/node/accelerator/duty_cycle
  • kubernetes.io/node/accelerator/memory_used
  • kubernetes.io/node/accelerator/memory_total

Metriche host

In GKE versione 1.28.1-gke.1066000 o successive, le VM TPU esportano le metriche di utilizzo delle TPU come metriche di sistema di GKE. In Cloud Monitoring sono disponibili le seguenti metriche per monitorare le prestazioni dell'host TPU:

  • Utilizzo TensorCore: percentuale attuale di TensorCore utilizzato. Il valore TensorCore equivale alla somma delle unità di moltiplicazione della matrice (MXU) più l'unità vettoriale. Il valore di utilizzo di TensorCore è la divisione delle operazioni TensorCore eseguite nell'ultimo periodo di campionamento (60 secondi) in base al numero supportato di operazioni TensorCore nello stesso periodo. Maggiore è il valore, maggiore è l'utilizzo.
  • Utilizzo larghezza di banda di memoria: percentuale attuale della larghezza di banda della memoria dell'acceleratore in uso. È calcolata dividendo la larghezza di banda della memoria utilizzata in un periodo di campionamento (60 secondi) per la larghezza di banda massima supportata nello stesso periodo di campionamento.

Queste metriche si trovano nello schema del nodo Kubernetes (k8s_node) e del container Kubernetes (k8s_container).

Container Kubernetes:

  • kubernetes.io/container/accelerator/tensorcore_utilization
  • kubernetes.io/container/accelerator/memory_bandwidth_utilization

Nodo Kubernetes:

  • kubernetes.io/container/node/tensorcore_utilization
  • kubernetes.io/container/node/memory_bandwidth_utilization

Per saperne di più, consulta Metriche Kubernetes e Metriche di sistema GKE.

Logging

I log emessi dai container in esecuzione sui nodi GKE, incluse le VM TPU, vengono raccolti dall'agente di logging GKE, inviati a Logging e sono visibili in Logging.

Suggerimenti per i carichi di lavoro TPU in Autopilot

I seguenti suggerimenti potrebbero migliorare l'efficienza dei carichi di lavoro TPU:

  • Utilizza pod di tempo di esecuzione esteso per un periodo di tolleranza fino a sette giorni prima che GKE arresti i tuoi pod per gli scale down o gli upgrade dei nodi. Puoi utilizzare periodi di manutenzione ed esclusioni con pod con tempo di esecuzione esteso per ritardare ulteriormente gli upgrade automatici dei nodi.
  • Utilizza le prenotazioni di capacità per assicurarti che i carichi di lavoro ricevano le TPU richieste senza essere messi in coda per verificare la disponibilità.