Utilizza un'immagine sistema operativo personalizzata

Puoi utilizzare un'immagine sistema operativo personalizzata per le tue 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'apporto di modifiche specifiche al 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 LTS (Long Term Support) predefinite ottimizzate per TPU. Le modifiche al sistema operativo elencate in questa pagina sono convalidate solo per le immagini Ubuntu LTS supportate da Google e ottimizzate per TPU.
  • È tua responsabilità estrapolare le modifiche richieste al sistema operativo per qualsiasi altra distribuzione del sistema operativo o per le immagini personalizzate. 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 sistema operativo diverse dalle immagini Ubuntu LTS predefinite ottimizzate per TPU. Devi creare e testare l'immagine sistema operativo personalizzata.

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

Prerequisiti

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

  • Docker
  • ethtool, versione 6.7 o successive se utilizzi TPU v6e o TPU7x
  • Python 3 con i pacchetti PyYAML e requests, disponibili in Ubuntu come python3-yaml e python3-requests
  • gcloud CLI

Per le VM TPU7x, l'immagine di base deve utilizzare una versione minima del kernel Linux del sistema operativo guest pari a 6.8.

Apporta modifiche durante la creazione dell'immagine

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

Disattiva gli upgrade automatici

Disattiva e rimuovi unattended-upgrades in modo che gli aggiornamenti dei pacchetti in background non blocchino dpkg o interrompano i workload TPU:

echo 'APT::Periodic::Unattended-Upgrade "0";' > /etc/apt/apt.conf.d/99cloudtpuconf
DEBIAN_FRONTEND=noninteractive apt-get -y purge unattended-upgrades

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

# 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

Modifica l'immagine per migliorare il rendimento

Per garantire prestazioni ottimali, regola 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 di 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.
  • transparent_hugepage=always: attiva le Transparent Huge Pages (THP).
  • iommu.passthrough=0: disattiva la modalità passthrough dell'unità di gestione della memoria di input/output (IOMMU).
  • panic=-1: configura il kernel in modo che si riavvii immediatamente se si verifica un kernel panic.
  • intel_iommu=on,sm_on: attiva Intel IOMMU in modalità scalabile. Obbligatorio per le architetture TPU7x e v5p.

I seguenti passaggi mostrano come aggiornare questi parametri del kernel:

  1. Imposta la variabile della riga di comando del kernel con i parametri comuni:

    kernel_cmdline="idle=poll transparent_hugepage=always iommu.passthrough=0 panic=-1"

  2. Se stai creando un'immagine per TPU7x o TPU v5p, aggiungi il parametro Intel IOMMU:

    kernel_cmdline="${kernel_cmdline} intel_iommu=on,sm_on"

  3. Aggiorna GRUB_CMDLINE_LINUX in /etc/default/grub e rigenera la configurazione GRUB:

    sed -i "s/^GRUB_CMDLINE_LINUX=\".*\"/GRUB_CMDLINE_LINUX=\"${kernel_cmdline}\"/" /etc/default/grub
echo "Status: New kernel cmdline: $(grep -e '^GRUB_CMDLINE_LINUX=' /etc/default/grub)"

update-grub

Installa l'agente vBar

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

Per eseguire il pull dell'immagine container dell'agente vBar durante la creazione dell'immagine, esegui i seguenti comandi:

  1. Autentica Docker con gcr.io:

    gcloud auth configure-docker gcr.io

  2. Esegui il pull dell'immagine Docker dell'agente vBar:

    docker pull gcr.io/cloud-tpu-v2-images/vbar_control_agent:0.0.3

Per scoprire come configurare l'immagine in modo da avviare il container dell'agente vBar all'avvio, consulta Avvia l'agente vBar.

(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 stabilire gli intervalli di raccolta, attivare o disattivare metriche specifiche o modificare le destinazioni di esportazione.

Per installare ed eseguire AI Telemetry Collector, esegui questi comandi:

  1. Autentica Docker con gcr.io:

    gcloud auth configure-docker gcr.io

  2. Esegui il pull dell'immagine Docker di AI Telemetry Collector:

    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 -d --restart=always --privileged --net=host gcr.io/cloud-tpu-v2-images/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.

Apporta modifiche al tempo di avvio

Configura l'immagine in modo che esegua le attività descritte nelle sezioni seguenti ogni volta che viene avviata una VM. 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: della configurazione cloud-init.

Avvia l'agente vBar

Avvia il container dell'agente di controllo vBar all'avvio:

docker run -d --restart=always --privileged --net=host \
  -e HOME=/home/tpu-runtime --name=vbar_control_agent \
  gcr.io/cloud-tpu-v2-images/vbar_control_agent:0.0.3 \
  vbar_control_agent_files/bin/vbar_control_agent --logtostderr

Configura le variabili di ambiente

Per assicurarti che l'ambiente sia inizializzato correttamente per i workload TPU, devi recuperare le variabili di configurazione del runtime dal server dei 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 dei metadati tpu-env e li salva in /home/tpu-runtime/tpu-env per l'utilizzo da parte dello stack software TPU.

write_files:
  - path: /var/scripts/configure-env-vars.sh
    permissions: '0755'
    owner: root
    content: |
      #!/bin/bash
      grep -q CLOUDSDK_PYTHON /etc/environment || echo "CLOUDSDK_PYTHON=/usr/bin/python3" >> /etc/environment
      export HOME=/home/tpu-runtime
      mkdir -p "${HOME}"
      curl -s 'http://metadata.google.internal/computeMetadata/v1/instance/attributes/tpu-env' -H 'Metadata-Flavor: Google' > /tmp/tpu-env.yaml
      python3 -c '
      import yaml
      with open("/tmp/tpu-env.yaml", "r") as stream_in:
        for k, v in (yaml.safe_load(stream_in) or {}).items():
          print("{var}={value}".format(var=k, value=str(v)))
      ' > "${HOME}/tpu-env"
      rm -f /tmp/tpu-env.yaml
      set -a
      source "${HOME}/tpu-env"
      set +a

Recupera i metadati della VM

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

write_files:
  - path: /var/scripts/get-vm-metadata.py
    permissions: '0755'
    owner: root
    content: |
      #!/usr/bin/env python3
      import os
      import sys
      import requests

      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.text)
      elif request.status_code in (403, 404):
        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 dei timeout aggiungendo 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.

write_files:
  - path: /var/scripts/configure-gcs-timeouts.sh
    permissions: '0755'
    owner: root
    content: |
      #!/bin/bash
      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

Configura le impostazioni di rete

Per ottimizzare il throughput e la latenza di rete sulle VM TPU, aggiungi i seguenti script alla sezione write_files: della configurazione cloud-init. Sulle VM TPU v6e e TPU7x, questi script abilitano anche la suddivisione dei dati TCP, l'aumento dei limiti del buffer di ricezione TCP e la distribuzione degli interrupt NIC tra i core NUMA. Questi script richiedono lo script /var/scripts/get-vm-metadata.py della sezione Recupera i metadati della VM.

write_files:
  - path: /var/scripts/configure-networking.sh
    permissions: '0755'
    owner: root
    content: |
      #!/bin/bash
      ACCELERATOR_TYPE=$(/usr/bin/python3 /var/scripts/get-vm-metadata.py accelerator-type)
      if [[ "$ACCELERATOR_TYPE" =~ ^v6e ]] || [[ "$ACCELERATOR_TYPE" =~ ^tpu7 ]]; then
        echo "Enabling max rx buffer size for v6e, tpu7x & tpu7"
        for port in $(ls /sys/class/net | grep ens); do
          dhclient -r "$port"
          if ethtool -h | grep -q tcp-data-split; then
            if ethtool -G "$port" tcp-data-split on; then
              echo "Max RX buffer size and TCP data split enabled for $port"
            else
              echo "Unable to enable max RX buffer size and TCP data split for $port"
            fi
          else
            if ethtool --set-priv-flags "$port" enable-max-rx-buffer-size on; then
              echo "Max RX buffer size enabled for $port"
            else
              echo "Unable to enable max RX buffer size for $port"
            fi
          fi
          dhclient "$port"
        done
      else
        echo "Skipping max rx buffer size for TPU generation $ACCELERATOR_TYPE before v6e"
      fi

  - path: /var/scripts/set_tcp_config.sh
    permissions: '0755'
    owner: root
    content: |
      #!/bin/bash
      echo "Setting TCP config"

      # Disable metrics cache
      sysctl -w net.ipv4.tcp_no_metrics_save=1

      # Disable slow start after idle
      sysctl -w net.ipv4.tcp_slow_start_after_idle=0

      # Set rto_min 5ms and enable quickack
      sysctl_rto_min_exists=$(sysctl -a | grep tcp_rto_min_us)
      if [[ -z "$sysctl_rto_min_exists" ]]; then
        ip route show | while IFS= read -r route; do
          if ! echo "$route" | grep -q "linkdown"; then
            ip route change ${route/lock/} rto_min 5ms quickack 1
          fi
        done
      else
        sysctl -w net.ipv4.tcp_rto_min_us=5000
        ip route show | while IFS= read -r route; do
          if ! echo "$route" | grep -q "linkdown"; then
            ip route change ${route/lock/} quickack 1
          fi
        done
      fi

      # Increase TCP zerocopy control memory
      sysctl -w net.core.optmem_max=131072

      # Disable Cubic Hystart Ack-Train
      echo 2 > /sys/module/tcp_cubic/parameters/hystart_detect

      ACCELERATOR_TYPE=$(/usr/bin/python3 /var/scripts/get-vm-metadata.py accelerator-type)
      if [[ "$ACCELERATOR_TYPE" =~ ^v6e ]] || [[ "$ACCELERATOR_TYPE" =~ ^tpu7 ]]; then
        # Increase TCP receive buffer limits
        echo "4096 41943040 314572800" > /proc/sys/net/ipv4/tcp_rmem
      fi

  - path: /var/scripts/spread-nic-irqs.sh
    permissions: '0755'
    owner: root
    content: |
      #!/bin/bash
      ACCELERATOR_TYPE=$(/usr/bin/python3 /var/scripts/get-vm-metadata.py accelerator-type)
      if [[ "$ACCELERATOR_TYPE" == v6e* ]] || [[ "$ACCELERATOR_TYPE" == tpu7* ]]; then
        node0_cores=$(echo /sys/devices/system/node/node0/cpu[0-9]* | wc -w)
        ht_buddies=$(cat /sys/devices/system/cpu/cpu0/topology/core_cpus_list | tr ',' ' ' | wc -w)
        total_schedulable_cores=$((node0_cores / ht_buddies))
        second_numa_starting_core=${total_schedulable_cores}

        first_numa_core=0
        second_numa_core=${second_numa_starting_core}
        for nic in $(ls -1 /sys/class/net); do
          if [[ -d "/sys/class/net/$nic/device" ]]; then
            NUM_QUEUES=$(echo /sys/class/net/"$nic"/queues/tx* | wc -w)
            base_ntfy_block=$(ls -1 /sys/class/net/"${nic}"/device/msi_irqs | sort -n | head -n 1)
            num_ntfy_blocks=$(ls /sys/class/net/"${nic}"/device/msi_irqs/ | wc -w)
            tx_irq_base_directory=$((base_ntfy_block))
            rx_irq_base_directory=$((base_ntfy_block + (num_ntfy_blocks / 2)))

            core=0
            max_core=${total_schedulable_cores}
            nic_numa=$(cat /sys/class/net/"${nic}"/device/numa_node)
            if [[ "${nic_numa}" -eq "1" ]]; then
              core=${second_numa_core}
              max_core=$((second_numa_starting_core + total_schedulable_cores))
            else
              core=${first_numa_core}
            fi

            for ((queue = 0; queue < $NUM_QUEUES; queue++)); do
              echo $core > /proc/irq/$((tx_irq_base_directory + $queue))/smp_affinity_list
              echo $core > /proc/irq/$((rx_irq_base_directory + $queue))/smp_affinity_list
              cp /proc/irq/$((tx_irq_base_directory + $queue))/smp_affinity /sys/class/net/"$nic"/queues/tx-"$queue"/xps_cpus
              core=$((core + 1))
              if [[ $core -ge $max_core ]]; then
                if [[ "${nic_numa}" -eq "1" ]]; then
                  core=${second_numa_starting_core}
                else
                  core=0
                fi
              fi
            done

            if [[ "${nic_numa}" -eq "1" ]]; then
              second_numa_core=${core}
            else
              first_numa_core=${core}
            fi
          fi
        done
      else
        echo "Skipping NIC IRQ spreading for TPU generation $ACCELERATOR_TYPE before v6e"
      fi

Esegui gli script di avvio

Per eseguire automaticamente gli script di configurazione all'avvio della VM, aggiungi una sezione runcmd: alla configurazione cloud-init:

runcmd:
  - /var/scripts/configure-env-vars.sh
  - /var/scripts/configure-gcs-timeouts.sh
  - /var/scripts/configure-networking.sh
  - /var/scripts/set_tcp_config.sh
  - /var/scripts/spread-nic-irqs.sh

Passaggi successivi