Usar uma imagem do SO personalizada

É possível usar uma imagem do SO personalizada para as VMs de TPU para pré-carregar software, usar uma distribuição específica do SO ou aplicar modificações personalizadas do kernel. Para criar uma imagem personalizada, é necessário fazer modificações específicas no sistema durante o processo de criação da imagem e configurar a imagem para processar as tarefas de inicialização necessárias para a funcionalidade da TPU.

Lembre-se das seguintes exonerações de responsabilidade se você usar uma imagem de SO personalizada com TPUs:

  • O Google fornece imagens padrão do Ubuntu com suporte de longo prazo (LTS) otimizadas para TPU. As mudanças no SO listadas nesta página são validadas apenas para as imagens do Ubuntu LTS otimizadas para TPU e compatíveis com o Google.
  • Você é responsável por extrapolar as mudanças necessárias no SO para qualquer outra distribuição de SO ou imagens personalizadas. O Google não garante que as modificações para o Ubuntu listadas nesta página funcionem com outras distribuições de SO ou outra imagem do Ubuntu com um kernel personalizado.
  • O Google não cria nem fornece testes para imagens do SO que não sejam as imagens padrão do Ubuntu LTS otimizadas para TPU. É preciso criar e testar a imagem do SO personalizada.

Para mais informações sobre as imagens padrão do Ubuntu LTS otimizadas para TPU, consulte Imagens do SO de TPU.

Pré-requisitos

A imagem de base precisa ter os seguintes componentes instalados:

  • Docker
  • ethtool, versão 6.7 ou mais recente se você usa a TPU v6e ou a TPU7x
  • Python 3 com os pacotes PyYAML e requests, disponíveis no Ubuntu como python3-yaml e python3-requests
  • CLI da gcloud

Para VMs TPU7x, a imagem de base precisa usar uma versão mínima do kernel do Linux do SO convidado de 6.8.

Fazer modificações durante a criação da imagem

Aplique as seguintes modificações ao criar sua imagem personalizada do Ubuntu.

Desativar upgrades automáticos

Desative e remova unattended-upgrades para que as atualizações de pacotes em segundo plano não mantenham o bloqueio dpkg nem interrompam as cargas de trabalho da TPU:

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

Vincular dispositivos de TPU ao VFIO

Para permitir que o SO convidado acesse o hardware da TPU, vincule os dispositivos TPU ao driver vfio-pci.

  1. Crie um arquivo de regras udev chamado 99-tpu-vfiopci.rules em /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. Crie um script chamado bind_to_vfio_pci.sh em /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. Torne o script executável:

    chmod +x /lib/udev/bind_to_vfio_pci.sh
    
  4. Conceda a todos os usuários do sistema acesso ao dispositivo de TPU:

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

Modificar a imagem para melhorar a performance

Para garantir o desempenho ideal, ajuste os seguintes limites e parâmetros do sistema.

Limites de memória

Permita que um único processo bloqueie memória ilimitada atualizando /etc/security/limits.conf:

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

Limites de arquivos

Aumente o número de arquivos abertos atualizando /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 do kernel

Atualize a configuração do GRUB (geralmente em /etc/default/grub) para incluir os seguintes parâmetros em GRUB_CMDLINE_LINUX:

  • idle=poll: impede que a CPU entre em estados de inatividade de baixo consumo de energia.
  • transparent_hugepage=always: ativa as Transparent Huge Pages (THP).
  • iommu.passthrough=0: desativa o modo de passagem da unidade de gerenciamento de memória de entrada e saída (IOMMU, na sigla em inglês).
  • panic=-1: configura o kernel para reiniciar imediatamente se ocorrer um kernel panic.
  • intel_iommu=on,sm_on: ativa a IOMMU da Intel no modo escalonável. Obrigatório para arquiteturas TPU7x e v5p.

As etapas a seguir mostram como atualizar esses parâmetros do kernel:

  1. Defina a variável de linha de comando do kernel com os parâmetros comuns:

    kernel_cmdline="idle=poll transparent_hugepage=always iommu.passthrough=0 panic=-1"
    
  2. Se você estiver criando uma imagem para TPU7x ou TPU v5p, anexe o parâmetro IOMMU da Intel:

    kernel_cmdline="${kernel_cmdline} intel_iommu=on,sm_on"
    
  3. Atualize GRUB_CMDLINE_LINUX em /etc/default/grub e gere novamente a configuração do 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
    

Instalar o agente vBar

O agente vBar é necessário para que a rede de interconexão entre chips (ICI, na sigla em inglês) funcione.

Para extrair a imagem do contêiner do agente do vBar durante a criação da imagem, execute os seguintes comandos:

  1. Autentique o Docker com gcr.io:

    gcloud auth configure-docker gcr.io
    
  2. Extraia a imagem Docker do agente do vBar:

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

Para saber como configurar a imagem para iniciar o contêiner do agente do vBar no momento da inicialização, consulte Iniciar o agente do vBar.

Opcional: instalar e executar o coletor de telemetria de IA

O AI Telemetry Collector é executado na VM de TPU e permite acessar métricas de ambiente de execução e infraestrutura pelo Cloud Monitoring ou pelo seu próprio pipeline de monitoramento baseado no Prometheus. É possível usar o AI Telemetry Collector com um SO personalizado usando a imagem Docker ai-telemetry-collector. É possível instalar a imagem no SO personalizado e usar um arquivo config.yaml para ditar os intervalos de coleta, ativar ou desativar métricas específicas ou alterar os destinos de exportação.

Para instalar e executar o AI Telemetry Collector, execute os seguintes comandos:

  1. Autentique o Docker com gcr.io:

    gcloud auth configure-docker gcr.io
    
  2. Extraia a imagem Docker do AI Telemetry Collector:

    docker pull gcr.io/cloud-tpu-v2-images/ai-telemetry-collector:latest
    
  3. Execute um contêiner usando a imagem do AI Telemetry Collector com a configuração padrão:

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

    Para saber como usar um arquivo de configuração personalizado ou adicionar outros arquivos de configuração, consulte Coletor de telemetria de IA.

Fazer modificações no tempo de inicialização

Configure a imagem para executar as tarefas nas seções a seguir sempre que uma VM for inicializada. É possível usar a ferramenta cloud-init para configurar tarefas de tempo de inicialização transmitindo metadados para as instâncias. As configurações nas seções a seguir usam módulos como write_files e runcmd. Os snippets que definem os arquivos a serem gravados precisam ser incluídos na chave write_files:, e os comandos que precisam ser executados no momento da inicialização precisam ser incluídos na chave runcmd: na configuração cloud-init.

Iniciar o agente vBar

Inicie o contêiner do agente de controle do vBar no momento da inicialização:

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 as variáveis de ambiente

Para garantir que o ambiente seja inicializado corretamente para cargas de trabalho de TPU, é necessário recuperar as variáveis de configuração do ambiente de execução do servidor de metadados do Compute Engine durante o processo de inicialização do sistema. Para fazer isso, adicione o snippet a seguir à seção write_files: da configuração cloud-init, que cria um script chamado /var/scripts/configure-env-vars.sh. Esse script automatiza a recuperação de atributos da chave de metadados tpu-env e os salva em /home/tpu-runtime/tpu-env para serem usados pela pilha de software da 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

Acessar metadados da VM

O snippet a seguir cria um script chamado /var/scripts/get-vm-metadata.py, um utilitário Python para consultar de forma programática o servidor de metadados em busca de atributos de instância específicos e tags de metadados personalizadas. Adicione o seguinte à seção write_files: da configuração 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))

Aumentar os tempos limite do Cloud Storage

Se a carga de trabalho interagir com o Cloud Storage, aumente as durações de tempo limite adicionando valores de tempo limite a /etc/environment. Para fazer isso, adicione o snippet abaixo à seção write_files: da configuração do cloud-init, que cria um script chamado /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

Definir as configurações de rede

Para otimizar a capacidade de processamento e a latência da rede em VMs de TPU, adicione os seguintes scripts à seção write_files: da configuração cloud-init. Nas VMs de TPU v6e e TPU7x, esses scripts também permitem a divisão de dados TCP, o aumento dos limites de buffer de recebimento TCP e a distribuição de interrupções de NIC entre núcleos NUMA. Esses scripts exigem o script /var/scripts/get-vm-metadata.py da seção Receber metadados da 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

Executar os scripts de inicialização

Para executar os scripts de configuração automaticamente quando a VM for inicializada, adicione uma seção runcmd: à configuração 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

A seguir