Utilizzare un'immagine sistema operativo personalizzata

Puoi utilizzare un'immagine sistema operativo personalizzata per le VM TPU per precaricare il software, utilizzare una distribuzione del sistema operativo specifica o applicare modifiche personalizzate al kernel. La creazione di un'immagine personalizzata comporta l'esecuzione di 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 le seguenti esclusioni di responsabilità se utilizzi un'immagine sistema operativo personalizzata con le TPU:

  • Google fornisce immagini Ubuntu con assistenza a lungo termine (LTS) predefinite ottimizzate per le TPU. Le modifiche del 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 del sistema operativo richieste 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 un'altra immagine Ubuntu con un kernel personalizzato.
  • Google non crea né fornisce test per immagini sistema operativo diverse dalle immagini Ubuntu LTS predefinite ottimizzate per le TPU. Devi creare e testare l'immagine sistema operativo personalizzata.

Per ulteriori informazioni sulle immagini Ubuntu LTS predefinite ottimizzate per le TPU, consulta Immagini sistema operativo TPU.

Prerequisiti

L'immagine di base deve avere i seguenti componenti installati:

  • Python 3
  • Gcloud CLI

Apportare modifiche durante la creazione dell'immagine

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

Associare i dispositivi TPU a VFIO

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

  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 un rendimento ottimale, modifica i seguenti limiti e parametri di sistema.

Limiti di memoria

Consenti a un singolo processo di bloccare la memoria illimitata 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.
  • intel_iommu=on,sm_on: abilita l'unità di gestione della memoria di input/output (IOMMU) Intel. Obbligatorio per le architetture TPU7x e v5p.
  • transparent_hugepage=always: abilita 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. Abilita 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
    

Installare l'agente vBar

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

Per installare l'agente vBar, esegui i seguenti 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) Installare ed eseguire l'agente di raccolta di telemetria AI

Il raccoglitore di telemetria AI 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 il raccoglitore di telemetria AI con un sistema operativo personalizzato utilizzando l'immagine Docker ai-telemetry-collector. Puoi installare l'immagine sul 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 il raccoglitore di telemetria AI, esegui i seguenti 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 del raccoglitore di telemetria AI 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 file di configurazione aggiuntivi, consulta Raccoglitore di telemetria AI.

Apportare modifiche all'avvio

Configura l'immagine in modo che esegua le attività nelle sezioni seguenti ogni volta che viene avviata una VM. Puoi utilizzare lo cloud-init strumento 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: e i comandi da eseguire all'avvio devono essere inclusi nella chiave runcmd: nella configurazione cloud-init.

Avviare l'agente vBar

Avvia l'agente di controllo vBar con gli ID utente e gruppo appropriati:

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

Configurare le variabili di ambiente

Per assicurarti che l'ambiente sia inizializzato correttamente per i carichi di lavoro TPU, devi recuperare le variabili di configurazione di runtime dal server dei metadati 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 dei 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

Recuperare i metadati della VM

Il seguente snippet 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 di istanza specifici e tag di metadati personalizzati. Aggiungi quanto segue alla sezione write_files: della configurazione 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))

Aumentare i timeout di Cloud Storage

Se il carico di lavoro interagisce con Cloud Storage, aumenta le durate dei 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