Utiliser une image d'OS personnalisée

Vous pouvez utiliser une image d'OS personnalisée pour vos VM TPU afin de précharger des logiciels, d'utiliser une distribution d'OS spécifique ou d'appliquer des modifications de kernel personnalisées. La création d'une image personnalisée implique d'apporter des modifications spécifiques au système lors du processus de création de l'image et de configurer l'image pour qu'elle gère les tâches de démarrage requises pour la fonctionnalité TPU.

Tenez compte des clauses de non-responsabilité suivantes si vous utilisez une image d'OS personnalisée avec des TPU :

  • Google fournit des images Ubuntu LTS (support à long terme) optimisées pour les TPU par défaut. Les modifications d'OS listées sur cette page ne sont validées que pour les images Ubuntu LTS compatibles avec Google et optimisées pour les TPU.
  • Il vous incombe d'extrapoler les modifications d'OS requises pour toute autre distribution d'OS ou image personnalisée. Google ne garantit pas que les modifications apportées à Ubuntu répertoriées sur cette page fonctionnent avec d'autres distributions d'OS ou une autre image Ubuntu avec un noyau personnalisé.
  • Google ne crée ni ne fournit de tests pour les images de l'OS autres que les images Ubuntu LTS par défaut optimisées pour les TPU. Vous devez créer et tester votre image d'OS personnalisée.

Pour en savoir plus sur les images Ubuntu LTS optimisées pour les TPU par défaut, consultez images de l'OS TPU.

Prérequis

Les composants suivants doivent être installés sur votre image de base :

  • Docker
  • ethtool, version 6.7 ou ultérieure si vous utilisez des TPU v6e ou TPU7x
  • Python 3 avec les packages PyYAML et requests, disponibles dans Ubuntu sous les noms python3-yaml et python3-requests
  • gcloud CLI

Pour les VM TPU7x, votre image de base doit utiliser une version minimale du noyau Linux de l'OS invité de 6.8.

Apporter des modifications lors de la création de l'image

Appliquez les modifications suivantes lors de la création de votre image Ubuntu personnalisée.

Désactiver les mises à niveau sans surveillance

Désactivez et supprimez unattended-upgrades afin que les mises à jour de packages en arrière-plan ne maintiennent pas le verrouillage dpkg ou n'interrompent pas les charges de travail TPU :

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

Lier des appareils TPU à VFIO

Pour permettre à l'OS invité d'accéder au matériel TPU, vous devez lier les appareils TPU au pilote vfio-pci.

  1. Créez un fichier de règles udev nommé 99-tpu-vfiopci.rules dans /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. Créez un script nommé bind_to_vfio_pci.sh dans /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. Rendez le script exécutable :

    chmod +x /lib/udev/bind_to_vfio_pci.sh

  4. Accordez à tous les utilisateurs du système l'accès à l'appareil TPU :

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

Modifier l'image pour améliorer les performances

Pour garantir des performances optimales, ajustez les limites et paramètres système suivants.

Limites de mémoire

Autorisez un seul processus à verrouiller une mémoire illimitée en mettant à jour /etc/security/limits.conf :

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

Limites de fichiers

Augmentez le nombre de fichiers ouverts en mettant à jour /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

Paramètres du kernel

Mettez à jour votre configuration GRUB (généralement dans /etc/default/grub) pour inclure les paramètres suivants dans GRUB_CMDLINE_LINUX :

  • idle=poll : empêche le processeur de passer à un état d'inactivité à faible consommation d'énergie.
  • transparent_hugepage=always : active les pages volumineuses transparentes (THP).
  • iommu.passthrough=0 : désactive le mode de transfert de l'unité de gestion de mémoire d'entrée-sortie (IOMMU).
  • panic=-1 : configure le noyau pour qu'il redémarre immédiatement en cas de panique du noyau.
  • intel_iommu=on,sm_on : active Intel IOMMU en mode évolutif. Obligatoire pour les architectures TPU7x et v5p.

Les étapes suivantes montrent comment mettre à jour ces paramètres de kernel :

  1. Définissez la variable de ligne de commande du noyau avec les paramètres courants :

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

  2. Si vous créez une image pour TPU7x ou TPU v5p, ajoutez le paramètre Intel IOMMU :

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

  3. Mettez à jour GRUB_CMDLINE_LINUX dans /etc/default/grub et regénérez la configuration 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

Installer l'agent vBar

L'agent vBar est requis pour que le réseau d'interconnexion inter-puces (ICI) fonctionne.

Pour extraire l'image de conteneur de l'agent vBar lors de la création de l'image, exécutez les commandes suivantes :

  1. Authentifiez Docker avec gcr.io :

    gcloud auth configure-docker gcr.io

  2. Extrayez l'image Docker de l'agent vBar :

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

Pour savoir comment configurer l'image afin de démarrer le conteneur de l'agent vBar au démarrage, consultez Démarrer l'agent vBar.

Facultatif : Installer et exécuter le collecteur de télémétrie IA

Le collecteur de télémétrie IA s'exécute dans la VM TPU et vous permet d'accéder aux métriques d'exécution et d'infrastructure via Cloud Monitoring ou votre propre pipeline de surveillance basé sur Prometheus. Vous pouvez utiliser AI Telemetry Collector avec un OS personnalisé à l'aide de l'image Docker ai-telemetry-collector. Vous pouvez installer l'image sur votre OS personnalisé et utiliser un fichier config.yaml pour définir les intervalles de collecte, activer ou désactiver des métriques spécifiques, ou modifier les destinations d'exportation.

Pour installer et exécuter le collecteur de télémétrie IA, exécutez les commandes suivantes :

  1. Authentifiez Docker avec gcr.io :

    gcloud auth configure-docker gcr.io

  2. Extrayez l'image Docker d'AI Telemetry Collector :

    docker pull gcr.io/cloud-tpu-v2-images/ai-telemetry-collector:latest

  3. Exécutez un conteneur à l'aide de l'image AI Telemetry Collector avec la configuration par défaut :

    docker run -d --restart=always --privileged --net=host gcr.io/cloud-tpu-v2-images/ai-telemetry-collector:latest

    Pour en savoir plus sur l'utilisation d'un fichier de configuration personnalisé ou l'ajout de fichiers de configuration supplémentaires, consultez AI Telemetry Collector.

Apporter des modifications au moment du démarrage

Configurez votre image pour qu'elle effectue les tâches décrites dans les sections suivantes chaque fois qu'une VM démarre. Vous pouvez utiliser l'outil cloud-init pour configurer les tâches de démarrage en transmettant des métadonnées à vos instances. Les configurations des sections suivantes utilisent des modules tels que write_files et runcmd. Les extraits qui définissent les fichiers à écrire doivent être inclus sous la clé write_files:, et les commandes à exécuter au démarrage doivent être incluses sous la clé runcmd: dans votre configuration cloud-init.

Démarrer l'agent vBar

GKE exécute et gère l'agent vBar pour vous.

Démarrez le conteneur de l'agent de contrôle vBar au démarrage :

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

Configurer les variables d'environnement

Pour vous assurer que votre environnement est correctement initialisé pour les charges de travail TPU, vous devez récupérer les variables de configuration d'exécution à partir du serveur de métadonnées Compute Engine pendant le processus de démarrage du système. Pour ce faire, ajoutez l'extrait suivant à la section write_files: de votre configuration cloud-init, qui crée un script nommé /var/scripts/configure-env-vars.sh. Ce script automatise la récupération des attributs à partir de la clé de métadonnées tpu-env et les enregistre dans /home/tpu-runtime/tpu-env pour qu'ils soient utilisés par la pile logicielle 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

Obtenir les métadonnées d'une VM

L'extrait de code suivant crée un script nommé /var/scripts/get-vm-metadata.py, un utilitaire Python permettant d'interroger de manière programmatique le serveur de métadonnées pour obtenir des attributs d'instance spécifiques et des tags de métadonnées personnalisés. Ajoutez ce qui suit à la section write_files: de votre configuration 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))

Augmenter les délais avant expiration de Cloud Storage

Si votre charge de travail interagit avec Cloud Storage, augmentez les délais avant expiration en ajoutant des valeurs de délai avant expiration à /etc/environment. Pour ce faire, ajoutez l'extrait suivant à la section write_files: de votre configuration cloud-init, qui crée un script nommé /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

Configuration des paramètres réseau

Pour optimiser le débit et la latence du réseau sur les VM TPU, ajoutez les scripts suivants à la section write_files: de votre configuration cloud-init. Sur les VM TPU v6e et TPU7x, ces scripts permettent également de diviser les données TCP, d'augmenter les limites de la mémoire tampon de réception TCP et de répartir les interruptions de la carte d'interface réseau sur les cœurs NUMA. Ces scripts nécessitent le script /var/scripts/get-vm-metadata.py de la section Obtenir les métadonnées de 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

Exécuter les scripts de démarrage

Pour exécuter automatiquement les scripts de configuration au démarrage de la VM, ajoutez une section runcmd: à votre configuration 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

Étapes suivantes