Visualiza las métricas de TPU en el panel de Ray

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

Antes de comenzar

Antes de comenzar, asegúrate de haber realizado las siguientes tareas:

  • Habilita la API de Google Kubernetes Engine.
    • Habilitar la API de Google Kubernetes Engine
  • Si deseas usar Google Cloud CLI para esta tarea, instala y, luego, inicializa gcloud CLI. Si ya instalaste gcloud CLI, ejecuta el comando gcloud components update para obtener la versión más reciente. Es posible que las versiones anteriores de gcloud CLI no admitan la ejecución de los comandos que se describen en este documento.

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

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

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

Prepara el entorno

En este instructivo, usarás Cloud Shell, que es un entorno de shell para administrar recursos alojados enGoogle Cloud.

Cloud Shell ya viene instalado en Google Cloud CLI y la herramienta de línea de comandos de kubectl. Gcloud CLI proporciona la interfaz de línea de comandos principal para Google Cloud, y kubectl proporciona la interfaz de línea de comandos principal para ejecutar comandos en clústeres de Kubernetes.

Iniciar Cloud Shell:

  1. Ve a la consola de Google Cloud .

    Google Cloud console

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

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

gcloud config set project PROJECT_ID

Reemplaza PROJECT_ID con el ID del proyecto.

Visualiza las métricas de TPU en el panel de Ray

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

  1. En tu shell, habilita las métricas en el panel 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 es similar a lo 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 agregar las definiciones de recursos personalizados (CRD), un Pod Monitor y un Service Monitor para los Pods de Ray a tu clúster.

  2. Instala el operador de KubeRay v1.5.0 o una versión posterior con la recopilació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 v1.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 establece una variable de entorno TPU_DEVICE_PLUGIN_HOST_IP que Ray usa para sondear las métricas de TPU desde el 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, configura los siguientes puertos y variables de entorno dentro de 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 usando Ray. Puedes usar una carga de trabajo de TPU de ejemplo, como Entrena un LLM con JAX, Ray Train y TPU Trillium en GKE o Entrega un LLM con TPU en GKE con KubeRay.

Consulta las métricas en el panel de Ray o Grafana

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

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

    Reemplaza RAY_CLUSTER_NAME por el nombre de tu RayCluster. Si tu carga de trabajo proviene del instructivo Entrena un LLM con JAX, Ray Train y TPU Trillium en GKE, este valor es maxtext-tpu-cluster.

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

  3. Redirige el puerto de la IU web de Grafana. Debes redireccionar puertos de Grafana para ver las métricas de TPU que se muestran en los paneles de Grafana integrados en el panel de Ray.

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

  4. En el panel de Ray, abre la pestaña Metrics y busca la pestaña TPU Utilization by Node. Cuando las TPU están en funcionamiento, las métricas por nodo para el uso de Tensor Core, el uso de HBM, el ciclo de trabajo de la TPU y el uso de memoria se transmiten al panel. Para ver estas métricas en Grafana, navega a View Tab in Grafana.

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

  5. Para ver los registros de libtpu, navega a la pestaña Registros en el panel de Ray y selecciona el nodo 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 TPU, consulta Depura los registros de la VM de TPU.

¿Qué sigue?