使用自訂 OS 映像檔

您可以為 TPU VM 使用自訂 OS 映像檔，預先載入軟體、使用特定 OS 發行版本，或套用自訂核心修改項目。建立自訂映像檔時，您需要在映像檔建立過程中進行特定系統修改，並設定映像檔來處理 TPU 功能所需的開機時間工作。

重要事項： 本文列出的系統修改、指令碼和設定步驟，僅適用於 Google 認證的特定 Ubuntu 版本。如需支援的版本清單，請參閱「Google 代管的 OS 映像檔」。這些指示僅適用於 Ubuntu，不適用於其他作業系統發行版本或其他 Ubuntu 版本。

如果使用自訂 OS 映像檔搭配 TPU，請注意下列免責事項：

• Google 提供預設的 TPU 最佳化 Ubuntu 長期支援 (LTS) 映像檔。本頁列出的 OS 變更僅適用於 Google 支援的 TPU 最佳化 Ubuntu LTS 映像檔。
• 您有責任推斷其他 OS 發行版本或自訂映像檔所需的 OS 變更。Google 不保證本頁列出的 Ubuntu 修改項目適用於其他 OS 發行版本，或使用自訂核心的其他 Ubuntu 映像檔。
• Google 不會建構或提供任何 OS 映像檔的測試，預設的 TPU 最佳化 Ubuntu LTS 映像檔除外。您必須建構及測試自訂 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：啟用透明巨頁 (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 遙測收集器

AI 遙測收集器會在 TPU VM 內執行，讓您透過 Cloud Monitoring 或自己的 Prometheus 型監控管道，存取執行階段和基礎架構指標。您可以使用 ai-telemetry-collector Docker 映像檔，搭配自訂 OS 使用 AI 遙測收集器。您可以將映像檔安裝到自訂 OS，並使用 config.yaml 檔案指定收集間隔、啟用或停用特定指標，或是變更匯出目的地。

如要安裝 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 等模組。定義要寫入檔案的程式碼片段應包含在 write_files: 鍵底下，而開機時應執行的指令則應包含在 cloud-init 設定中的 runcmd: 鍵底下。

啟動 vBar 代理程式

使用適當的使用者和群組 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

取得 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

後續步驟

• 查看可用的 TPU OS 映像檔。
• 瞭解如何管理 TPU VM。