DWS 및 Kueue를 사용하여 TPU에서 AI 학습 최적화

이 문서에서는 Tensor Processing Unit (TPU)에서 AI 학습 가용성을 극대화하도록 GKE 클러스터를 구성하는 방법을 설명합니다. Kueue라는 오픈소스 작업 대기열 도구와 Google Cloud의 동적 워크로드 스케줄러(DWS)를 사용하여 자동 장애 복구 시스템을 구성합니다.

이 문서에서 정의하는 구성은 기본 노드 풀과 백업 노드 풀을 설정합니다.

  • 온디맨드 (계획 A): 이 노드 풀은 노드에 가장 적합합니다. 시스템은 먼저 이러한 주문형 머신에서 작업을 예약하려고 시도합니다. 온디맨드라는 용어는 이러한 머신이 실행되기 시작하면 중단 없이 안정적으로 액세스할 수 있음을 의미합니다.
  • DWS flex-start (계획 B): 백업 노드 풀입니다. 계획 A 머신을 사용할 수 없는 경우 Kueue 스케줄링 프로그램은 자동으로 이 계획 B 풀에 작업을 할당합니다. 그런 다음 DWS는 계획 B 하드웨어를 검색하지만 해당 하드웨어도 사용 불가능할 수 있으므로 즉시 액세스할 수 있다고 보장하지는 않습니다. 하지만 DWS는 포기하지 않습니다. 최대 7일 동안 요청을 대기열에 보관하고 머신이 제공되는 즉시 자동으로 제공합니다.

이 접근 방식은 작업이 대기열에서 기다리는 시간을 최소화합니다. 즉, 사용 가능한 리소스를 수동으로 확인하거나 다른 머신에 맞게 스크립트를 다시 작성할 필요가 없습니다.

구성 단계 개요

자동 대체 시스템을 구성하려면 여러 구성 단계를 완료해야 합니다. 이 구성을 다음 두 가지 카테고리로 나누는 것이 좋습니다.

  • 클러스터 관리자 작업: GKE 클러스터 생성, 노드 풀 프로비저닝, Kueue 스케줄링 컨트롤러 설치와 같은 일회성 인프라 구성
  • AI 개발자 작업: 학습 작업 요구사항 정의, 워크로드 제출과 같은 반복적인 일일 워크플로

이러한 단계를 모두 직접 수행하더라도 이 구분을 염두에 두면 전체 프로세스를 명확히 하는 데 도움이 됩니다.

시스템을 구성하기 전에 실행할 구성 단계를 검토하세요.

주제 작업
인프라 구성 (클러스터 관리자) 1. GKE 클러스터 만들기
2 노드 풀 만들기
3 노드 풀에서 flex-start의 상태 확인
Kueue 설치 및 구성 (클러스터 관리자) 1. Kueue 설치
2 구성 규칙 정의
학습 작업 실행 (AI 개발자) 1. ConfigMap 만들기
2. RayJob 매니페스트 정의
3. 워크로드를 제출합니다
4. RayJob에 연결
5 로그 확인하기

주요 개념

  • 온디맨드 (계획 A) 노드 풀: 기본 우선순위가 높은 노드 풀입니다. 작업은 항상 이 풀을 먼저 사용하려고 합니다.
  • DWS flex-start (계획 B) 노드 풀: 백업 노드 풀입니다. 기본 풀의 머신을 사용할 수 없는 경우 시스템은 이 풀을 자동으로 사용하여 사용 가능한 하드웨어를 검색합니다.
  • Kueue: 작업 대기열을 관리하는 예약 프로그램입니다. 작업 요청을 가로채서 사용할 노드 풀 (계획 A 또는 계획 B)을 결정합니다.
  • 작업: 실행하려는 AI 학습 워크로드입니다. 이 문서에서는 RayJob 매니페스트를 사용하여 이를 정의합니다.

시작하기 전에

  1. Google Cloud 콘솔의 프로젝트 선택기 페이지에서 Google Cloud 프로젝트를 선택하거나 만듭니다.

    프로젝트를 선택하거나 만드는 데 필요한 역할

    • 프로젝트 선택: 프로젝트를 선택하는 데는 특정 IAM 역할이 필요하지 않습니다. 역할이 부여된 프로젝트를 선택하면 됩니다.
    • 프로젝트 만들기: 프로젝트를 만들려면 resourcemanager.projects.create 권한이 포함된 프로젝트 생성자 역할(roles/resourcemanager.projectCreator)이 필요합니다. 역할 부여 방법 알아보기

    프로젝트 선택기로 이동

  2. Google Cloud 프로젝트에 결제가 사용 설정되어 있는지 확인합니다.

  3. Google Kubernetes Engine, Cloud TPU API API를 사용 설정합니다.

    API 사용 설정에 필요한 역할

    API를 사용 설정하려면 serviceusage.services.enable 권한이 포함된 서비스 사용량 관리자 IAM 역할(roles/serviceusage.serviceUsageAdmin)이 필요합니다. 역할 부여 방법 알아보기

    API 사용 설정

  4. Google Cloud 콘솔에서 Cloud Shell을 활성화합니다.

    Cloud Shell 활성화

  5. TPU flex-start VM을 사용하기에 충분한 선점형 할당량이 있는지 확인합니다. 기본 할당량이 필요에 충분하지 않으면 더 높은 할당량을 요청하세요. 자세한 내용은 Cloud TPU 할당량 및 Cloud TPU 환경 설정을 참고하세요.

환경 변수 정의

이 문서에서 실행하는 명령어를 간소화하려면 Cloud Shell에서 환경 변수를 설정하면 됩니다. 이러한 변수는 Google Cloud 프로젝트의 ID, 노드 풀의 이름, GKE 클러스터의 위치와 같은 값을 저장합니다.

이러한 변수를 정의한 후에는 매번 값을 다시 입력하거나 바꾸는 대신 변수 이름 (예: $CLUSTER_NAME)을 참조하여 여러 명령어에서 재사용할 수 있습니다. 이 접근 방식을 사용하면 프로세스를 더 쉽게 따라갈 수 있으므로 오류 위험이 줄어듭니다.

Cloud Shell에서 다음과 같은 유용한 환경 변수를 정의하려면 다음 명령어를 실행하세요.

export PROJECT_ID=$(gcloud config get project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export ZONE="us-east5-b"
export REGION="us-east5"
export CLUSTER_NAME="tpu-cluster"
export GKE_VERSION="1.34"
export ONDEMAND_NODEPOOL="on-demand-pool"
export DWS_NODEPOOL="dws-pool"

다음은 이러한 환경 변수에 대한 설명입니다.

  • PROJECT_ID: Google Cloud 프로젝트의 ID입니다.
  • PROJECT_NUMBER: 프로젝트의 고유 식별자 번호입니다 (예: 123456789012).
  • ZONE: 클러스터의 컴퓨팅 영역입니다 (예: us-east5-b). 선택한 가속기 유형을 사용할 수 있는 영역을 선택합니다. 가용성 정보는 Cloud TPU 할당량 또는 GPU 할당량을 참고하세요.
  • REGION: 클러스터 리소스를 만드는 리전입니다 (예: us-east5).
  • CLUSTER_NAME: GKE 클러스터에 선택한 이름입니다.
  • GKE_VERSION: 클러스터의 GKE 버전입니다. 버전 1.34 이상을 사용합니다.
  • ONDEMAND_NODEPOOL: 표준 주문형 노드 풀의 이름입니다. (이 노드 풀이 계획 A 노드 풀입니다.)
  • DWS_NODEPOOL: DWS 플렉스 시작 노드 풀의 이름입니다. (이것이 계획 B 노드 풀입니다.)

인프라 구성 (클러스터 관리자)

클러스터 관리자는 대체 메커니즘을 지원하도록 GKE 클러스터와 노드 풀을 구성합니다.

GKE 클러스터 만들기

먼저 GKE 클러스터를 만듭니다. 이 클러스터는 Kueue 컨트롤러를 설치하고, 노드 풀을 구성하고, AI 학습 작업을 실행하는 환경입니다. 클러스터를 만들고 연결하려면 다음 단계를 따르세요.

  1. 클러스터를 만듭니다.

    gcloud container clusters create ${CLUSTER_NAME} \
  --cluster-version=${GKE_VERSION} \
  --machine-type=n2-standard-16 \
  --location=${ZONE} \
  --enable-image-streaming \
  --addons=RayOperator \
  --project=${PROJECT_ID}

    이 명령어는 다음 주요 플래그를 사용합니다.

    • --addons=RayOperator: 클러스터에 Ray Operator를 설치합니다. 이 연산자는 이 문서의 뒷부분에서 제출하는 RayJob 워크로드를 관리하는 데 필요합니다.
    • --enable-image-streaming: 클러스터에서 컨테이너 이미지를 더 빠르게 가져올 수 있습니다. 이 기능은 대규모 AI 컨테이너 이미지가 실행되는 데 걸리는 시간을 크게 줄여줍니다.

  2. kubectl CLI가 클러스터에 연결할 수 있도록 클러스터의 사용자 인증 정보를 검색합니다. 이 명령어는 기본적으로 ~/.kube/config 디렉터리에 저장된 Kubernetes 구성 파일을 업데이트합니다.

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --location=${ZONE} \
  --project=${PROJECT_ID}

노드 풀 만들기

환경의 기본 및 백업 노드 풀(온디맨드 노드 풀(계획 A) 및 DWS flex-start 노드 풀(계획 B))을 만듭니다.

  1. 온디맨드 노드 풀 만들기: 이 풀은 학습 작업의 기본 리소스로 사용됩니다.

    gcloud container node-pools create ${ONDEMAND_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    이 예에서 첫 번째 선택 머신은 TPU v6e 가속기입니다. --machine-type=ct6e-standard-4t 플래그를 사용하여 이 하드웨어를 지정합니다. 이 머신 유형을 AI 모델에 원하는 하드웨어(예: GPU 또는 다양한 TPU)에 맞게 변경할 수 있습니다.

  2. DWS flex-start 노드 풀 만들기: 이 예에서는 기본 주문형 풀에 선택한 것과 동일한 머신 유형 (--machine-type=ct6e-standard-4t)을 선택합니다. 계획 B 노드 풀에서 다른 머신 유형을 사용하지 않아도 됩니다. 이 특정 하드웨어를 정말로 원한다면 플랜 B로 선택하는 것은 즉시 사용할 수 없는 경우 다른 획득 방법으로 전환한다는 의미입니다. 이 대체 방법은 최대 7일 동안 사용 가능한 하드웨어를 지속적으로 검색하는 DWS를 사용합니다.

    gcloud container node-pools create ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --enable-queued-provisioning \
  --flex-start \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    이러한 명령어는 다음 주요 플래그를 사용합니다.

    • --num-nodes=0, --min-nodes=0, --max-nodes=4, --enable-autoscaling: 이 조합을 사용하면 작업에 노드가 필요할 때 노드 풀을 0개에서 확장하고 유휴 상태일 때 다시 축소할 수 있으므로 비용을 절감할 수 있습니다.
    • --tpu-topology: TPU 칩의 물리적 배열을 정의합니다. 칩의 물리적 배열이 분산 학습 작업의 실행 속도에 영향을 미치기 때문에 이 레이아웃을 지정합니다.
    • --reservation-affinity=none: 노드 풀이 미리 예약된 하드웨어를 사용하지 않도록 합니다. Google Cloud 를 사용하면 가용성을 보장하기 위해 특정 머신을 예약할 수 있습니다. 이 플래그를 none로 설정하면 시스템에서 이러한 예약을 우회하고 예약되지 않은 머신을 대신 동적으로 요청합니다.
    • --enable-queued-provisioning 및 --flex-start: (Plan B 풀만 해당) 이러한 플래그를 사용하면 DWS가 유연한 용량을 사용할 수 있게 될 때 Plan B 풀의 노드를 프로비저닝할 수 있습니다.

노드 풀에서 flex-start의 상태 확인

DWS flex-start 노드 풀을 검사하고 flex-start가 사용 설정되어 있는지 확인합니다.

gcloud container node-pools describe ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --format="get(config.flexStart)"

flex-start가 사용 설정되면 출력은 True입니다.

Kueue 설치 및 구성 (클러스터 관리자)

이 섹션에서는 클러스터에 Kueue 컨트롤러를 설치합니다. Kueue는 작업 대기열을 관리하는 스케줄링 프로그램입니다. 작업 요청을 가로채서 사용할 노드 풀 (온디맨드 또는 DWS 플렉스 시작)을 결정한 다음 작업을 할당합니다.

Kueue 설치

다음 명령어를 실행하여 Kueue를 설치합니다. 이 명령어는 공식 저장소에서 설치 매니페스트를 다운로드하여 클러스터에 적용합니다.

helm install kueue oci://registry.k8s.io/kueue/charts/kueue \
  --namespace kueue-system \
  --create-namespace \
  --set "controllerManager.featureGates[0].name=ElasticJobsViaWorkloadSlices" \
  --set "controllerManager.featureGates[0].enabled=true"

구성 규칙 정의

우선순위 규칙을 정의하는 YAML 매니페스트를 만듭니다. 이 규칙은 Kueue에 온디맨드 풀을 먼저 사용하고 DWS flex-start 풀을 두 번째로 사용하도록 지시합니다.

  1. 다음 콘텐츠를 포함한 dws-tpu-queue.yaml 파일을 만듭니다. 이 파일은 두 가지 리소스 플레이버 (온디맨드 및 DWS 플렉스 시작)와 우선순위를 지정하는 클러스터 대기열을 정의합니다. 이 구성 파일은 Kueue가 작업을 처리하는 데 사용하는 로직을 정의합니다.

    • ResourceFlavor: 이 문서의 시작 부분에서 두 개의 노드 풀을 만들고 환경 변수 ${ONDEMAND_NODEPOOL} 및 ${DWS_NODEPOOL}를 사용하여 이름을 할당했습니다. 이러한 노드 풀을 만들 때 GKE는 이러한 풀의 모든 노드에 환경 변수에 대해 선택한 이름으로 자동으로 라벨을 지정했습니다. ResourceFlavor 섹션은 Kueue에 이러한 라벨이 있는 노드를 찾도록 지시합니다.
    • ClusterQueue: 매니페스트의 이 섹션은 우선순위 규칙을 정의합니다. 주문형 버전이 먼저 나열되므로 Kueue는 주문형 머신을 먼저 프로비저닝하려고 시도합니다. Kueue가 이러한 머신을 가져올 수 없는 경우 대신 DWS 플렉스 시작 머신을 프로비저닝하려고 시도합니다.
    • Quotas: 파일은 주문형 노드 풀에서 작업이 언제든지 사용할 수 있는 총 리소스 (예: CPU, 메모리, TPU 칩)에 대한 한도인 할당량을 설정합니다. 작업이 이 한도에 도달하면 Kueue는 dws-tpu-queue.yaml에서 훨씬 높은 할당량 한도로 구성한 DWS 플렉스 시작 머신 (계획 B 머신)을 자동으로 프로비저닝하려고 시도합니다. 
      apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "default-cpu"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: default-pool
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "on-demand"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${ONDEMAND_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "dws"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${DWS_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: "cluster-queue"
spec:
  namespaceSelector: {}
  resourceGroups:
    - coveredResources: ["cpu", "memory", "google.com/tpu"]
      flavors:
        - name: "default-cpu" # Used for Ray Head Pod.
          resources:
            - name: "cpu"
              nominalQuota: 10
            - name: "memory"
              nominalQuota: 20Gi
            - name: "google.com/tpu"
              nominalQuota: 0
        - name: "on-demand" # First choice: on-demand node-pool.
          resources:
            - name: "cpu"
              nominalQuota: 40
            - name: "memory"
              nominalQuota: 75Gi
            - name: "google.com/tpu"
              nominalQuota: 16
        - name: "dws" # If on-demand is unavailable, fallback to DWS.
          resources:
            - name: "cpu"
              nominalQuota: 1000000000
            - name: "memory"
              nominalQuota: 1000000000Gi
            - name: "google.com/tpu"
              nominalQuota: 1000000000 # "Infinite" quota
  admissionChecksStrategy:
    admissionChecks:
      - name: "dws-prov"
        onFlavors: [dws]
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  namespace: "default"
  name: "user-queue"
spec:
  clusterQueue: "cluster-queue"
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: AdmissionCheck
metadata:
  name: dws-prov
spec:
  controllerName: kueue.x-k8s.io/provisioning-request
  parameters:
    apiGroup: kueue.x-k8s.io
    kind: ProvisioningRequestConfig
    name: dws-config
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ProvisioningRequestConfig
metadata:
  name: dws-config
spec:
  provisioningClassName: queued-provisioning.gke.io
  managedResources:
    - google.com/tpu

  2. 클러스터에 구성을 적용합니다. 다음 명령어는 envsubst라는 명령줄 도구를 사용하여 dws-tpu-queue.yaml 파일에 표시되는 자리표시자 변수를 대체합니다. envsubst는 자리표시자를 이전에 정의한 환경 변수의 값으로 바꿉니다.

    envsubst < dws-tpu-queue.yaml | kubectl apply -f -

학습 작업 실행 (AI 개발자)

AI 개발자는 RayJob 매니페스트를 만들어 학습 워크로드를 정의하고 제출합니다. 이 매니페스트에서 리소스 요구사항을 지정하면 클러스터 관리자가 이전에 Kueue 및 DWS로 구성한 자동 대체 시스템이 기본 노드 풀을 처리합니다.

이 섹션에서는 다음 단계를 수행합니다.

  • Python 학습 스크립트를 만듭니다.
  • 해당 스크립트를 Kubernetes ConfigMap에 저장합니다.
  • 학습 스크립트가 노드에서 실행될 수 있도록 ConfigMap을 볼륨으로 마운트하는 RayJob을 배포합니다.

이 단계를 수행하면 Ray Train이 노드에 JAX 워크로드를 자동으로 분산하고 Kueue가 필요한 머신을 가져오는 작업을 처리합니다.

학습 스크립트

다음 Python 스크립트를 train.py 파일에 복사하여 붙여넣습니다.

import time
import ray
from ray import train
from ray.train import ScalingConfig
from ray.train.v2.jax import JaxTrainer
import jax
import jax.numpy as jnp

def train_func():
    # JaxTrainer handles JAX distributed setup.
    print(f"Local Devices: {jax.local_devices()}")

    # Simple Linear Regression Training Loop.
    key = jax.random.PRNGKey(0)
    x = jax.random.normal(key, (1000, 10))
    w_true = jax.random.normal(key, (10, 1))
    y = jnp.dot(x, w_true)

    # Initialize weights
    w = jnp.zeros((10, 1))
    learning_rate = 0.1

    @jax.jit
    def update(w, x, y):
        y_pred = jnp.dot(x, w)
        loss = jnp.mean((y_pred - y) ** 2)
        grad = jax.grad(lambda w: jnp.mean((jnp.dot(x, w) - y) ** 2))(w)
        return w - learning_rate * grad, loss

    # Training loop
    print("Starting training...")
    for epoch in range(50):
        w, loss = update(w, x, y)
        if epoch % 10 == 0:
            train.report({"loss": loss.item(), "epoch": epoch})
            print(f"Epoch {epoch}: Loss {loss:.4f}")

    print("Training Complete!")
    # Allow metrics to sync before closing
    time.sleep(3)

def main():
    scaling_config = ScalingConfig(
        num_workers=4,
        resources_per_worker={"TPU": 4},
        use_tpu=True,
        topology="4x4",
        accelerator_type="TPU-V6E"
    )

    trainer = JaxTrainer(
        train_loop_per_worker=train_func,
        scaling_config=scaling_config
    )

    result = trainer.fit()
    print(f"Run Result: {result.metrics}")

if __name__ == "__main__":
    main()

학습 스크립트는 고성능 수치 계산을 위한 Python 라이브러리인 JAX를 사용하여 선형 회귀 모델을 학습시킵니다. 이 스크립트는 자동 대체에 DWS와 Kueue를 사용하는 방법을 보여주기 위해 설계된 단순화된 예시이며 데이터 병렬 처리 또는 모델 병렬 처리를 실행하지 않습니다.

학습 스크립트의 ScalingConfig 섹션은 학습 작업의 하드웨어 요구사항을 정의합니다. 이 섹션에서는 이전에 구성한 노드 풀의 실제 레이아웃과 일치하는 4x4 TPU 토폴로지를 요청합니다.

ConfigMap 만들기

train.py 스크립트의 콘텐츠를 Kubernetes ConfigMap 객체에 업로드합니다. 이렇게 하면 클러스터가 스크립트를 저장하고 RayJob에서 사용할 수 있습니다.

kubectl create configmap jax-train-script --from-file=train.py

다음 섹션에서 정의하는 RayJob은 이 ConfigMap을 볼륨으로 마운트합니다. 이렇게 하면 Ray 소프트웨어가 스크립트 파일을 찾아 실행할 수 있도록 Ray 컨테이너 내에 스크립트 파일이 표시됩니다.

RayJob 매니페스트 적용

다음 콘텐츠를 포함한 rayjob-tpu-v6e-dws.yaml 파일을 만듭니다. 이 매니페스트는 학습 작업을 정의하고 시스템에 라우팅 방법을 알려줍니다.

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: rayjob-tpu-v6e-dws-${JOB_ID}
  labels:
    kueue.x-k8s.io/queue-name: user-queue
  annotations:
    kueue.x-k8s.io/elastic-job: "true"
spec:
  shutdownAfterJobFinishes: true
  entrypoint: python /app/train.py
  runtimeEnvYAML: |
    pip:
      - jax[tpu]==0.8.2
      - pandas==2.3.3
  rayClusterSpec:
    enableInTreeAutoscaling: true
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          nodeSelector:
            cloud.google.com/gke-nodepool: default-pool
          containers:
            - name: ray-head
              image: rayproject/ray:2.53.0-py311
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "2"
                  memory: "4Gi"
                requests:
                  cpu: "2"
                  memory: "4Gi"
              volumeMounts:
                - mountPath: /app
                  name: train-script-volume
          volumes:
            - name: train-script-volume
              configMap:
                name: jax-train-script
    workerGroupSpecs:
      - replicas: 1
        minReplicas: 1
        maxReplicas: 2
        numOfHosts: 4
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            tolerations:
              - key: "google.com/tpu"
                operator: "Exists"
                effect: "NoSchedule"
              - key: "cloud.google.com/gke-queued"
                operator: "Exists"
                effect: "NoSchedule"
            containers:
              - name: ray-worker
                image: rayproject/ray:2.53.0-py311
                resources:
                  limits:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                  requests:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                volumeMounts:
                  - mountPath: /app
                    name: train-script-volume
            volumes:
              - name: train-script-volume
                configMap:
                  name: jax-train-script
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 4x4

이 매니페스트에는 대체 시스템이 작동하도록 하는 세 가지 구성이 포함되어 있습니다.

  • 특정 하드웨어 요청: nodeSelector 섹션은 스크립트에 필요한 하드웨어를 지정합니다 (이 예에서는 4x4 토폴로지가 있는 tpu-v6e-slice).
  • 대기열 선택: kueue.x-k8s.io/queue-name 라벨은 작업을 Kueue로 직접 라우팅합니다. 이렇게 하면 자동 대체 로직이 사용 설정됩니다.
  • DWS flex-start 노드 허용: tolerations 섹션을 사용하면 Plan B 노드 풀에서 작업을 실행할 수 있습니다. DWS 유연한 시작 노드는 일반 워크로드가 실수로 실행되지 않도록 GKE에서 특별히 표시 (taint)하므로 작업에서 cloud.google.com/gke-queued taint를 명시적으로 허용해야 합니다.

워크로드 제출

대체 시스템이 작동하는지 증명하려면 두 개의 작업을 제출해야 합니다. 첫 번째 작업은 계획 A 주문형 용량을 소비하므로 두 번째 작업은 계획 B DWS flex-start 용량으로 대체됩니다.

다음 명령어를 실행하여 두 작업을 제출합니다. 이 명령어는 for 루프와 envsubst를 사용하여 실행마다 고유한 작업 ID를 매니페스트에 삽입합니다.

for i in 1 2; do
  export JOB_ID=$i
  envsubst < rayjob-tpu-v6e-dws.yaml | kubectl apply -f -
  echo "Submitted Job $i"
  sleep 2
done

작업을 제출하면 시스템에서 다음과 같이 워크로드를 처리합니다.

  1. 인터셉트: Kueue는 대기열 라벨을 사용하여 작업을 감지하고 일시적으로 일시중지합니다.
  2. 결정: Kueue는 관리자의 규칙에 따라 리소스 가용성을 평가합니다. 먼저 Plan A 풀을 확인합니다.
  3. 과제:
    • 계획 A 리소스는 첫 번째 작업에 사용할 수 있으므로 Kueue는 작업 1을 계획 A에 할당합니다.
    • 작업 1이 계획 A 리소스를 사용하므로 Kueue는 자동으로 작업 2를 계획 B (DWS 플렉스 시작) 풀에 할당합니다.
  4. 시작: Kueue가 작업을 일시중지 해제합니다. 이 작업은 GKE 클러스터 자동 확장 처리를 트리거하여 노드를 프로비저닝하고 학습 스크립트를 시작합니다.

RayJob에 연결

최종 확인 단계로 kubectl port-forward 명령어를 사용하여 Ray 대시보드에 연결하고 작업이 실행되는 것을 확인할 수 있습니다.

첫 번째 작업의 상태를 확인하려면 다음 명령어를 실행합니다.

kubectl port-forward service/rayjob-tpu-v6e-dws-1-head-svc 8265:8265 &

이 명령어를 실행한 후 웹브라우저를 열고 http://localhost:8265로 이동합니다. Ray 대시보드에서 작업 상태와 보고된 측정항목을 확인하여 두 작업이 각각의 노드 풀에서 성공적으로 완료되었는지 확인할 수 있습니다.

다음 명령어를 실행하여 첫 번째 작업의 로그를 볼 수도 있습니다.

kubectl logs job/rayjob-tpu-v6e-dws-1

잘린 학습 스크립트 출력은 다음과 유사해야 합니다. 출력 끝부분에 Training Complete! 및 Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded 메시지가 표시됩니다.

(pid=, ip=10.68.3.4) 5] XLA::TPU program HBM usage: 52.5K / 31.25G
(pid=, ip=10.68.9.4) :2152] XLA::TPU program VMEM usage: 141.0K / 128.00M [repeated 5x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722540     855 deepsea_compiler_backend.cc:2163] Total hbm usage >= 260.14M: [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167]     reserved           204B [repeated 19x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722542     855 deepsea_compiler_backend.cc:2163]     program           70.0K  [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     arguments            0B  [repeated 12x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777627     888 deepsea_compiler_backend.cc:2163] Output size 0B; shares 0B with arguments. [repeated 14x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777625     888 deepsea_compiler_backend.cc:2163] Total host usage >= 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     program         unknown size  [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167] Program sflag requirement 224B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167]     scoped              40B [repeated 21x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167] Program vmem requirement 141.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program smem requirement 40B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program host requirement 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program hbm requirement 70.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2167]     overlays          70.0K [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2175] XLA::TPU program SMEM usage: 1.9K / 1.00M (3 parameters) [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167]     HLO temp          76.0K (0.0% utilization: Unpadded (0B) Padded (0B), 100.0% fragmentation (76.0K)) [repeated 14x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Training Complete! [repeated 3x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Epoch 40: Loss 0.0000 [repeated 3x across cluster]
2026-03-20 03:59:51,008 SUCC cli.py:65 -- ------------------------------------------
2026-03-20 03:59:51,008 SUCC cli.py:66 -- Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded
2026-03-20 03:59:51,008 SUCC cli.py:67 -- ------------------------------------------

삭제

이 문서에서 사용된 리소스 비용이 Google Cloud 계정에 청구되지 않도록 하려면 리소스가 포함된 프로젝트를 삭제하거나 프로젝트는 유지하되 개별 리소스를 삭제하세요.

프로젝트 삭제

  • Google Cloud 콘솔에서 리소스 관리 페이지로 이동합니다.

    리소스 관리로 이동

  • 프로젝트 목록에서 삭제할 프로젝트를 선택하고 삭제를 클릭합니다.
  • 대화상자에서 프로젝트 ID를 입력한 후 종료를 클릭하여 프로젝트를 삭제합니다.

    • 개별 리소스 삭제

    이 문서에서 사용한 GGoogle Cloud 프로젝트를 유지하려면 다음 명령어를 실행하여 클러스터를 삭제합니다.

    gcloud container clusters delete ${CLUSTER_NAME} \
    --location=${ZONE} \
    --project=${PROJECT_ID} \
    --quiet

    요약

    이 문서에서는 Ray 학습 환경을 구성하고 테스트했습니다. 이 환경에서는 기본 노드 풀과 백업 DWS 풀을 사용하여 하드웨어 가용성을 극대화합니다. 기본 머신을 사용할 수 없을 때 DWS로 자동 대체하여 학습 작업이 대기열에서 대기하는 시간을 최소화했습니다.

    이 문제를 해결하기 위해 다음 단계를 수행했습니다.

    1. GKE 클러스터 생성: 노드 풀 및 예약 도구를 호스팅할 환경을 설정했습니다.
    2. 노드 풀 구성: 주문형 노드 풀 (계획 A)과 DWS 노드 풀 (계획 B)을 만들었습니다.
    3. Kueue 설치 및 구성: Kueue 컨트롤러를 배포하고 시스템이 먼저 계획 A를 시도한 후 계획 B로 대체하도록 지시하는 우선순위 규칙을 적용했습니다.
    4. ConfigMap 생성: 테스트 워크로드로 사용할 단순화된 JAX 학습 스크립트를 클러스터에 배포했습니다.
    5. RayJob 매니페스트 정의: 특정 하드웨어를 요청하고, Kueue 컨트롤러로 라우팅하고, DWS 노드를 허용하도록 작업을 구성했습니다.
    6. 워크로드 제출: 계획 A 리소스가 사용될 때 Kueue가 두 번째 작업을 계획 B로 자동 라우팅하도록 두 작업을 제출했습니다.
    7. 결과 확인: 포트 포워딩을 사용하여 Ray 대시보드에 연결하고 두 작업이 모두 성공적으로 실행되었는지 확인했습니다.

    다음 단계