使用自定义操作系统映像

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

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

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

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

前提条件

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

  • Docker
  • ethtool，如果您使用 TPU v6e 或 TPU7x，则为 6.7 或更高版本
  • Python 3 以及 PyYAML 和 requests 软件包，在 Ubuntu 中以 python3-yaml 和 python3-requests 的形式提供
  • gcloud CLI

对于 TPU7x 虚拟机，您的基础映像必须使用 6.8 或更高版本的客机操作系统 Linux 内核。

在创建映像期间进行修改

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

停用无人值守升级

停用并移除 unattended-upgrades，以便后台软件包更新不会保留 dpkg 锁定或中断 TPU 工作负载：

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

将 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

# 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 进入低功耗空闲状态。
  • transparent_hugepage=always：启用透明大内存页 (THP)。
  • iommu.passthrough=0：停用输入输出内存管理单元 (IOMMU) 直通模式。
  • panic=-1：将内核配置为在发生内核崩溃时立即重新启动。
  • intel_iommu=on,sm_on：在可伸缩模式下启用 Intel IOMMU。TPU7x 和 v5p 架构需要此参数。

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

  1. 使用常用参数设置内核命令行变量：

    kernel_cmdline="idle=poll transparent_hugepage=always iommu.passthrough=0 panic=-1"

  2. 如果您要为 TPU7x 或 TPU v5p 构建映像，请附加 Intel IOMMU 参数：

    kernel_cmdline="${kernel_cmdline} intel_iommu=on,sm_on"

  3. 更新 /etc/default/grub 中的 GRUB_CMDLINE_LINUX 并重新生成 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

安装 vBar 代理

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

如需在创建映像期间拉取 vBar 代理容器映像，请运行以下命令：

  1. 使用 gcr.io 对 Docker 进行身份验证：

    gcloud auth configure-docker gcr.io

  2. 拉取 vBar 代理 Docker 映像：

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

如需了解如何配置映像以在启动时启动 vBar 代理容器，请参阅启动 vBar 代理。

可选：安装并运行 AI Telemetry Collector

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

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

  1. 使用 gcr.io 对 Docker 进行身份验证：

    gcloud auth configure-docker gcr.io

  2. 拉取 AI Telemetry Collector Docker 映像：

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

  3. 使用 AI Telemetry Collector 映像运行具有默认配置的容器：

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

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

在启动时进行修改

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

启动 vBar 代理

在启动时启动 vBar 控制代理容器：

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

配置环境变量

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

获取虚拟机元数据

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

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

增加 Cloud Storage 超时时间

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

配置网络设置

如需优化 TPU 虚拟机上的网络吞吐量和延迟时间，请将以下脚本添加到 cloud-init 配置的 write_files: 部分。在 TPU v6e 和 TPU7x 虚拟机上，这些脚本还支持 TCP 数据拆分、提高 TCP 接收缓冲区限制，以及在 NUMA 核心之间分散 NIC 中断。这些脚本需要获取虚拟机元数据部分中的 /var/scripts/get-vm-metadata.py 脚本。

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

运行启动脚本

如需在虚拟机启动时自动运行配置脚本，请在 cloud-init 配置中添加 runcmd: 部分：

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

后续步骤