Veja as métricas da TPU no painel de controlo do Ray

Este documento mostra como ver as métricas da TPU no painel de controlo do Ray com o KubeRay no Google Kubernetes Engine (GKE). Num cluster do GKE com o suplemento do GKE do Ray, as métricas de TPU estão disponíveis no Cloud Monitoring.

Antes de começar

Antes de começar, certifique-se de que realizou as seguintes tarefas:

  • Ative a API Google Kubernetes Engine.
  • Ative a API Google Kubernetes Engine
  • Se quiser usar a CLI gcloud para esta tarefa, instale-a e, em seguida, inicialize a CLI gcloud. Se instalou anteriormente a CLI gcloud, execute o comando gcloud components update para obter a versão mais recente. As versões anteriores da CLI gcloud podem não suportar a execução dos comandos neste documento.
  • Tem um cluster do GKE que cumpre as seguintes condições:

    • Um node pool configurado para usar TPUs.
    • O suplemento Ray on GKE não pode estar ativado neste cluster.
    • Uma carga de trabalho de serviço ou de preparação no GKE com o Ray e as TPUs.

    Se precisar de criar um cluster com estas condições, escolha um exemplo de carga de trabalho de TPU, como Forme um MDG com JAX, Ray Train e TPU Trillium no GKE ou Publique um MDG com TPUs no GKE com KubeRay e siga os passos de configuração para configurar um cluster.

Prepare o seu ambiente

Neste tutorial, vai usar o Cloud Shell, que é um ambiente de shell para gerir recursos alojados no Google Cloud.

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

Inicie o Cloud Shell:

  1. Aceda à Google Cloud consola.

    Google Cloud consola

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

É aberta uma sessão do Cloud Shell num frame na parte inferior da consola. Use esta shell para executar os comandos gcloud e kubectl. Antes de executar comandos, defina o projeto predefinido na CLI do Google Cloud com o seguinte comando:

gcloud config set project PROJECT_ID

Substitua PROJECT_ID pelo ID do seu projeto.

Veja as métricas da TPU no painel de controlo do Ray

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

  1. Na shell, ative as métricas no painel de controlo do Ray instalando o Kubernetes Prometheus Stack.

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

    O resultado é semelhante ao 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
    

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

  2. Instale o operador KubeRay v1.5.0 ou posterior com a extração de métricas ativada. Para ver as métricas da TPU no Grafana com os painéis de controlo predefinidos fornecidos no KubeRay, tem de instalar o KubeRay v1.5.0 ou 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. Instale o webhook da TPU do KubeRay. Este webhook define uma variável de ambiente TPU_DEVICE_PLUGIN_HOST_IP que o Ray usa para sondar as métricas da TPU a partir 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 a carga de trabalho para registar 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 através do Ray. Pode usar um exemplo de carga de trabalho de TPU, como Forme um MDG com JAX, Ray Train e TPU Trillium no GKE ou Publique um MDG com TPUs no GKE com KubeRay.

Veja as métricas no painel de controlo do Ray ou no Grafana

  1. Para estabelecer ligação 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 sua carga de trabalho for do tutorial Forme um MDG com JAX, Ray Train e TPU Trillium no GKE, este valor é maxtext-tpu-cluster.

  2. Para ver o painel de controlo do Ray, navegue para http://localhost:8265/ na sua máquina local.

  3. Encaminhe a porta da IU da Web do Grafana. Tem de encaminhar portas do Grafana para ver as métricas da TPU apresentadas nos painéis de controlo do Grafana incorporados no painel de controlo do Ray.

    kubectl port-forward -n prometheus-system service/prometheus-grafana 3000:http-web
    
  4. No painel de controlo do Ray, abra o separador Métricas e encontre o separador Utilização da TPU por nó. Quando as TPUs estão em execução, as métricas por nó para a utilização do Tensor Core, a utilização da HBM, o ciclo de trabalho da TPU e a utilização da memória são transmitidas para o painel de controlo. Pode ver estas métricas no Grafana navegando para o separador Vista no Grafana.

    Gráfico que mostra métricas no painel de controlo do Ray ou no Grafana

  5. Para ver os registos libtpu, navegue para o separador Registos no painel de controlo do Ray e selecione o nó da TPU. Os registos da libtpu são escritos no diretório /tpu_logs.

Para mais informações sobre a configuração do registo da TPU, consulte o artigo Depurar registos de VMs da TPU.

O que se segue?