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

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

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

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

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

前提条件

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

  • Python 3
  • gcloud CLI

イメージの作成中に変更を行う

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

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
#!/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: Transparent Huge Pages（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. Transparent Huge Pages（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 イメージを pull します。

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

AI テレメトリー コレクタをインストールするには、次のコマンドを実行します。

  1. Artifact Registry で Docker を認証します。

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

  2. Artifact Registry から Docker イメージを pull します。

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

  3. デフォルト構成で AI テレメトリー コレクタ イメージを使用してコンテナを実行します。

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

    カスタム構成ファイルの使用や追加の構成ファイルの追加については、AI テレメトリー コレクタをご覧ください。

起動時間の変更

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

 - 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

VM メタデータを取得する

次のスニペットは、/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

次のステップ