Usa una imagen de SO personalizada

Puedes usar una imagen de SO personalizada para tus VMs de TPU y precargar software, usar una distribución de SO específica o aplicar modificaciones personalizadas del kernel. Crear una imagen personalizada implica realizar modificaciones específicas del sistema durante el proceso de creación de la imagen y configurar la imagen para que controle las tareas de tiempo de arranque necesarias para la funcionalidad de la TPU.

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

  • Google proporciona imágenes predeterminadas de Ubuntu con asistencia a largo plazo (LTS) optimizadas para TPU. Los cambios en el 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 del 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 optimizadas para TPU predeterminadas, consulta Imágenes del SO de TPU.

Requisitos previos

Tu imagen base debe tener instalados los siguientes componentes:

  • Python 3
  • gcloud CLI

Realiza modificaciones durante la creación de imágenes

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

Vincula dispositivos de TPU a VFIO

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

  1. Crea un archivo de reglas 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
#!/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. Haz que la secuencia de comandos sea ejecutable:

    chmod +x /lib/udev/bind_to_vfio_pci.sh

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

    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

Para permitir que un solo proceso bloquee una cantidad ilimitada de memoria, actualiza /etc/security/limits.conf de la siguiente manera:

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

Límites de archivos

Actualiza /etc/security/limits.conf para aumentar la cantidad de archivos abiertos:

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 la configuración de GRUB (por lo general, en /etc/default/grub) para incluir los siguientes parámetros en GRUB_CMDLINE_LINUX:

  • idle=poll: Evita que la CPU ingrese en estados de inactividad de bajo consumo.
  • intel_iommu=on,sm_on: Habilita la unidad de administración de memoria de entrada/salida (IOMMU) de Intel. Se requiere para las arquitecturas de TPU7x y v5p.
  • transparent_hugepage=always: Habilita las páginas enormes transparentes (THP).

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

  1. Para evitar que la CPU pase a un estado de inactividad de bajo consumo, configura la siguiente variable, que usarás en el siguiente paso.

    kernel_cmdline="idle=poll"

  2. Habilita la unidad de administración de memoria de entrada y salida (IOMMU) de Intel. Este paso es obligatorio para TPU7x y 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. Habilita las páginas enormes transparentes (THP):

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

update-grub

Instala el agente de vBar

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

Para instalar el agente de vBar, ejecuta los siguientes comandos:

  1. Autentica Docker con Artifact Registry:

    gcloud auth configure-docker us-docker.pkg.dev

  2. Extrae la imagen de Docker de Artifact Registry:

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

  3. Ejecuta un contenedor con la imagen del agente vBar:

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

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 de tiempo 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 a través de la imagen de Docker ai-telemetry-collector. Puedes instalar la imagen en tu SO personalizado y usar un archivo config.yaml para indicar los intervalos de recopilación, habilitar o inhabilitar métricas específicas, o cambiar los destinos de exportación.

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

  1. Autentica Docker con Artifact Registry:

    gcloud auth configure-docker us-docker.pkg.dev

  2. Extrae la imagen de Docker de Artifact Registry:

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

  3. Ejecuta un contenedor con la imagen del recopilador de telemetría de IA con la configuración predeterminada:

    docker run --privileged --net=host 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 arranque 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 durante el inicio deben incluirse en la clave runcmd: en tu configuración de cloud-init.

Inicia el agente de vBar

Inicia el agente de control de vBar con los IDs de usuario y grupo adecuados:

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

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 tiempo 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, lo que crea un script llamado /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-env para que los use la pila de software de la 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

Obtén metadatos de la VM

El siguiente fragmento 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:

 - 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 los tiempos de espera de Cloud Storage

Si tu carga de trabajo interactúa con Cloud Storage, aumenta la duración del tiempo 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.

 - 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

¿Qué sigue?