Ray ダッシュボードで TPU 指標を表示する

このドキュメントでは、Google Kubernetes Engine（GKE）で KubeRay を使用して Ray ダッシュボードで TPU 指標を表示する方法について説明します。Ray GKE アドオンを使用する GKE クラスタでは、TPU 指標が Cloud Monitoring で使用可能です。

始める前に

作業を始める前に、次のタスクが完了していることを確認してください。

  • Google Kubernetes Engine API を有効にする。
    • Google Kubernetes Engine API の有効化
  • このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。gcloud CLI をインストール済みの場合は、gcloud components update コマンドを実行して最新のバージョンを取得します。以前のバージョンの gcloud CLI では、このドキュメントのコマンドを実行できない場合があります。

環境を準備する

このチュートリアルでは、Google Cloudでホストされているリソースを管理するためのシェル環境である Cloud Shell を使用します。

Cloud Shell には、Google Cloud CLI と kubectl コマンドライン ツールがプリインストールされています。gcloud CLI は Google Cloudへの主要なコマンドライン インターフェースを提供し、kubectl は Kubernetes クラスタにコマンドを実行するためのコマンドライン インターフェースを提供します。

Cloud Shell を起動します。

  1. Google Cloud コンソールに移動します。

    Google Cloud コンソール

  2. コンソールの右上隅にある [Cloud Shell をアクティブにする] ボタン をクリックします。

コンソールの下部にあるフレーム内で Cloud Shell セッションが開きます。このシェルで gcloud コマンドと kubectl コマンドを実行します。コマンドを実行する前に、Google Cloud CLI で次のコマンドを使用してデフォルト プロジェクトを設定します。

gcloud config set project PROJECT_ID

PROJECT_ID は、実際のプロジェクト ID に置き換えます。

Ray ダッシュボードで TPU 指標を表示する

Ray ダッシュボードで TPU 指標を表示する手順は次のとおりです。

  1. シェルで、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 オペレーター v1.5.0 以降をインストールします。KubeRay で提供されるデフォルトのダッシュボードを使用して Grafana で TPU 指標を表示するには、KubeRay v1.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 は、Ray が tpu-device-plugin DaemonSet から TPU 指標をポーリングするために使用する環境変数 TPU_DEVICE_PLUGIN_HOST_IP を設定します。

    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 でワークロードを実行します。GKE で JAX、Ray Train、TPU Trillium を使用して LLM をトレーニングするや、KubeRay を使用して、GKE 上で TPU を使用する LLM をサービングするなどの TPU ワークロードの例を使用できます。

Ray ダッシュボードまたは Grafana で指標を表示する

  1. RayCluster に接続するには、Ray ヘッドサービスをポート転送します。

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

    RAY_CLUSTER_NAME は、RayCluster の名前に置き換えます。ワークロードが GKE で JAX、Ray Train、TPU Trillium を使用して LLM をトレーニングするチュートリアルのものである場合、この値は maxtext-tpu-cluster です。

  2. Ray ダッシュボードを表示するには、ローカルマシンの http://localhost:8265/ に移動します。

  3. Grafana ウェブ UI をポート転送します。Ray ダッシュボードに埋め込まれた Grafana ダッシュボードに表示される TPU 指標を表示するには、Grafana をポート転送する必要があります。

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

  4. Ray ダッシュボードで、[Metrics] タブを開き、[TPU Utilization by Node] タブを見つけます。TPU の実行中、Tensor Core 使用率、HBM 使用率、TPU デューティ サイクル、メモリ使用量に関するノードごとの指標がダッシュボードにストリーミングされます。これらの指標を Grafana で表示するには、Grafana の [View] タブに移動します。

    Ray ダッシュボードまたは Grafana の指標を示すグラフ

  5. libtpu ログを表示するには、Ray ダッシュボードの [Logs] タブに移動して、TPU ノードを選択します。Libtpu ログは /tpu_logs ディレクトリに書き込まれます。

TPU ロギングの構成の詳細については、TPU VM ログのデバッグをご覧ください。

次のステップ