Conferir métricas de TPU no painel do Ray

Neste documento, mostramos como visualizar métricas de TPU no painel do Ray com o KubeRay no Google Kubernetes Engine (GKE). Em um cluster do GKE com o complemento do Ray GKE, as métricas de TPU estão disponíveis no Cloud Monitoring.

Antes de começar

Antes de começar, verifique se você realizou as tarefas a seguir:

  • Ativar a API Google Kubernetes Engine.
    • Ativar a API Google Kubernetes Engine
  • Se você quiser usar a CLI do Google Cloud para essa tarefa, instale e inicialize a gcloud CLI. Se você instalou a gcloud CLI anteriormente, instale a versão mais recente executando o comando gcloud components update. Talvez as versões anteriores da gcloud CLI não sejam compatíveis com a execução dos comandos neste documento.

  • Você tem um cluster do GKE que atende às seguintes condições:

    • Um pool de nós configurado para usar TPUs.
    • O complemento Ray no GKE não pode estar ativado nesse cluster.
    • Uma carga de trabalho de disponibilização ou treinamento no GKE usando o Ray e as TPUs.

    Se você precisar criar um cluster com essas condições, escolha um exemplo de carga de trabalho de TPU, como Treinar um LLM usando JAX, Ray Train e TPU Trillium no GKE ou Disponibilizar um LLM usando TPUs no GKE com KubeRay e siga as etapas de configuração para configurar um cluster.

Preparar o ambiente

Neste tutorial, você vai usar o Cloud Shell, um ambiente shell para gerenciar recursos hospedados noGoogle Cloud.

O Cloud Shell vem pré-instalado com a Google Cloud CLI e a ferramenta de linha de comando kubectl. A CLI gcloud fornece a principal interface de linha de comando do Google Cloud, e kubectl fornece a principal interface de linha de comando para executar comandos em clusters do Kubernetes.

Inicie o Cloud Shell

  1. Acesse o console do Google Cloud .

    Google Cloud console

  2. No canto superior direito do console, clique no botão Ativar Cloud Shell:

Uma sessão do Cloud Shell é aberta dentro de um quadro inferior no console. Use esse shell para executar os comandos gcloud e kubectl. Antes de executar comandos, defina o projeto padrão na Google Cloud CLI usando o seguinte comando:

gcloud config set project PROJECT_ID

Substitua PROJECT_ID pela ID do seu projeto.

Conferir métricas de TPU no painel do Ray

Para conferir as métricas de TPU no painel do Ray, faça o seguinte:

  1. No shell, ative as métricas no painel do Ray instalando a pilha do Prometheus do Kubernetes.

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

    O resultado será o seguinte:

    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

    Essa instalação usa um gráfico do Helm para adicionar as definições de recursos personalizados (CRDs), um monitor de pod e um monitor de serviço para pods do Ray ao cluster.

  2. Instale o operador KubeRay v1.5.0 ou mais recente com a raspagem de métricas ativada. Para conferir as métricas de TPU no Grafana com os painéis padrão fornecidos no KubeRay, instale o KubeRay v1.5.0 ou mais recente.

    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. Instale o webhook de TPU do KubeRay. Esse webhook define uma variável de ambiente TPU_DEVICE_PLUGIN_HOST_IP que o Ray usa para consultar as métricas da TPU do 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 sua carga de trabalho para registrar as métricas, defina as seguintes portas e variáveis de ambiente na especificação do 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. Execute uma carga de trabalho com TPUs usando o Ray. Você pode usar um exemplo de carga de trabalho da TPU, como Treinar um LLM usando JAX, Ray Train e TPU Trillium no GKE ou Disponibilizar um LLM usando TPUs no GKE com o KubeRay.

Ver as métricas no painel do Ray ou no Grafana

  1. Para se conectar ao RayCluster, encaminhe a porta do serviço principal do Ray:

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

    Substitua RAY_CLUSTER_NAME pelo nome do seu RayCluster. Se a carga de trabalho for do tutorial Treinar um LLM usando JAX, Ray Train e TPU Trillium no GKE, esse valor será maxtext-tpu-cluster.

  2. Para acessar o painel do Ray, navegue até http://localhost:8265/ na sua máquina local.

  3. Encaminhe a porta da UI da Web do Grafana. É necessário encaminhar a porta do Grafana para visualizar as métricas da TPU, que são mostradas nos painéis do Grafana incorporados ao painel do Ray.

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

  4. No painel do Ray, abra a guia Métricas e encontre a guia Utilização da TPU por nó. Quando as TPUs estão em execução, as métricas por nó de utilização do Tensor Core, utilização da HBM, ciclo de trabalho da TPU e uso de memória são transmitidas para o painel. Para conferir essas métricas no Grafana, navegue até a guia "Visualizar" no Grafana.

    Gráfico mostrando métricas no painel do Ray ou no Grafana

  5. Para conferir os registros do libtpu, acesse a guia Registros no painel do Ray e selecione o nó da TPU. Os registros do Libtpu são gravados no diretório /tpu_logs.

Para mais informações sobre como configurar o registro em registros da TPU, consulte Como depurar registros de VMs de TPU.

A seguir