이 페이지는 Cloud Translation API를 통해 번역되었습니다.

TPU Trillium으로 KubeRay 구성

Autopilot Standard

이 튜토리얼에서는 Google Kubernetes Engine (GKE)에서 TPU Trillium으로 KubeRay를 구성하는 방법을 보여줍니다. TPU Trillium에 필요한 환경 변수와 포드 사양을 비롯해 단일 호스트 및 멀티 호스트 TPU 구성을 설정하는 방법을 알아봅니다.

이 튜토리얼은 단일 호스트 및 멀티 호스트 노드 풀에 KubeRay를 사용하여 TPU Trillium 초기화를 구성하는 방법을 알아보려는 플랫폼 관리자 및 운영자, 데이터 및 AI 전문가를 대상으로 합니다. 이 튜토리얼에서는 TPU 초기화가 성공했는지 확인하는 Jax 스크립트를 실행하는 방법을 보여줍니다. 이 튜토리얼에서는 모델을 배포하지 않습니다.

GKE에서 KubeRay를 구성하기 전에 GKE의 Ray 정의 및 용어를 숙지해야 합니다.

개요

이 튜토리얼에서는 KubeRay를 사용하여 TPU Trillium 초기화가 성공했는지 확인하는 Jax로 Python 스크립트를 실행하는 방법을 보여줍니다. Jax는 머신러닝 워크로드를 지원하는 고성능 수치 연산 라이브러리입니다. KubeRay는 Kubernetes에서 Ray 애플리케이션을 배포, 관리, 모니터링하는 통합 방법을 제공하는 Kubernetes 연산자입니다.

Trillium TPU (v6e)에는 이전 TPU 세대와 다른 특정 환경 변수와 Pod 사양이 필요합니다. 이 튜토리얼에서는 Trillium TPU에서 KubeRay를 사용하여 워크로드를 성공적으로 배포하는 데 필요한 구성을 제공합니다.

시작하기 전에

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

Google Kubernetes Engine API를 사용 설정합니다.

Google Kubernetes Engine API 사용 설정

이 태스크에 Google Cloud CLI를 사용하려면 gcloud CLI를 설치한 후 초기화하세요. 이전에 gcloud CLI를 설치한 경우 gcloud components update를 실행하여 최신 버전을 가져옵니다.
참고: 기존 gcloud CLI 설치의 경우 compute/region 속성을 설정해야 합니다. 주로 영역 클러스터를 사용하는 경우에는 대신 compute/zone을 설정합니다. 기본 위치를 설정하면 gcloud CLI에서 One of [--zone, --region] must be supplied: Please specify location과 같은 오류를 방지할 수 있습니다. 클러스터의 위치가 설정한 기본값과 다른 경우 특정 명령어에서 위치를 지정해야 할 수 있습니다.

Ray CLI(버전 2.37.0)가 설치되어 있는지 확인합니다.

Cloud Shell 활성화

Cloud Shell에는 이 튜토리얼에서 사용되는 gcloud, helm, kubectl 명령줄 도구가 사전 설치되어 있습니다.

Google Cloud 콘솔로 이동합니다.
Google Cloud 콘솔 창 상단에서 Cloud Shell 활성화 버튼을 클릭합니다.

Google Cloud 콘솔의 새 프레임 내에 Cloud Shell 세션이 열리면서 명령줄 프롬프트가 표시됩니다.

GKE 클러스터 및 노드 풀 만들기

GKE Autopilot 또는 Standard 클러스터의 TPU에서 KubeRay를 구성할 수 있습니다. 완전 관리형 Kubernetes 환경을 위해서는 Autopilot을 사용하는 것이 좋습니다. 워크로드에 가장 적합한 GKE 작업 모드를 선택하려면 GKE 작업 모드 정보를 참고하세요.

Autopilot

Cloud Shell에서 다음 명령어를 실행합니다.
```
gcloud container clusters create-auto CLUSTER_NAME \
    --enable-ray-operator \
    --release-channel=rapid \
    --location=LOCATION
```
다음을 바꿉니다.
- CLUSTER_NAME: 새 클러스터의 이름
- LOCATION: TPU Trillium 용량을 사용할 수 있는 리전입니다. 자세한 내용은 GKE의 TPU 가용성을 참조하세요.
GKE는 Ray 연산자 부가기능이 사용 설정된 Autopilot 클러스터를 만듭니다. 부가기능은 클러스터 컨트롤 플레인에 Ray TPU 웹훅을 자동으로 설치합니다.

클러스터와 통신하려면 kubectl을 구성하세요.

gcloud container clusters get-credentials CLUSTER_NAME --location=LOCATION

표준

Cloud Shell에서 다음 명령어를 실행하여 Ray 연산자 부가기능을 사용 설정하는 Standard 클러스터를 만듭니다.
```
gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --addons=RayOperator \
  --cluster-version=1.33 \
  --machine-type=n1-standard-16
```
다음을 바꿉니다.
- CLUSTER_NAME: 새 클러스터의 이름
- LOCATION: TPU Trillium 용량을 사용할 수 있는 리전입니다. 자세한 내용은 GKE의 TPU 가용성을 참고하세요.
클러스터 만들기는 몇 분 정도 걸릴 수 있습니다.

클러스터와 통신하려면 kubectl을 구성하세요.

gcloud container clusters get-credentials CLUSTER_NAME --location=LOCATION

단일 호스트 또는 멀티 호스트 TPU 슬라이스 노드 풀을 만들 수 있습니다.

단일 호스트

Cloud Shell에서 다음 명령어를 실행합니다.

gcloud container node-pools create v6e-4 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=1 \
    --threads-per-core=1 \
    --tpu-topology=2x2

멀티 호스트

Cloud Shell에서 다음 명령어를 실행합니다.

gcloud container node-pools create v6e-16 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=4 \
    --threads-per-core=1 \
    --tpu-topology=4x4

RayJob 커스텀 리소스 실행

RayJob 매니페스트를 정의하면 KubeRay에 다음 작업을 수행하도록 지시할 수 있습니다.

RayCluster 생성: RayJob 사양에는 원하는 Ray 클러스터 구성 (헤드 및 작업자 그룹)을 정의하는 rayClusterSpec이 포함됩니다.
특정 작업 실행: RayJob 내의 entrypoint 필드는 생성된 Ray 클러스터 내에서 실행할 명령어 또는 스크립트를 지정합니다. 이 튜토리얼에서 entrypoint은 TPU Trillium 초기화를 확인하도록 설계된 Python 스크립트 (tpu_list_devices.py)입니다.

RayJob 커스텀 리소스를 만들려면 다음 단계를 완료하세요.

단일 호스트

다음 ray-job.tpu-v6e-singlehost.yaml 매니페스트를 만듭니다.

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: v6e-4-job
spec:
  entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
  runtimeEnvYAML: |
    working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
    pip:
      - jax[tpu]==0.4.33
      - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
  rayClusterSpec:
    rayVersion: '2.43.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
          -   name: ray-head
              image: rayproject/ray:2.43.0-py310
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "8"
                  memory: 40G
                requests:
                  cpu: "8"
                  memory: 40G
    workerGroupSpecs:
    -   replicas: 1
        minReplicas: 1
        maxReplicas: 1
        numOfHosts: 1
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            containers:
            -   name: ray-worker
                image: rayproject/ray:2.43.0-py310
                resources:
                  limits:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
                  requests:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 2x2

매니페스트를 적용합니다.

kubectl apply -f ray-job.tpu-v6e-singlehost.yaml

RayJob이 생성되어 실행 중인지 확인합니다.

kubectl get rayjobs v6e-4-job

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

NAME      JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME       START TIME  END TIME   AGE
v6e-4-job PENDING      Running             v6e-4-job-raycluster   2024-10-15T23:15:22Z  20s

RayJob의 출력을 인쇄합니다.

kubectl logs -l=job-name=v6e-4-job

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

2024-10-15 16:15:40,222 INFO cli.py:300 -- ray job stop v6e-4-job-hzq5q
2024-10-15 16:15:40,246 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
2024-10-15 16:15:40,112 INFO job_manager.py:528 -- Runtime env is setting up.
2024-10-15 16:15:50,181 INFO worker.py:1461 -- Using address 10.84.1.25:6379 set in the environment variable RAY_ADDRESS
2024-10-15 16:15:50,181 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.25:6379...
2024-10-15 16:15:50,186 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.25:8265
['TPU cores:4']
2024-10-15 16:16:12,349 SUCC cli.py:63 -- -------------------------------------
2024-10-15 16:16:12,349 SUCC cli.py:64 -- Job 'v6e-4-job-hzq5q' succeeded
2024-10-15 16:16:12,349 SUCC cli.py:65 -- -------------------------------------

멀티 호스트

다음 ray-job.tpu-v6e-multihost.yaml 매니페스트를 만듭니다.

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: v6e-16-job
spec:
  entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
  runtimeEnvYAML: |
    working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
    pip:
      - jax[tpu]==0.4.33
      - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
  rayClusterSpec:
    rayVersion: '2.43.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
          -   name: ray-head
              image: rayproject/ray:2.43.0-py310
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "8"
                  memory: 40G
                requests:
                  cpu: "8"
                  memory: 40G
    workerGroupSpecs:
      - replicas: 1
        minReplicas: 1
        maxReplicas: 1
        numOfHosts: 4
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            containers:
              - name: ray-worker
                image: rayproject/ray:2.43.0-py310
                resources:
                  limits:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
                  requests:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
                env:
                - name: NODE_IP
                  valueFrom:
                    fieldRef:
                      fieldPath: status.hostIP
                - name: VBAR_CONTROL_SERVICE_URL
                  value: $(NODE_IP):8353
                - name: JAX_PLATFORMS
                  value: tpu,cpu
                - name: ENABLE_PJRT_COMPATIBILITY
                  value: "true"
                ports:
                - containerPort: 8081
                  name: mxla
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 4x4

매니페스트를 적용합니다.

kubectl apply -f ray-job.tpu-v6e-multihost.yaml

v6e-16 RayJob이 생성되고 실행 중인지 확인합니다.

kubectl get rayjobs v6e-16-job

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

NAME         JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME              START TIME             END TIME   AGE
v6e-16-job                Running             v6e-16-job-raycluster-qr6vk   2024-10-16T19:28:19Z              66s

v6e-16 RayJob의 출력을 인쇄합니다.

kubectl logs -l=job-name=v6e-16-job

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

2024-10-16 12:21:33,986 INFO cli.py:300 -- ray job stop v6e-16-job-z44s7
2024-10-16 12:21:34,011 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
2024-10-16 12:21:33,826 INFO job_manager.py:528 -- Runtime env is setting up.
2024-10-16 12:21:46,327 INFO worker.py:1461 -- Using address 10.84.1.61:6379 set in the environment variable RAY_ADDRESS
2024-10-16 12:21:46,327 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.61:6379...
2024-10-16 12:21:46,333 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.61:8265
['TPU cores:16', 'TPU cores:16', 'TPU cores:16', 'TPU cores:16']
2024-10-16 12:22:12,156 SUCC cli.py:63 -- ---------------------------------
2024-10-16 12:22:12,156 SUCC cli.py:64 -- Job 'v6e-16-job-z44s7' succeeded
2024-10-16 12:22:12,156 SUCC cli.py:65 -- ---------------------------------

Ray 대시보드에서 RayJob 보기

GKE에서 RayCluster 서비스를 생성했는지 확인하고 RayCluster 인스턴스에 연결합니다.

단일 호스트

RayJob에 대해 생성된 RayCluster의 이름을 가져옵니다.

export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-4-job -o jsonpath='{.status.rayClusterName}')

RayCluster 헤드 서비스의 이름을 검색합니다.

export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')

헤드 서비스를 포트 전달하여 Ray 대시보드에 연결합니다.
```
kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
```
웹브라우저를 열고 다음 URL을 입력합니다.
```
http://localhost:8265/#/jobs
```
RayJob 상태 및 관련 로그를 확인합니다.

멀티 호스트

RayJob에 대해 생성된 RayCluster의 이름을 가져옵니다.

export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-16-job -o jsonpath='{.status.rayClusterName}')

RayCluster 헤드 서비스의 이름을 검색합니다.

export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')

헤드 서비스를 포트 전달하여 Ray 대시보드에 연결합니다.
```
kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
```
웹브라우저를 열고 다음 URL을 입력합니다.
```
http://localhost:8265/#/jobs
```
RayJob 상태 및 관련 로그를 확인합니다.

Ray는 TPU_WORKER_ID=0 값에 해당하는 Ray 작업자 노드를 식별하기 위해 TPU-{accelerator}-Head 리소스를 설정합니다. 멀티 호스트 TPU 그룹에서 TPU_WORKER_ID=0가 있는 Ray 노드에는 리소스에 TPU-v6e-16-head: 1.0이 설정되어 있습니다. 이 TPU_WORKER_ID 환경 변수는 KubeRay용 변경 GKE 웹훅에 의해 설정됩니다.

삭제

튜토리얼을 완료한 후 계정에 원치 않는 요금이 청구되지 않도록 RayJob을 삭제합니다.

단일 호스트

kubectl delete rayjobs v6e-4-job

멀티 호스트

kubectl delete rayjobs v6e-16-job

다음 단계

Kubernetes의 Ray 알아보기
GKE에서 TPU로 vLLM을 서빙하는 방법을 알아보세요.
GKE에서 TPU를 사용하여 SDXL을 제공하는 방법을 알아보세요.
GKE의 TPU 정보 자세히 알아보기

TPU Trillium으로 KubeRay 구성 컬렉션을 사용해 정리하기 내 환경설정을 기준으로 콘텐츠를 저장하고 분류하세요.

개요

시작하기 전에

Cloud Shell 활성화

GKE 클러스터 및 노드 풀 만들기

Autopilot

표준

단일 호스트

멀티 호스트

RayJob 커스텀 리소스 실행

단일 호스트

멀티 호스트

Ray 대시보드에서 RayJob 보기

단일 호스트

멀티 호스트

삭제

단일 호스트

멀티 호스트

다음 단계

TPU Trillium으로 KubeRay 구성