Visualizza le metriche TPU nella dashboard Ray

Questo documento mostra come visualizzare le metriche TPU nella dashboard di Ray con KubeRay su Google Kubernetes Engine (GKE). In un cluster GKE con il componente aggiuntivo Ray GKE, le metriche TPU sono disponibili in Cloud Monitoring.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti operazioni:

  • Attiva l'API Google Kubernetes Engine.
    • Attiva l'API Google Kubernetes Engine
  • Se vuoi utilizzare Google Cloud CLI per questa attività, installala e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, scarica l'ultima versione eseguendo il comando gcloud components update. Le versioni precedenti di gcloud CLI potrebbero non supportare l'esecuzione dei comandi in questo documento.

prepara l'ambiente

In questo tutorial utilizzerai Cloud Shell, un ambiente shell per la gestione delle risorse ospitate su Google Cloud.

Cloud Shell include Google Cloud CLI e lo strumento a riga di comando kubectl preinstallati. gcloud CLI fornisce l'interfaccia a riga di comando principale per Google Cloud, mentre kubectl fornisce l'interfaccia a riga di comando principale per l'esecuzione di comandi sui cluster Kubernetes.

Avvia Cloud Shell:

  1. Vai alla console Google Cloud .

    Google Cloud console

  2. Nell'angolo in alto a destra della console, fai clic sul pulsante Attiva Cloud Shell:

Una sessione di Cloud Shell si apre in un frame nella parte inferiore della console. Utilizza questa shell per eseguire i comandi gcloud e kubectl. Prima di eseguire i comandi, imposta il progetto predefinito in Google Cloud CLI utilizzando questo comando:

gcloud config set project PROJECT_ID

Sostituisci PROJECT_ID con l'ID progetto.

Visualizza le metriche TPU nella dashboard Ray

Per visualizzare le metriche TPU nella dashboard di Ray:

  1. Nella shell, abilita le metriche nella dashboard Ray installando lo stack Kubernetes Prometheus.

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

    L'output è simile al seguente:

    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

    Questa installazione utilizza un grafico Helm per aggiungere le definizioni di risorse personalizzate (CRD), un monitor pod e un monitor servizi per i pod Ray al cluster.

  2. Installa l'operatore KubeRay v1.5.0 o versioni successive con lo scraping delle metriche abilitato. Per visualizzare le metriche TPU in Grafana con le dashboard predefinite fornite in KubeRay, devi installare KubeRay v1.5.0 o versioni successive.

    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. Installa l'webhook TPU di KubeRay. Questo webhook imposta una variabile di ambiente TPU_DEVICE_PLUGIN_HOST_IP che Ray utilizza per eseguire il polling delle metriche TPU dal DaemonSet tpu-device-plugin.

    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. Per preparare il carico di lavoro a registrare le metriche, imposta le seguenti porte e variabili di ambiente all'interno della specifica 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. Esegui un workload con le TPU utilizzando Ray. Puoi utilizzare un carico di lavoro TPU di esempio, ad esempio Addestra un LLM utilizzando JAX, Ray Train e TPU Trillium su GKE o Gestisci un LLM utilizzando le TPU su GKE con KubeRay.

Visualizza le metriche nella dashboard Ray o in Grafana

  1. Per connetterti a RayCluster, inoltra la porta del servizio head di Ray:

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

    Sostituisci RAY_CLUSTER_NAME con il nome del tuo RayCluster. Se il tuo workload proviene dal tutorial Addestra un LLM utilizzando JAX, Ray Train e TPU Trillium su GKE, questo valore è maxtext-tpu-cluster.

  2. Per visualizzare la dashboard di Ray, vai a http://localhost:8265/ sulla tua macchina locale.

  3. Esegui il port forwarding della GUI web di Grafana. Devi eseguire il port forwarding di Grafana per visualizzare le metriche TPU visualizzate nelle dashboard Grafana incorporate nella dashboard Ray.

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

  4. Nella dashboard Ray, apri la scheda Metriche e trova la scheda Utilizzo TPU per nodo. Quando le TPU sono in esecuzione, le metriche per nodo relative all'utilizzo di Tensor Core, all'utilizzo di HBM, al ciclo di lavoro della TPU e all'utilizzo della memoria vengono trasmesse in streaming alla dashboard. Queste metriche sono visualizzabili in Grafana accedendo alla scheda Visualizza in Grafana.

    Grafico che mostra le metriche nella dashboard Ray o in Grafana

  5. Per visualizzare i log di libtpu, vai alla scheda Log nella dashboard di Ray e seleziona il nodo TPU. I log di libtpu vengono scritti nella directory /tpu_logs.

Per saperne di più sulla configurazione della registrazione TPU, consulta Debug dei log delle VM TPU.

Passaggi successivi