Utilizza un'immagine del sistema operativo personalizzata

Puoi utilizzare un'immagine sistema operativo personalizzata per le tue VM TPU per precaricare software, utilizzare una distribuzione del sistema operativo specifica o applicare modifiche personalizzate al kernel. La creazione di un'immagine personalizzata comporta modifiche specifiche del sistema durante il processo di creazione dell'immagine e la configurazione dell'immagine per gestire le attività di avvio richieste per la funzionalità TPU.

Tieni presente i seguenti disclaimer se utilizzi un'immagine sistema operativo personalizzata con le TPU:

  • Google fornisce immagini Ubuntu con assistenza a lungo termine (LTS) ottimizzate per TPU predefinite. Le modifiche al sistema operativo elencate in questa pagina vengono convalidate solo per le immagini Ubuntu LTS ottimizzate per le TPU supportate da Google.
  • È tua responsabilità estrapolare le modifiche necessarie al sistema operativo per qualsiasi altra distribuzione del sistema operativo o immagine personalizzata. Google non garantisce che le modifiche per Ubuntu elencate in questa pagina funzionino con altre distribuzioni del sistema operativo o con un'altra immagine Ubuntu con un kernel personalizzato.
  • Google non crea né fornisce test per immagini del sistema operativo diverse dalle immagini Ubuntu LTS ottimizzate per TPU predefinite. Devi creare e testare l'immagine del sistema operativo personalizzato.

Per saperne di più sulle immagini Ubuntu LTS predefinite ottimizzate per TPU, consulta Immagini del sistema operativo TPU.

Prerequisiti

Nell'immagine di base devono essere installati i seguenti componenti:

  • Python 3
  • gcloud CLI

Apportare modifiche durante la creazione dell'immagine

Applica le seguenti modifiche durante la creazione dell'immagine Ubuntu personalizzata.

Associa i dispositivi TPU a VFIO

Per consentire al sistema operativo guest di accedere all'hardware TPU, devi associare i dispositivi TPU al driver vfio-pci.

  1. Crea un file di regole udev denominato 99-tpu-vfiopci.rules in /etc/udev/rules.d/:

    # Rules for binding vfio-enabled TPU devices to vfio-pci.
    
    # v5p
    SUBSYSTEM=="pci", ACTION=="add", ATTRS{vendor}=="0x1ae0", ATTRS{device}=="0x0062", ATTRS{subsystem_vendor}=="0x1ae0", ATTRS{subsystem_device}=="0x00ad", DRIVER!="vfio-pci", TAG+="bind_to_vfio_pci"
    
    # v6e
    SUBSYSTEM=="pci", ACTION=="add", ATTRS{vendor}=="0x1ae0", ATTRS{device}=="0x006f", ATTRS{subsystem_vendor}=="0x1ae0", ATTRS{subsystem_device}=="0x00d1", DRIVER!="vfio-pci", TAG+="bind_to_vfio_pci"
    
    # TPU7x
    SUBSYSTEM=="pci", ACTION=="add", ATTRS{vendor}=="0x1ae0", ATTRS{device}=="0x0076", ATTRS{subsystem_vendor}=="0x1ae0", ATTRS{subsystem_device}=="0x00f2", DRIVER!="vfio-pci", TAG+="bind_to_vfio_pci"
    
    # Bind all 'bind_to_vfio_pci' tagged devices to vfio-pci.
    TAG=="bind_to_vfio_pci", RUN+="/lib/udev/bind_to_vfio_pci.sh $kernel"
    
  2. Crea uno script denominato bind_to_vfio_pci.sh in /lib/udev/:

    #!/bin/bash
    #!/usr/bin/env bash
    
    # Run ./bind_to_vfio_pci.sh <DBDF>
    # Binds the device at <DBDF> to vfio-pci.
    # If the device is already bound to a driver, unbinds it first.
    
    # Load the vfio-pci module into the kernel. No-op if already loaded.
    modprobe vfio-pci
    
    DBDF_REGEX="^[[:xdigit:]]{4}:[[:xdigit:]]{2}:[[:xdigit:]]{2}.[[:xdigit:]]$"
    
    unset BDF
    if [[ $1 =~ $DBDF_REGEX ]]; then
        BDF=$1
    else
        echo "Error: BDF arg ($1) is not in form dddd:bb:dd.f"
        exit 1
    fi
    
    PCI_PATH="/sys/bus/pci/devices/$BDF"
    
    echo "vfio-pci" > "$PCI_PATH/driver_override"
    
    PCI_DRIVER_PATH="$PCI_PATH/driver"
    if [[ -d "$PCI_DRIVER_PATH" ]]; then
        curr_driver=$(readlink "$PCI_DRIVER_PATH")
            curr_driver=${curr_driver##*/}
        if [[ $curr_driver == "vfio-pci" ]]; then
            echo "$BDF already bound to vfio-pci"
            exit 0
        else
            echo "$BDF" > "$PCI_DRIVER_PATH/unbind"
            if [[ -d "$PCI_DRIVER_PATH" ]]; then
                echo "Error: Unable to unbind $PCI_DRIVER_PATH"
                exit 1
            fi
            echo "Unbound $BDF from driver $curr_driver"
        fi
    fi
    echo "$BDF" > /sys/bus/pci/drivers_probe
    echo "Bound $BDF to vfio-pci"
    
    # Grant read/write access on VFIO device to all users
    IOMMU_GROUP=$(readlink "$PCI_PATH/iommu_group" | xargs basename)
    VFIO_DEV="/dev/vfio/$IOMMU_GROUP"
    if [[ -c "$VFIO_DEV" ]]; then
        chmod 0666 "$VFIO_DEV"
    else
        echo "$VFIO_DEV not found"
        exit 1
    fi
    
    # Set allow_unsafe_interrupts for x86 platforms.
    (uname -a | grep -q x86_64) && echo 1 > /sys/module/vfio_iommu_type1/parameters/allow_unsafe_interrupts
    
    # This is only needed to avoid non-zero exit code from previous command.
    echo "All Done!"
    
  3. Rendi eseguibile lo script:

    chmod +x /lib/udev/bind_to_vfio_pci.sh
    
  4. Concedi a tutti gli utenti del sistema l'accesso al dispositivo TPU:

    echo 'KERNEL=="accel*" MODE="0666"' >> /etc/udev/rules.d/99-tpu.rules
    

Modificare l'immagine per migliorare il rendimento

Per garantire prestazioni ottimali, modifica i seguenti limiti e parametri di sistema.

Limiti di memoria

Consenti a un singolo processo di bloccare una quantità illimitata di memoria aggiornando /etc/security/limits.conf:

echo '*  hard  memlock  unlimited' >> /etc/security/limits.conf
echo '*  soft  memlock  unlimited' >> /etc/security/limits.conf

Limiti dei file

Aumenta il numero di file aperti aggiornando /etc/security/limits.conf:

echo "*    soft    nofile       100000" >> /etc/security/limits.conf
echo "*    hard    nofile       100000" >> /etc/security/limits.conf
echo "root soft    nofile       100000" >> /etc/security/limits.conf
echo "root hard    nofile       100000" >> /etc/security/limits.conf

Parametri del kernel

Aggiorna la configurazione GRUB (in genere in /etc/default/grub) per includere i seguenti parametri in GRUB_CMDLINE_LINUX:

  • idle=poll: impedisce alla CPU di entrare in stati di inattività a basso consumo energetico.
  • intel_iommu=on,sm_on: attiva l'unità di gestione della memoria di input/output (IOMMU) di Intel. Obbligatorio per le architetture TPU7x e v5p.
  • transparent_hugepage=always: attiva le pagine enormi trasparenti (THP).

I seguenti passaggi mostrano come aggiornare questi parametri del kernel:

  1. Impedisci alla CPU di passare a uno stato di inattività a basso consumo impostando la seguente variabile, che utilizzerai nel passaggio successivo.

    kernel_cmdline="idle=poll"
    
  2. Attiva l'unità di gestione della memoria di input/output (IOMMU) Intel. Questo passaggio è obbligatorio per TPU7x e TPU v5p.

    kernel_cmdline="${kernel_cmdline} intel_iommu=on,sm_on";
    sed -i "s/GRUB_CMDLINE_LINUX=\"\"/GRUB_CMDLINE_LINUX=\"${kernel_cmdline}\"/" /etc/default/grub
    echo "Status: New kernel cmdline: $(cat /etc/default/grub | grep -e '^GRUB_CMDLINE_LINUX=')"
    
    update-grub
    
  3. Abilita le pagine enormi trasparenti (THP):

    echo "Status: Enabling THP"
    sed -i -r 's/GRUB_CMDLINE_LINUX="[a-zA-Z0-9_= ]*/& transparent_hugepage=always/' /etc/default/grub
    
    update-grub
    

Installa l'agente vBar

L'agente vBar è necessario per il funzionamento della rete di interconnessione tra chip (ICI).

Per installare l'agente vBar, esegui questi comandi:

  1. Autentica Docker con Artifact Registry:

    gcloud auth configure-docker us-docker.pkg.dev
    
  2. Estrai l'immagine Docker da Artifact Registry:

    docker pull gcr.io/cloud-tpu-v2-images/vbar_control_agent:0.0.1
    
  3. Esegui un container utilizzando l'immagine dell'agente vBar:

    docker run --privileged --net=host vbar_control_agent:0.0.1
    

(Facoltativo) Installa ed esegui AI Telemetry Collector

AI Telemetry Collector viene eseguito all'interno della VM TPU e ti consente di accedere alle metriche di runtime e dell'infrastruttura tramite Cloud Monitoring o tramite la tua pipeline di monitoraggio basata su Prometheus. Puoi utilizzare AI Telemetry Collector con un sistema operativo personalizzato utilizzando l'immagine Docker ai-telemetry-collector. Puoi installare l'immagine sul tuo sistema operativo personalizzato e utilizzare un file config.yaml per definire gli intervalli di raccolta, attivare o disattivare metriche specifiche o modificare le destinazioni di esportazione.

Per installare AI Telemetry Collector, esegui questi comandi:

  1. Autentica Docker con Artifact Registry:

    gcloud auth configure-docker us-docker.pkg.dev
    
  2. Estrai l'immagine Docker da Artifact Registry:

    docker pull gcr.io/cloud-tpu-v2-images/ai-telemetry-collector:latest
    
  3. Esegui un container utilizzando l'immagine AI Telemetry Collector con la configurazione predefinita:

    docker run --privileged --net=host ai-telemetry-collector:latest
    

    Per informazioni sull'utilizzo di un file di configurazione personalizzato o sull'aggiunta di altri file di configurazione, consulta AI Telemetry Collector.

Apportare modifiche al tempo di avvio

Configura l'immagine per eseguire le attività nelle sezioni seguenti ogni volta che una VM viene avviata. Puoi utilizzare lo strumento cloud-init per configurare le attività di avvio passando i metadati alle istanze. Le configurazioni nelle sezioni seguenti utilizzano moduli come write_files e runcmd. Gli snippet che definiscono i file da scrivere devono essere inclusi nella chiave write_files:, mentre i comandi da eseguire all'avvio devono essere inclusi nella chiave runcmd: nella configurazione cloud-init.

Avvia l'agente vBar

Avvia l'agente di controllo della barra verticale con gli ID utente e gruppo appropriati:

vbar_control_agent --logtostderr --gid= --uid=  --chroot= --census_enabled=false --loas_pwd_fallback_in_corp

Configura le variabili di ambiente

Per assicurarti che l'ambiente venga inizializzato correttamente per i workload TPU, devi recuperare le variabili di configurazione di runtime dal server di metadati di Compute Engine durante il processo di avvio del sistema. Per farlo, aggiungi il seguente snippet alla sezione write_files: della configurazione cloud-init, che crea uno script denominato /var/scripts/configure-env-vars.sh. Questo script automatizza il recupero degli attributi dalla chiave di metadati tpu-env e li salva in /${HOME}/tpu-env per essere utilizzati dallo stack software TPU.

 - path: /var/scripts/configure-env-vars.sh
    permissions: 0444
    owner: root
    content: |
      grep -q CLOUDSDK_PYTHON /etc/environment || echo "CLOUDSDK_PYTHON=/usr/bin/python3" >> /etc/environment

      export HOME=/home/tpu-runtime
      curl -s 'http://metadata.google.internal/computeMetadata/v1/instance/attributes/tpu-env' -H 'Metadata-Flavor: Google' > /tmp/tpu-env.yaml

      eval $(python3 -c '''
      import yaml
      stream_in=open("/tmp/tpu-env.yaml", "r")
      for k,v in yaml.safe_load(stream_in).items():
        print("{var}=\"{value}\"".format(var = k, value = str(v)))
      ''' > "/${HOME}/tpu-env"
      )

      rm -f "/tmp/tpu-env.yaml"

      printenv
      cat ${HOME}/tpu-env

Recupera i metadati della VM

Lo snippet seguente crea uno script denominato /var/scripts/get-vm-metadata.py, un'utilità Python per eseguire query a livello di programmazione sul server dei metadati per attributi specifici dell'istanza e tag di metadati personalizzati. Aggiungi quanto segue alla sezione write_files: della configurazione di cloud-init:

 - path: /var/scripts/get-vm-metadata.py
    permissions: 0444
    owner: root
    content: |
      import sys, requests, os

      if len(sys.argv) < 2:
        sys.stderr.write('Must provide key')
        os._exit(1)

      key = sys.argv[1]
      default = None
      if len(sys.argv) > 2:
        default = sys.argv[2]

      attribute_type = 'attributes'
      if len(sys.argv) > 3:
        attribute_type = sys.argv[3]

      request = requests.get("http://metadata.google.internal/computeMetadata/v1/instance/{}/{}".format(attribute_type, key), headers={'Metadata-Flavor': 'Google'})
      if request.status_code == 200:
        print(request.content)
      elif request.status_code == 404 or request.status_code == '403':
        sys.stderr.write('Metadata key: {} does not exist\n'.format(key))
        if default:
          print(default)
      else:
        sys.stderr.write('Lookup failed with: {}'.format(request))

Aumenta i timeout di Cloud Storage

Se il tuo workload interagisce con Cloud Storage, aumenta la durata del timeout aggiungendo i valori di timeout a /etc/environment. Per farlo, aggiungi il seguente snippet alla sezione write_files: della configurazione cloud-init, che crea uno script denominato /var/scripts/configure-gcs-timeouts.sh.

 - path: /var/scripts/configure-gcs-timeouts.sh
    permissions: 0444
    owner: root
    content: |
      echo "GCS_RESOLVE_REFRESH_SECS=60" >> /etc/environment
      echo "GCS_REQUEST_CONNECTION_TIMEOUT_SECS=300" >> /etc/environment
      echo "GCS_METADATA_REQUEST_TIMEOUT_SECS=300" >> /etc/environment
      echo "GCS_READ_REQUEST_TIMEOUT_SECS=300" >> /etc/environment
      echo "GCS_WRITE_REQUEST_TIMEOUT_SECS=600" >> /etc/environment

Passaggi successivi