Usar uma imagem personalizada do SO

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

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

Lembre-se das seguintes isenções de responsabilidade se você usar uma imagem personalizada do SO 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ão validadas apenas para as imagens do Ubuntu LTS otimizadas para TPU com suporte do Google.
• Você é responsável por extrapolar as mudanças necessárias do SO para qualquer outra distribuição de SO ou imagens personalizadas. O Google não garante que as modificações do 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. Você precisa criar e testar sua imagem personalizada do SO.

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

Pré-requisitos

Sua imagem de base precisa ter os seguintes componentes instalados:

• Python 3
• CLI gcloud

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

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, é necessário vincular dispositivos de TPU ao vfio-pci driver.

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 acesso ao dispositivo de TPU a todos os usuários do sistema:

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

Modificar a imagem para melhorar a performance

Para garantir a performance 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 arquivo

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 baixa energia.
• intel_iommu=on,sm_on: ativa a unidade de gerenciamento de memória de entrada/saída (IOMMU) da Intel. Necessário para arquiteturas TPU7x e v5p.
• transparent_hugepage=always: ativa páginas enormes transparentes (THP).

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

1. Impeça que a CPU entre em um estado de inatividade de baixa energia definindo 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 é necessá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 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.

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: instalar e executar o coletor de telemetria de IA

O coletor de telemetria de IA é executado na VM de TPU e permite acessar métricas de infraestrutura e de ambiente 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. Você pode instalar a imagem no 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.

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 de configuração, consulte Coletor de telemetria de IA.

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

Configure sua imagem para executar as tarefas nas seções a seguir sempre que uma VM for inicializada. É possível usar a cloud-init ferramenta para configurar tarefas de inicialização transmitindo metadados para suas instâncias. As configurações nas seções a seguir usam módulos como write_files e runcmd. Os snippets que definem 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 agente de controle vBar com os IDs de usuário e grupo apropriados:

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

Configurar 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 tpu-env chave de metadados 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 para 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 do tempo limite adicionando valores de tempo limite a /etc/environment. 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-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 de SO da TPU disponíveis.
• Saiba como gerenciar VMs de TPU.