使用自定义操作系统映像

您可以为 TPU 虚拟机使用自定义操作系统映像，以预加载软件、使用特定的操作系统分发版本或应用自定义内核修改。创建自定义映像需要在映像创建过程中进行特定的系统修改，并配置映像以处理 TPU 功能所需的启动时任务。

重要提示： 本文档中列出的系统修改、脚本和配置步骤仅针对经过 Google 认证的特定 Ubuntu 版本进行了验证。如需查看支持的版本的列表，请参阅 Google 管理的操作系统映像。这些说明是 Ubuntu 专用的，未针对其他操作系统分发版或其他 Ubuntu 版本进行验证。

如果您将自定义操作系统映像与 TPU 搭配使用，请注意以下免责声明：

• Google 提供默认的已针对 TPU 优化的 Ubuntu 长期支持 (LTS) 映像。此页面上列出的操作系统变更仅针对 Google 支持的、经过 TPU 优化的 Ubuntu LTS 映像进行了验证。
• 您有责任针对任何其他操作系统分发或自定义映像推断所需的操作系统更改。Google 无法保证此页面上列出的 Ubuntu 修改适用于其他操作系统分发版本或具有自定义内核的其他 Ubuntu 映像。
• Google 不会构建或提供除默认 TPU 优化型 Ubuntu LTS 映像之外的任何操作系统映像的测试。您必须构建并测试自定义操作系统映像。

如需详细了解默认的 TPU 优化 Ubuntu LTS 映像，请参阅 TPU 操作系统映像。

前提条件

您的基础映像必须安装以下组件：

• Python 3
• gcloud CLI

在创建图片期间进行修改

在构建自定义 Ubuntu 映像时，应用以下修改。

将 TPU 设备绑定到 VFIO

如需允许 Guest 操作系统访问 TPU 硬件，您必须将 TPU 设备绑定到 vfio-pci 驱动程序。

1. 在 /etc/udev/rules.d/ 中创建一个名为 99-tpu-vfiopci.rules 的 udev 规则文件：

   # 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. 在 /lib/udev/ 中创建一个名为 bind_to_vfio_pci.sh 的脚本：

   #!/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. 让该脚本可执行：

   chmod +x /lib/udev/bind_to_vfio_pci.sh
   

4. 向系统中的所有用户授予对 TPU 设备的访问权限：

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

修改图片以提升效果

为确保最佳性能，请调整以下系统限制和参数。

内存限制

通过更新 /etc/security/limits.conf 允许单个进程锁定无限内存：

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

文件限制

通过更新 /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

内核参数

更新 GRUB 配置（通常位于 /etc/default/grub 中），以在 GRUB_CMDLINE_LINUX 中包含以下参数：

• idle=poll：防止 CPU 进入低功耗空闲状态。
• intel_iommu=on,sm_on：启用 Intel 输入-输出内存管理单元 (IOMMU)。对于 TPU7x 和 v5p 架构是必需的。
• transparent_hugepage=always：启用透明大内存页 (THP)。

以下步骤展示了如何更新这些内核参数：

1. 通过设置以下变量来防止 CPU 进入低功耗空闲状态，您将在下一步中使用该变量。

   kernel_cmdline="idle=poll"
   

2. 启用 Intel 输入/输出内存管理单元 (IOMMU)。对于 TPU7x 和 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. 启用透明巨型页面 (THP)：

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

安装 vBar 代理

芯片间互连 (ICI) 网络需要 vBar 代理才能正常运行。

注意： 如果您要构建自定义操作系统映像以用于 Google Kubernetes Engine，请跳过本部分。GKE 会为您安装和管理 vBar 代理，

在映像中安装该代理可能会与 GKE 提供的版本发生冲突。

如需安装 vBar 代理，请运行以下命令：

1. 使用 Artifact Registry 对 Docker 进行身份验证：

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

2. 从 Artifact Registry 中拉取 Docker 映像：

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

3. 使用 vBar 代理映像运行容器：

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

可选：安装并运行 AI 遥测数据收集器

AI 遥测收集器在 TPU 虚拟机内运行，可让您通过 Cloud Monitoring 或您自己的基于 Prometheus 的监控流水线访问运行时和基础设施指标。您可以使用 ai-telemetry-collector Docker 映像将 AI 遥测收集器与自定义操作系统搭配使用。您可以将映像安装到自定义操作系统上，并使用 config.yaml 文件来指定收集间隔、启用或停用特定指标，或更改导出目的地。

注意： 如果您要构建自定义操作系统映像以用于 Google Kubernetes Engine，请跳过本部分。GKE 会为您安装和管理 AI 遥测收集器。

如需安装 AI 遥测收集器，请运行以下命令：

1. 使用 Artifact Registry 对 Docker 进行身份验证：

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

2. 从 Artifact Registry 中拉取 Docker 映像：

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

3. 使用 AI 遥测收集器映像运行容器，并采用默认配置：

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

   如需了解如何使用自定义配置文件或添加其他配置文件，请参阅 AI 遥测收集器。

修改启动时间

将映像配置为在每次虚拟机启动时执行以下部分中的任务。您可以使用 cloud-init 工具通过向实例传递元数据来配置启动时任务。以下部分中的配置使用 write_files 和 runcmd 等模块。定义要写入的文件的代码段应包含在 cloud-init 配置中的 write_files: 键下，而应在启动时运行的命令应包含在 runcmd: 键下。

启动 vBar 代理

注意： 如果您要构建自定义操作系统映像以用于 Google Kubernetes Engine，请跳过本部分。GKE 会为您运行和管理 vBar 代理。

使用适当的用户 ID 和群组 ID 启动 vBar 控制代理：

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

配置环境变量

为确保您的环境已针对 TPU 工作负载正确初始化，您必须在系统启动过程中从 Compute Engine 元数据服务器检索运行时配置变量。为此，请将以下代码段添加到 cloud-init 配置的 write_files: 部分，以创建一个名为 /var/scripts/configure-env-vars.sh 的脚本。此脚本可自动从 tpu-env 元数据密钥中检索属性，并将其保存在 /${HOME}/tpu-env 中，以供 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

获取虚拟机元数据

以下代码段创建了一个名为 /var/scripts/get-vm-metadata.py 的脚本，这是一个 Python 实用程序，用于以编程方式查询元数据服务器以获取特定实例属性和自定义元数据标记。将以下内容添加到 cloud-init 配置的 write_files: 部分：

 - 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))

增加 Cloud Storage 超时时间

如果您的工作负载与 Cloud Storage 进行交互，请通过向 /etc/environment 添加超时值来增加超时时长。为此，请将以下代码段添加到 cloud-init 配置的 write_files: 部分，该代码段会创建一个名为 /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

后续步骤

• 查看可用的 TPU 操作系统映像。
• 了解如何管理 TPU 虚拟机。