Ray 대시보드에서 TPU 측정항목 보기

이 문서에서는 Google Kubernetes Engine(GKE)에서 KubeRay를 사용하여 Ray 대시보드에서 TPU 측정항목을 보는 방법을 보여줍니다. Ray GKE 부가기능이 있는 GKE 클러스터에서는 Cloud Monitoring에서 TPU 측정항목을 사용할 수 있습니다.

시작하기 전에

시작하기 전에 다음 태스크를 수행했는지 확인합니다.

  • Google Kubernetes Engine API를 사용 설정합니다.
    • Google Kubernetes Engine API 사용 설정
  • 이 태스크에 Google Cloud CLI를 사용하려면 gcloud CLI를 설치한 후 초기화합니다. 이전에 gcloud CLI를 설치한 경우 gcloud components update 명령어를 실행하여 최신 버전을 가져옵니다. 이전 gcloud CLI 버전에서는 이 문서의 명령어를 실행하지 못할 수 있습니다.

개발 환경 준비

이 튜토리얼에서는Google Cloud에서 호스팅되는 리소스를 관리하기 위한 셸 환경인 Cloud Shell을 사용합니다.

Cloud Shell에는 Google Cloud CLI 및 kubectl 명령줄 도구가 사전 설치되어 있습니다. gcloud CLI는 Google Cloud의 기본 명령줄 인터페이스를 제공하며 kubectl은 Kubernetes 클러스터를 대상으로 명령어를 실행할 수 있는 기본 명령줄 인터페이스를 제공합니다.

Cloud Shell 시작:

  1. Google Cloud 콘솔로 이동합니다.

    Google Cloud 콘솔

  2. 콘솔의 오른쪽 상단에서 Cloud Shell 활성화 버튼 을 클릭합니다.

콘솔 하단의 프레임에서 Cloud Shell 세션이 열립니다. 이 셸을 사용하여 gcloud 및 kubectl 명령어를 실행합니다. 명령어를 실행하기 전에 다음 명령어를 사용하여 Google Cloud CLI에서 기본 프로젝트를 설정합니다.

gcloud config set project PROJECT_ID

PROJECT_ID를 프로젝트 ID로 바꿉니다.

Ray 대시보드에서 TPU 측정항목 보기

Ray 대시보드에서 TPU 측정항목을 보려면 다음 단계를 수행하세요.

  1. 셸에서 Kubernetes Prometheus 스택을 설치하여 Ray 대시보드에서 측정항목을 사용 설정합니다.

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

    출력은 다음과 비슷합니다.

    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

    이 설치에서는 Helm 차트를 사용하여 Ray 포드의 커스텀 리소스 정의(CRD), 포드 모니터, 서비스 모니터를 클러스터에 추가합니다.

  2. 측정항목 스크래핑이 사용 설정된 KubeRay 연산자 v1.5.0 이상을 설치합니다. KubeRay에 제공된 기본 대시보드를 사용하여 Grafana에서 TPU 측정항목을 보려면 KubeRay v1.5.0 이상을 설치해야 합니다.

    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. KubeRay TPU 웹훅을 설치합니다. 이 웹훅은 Ray가 tpu-device-plugin DaemonSet에서 TPU 측정항목을 폴링하는 데 사용하는 환경 변수 TPU_DEVICE_PLUGIN_HOST_IP를 설정합니다.

    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. 측정항목을 로깅하도록 워크로드를 준비하려면 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. Ray를 사용하여 TPU로 워크로드를 실행합니다. GKE에서 JAX, Ray Train, TPU Trillium을 사용하여 LLM 학습 또는 KubeRay를 사용하여 GKE에서 TPU로 LLM 제공과 같은 예시 TPU 워크로드를 사용할 수 있습니다.

Ray 대시보드 또는 Grafana에서 측정항목 보기

  1. RayCluster에 연결하려면 Ray 헤드 서비스를 포트 전달하세요.

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

    RAY_CLUSTER_NAME을 RayCluster 이름으로 바꿉니다. 워크로드를 GKE에서 JAX, Ray Train, TPU Trillium을 사용하여 LLM 학습 튜토리얼에서 가져온 경우 이 값은 maxtext-tpu-cluster입니다.

  2. Ray 대시보드를 보려면 로컬 머신에서 http://localhost:8265/로 이동합니다.

  3. Grafana 웹 UI를 포트 전달합니다. Ray 대시보드에 삽입된 Grafana 대시보드에 표시되는 TPU 측정항목을 보려면 Grafana를 포트 전달해야 합니다.

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

  4. Ray 대시보드에서 측정항목 탭을 열고 노드별 TPU 사용률 탭을 찾습니다. TPU가 실행되면 텐서 코어 사용률, HBM 사용률, TPU 작업 주기, 메모리 사용량에 관한 노드별 측정항목이 대시보드로 스트리밍됩니다. 이러한 측정항목은 Grafana에서 Grafana의 보기 탭으로 이동하여 볼 수 있습니다.

    Ray 대시보드 또는 Grafana의 측정항목을 보여주는 그래프

  5. libtpu 로그를 보려면 Ray 대시보드의 로그 탭으로 이동하여 TPU 노드를 선택합니다. Libtpu 로그는 /tpu_logs 디렉터리에 기록됩니다.

TPU 로깅 구성에 대한 자세한 내용은 TPU VM 로그 디버깅을 참조하세요.

다음 단계