在 Ray 資訊主頁上查看 TPU 指標

本文說明如何透過 Google Kubernetes Engine (GKE) 上的 KubeRay，在 Ray 資訊主頁上查看 TPU 指標。在搭載 Ray GKE 外掛程式的 GKE 叢集中，Cloud Monitoring 會提供 TPU 指標。

事前準備

開始之前，請確認您已完成下列工作：

  • 啟用 Google Kubernetes Engine API。
    • 啟用 Google Kubernetes Engine API
  • 如要使用 Google Cloud CLI 執行這項工作，請安裝並初始化 gcloud CLI。如果您先前已安裝 gcloud CLI，請執行 gcloud components update 指令，取得最新版本。較舊的 gcloud CLI 版本可能不支援執行本文中的指令。

準備環境

在本教學課程中，您將使用 Cloud Shell，這是用來管理Google Cloud上託管資源的殼層環境。

Cloud Shell 已預先安裝 Google Cloud CLI 和 kubectl 指令列工具。gcloud CLI 是 Google Cloud的主要指令列介面，而 kubectl 是用於執行 Kubernetes 叢集相關指令的主要指令列介面。

啟動 Cloud Shell：

  1. 前往 Google Cloud 控制台。

    Google Cloud console

  2. 按一下主控台右上角的「啟用 Cloud Shell」按鈕：

此時 Cloud Shell 工作階段會在主控台底部的頁框中開啟，您可以使用這個殼層來執行 gcloud 和 kubectl 指令。執行指令前，請先使用下列指令，在 Google Cloud CLI 中設定預設專案：

gcloud config set project PROJECT_ID

將 PROJECT_ID 替換為專案 ID。

在 Ray 資訊主頁上查看 TPU 指標

如要在 Ray 資訊主頁上查看 TPU 指標，請按照下列步驟操作：

  1. 在 Shell 中安裝 Kubernetes Prometheus Stack，即可在 Ray 資訊主頁中啟用指標。

    ./install/prometheus/install.sh --auto-load-dashboard true
kubectl get all -n prometheus-system

    輸出結果會與下列內容相似：

    NAME                                                  READY   UP-TO-DATE   AVAILABLE   AGE
deployment.apps/prometheus-grafana                    1/1     1            1           46s
deployment.apps/prometheus-kube-prometheus-operator   1/1     1            1           46s
deployment.apps/prometheus-kube-state-metrics         1/1     1            1           46s

    這個安裝程序會使用 Helm 資訊圖表，將自訂資源定義 (CRD)、Pod 監控器和 Ray Pod 的服務監控器新增至叢集。

  2. 安裝 KubeRay 運算子 1.5.0 以上版本，並啟用指標擷取功能。如要在 Grafana 中查看 TPU 指標，並使用 KubeRay 提供的預設資訊主頁，您必須安裝 KubeRay 1.5.0 以上版本。

    helm repo add kuberay https://ray-project.github.io/kuberay-helm/
helm install kuberay-operator kuberay/kuberay-operator --version 1.5.0-rc.0 \
    --set metrics.serviceMonitor.enabled=true \
    --set metrics.serviceMonitor.selector.release=prometheus

  3. 安裝 KubeRay TPU Webhook。這個 Webhook 會設定環境變數 TPU_DEVICE_PLUGIN_HOST_IP，供 Ray 從 tpu-device-plugin DaemonSet 輪詢 TPU 指標。

    helm install kuberay-tpu-webhook \
oci://us-docker.pkg.dev/ai-on-gke/kuberay-tpu-webhook-helm/kuberay-tpu-webhook \
    --set tpuWebhook.image.tag=v1.2.6-gke.0

  4. 如要準備工作負載來記錄指標，請在 RayCluster 規格中設定下列連接埠和環境變數：

    headGroupSpec:
...
ports:
...
- containerPort: 44217
  name: as-metrics
- containerPort: 44227
  name: dash-metrics
env:
- name: RAY_GRAFANA_IFRAME_HOST
  value: http://127.0.0.1:3000
- name: RAY_GRAFANA_HOST
  value: http://prometheus-grafana.prometheus-system.svc:80
- name: RAY_PROMETHEUS_HOST
  value: http://prometheus-kube-prometheus-prometheus.prometheus-system.svc:9090

  5. 使用 Ray 執行含有 TPU 的工作負載。您可以採用 TPU 工作負載範例，例如「使用 JAX、Ray Train 和 TPU Trillium 在 GKE 上訓練 LLM」或「透過 KubeRay 在 GKE 上使用 TPU 提供 LLM」。

在 Ray 資訊主頁或 Grafana 中查看指標

  1. 如要連線至 RayCluster，請轉送 Ray head 服務的通訊埠：

    kubectl port-forward service/RAY_CLUSTER_NAME-head-svc 8265:8265

    將 RAY_CLUSTER_NAME 替換為 RayCluster 的名稱。如果您的工作負載來自「使用 JAX、Ray Train 和 TPU Trillium 在 GKE 上訓練 LLM」教學課程，則這個值為 maxtext-tpu-cluster。

  2. 如要查看 Ray 資訊主頁，請在本機電腦上前往 http://localhost:8265/。

  3. 轉送 Grafana 網頁版 UI 的通訊埠。您需要轉送 Grafana 通訊埠，才能查看顯示在 Ray 資訊主頁內嵌 Grafana 資訊主頁上的 TPU 指標。

    kubectl port-forward -n prometheus-system service/prometheus-grafana 3000:http-web

  4. 在 Ray 資訊主頁中，開啟「指標」分頁，然後找到「依節點劃分的 TPU 使用率」分頁。TPU 執行時，系統會將 Tensor Core 使用率、HBM 使用率、TPU 任務週期和記憶體用量的節點指標，串流至資訊主頁。如要在 Grafana 中查看這些指標，請前往 Grafana 中的「View」分頁。

    顯示 Ray 資訊主頁或 Grafana 中指標的圖表

  5. 如要查看 libtpu 記錄，請前往 Ray 資訊主頁的「Logs」(記錄) 分頁，然後選取 TPU 節點。Libtpu 記錄會寫入 /tpu_logs 目錄。

如要進一步瞭解如何設定 TPU 記錄，請參閱「偵錯 TPU VM 記錄」。

後續步驟