Usa una imagen de SO personalizada

Puedes usar una imagen de SO personalizada para tus VMs de TPU con el objetivo de precargar software, usar una distribución de SO específica o aplicar modificaciones personalizadas del kernel. Para crear una imagen personalizada, se deben realizar modificaciones específicas en el sistema durante el proceso de creación de la imagen y configurarla para que controle las tareas de tiempo de inicio necesarias para la funcionalidad de la TPU.

Ten en cuenta las siguientes exenciones de responsabilidad si usas una imagen de SO personalizada con TPU:

  • Google proporciona imágenes predeterminadas de asistencia a largo plazo (LTS) de Ubuntu optimizadas para TPU. Los cambios del SO que se indican en esta página solo se validan para las imágenes de Ubuntu LTS optimizadas para TPU y compatibles con Google.
  • Eres responsable de extrapolar los cambios necesarios del SO para cualquier otra distribución del SO o imágenes personalizadas. Google no garantiza que las modificaciones para Ubuntu que se indican en esta página funcionen con otras distribuciones de SO o con otra imagen de Ubuntu con un kernel personalizado.
  • Google no compila ni proporciona pruebas para ninguna imagen de SO que no sean las imágenes predeterminadas de Ubuntu LTS optimizadas para TPU. Debes compilar y probar tu imagen de SO personalizada.

Para obtener más información sobre las imágenes de Ubuntu LTS predeterminadas optimizadas para TPU, consulta Imágenes de SO de TPU.

Requisitos previos

La imagen base debe tener instalados los siguientes componentes:

  • Docker
  • ethtool, versión 6.7 o posterior si usas TPU v6e o TPU7x
  • Python 3 con los paquetes PyYAML y requests, disponibles en Ubuntu como python3-yaml y python3-requests
  • gcloud CLI

Para las VMs de TPU7x, tu imagen base debe usar una versión mínima de kernel de Linux del SO invitado de 6.8.

Realiza modificaciones durante la creación de imágenes

Aplica las siguientes modificaciones mientras compilas tu imagen de Ubuntu personalizada.

Inhabilitar las actualizaciones automáticas

Inhabilita y quita unattended-upgrades para que las actualizaciones de paquetes en segundo plano no retengan el bloqueo de dpkg ni interrumpan las cargas de trabajo de TPU:

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

Vincula dispositivos TPU a VFIO

Para permitir que el SO invitado acceda al hardware de TPU, debes vincular los dispositivos TPU al controlador vfio-pci.

  1. Crea un archivo de reglas de udev llamado 99-tpu-vfiopci.rules en /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 una secuencia de comandos llamada bind_to_vfio_pci.sh en /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. Haz que la secuencia de comandos sea ejecutable:

    chmod +x /lib/udev/bind_to_vfio_pci.sh

  4. Otorga acceso al dispositivo TPU a todos los usuarios del sistema:

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

Modifica la imagen para mejorar el rendimiento

Para garantizar un rendimiento óptimo, ajusta los siguientes límites y parámetros del sistema.

Límites de memoria

Permite que un solo proceso bloquee memoria ilimitada actualizando /etc/security/limits.conf:

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

Límites de archivos

Aumenta la cantidad de archivos abiertos actualizando /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

Parámetros del kernel

Actualiza tu configuración de GRUB (por lo general, en /etc/default/grub) para incluir los siguientes parámetros en GRUB_CMDLINE_LINUX:

  • idle=poll: Impide que la CPU entre en estados de inactividad de bajo consumo.
  • transparent_hugepage=always: Habilita las páginas enormes transparentes (THP).
  • iommu.passthrough=0: Inhabilita el modo de transferencia de la unidad de administración de la memoria de entrada y salida (IOMMU).
  • panic=-1: Configura el kernel para que se reinicie de inmediato si se produce un error irrecuperable del kernel.
  • intel_iommu=on,sm_on: Habilita la IOMMU de Intel en modo escalable. Es obligatorio para las arquitecturas TPU7x y v5p.

En los siguientes pasos, se muestra cómo actualizar estos parámetros del kernel:

  1. Establece la variable de línea de comandos del kernel con los parámetros comunes:

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

  2. Si compilas una imagen para TPU7x o TPU v5p, agrega el parámetro de IOMMU de Intel:

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

  3. Actualiza GRUB_CMDLINE_LINUX en /etc/default/grub y vuelve a generar la configuración de 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

Instala el agente de vBar

El agente de vBar es necesario para que funcione la red de interconexión entre chips (ICI).

Para extraer la imagen de contenedor del agente de vBar durante la creación de la imagen, ejecuta los siguientes comandos:

  1. Autentica Docker con gcr.io:

    gcloud auth configure-docker gcr.io

  2. Extrae la imagen de Docker del agente de vBar:

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

Para aprender a configurar la imagen para que inicie el contenedor del agente de vBar en el momento del inicio, consulta Inicia el agente de vBar.

Opcional: Instala y ejecuta el recopilador de telemetría de IA

El recopilador de telemetría de IA se ejecuta dentro de la VM de TPU y te permite acceder a las métricas de infraestructura y del entorno de ejecución a través de Cloud Monitoring o de tu propia canalización de supervisión basada en Prometheus. Puedes usar el recopilador de telemetría de IA con un SO personalizado usando la imagen de Docker ai-telemetry-collector. Puedes instalar la imagen en tu SO personalizado y usar un archivo config.yaml para dictar los intervalos de recopilación, habilitar o inhabilitar métricas específicas, o cambiar los destinos de exportación.

Para instalar y ejecutar el recopilador de telemetría de IA, ejecuta los siguientes comandos:

  1. Autentica Docker con gcr.io:

    gcloud auth configure-docker gcr.io

  2. Extrae la imagen de Docker del recopilador de telemetría de IA:

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

  3. Ejecuta un contenedor con la imagen de AI Telemetry Collector con la configuración predeterminada:

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

    Para obtener información sobre cómo usar un archivo de configuración personalizado o agregar archivos de configuración adicionales, consulta AI Telemetry Collector.

Realiza modificaciones en el tiempo de inicio

Configura tu imagen para que realice las tareas de las siguientes secciones cada vez que se inicie una VM. Puedes usar la herramienta cloud-init para configurar tareas de tiempo de inicio pasando metadatos a tus instancias. Las configuraciones de las siguientes secciones usan módulos como write_files y runcmd. Los fragmentos que definen los archivos que se escribirán deben incluirse en la clave write_files:, y los comandos que se deben ejecutar en el momento del inicio deben incluirse en la clave runcmd: de tu configuración de cloud-init.

Inicia el agente de vBar

Inicia el contenedor del agente de control de vBar en el momento del inicio:

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

Configure las variables de entorno

Para asegurarte de que tu entorno se inicialice correctamente para las cargas de trabajo de TPU, debes recuperar las variables de configuración del entorno de ejecución del servidor de metadatos de Compute Engine durante el proceso de inicio del sistema. Para ello, agrega el siguiente fragmento a la sección write_files: de tu configuración de cloud-init, que crea una secuencia de comandos llamada /var/scripts/configure-env-vars.sh. Esta secuencia de comandos automatiza la recuperación de atributos de la clave de metadatos tpu-env y los guarda en /home/tpu-runtime/tpu-env para que los use la pila de software de 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

Obtén metadatos de VM

En el siguiente fragmento, se crea una secuencia de comandos llamada /var/scripts/get-vm-metadata.py, una utilidad de Python para consultar de forma programática el servidor de metadatos en busca de atributos de instancia específicos y etiquetas de metadatos personalizadas. Agrega lo siguiente a la sección write_files: de tu configuración de 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 los tiempos de espera de Cloud Storage

Si tu carga de trabajo interactúa con Cloud Storage, aumenta la duración de los tiempos de espera agregando valores de tiempo de espera a /etc/environment. Para ello, agrega el siguiente fragmento a la sección write_files: de tu configuración de cloud-init, que crea una secuencia de comandos llamada /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

Establecer la configuración de red

Para optimizar la capacidad de procesamiento y la latencia de la red en las VMs de TPU, agrega las siguientes secuencias de comandos a la sección write_files: de tu configuración de cloud-init. En las VMs TPU v6e y TPU7x, estas secuencias de comandos también habilitan la división de datos de TCP, aumentan los límites del búfer de recepción de TCP y distribuyen las interrupciones de NIC en los núcleos de NUMA. Estas secuencias de comandos requieren la secuencia de comandos /var/scripts/get-vm-metadata.py de la sección Obtén metadatos 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

Ejecuta las secuencias de comandos de inicio

Para ejecutar las secuencias de comandos de configuración automáticamente cuando se inicia la VM, agrega una sección runcmd: a tu configuración de 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

¿Qué sigue?