使用自定义操作系统映像

您可以为 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

如需允许访客操作系统访问 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 代理才能正常运行。

如需安装 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 Telemetry Collector

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

如需安装 AI Telemetry Collector，请运行以下命令：

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 Telemetry Collector 映像和默认配置运行容器：

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

   如需了解如何使用自定义配置文件或添加其他 配置文件，请参阅 AI Telemetry Collector。

进行启动时间修改

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

启动 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 虚拟机。