Usar uma imagem do SO personalizada

É possível usar uma imagem de SO personalizada para suas VMs de TPU e pré-carregar software, usar uma distribuição de SO específica ou aplicar modificações personalizadas no 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.

Importante: as modificações do sistema, os scripts e as etapas de configuração listados neste documento são validados apenas para as versões específicas do Ubuntu certificadas pelo Google. Para uma lista de versões compatíveis, consulte Imagens de SO gerenciadas pelo Google. Estas instruções são específicas do Ubuntu e não foram validadas para outras distribuições de sistemas operacionais ou versões do Ubuntu.

Tenha em mente as seguintes exonerações de responsabilidade se você usar uma imagem de SO personalizada com TPUs:

• O Google oferece imagens padrão do Ubuntu com suporte de longo prazo (LTS) otimizadas para TPU. As mudanças de SO listadas nesta página só são validadas para as imagens do Ubuntu LTS otimizadas para TPU e com suporte do 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 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 oferece testes para imagens de SO que não sejam as imagens padrão do Ubuntu LTS otimizadas para TPU. É necessário criar e testar sua imagem de 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

Sua imagem de base precisa ter os seguintes componentes instalados:

• Python 3
• CLI da gcloud

Fazer modificações durante a criação de imagens

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

Vincular dispositivos de TPU ao VFIO

Para permitir que o SO convidado acesse o hardware da TPU, vincule dispositivos de 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
   #!/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. Torne o script executável:

   chmod +x /lib/udev/bind_to_vfio_pci.sh
   

4. Conceda a todos os usuários no 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 (normalmente 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.
• intel_iommu=on,sm_on: ativa a unidade de gerenciamento de memória de entrada/saída (IOMMU) da Intel. Obrigatório para arquiteturas TPU7x e v5p.
• transparent_hugepage=always: ativa páginas enormes transparentes (THP, na sigla em inglês).

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

1. Para evitar que a CPU entre em um estado de inatividade de baixo consumo de energia, defina a seguinte variável, que será usada na próxima etapa.

   kernel_cmdline="idle=poll"
   

2. Ative a unidade de gerenciamento de memória de entrada/saída (IOMMU) da Intel. Essa etapa é obrigatória para TPU7x e 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. Ative as 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
   

Instalar o agente vBar

O agente vBar é necessário para que a rede de interconexão entre chips (ICI) funcione.

Observação: pule esta seção se você estiver criando uma imagem de SO personalizada para uso com o Google Kubernetes Engine. O GKE instala e gerencia o agente vBar para você, e a instalação na imagem pode entrar em conflito com a versão fornecida pelo GKE.

Para instalar o agente vBar, execute os seguintes comandos:

1. Autentique o Docker com o Artifact Registry:

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

2. Extraia a imagem Docker do Artifact Registry:

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

3. Execute um contêiner usando a imagem do agente vBar:

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

Opcional: instale e execute o coletor de telemetria de IA

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

Observação: pule esta seção se você estiver criando uma imagem de SO personalizada para uso com o Google Kubernetes Engine. O GKE instala e gerencia o coletor de telemetria de IA para você.

Para instalar o coletor de telemetria de IA, execute os seguintes comandos:

1. Autentique o Docker com o Artifact Registry:

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

2. Extraia a imagem Docker do Artifact Registry:

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

3. Execute um contêiner usando a imagem do coletor de telemetria de IA com a configuração padrão:

   docker run --privileged --net=host ai-telemetry-collector:latest
   

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

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

Configure a imagem para realizar 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 às 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 na inicialização precisam ser incluídos na chave runcmd: na configuração cloud-init.

Iniciar o agente vBar

Observação: pule esta seção se você estiver criando uma imagem de SO personalizada para uso com o Google Kubernetes Engine. O GKE executa e gerencia o agente vBar para você.

Inicie o agente de controle vBar com os IDs de usuário e grupo adequados:

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

Configure as variáveis de ambiente

Para garantir que seu ambiente seja inicializado corretamente para cargas de trabalho de TPU, recupere as variáveis de configuração de tempo de execução do servidor de metadados do Compute Engine durante o processo de inicialização do sistema. Para fazer isso, adicione o seguinte snippet à 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-env para serem usados pela pilha de software da 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

Receber metadados da VM

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

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 seguinte snippet à seção write_files: da sua configuração cloud-init, que cria um script chamado /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

A seguir

• Analise as imagens do SO da TPU disponíveis.
• Saiba como gerenciar VMs de TPU.