Ver métricas de TPU en el panel de control de Ray

En este documento se explica cómo ver las métricas de TPU en el panel de control de Ray con KubeRay en Google Kubernetes Engine (GKE). En un clúster de GKE con el complemento Ray GKE, las métricas de TPU están disponibles en Cloud Monitoring.

Antes de empezar

Antes de empezar, asegúrate de que has realizado las siguientes tareas:

  • Habilita la API de Google Kubernetes Engine.
    • Habilitar la API de Google Kubernetes Engine
  • Si quieres usar Google Cloud CLI para esta tarea, instálala y, a continuación, inicialízala. Si ya has instalado la CLI de gcloud, obtén la versión más reciente ejecutando el comando gcloud components update. Es posible que las versiones anteriores de la interfaz de línea de comandos de gcloud no admitan la ejecución de los comandos de este documento.

  • Tienes un clúster de GKE que cumple las siguientes condiciones:

    • Un grupo de nodos configurado para usar TPUs.
    • El complemento Ray on GKE no debe estar habilitado en este clúster.
    • Una carga de trabajo de servicio o entrenamiento en GKE que usa Ray y TPUs.

    Si necesitas crear un clúster con estas condiciones, elige una carga de trabajo de TPU de ejemplo, como Entrenar un LLM con JAX, Ray Train y TPU Trillium en GKE o Servir un LLM con TPUs en GKE con KubeRay, y sigue los pasos de configuración para configurar un clúster.

Prepara tu entorno

En este tutorial, usarás Cloud Shell, un entorno de shell para gestionar recursos alojados enGoogle Cloud.

Cloud Shell tiene preinstaladas la CLI de Google Cloud y la herramienta de línea de comandos kubectl. La CLI de gcloud proporciona la interfaz de línea de comandos principal de Google Cloud, y kubectl proporciona la interfaz de línea de comandos principal para ejecutar comandos en clústeres de Kubernetes.

Abre Cloud Shell:

  1. Ve a la Google Cloud consola.

    Google Cloud console

  2. En la esquina superior derecha de la consola, haz clic en el botón Activar Cloud Shell:

Se abrirá una sesión de Cloud Shell en un marco de la parte inferior de la consola. Esta shell se usa para ejecutar comandos gcloud y kubectl. Antes de ejecutar comandos, define tu proyecto predeterminado en la CLI de Google Cloud con el siguiente comando:

gcloud config set project PROJECT_ID

Sustituye PROJECT_ID por el ID del proyecto.

Ver métricas de TPU en el panel de control de Ray

Para ver las métricas de TPU en el panel de control de Ray, haz lo siguiente:

  1. En tu shell, habilita las métricas en el panel de control de Ray instalando la pila de Prometheus de Kubernetes.

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

    El resultado debería ser similar al siguiente:

    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

    Esta instalación usa un gráfico de Helm para añadir las definiciones de recursos personalizados (CRDs), un monitor de pods y un monitor de servicios para pods de Ray a tu clúster.

  2. Instala el operador KubeRay v1.5.0 o una versión posterior con la función de recolección de métricas habilitada. Para ver las métricas de TPU en Grafana con los paneles predeterminados que se proporcionan en KubeRay, debes instalar KubeRay 1.5.0 o una versión posterior.

    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. Instala el webhook de TPU de KubeRay. Este webhook define una variable de entorno TPU_DEVICE_PLUGIN_HOST_IP que Ray usa para sondear las métricas de TPU del 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. Para preparar tu carga de trabajo para registrar las métricas, define los siguientes puertos y variables de entorno en la especificación de 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. Ejecuta una carga de trabajo con TPUs mediante Ray. Puedes usar una carga de trabajo de TPU de ejemplo, como Entrenar un LLM con JAX, Ray Train y TPU Trillium en GKE o Servir un LLM con TPUs en GKE con KubeRay.

Ver las métricas en el panel de control de Ray o en Grafana

  1. Para conectarte al RayCluster, reenvía el servicio principal de Ray:

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

    Sustituye RAY_CLUSTER_NAME por el nombre de tu RayCluster. Si tu carga de trabajo procede del tutorial Entrenar un LLM con JAX, Ray Train y TPU Trillium en GKE, este valor es maxtext-tpu-cluster.

  2. Para ver el panel de control de Ray, ve a http://localhost:8265/ en tu máquina local.

  3. Redirige los puertos de la interfaz web de Grafana. Debes reenviar el puerto de Grafana para ver las métricas de TPU, que se muestran en los paneles de control de Grafana insertados en el panel de control de Ray.

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

  4. En el panel de control de Ray, abre la pestaña Métricas y busca la pestaña Utilización de TPU por nodo. Cuando las TPUs están en funcionamiento, las métricas por nodo de uso de Tensor Core, uso de HBM, ciclo de trabajo de TPU y uso de memoria se envían al panel de control. Para ver estas métricas en Grafana, vaya a la pestaña Ver en Grafana.

    Gráfico que muestra las métricas en el panel de control de Ray o en Grafana

  5. Para ver los registros de libtpu, vaya a la pestaña Registros del panel de control de Ray y seleccione el nodo de TPU. Los registros de libtpu se escriben en el directorio /tpu_logs.

Para obtener más información sobre cómo configurar el registro de TPUs, consulta Depurar registros de VMs de TPU.

Siguientes pasos