Afficher les métriques TPU dans le tableau de bord Ray

Ce document explique comment afficher les métriques de TPU dans le tableau de bord Ray avec KubeRay sur Google Kubernetes Engine (GKE). Dans un cluster GKE avec le module complémentaire Ray GKE, les métriques TPU sont disponibles dans Cloud Monitoring.

Avant de commencer

Avant de commencer, effectuez les tâches suivantes :

  • Activez l'API Google Kubernetes Engine.
    • Activer l'API Google Kubernetes Engine
  • Si vous souhaitez utiliser Google Cloud CLI pour cette tâche, installez puis initialisez la gcloud CLI. Si vous avez déjà installé la gcloud CLI, obtenez la dernière version en exécutant la commande gcloud components update. Il est possible que les versions antérieures de la gcloud CLI ne permettent pas d'exécuter les commandes de ce document.

Préparer votre environnement

Dans ce tutoriel, vous allez utiliser Cloud Shell, un environnement shell permettant de gérer les ressources hébergées surGoogle Cloud.

Cloud Shell est préinstallé avec Google Cloud CLI et l'outil de ligne de commande kubectl. Gcloud CLI est l'interface de ligne de commande principale utilisée sur Google Cloud, et kubectl fournit l'interface de ligne de commande principale permettant d'exécuter des commandes sur les clusters Kubernetes.

Lancez Cloud Shell :

  1. Accédez à la console Google Cloud .

    ConsoleGoogle Cloud

  2. Dans l'angle supérieur droit de la console, cliquez sur le bouton Activer Cloud Shell :

Une session Cloud Shell s'ouvre dans un cadre situé en bas de la console. Ce shell vous permet d'exécuter les commandes gcloud et kubectl. Avant d'exécuter des commandes, définissez votre projet par défaut dans Google Cloud CLI à l'aide de la commande suivante :

gcloud config set project PROJECT_ID

Remplacez PROJECT_ID par l'ID du projet.

Afficher les métriques TPU dans le tableau de bord Ray

Pour afficher les métriques TPU dans le tableau de bord Ray :

  1. Dans votre shell, activez les métriques dans le tableau de bord Ray en installant la pile Kubernetes Prometheus.

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

    Le résultat ressemble à ce qui suit :

    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

    Cette installation utilise un graphique Helm pour ajouter les définitions de ressources personnalisées (CRD), un PodMonitor et un ServiceMonitor pour les pods Ray à votre cluster.

  2. Installez l'opérateur KubeRay v1.5.0 ou version ultérieure avec l'extraction des métriques activée. Pour afficher les métriques TPU dans Grafana avec les tableaux de bord par défaut fournis dans KubeRay, vous devez installer KubeRay v1.5.0 ou version ultérieure.

    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. Installez le webhook TPU KubeRay. Ce webhook définit une variable d'environnement TPU_DEVICE_PLUGIN_HOST_IP que Ray utilise pour interroger les métriques TPU à partir du 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. Pour préparer votre charge de travail à enregistrer les métriques, définissez les ports et les variables d'environnement suivants dans la spécification 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. Exécutez une charge de travail avec des TPU à l'aide de Ray. Vous pouvez utiliser un exemple de charge de travail TPU, comme Entraîner un LLM à l'aide de JAX, Ray Train et TPU Trillium sur GKE ou Mettre en service un LLM à l'aide de TPU sur GKE avec KubeRay.

Afficher les métriques dans le tableau de bord Ray ou Grafana

  1. Pour vous connecter au RayCluster, transférez le port du service principal Ray :

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

    Remplacez RAY_CLUSTER_NAME par le nom de votre RayCluster. Si votre charge de travail provient du tutoriel Entraîner un LLM à l'aide de JAX, Ray Train et TPU Trillium sur GKE, cette valeur est maxtext-tpu-cluster.

  2. Pour afficher le tableau de bord Ray, accédez à http://localhost:8265/ sur votre ordinateur local.

  3. Transférez le port de l'interface utilisateur Web de Grafana. Vous devez transférer le port Grafana pour afficher les métriques TPU qui sont affichées dans les tableaux de bord Grafana intégrés au tableau de bord Ray.

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

  4. Dans le tableau de bord Ray, ouvrez l'onglet Metrics (Métriques), puis recherchez l'onglet TPU Utilization by Node (Utilisation des TPU par nœud). Lorsque les TPU sont en cours d'exécution, les métriques par nœud concernant l'utilisation des Tensor Cores, l'utilisation de la mémoire HBM, le cycle d'utilisation des TPU et l'utilisation de la mémoire sont diffusées dans le tableau de bord. Pour afficher ces métriques dans Grafana, accédez à View Tab in Grafana (Onglet "Afficher" dans Grafana).

    Graphique affichant les métriques dans le tableau de bord Ray ou Grafana

  5. Pour afficher les journaux libtpu, accédez à l'onglet Logs (Journaux) du tableau de bord Ray, puis sélectionnez le nœud TPU. Les journaux Libtpu sont écrits dans le répertoire /tpu_logs.

Pour en savoir plus sur la configuration de la journalisation des TPU, consultez Déboguer les journaux des VM TPU.

Étapes suivantes