Esegui applicazioni Cloud TPU su GKE

Questa guida spiega come:

Per ulteriori informazioni sulle architetture delle VM TPU, consulta Architettura di sistema. Questa guida può essere utilizzata solo con l'architettura dei nodi TPU.

Vantaggi dell'esecuzione di applicazioni Cloud TPU su GKE

Le applicazioni di addestramento Cloud TPU possono essere configurate per l'esecuzione in container GKE all'interno dei pod GKE. Quando lo saranno, vedrai i seguenti vantaggi:

  • Configurazione e gestione del flusso di lavoro migliorate: GKE gestisce il ciclo di vita delle TPU. Dopo aver configurato l'inizializzazione e l'addestramento di Cloud TPU con GKE, i carichi di lavoro possono essere ripetuti e gestiti da GKE, incluso il ripristino da errori dei job.

  • Costo ottimizzato: paghi per la TPU solo quando il job è attivo. GKE crea ed elimina automaticamente le TPU in base ai requisiti delle risorse del pod.

  • Utilizzo flessibile: la richiesta di un acceleratore hardware diverso (CPU, GPU o TPU) è una piccola modifica alle specifiche del pod:

    kind: Pod
    metadata:
      name: example-tpu
      annotations:
        # The Cloud TPUs that will be created for this Job will support
        # TensorFlow 2.12.1. This version MUST match the
        # TensorFlow version that your model is built on.
        tf-version.cloud-tpus.google.com: "2.12.1"
    spec:
      containers:
      - name: example-container
        resources:
          limits:
            cloud-tpus.google.com/v2: 8
            # See the line above for TPU, or below for CPU / GPU.
            # cpu: 2
            # nvidia.com/gpu: 1
    
  • Scalabilità: GKE fornisce API (job e deployment) in grado di scalare fino a centinaia di pod e nodi TPU GKE.

  • Tolleranza di errore: l'API Job di GKE, insieme al meccanismo di checkpoint TensorFlow, fornisce la semantica del completamento dell'esecuzione. I job di addestramento verranno automaticamente eseguiti nuovamente con l'ultimo stato letto dal checkpoint se si verificano errori sulle istanze VM o sui nodi Cloud TPU.

Requisiti e limitazioni di configurazione di Cloud TPU e GKE

Tieni presente quanto segue quando definisci la configurazione GKE:

  • Cloud TPU non è supportato nei pool di nodi Windows Server.
  • Devi creare il tuo cluster GKE e i pool di nodi in una zona in cui è disponibile Cloud TPU. Devi inoltre creare i bucket Cloud Storage per conservare i dati e i modelli di addestramento nella stessa regione del cluster GKE. Consulta il documento su tipi e zone per un elenco delle zone disponibili.
  • Devi utilizzare indirizzi IP conformi alla specifica RFC 1918 per i cluster GKE. Per ulteriori informazioni, consulta Networking di GKE.
  • Ogni container può richiedere al massimo una Cloud TPU, ma più container in un pod possono richiedere una Cloud TPU ciascuno.

Prima di iniziare

  1. Accedi al tuo account Google Cloud. Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti gratuiti per l'esecuzione, il test e il deployment dei carichi di lavoro.
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.

  4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  5. Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.

  6. Quando utilizzi Cloud TPU con GKE, il tuo progetto utilizza componenti fatturabili di Google Cloud. Controlla i prezzi di Cloud TPU e i prezzi di GKE per stimare i tuoi costi e segui le istruzioni per ripulire le risorse quando hai finito.

  7. Abilita le seguenti API nella console Google Cloud:

Crea un nuovo cluster con il supporto di Cloud TPU

Utilizza le istruzioni seguenti per configurare il tuo ambiente e creare un cluster GKE con il supporto di Cloud TPU, utilizzando gcloud CLI:

  1. Installa i componenti gcloud necessari per eseguire GKE con Cloud TPU:

    $ gcloud components install kubectl 
  2. Configura gcloud con il tuo ID progetto Google Cloud:

    $ gcloud config set project project-name
    

    Sostituisci project-name con il nome del tuo progetto Google Cloud.

    La prima volta che esegui questo comando in una nuova VM Cloud Shell, viene visualizzata una pagina Authorize Cloud Shell. Fai clic su Authorize in fondo alla pagina per consentire a gcloud di effettuare chiamate API Google Cloud con le tue credenziali.

  3. Configura gcloud con la zona in cui prevedi di utilizzare una risorsa Cloud TPU. Questo esempio utilizza us-central1-b, ma puoi utilizzare una TPU in qualsiasi zona supportata.

    $ gcloud config set compute/zone us-central1-b
    
  4. Utilizza il comando gcloud container clusters create per creare un cluster su GKE che supporti Cloud TPU.

    $ gcloud container clusters create cluster-name \
      --release-channel=stable \
      --scopes=cloud-platform \
      --enable-ip-alias \
      --enable-tpu
    

    Descrizioni flag di comando

    canale di rilascio
    I canali di rilascio consentono di gestire gli upgrade automatici per i cluster. Quando crei un nuovo cluster, puoi scegliere il canale di rilascio. Verrà eseguito l'upgrade del cluster solo alle versioni offerte in quel canale.
    ambiti
    Garantisce che tutti i nodi nel cluster abbiano accesso al bucket Cloud Storage. Affinché questa operazione funzioni, il cluster e il bucket di archiviazione devono trovarsi nello stesso progetto. Tieni presente che i pod Kubernetes ereditano per impostazione predefinita gli ambiti dei nodi in cui viene eseguito il deployment. Di conseguenza, scopes=cloud-platform assegna a tutti i pod Kubernetes in esecuzione nel cluster l'ambito cloud-platform. Se vuoi limitare l'accesso in base ai singoli pod, consulta la guida di GKE per l'autenticazione con gli account di servizio.
    enable-ip-alias
    Indica che il cluster utilizza intervalli IP alias. Questa operazione è necessaria per l'utilizzo di Cloud TPU su GKE.
    enable-tpu
    Indica che il cluster deve supportare Cloud TPU.
    tpu-ipv4-cidr (facoltativo, non specificato sopra)
    Indica l'intervallo CIDR da utilizzare per Cloud TPU. Specifica IP_RANGE nel formato IP/20, ad esempio 10.100.0.0/20. Se non specifichi questo flag, viene allocato e assegnato automaticamente un intervallo CIDR con dimensioni pari a /20.

Una volta creato il cluster, dovresti visualizzare un messaggio simile al seguente:

NAME             LOCATION       MASTER_VERSION    MASTER_IP     MACHINE_TYPE   NODE_VERSION      NUM_NODES  STATUS
cluster-name  us-central1-b  1.16.15-gke.4901  34.71.245.25  n1-standard-1  1.16.15-gke.4901  3          RUNNING

Richiedi una Cloud TPU nelle specifiche del tuo pod Kubernetes

Nella specifica del pod Kubernetes:

  • Devi creare i modelli nei container utilizzando la stessa versione di TensorFlow. Vedi le versioni supportate.

  • Specifica la risorsa Cloud TPU nella sezione limits sotto il campo resource nella specifica del container.

    Tieni presente che l'unità della risorsa Cloud TPU è il numero di core Cloud TPU. La seguente tabella elenca esempi di richieste di risorse valide. Consulta Tipi e zone TPU per un elenco completo delle risorse TPU valide.

    Se la risorsa che deve essere utilizzata è un pod di Cloud TPU, richiedi la quota poiché la quota predefinita per il pod di Cloud TPU è pari a zero.

    Richiesta di risorse Tipo di Cloud TPU
    cloud-tpus.google.com/v2: 8 Un dispositivo Cloud TPU v2 (8 core)
    cloud-tpus.google.com/previous-v2: 8 Un dispositivo Cloud TPU v2 prerilasciabile (8 core)
    cloud-tpus.google.com/v3: 8 Un dispositivo Cloud TPU v3 (8 core)
    cloud-tpus.google.com/predict-v3: 8 Un dispositivo Cloud TPU v3 prerilasciabile (8 core)
    cloud-tpus.google.com/v2: 32 Un pod di Cloud TPU v2-32 (32 core)
    cloud-tpus.google.com/v3: 32 Un pod di Cloud TPU v3-32 (32 core)

    Per ulteriori informazioni su come specificare risorse e limiti nelle specifiche dei pod, consulta la documentazione di Kubernetes.

La seguente specifica del pod di esempio richiede una TPU Cloud TPU v2-8 prerilasciabile con TensorFlow 2.12.1.

La durata dei nodi Cloud TPU è associata ai pod Kubernetes che li richiedono. Cloud TPU viene creato on demand quando il pod Kubernetes viene pianificato, e riciclato quando il pod Kubernetes viene eliminato.

apiVersion: v1
kind: Pod
metadata:
  name: gke-tpu-pod
  annotations:
     # The Cloud TPUs that will be created for this Job will support
     # TensorFlow 2.12.1. This version MUST match the
     # TensorFlow version that your model is built on.
     tf-version.cloud-tpus.google.com: "2.12.1"
spec:
  restartPolicy: Never
  containers:
  - name: gke-tpu-container
    # The official TensorFlow 2.12.1 image.
    # https://hub.docker.com/r/tensorflow/tensorflow
    image: tensorflow/tensorflow:2.12.1
    command:
    - python
    - -c
    - |
      import tensorflow as tf
      print("Tensorflow version " + tf.__version__)

      tpu = tf.distribute.cluster_resolver.TPUClusterResolver('$(KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS)')
      print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

      tf.config.experimental_connect_to_cluster(tpu)
      tf.tpu.experimental.initialize_tpu_system(tpu)
      strategy = tf.distribute.TPUStrategy(tpu)

      @tf.function
      def add_fn(x,y):
          z = x + y
          return z

      x = tf.constant(1.)
      y = tf.constant(1.)
      z = strategy.run(add_fn, args=(x,y))
      print(z)
    resources:
      limits:
        # Request a single Preemptible v2-8 Cloud TPU device to train the model.
        cloud-tpus.google.com/preemptible-v2: 8

Creazione del job

Segui questi passaggi per creare il job nel cluster GKE e per installare kubectl

  1. Con un editor di testo, crea la specifica del pod example-job.yaml, quindi copia e incolla la specifica del pod mostrata in precedenza.

  2. Esegui il job:

    $ kubectl create -f example-job.yaml
    
    pod "gke-tpu-pod" created

    Questo comando crea il job che pianifica automaticamente il pod.

  3. Verifica che il pod GKE sia stato pianificato e che sia stato eseguito il provisioning dei nodi Cloud TPU. Un pod GKE che richiede nodi Cloud TPU può rimanere in attesa per 5 minuti prima dell'esecuzione. Vedrai un output simile al seguente finché il pod GKE non viene pianificato.

    $ kubectl get pods -w
    
    NAME          READY     STATUS    RESTARTS   AGE
    gke-tpu-pod   0/1       Pending   0          1m
    

    Dopo circa 5 minuti, il risultato sarà simile a questo:

    NAME          READY     STATUS              RESTARTS   AGE
    gke-tpu-pod   0/1       Pending             0          21s
    gke-tpu-pod   0/1       Pending             0          2m18s
    gke-tpu-pod   0/1       Pending             0          2m18s
    gke-tpu-pod   0/1       ContainerCreating   0          2m18s
    gke-tpu-pod   1/1       Running             0          2m48s
    gke-tpu-pod   0/1       Completed           0          3m8s
    

    Devi usare Ctrl-C per uscire dal comando "kubectl get".

    Puoi stampare informazioni di log e recuperare informazioni più dettagliate su ciascun pod GKE utilizzando i seguenti comandi kubectl. Ad esempio, per vedere l'output di log per il tuo pod GKE, usa:

    $ kubectl logs gke-tpu-pod

    Dovresti vedere un output simile al seguente:

    2021-09-24 18:55:25.400699: I tensorflow/core/platform/cpu_feature_guard.cc:142]
    This TensorFlow binary is optimized with oneAPI Deep Neural Network Library
    (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
    To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
    2021-09-24 18:55:25.405947: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
    Initialize GrpcChannelCache for job worker -> {0 -> 10.0.16.2:8470}
    2021-09-24 18:55:25.406058: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
    Initialize GrpcChannelCache for job localhost -> {0 -> localhost:32769}
    2021-09-24 18:55:28.091729: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
    Initialize GrpcChannelCache for job worker -> {0 -> 10.0.16.2:8470}
    2021-09-24 18:55:28.091896: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
    Initialize GrpcChannelCache for job localhost -> {0 -> localhost:32769}
    2021-09-24 18:55:28.092579: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:427]
    Started server with target: grpc://localhost:32769
    Tensorflow version 2.12.1
    Running on TPU  ['10.0.16.2:8470']
    PerReplica:{
      0: tf.Tensor(2.0, shape=(), dtype=float32),
      1: tf.Tensor(2.0, shape=(), dtype=float32),
      2: tf.Tensor(2.0, shape=(), dtype=float32),
      3: tf.Tensor(2.0, shape=(), dtype=float32),
      4: tf.Tensor(2.0, shape=(), dtype=float32),
      5: tf.Tensor(2.0, shape=(), dtype=float32),
      6: tf.Tensor(2.0, shape=(), dtype=float32),
      7: tf.Tensor(2.0, shape=(), dtype=float32)
    }
    

    Per vedere una descrizione completa del pod GKE, utilizza:

    $ kubectl describe pod gke-tpu-pod
    

    Per ulteriori dettagli, consulta Introspezione e debug delle applicazioni.

Crea e containerizza il tuo modello nell'immagine Docker

Per maggiori dettagli su questo processo, fai riferimento a Creare e containerizzare il tuo modello.

Abilita il supporto di Cloud TPU su un cluster esistente

Per abilitare il supporto di Cloud TPU su un cluster GKE esistente, esegui questi passaggi in Google Cloud CLI:

  1. Abilita il supporto di Cloud TPU:

    gcloud beta container clusters update cluster-name --enable-tpu
    

    Sostituisci cluster-name con il nome del tuo cluster.

  2. Aggiorna la voce kubeconfig:

    gcloud container clusters get-credentials cluster-name
    

Impostazione di un intervallo CIDR personalizzato

Per impostazione predefinita, GKE alloca un blocco CIDR con dimensioni pari a /20 per le TPU di cui è stato eseguito il provisioning dal cluster. Puoi specificare un intervallo CIDR personalizzato per Cloud TPU eseguendo questo comando:

gcloud beta container clusters update cluster-name \
  --enable-tpu \
  --tpu-ipv4-cidr 10.100.0.0/20

Sostituisci quanto segue:

  • cluster-name: il nome del cluster esistente.
  • 10.100.0.0/20: il tuo intervallo CIDR personalizzato.

Utilizzo di intervalli CIDR esistenti con VPC condiviso

Segui la guida sulle TPU nei cluster GKE che utilizzano un VPC condiviso per verificare la configurazione corretta per il tuo VPC condiviso.

Disabilitazione di Cloud TPU in un cluster

Per disabilitare il supporto di Cloud TPU su un cluster GKE esistente, esegui questi passaggi in Google Cloud CLI:

  1. Verifica che nessuno dei tuoi carichi di lavoro utilizzi Cloud TPU:

    $ kubectl get tpu
    
  2. Disabilita il supporto di Cloud TPU nel tuo cluster:

    $ gcloud beta container clusters update cluster-name --no-enable-tpu
    

    Sostituisci cluster-name con il nome del tuo cluster.

    Questa operazione richiede circa 5 minuti per i cluster di zona e circa 15 minuti per i cluster a livello di regione, a seconda della regione del cluster.

  3. Quando le operazioni sono state completate senza errori, puoi verificare che le TPU di cui è stato eseguito il provisioning dal cluster siano state rimosse:

    $ gcloud compute tpus list
    

    I nomi delle TPU create da Cloud TPU hanno il seguente formato:

    $ gke-cluster-name-cluster-id-tpu-tpu-id
    

    Sostituisci quanto segue:

    • cluster-name: il nome del cluster esistente.
    • cluster-id: l'ID del cluster esistente.
    • tpu-id: l'ID della Cloud TPU.

    Se vengono visualizzate TPU, puoi eliminarle manualmente eseguendo:

    $ gcloud compute tpus delete gke-cluster-name-cluster-id-tpu-tpu-id
    

Esegui la pulizia

Quando hai finito di utilizzare Cloud TPU su GKE, pulisci le risorse per evitare che al tuo account di fatturazione Cloud vengano addebitati costi aggiuntivi.

  1. Esegui questo comando per eliminare il cluster GKE, sostituendo cluster-name con il nome del cluster e project-name con il nome del tuo progetto Google Cloud:

    $ gcloud container clusters delete cluster-name \
    --project=project-name --zone=us-central1-b
    
  2. Quando hai finito di esaminare i dati, usa il comando gsutil per eliminare il bucket Cloud Storage che hai creato. Sostituisci bucket-name con il nome del tuo bucket Cloud Storage:

    $ gsutil rm -r gs://bucket-name