カスタム OS イメージを使用する

TPU VM にカスタム OS イメージを使用すると、ソフトウェアのプリロード、特定の OS ディストリビューションの使用、カスタム カーネルの変更の適用を行うことができます。カスタム イメージを作成するには、イメージ作成プロセス中に特定のシステム変更を行い、TPU 機能に必要な起動タスクを処理するようにイメージを構成します。

TPU でカスタム OS イメージを使用する場合は、次の免責事項に注意してください。

  • Google は、デフォルトの TPU に最適化された Ubuntu の長期サポート（LTS）イメージを提供しています。このページに記載されている OS の変更は、Google がサポートする TPU に最適化された Ubuntu LTS イメージでのみ検証されています。
  • 他の OS ディストリビューションまたはカスタム イメージに必要な OS の変更を推定する責任は、お客様にあります。このページに記載されている Ubuntu の変更が、他の OS ディストリビューションやカスタム カーネルを含む別の Ubuntu イメージで動作することは保証されません。
  • Google は、デフォルトの TPU に最適化された Ubuntu LTS イメージ以外の OS イメージのビルドやテストを行いません。カスタム OS イメージをビルドしてテストする必要があります。

デフォルトの TPU 最適化 Ubuntu LTS イメージの詳細については、TPU OS イメージをご覧ください。

前提条件

ベースイメージには、次のコンポーネントがインストールされている必要があります。

  • Docker
  • ethtool、TPU v6e または TPU7x を使用する場合はバージョン 6.7 以降
  • PyYAML パッケージと requests パッケージを含む Python 3（Ubuntu では python3-yaml と python3-requests として利用可能）
  • gcloud CLI

TPU7x VM の場合、ベースイメージは 6.8 以降のゲスト OS Linux カーネル バージョンを使用する必要があります。

イメージの作成時に変更を加える

カスタム Ubuntu イメージをビルドするときに、次の変更を適用します。

自動アップグレードを無効にする

バックグラウンド パッケージの更新が dpkg ロックを保持したり、TPU ワークロードを中断したりしないように、unattended-upgrades を無効にして削除します。

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

TPU デバイスを VFIO にバインドする

ゲスト OS が 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 を更新して、1 つのプロセスで無制限のメモリをロックできるようにします。

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: Transparent Huge Page（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 エージェント コンテナ イメージを pull するには、次のコマンドを実行します。

  1. gcr.io で Docker を認証します。

    gcloud auth configure-docker gcr.io

  2. vBar エージェントの Docker イメージを pull します。

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

起動時に vBar エージェント コンテナを起動するようにイメージを構成する方法については、vBar エージェントを起動するをご覧ください。

省略可: AI Telemetry Collector をインストールして実行する

AI Telemetry Collector は TPU VM 内で実行され、Cloud Monitoring または独自の Prometheus ベースのモニタリング パイプラインを介してランタイム指標とインフラストラクチャ指標にアクセスできます。ai-telemetry-collector Docker イメージを使用すると、カスタム OS で AI Telemetry Collector を使用できます。カスタム OS にイメージをインストールし、config.yaml ファイルを使用して、収集間隔の指定、特定の指標の有効化または無効化、エクスポート先の変更を行うことができます。

AI Telemetry Collector をインストールして実行するには、次のコマンドを実行します。

  1. gcr.io で Docker を認証します。

    gcloud auth configure-docker gcr.io

  2. AI Telemetry Collector の Docker イメージを pull します。

    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 をご覧ください。

起動時に変更を加える

VM が起動するたびに次のセクションのタスクを実行するようにイメージを構成します。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 メタデータキーから属性の取得を自動化し、TPU ソフトウェア スタックで使用される /home/tpu-runtime/tpu-env に保存します。

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

VM メタデータを取得する

次のスニペットは、/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 VM のネットワーク スループットとレイテンシを最適化するには、cloud-init 構成の write_files: セクションに次のスクリプトを追加します。TPU v6e VM と TPU7x VM では、これらのスクリプトにより、TCP データ分割、TCP 受信バッファの上限の引き上げ、NUMA コア間での NIC 割り込みの分散も有効になります。これらのスクリプトには、VM メタデータを取得するセクションの /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

ブート スクリプトを実行する

VM の起動時に構成スクリプトを自動的に実行するには、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

次のステップ