LabHub

블로그

Kubernetes AI 학습 파이프라인: Volcano, Training Operator, Kueue 분석

2026-03-01 · 31 분 읽기 · #kubernetes#ai#distributed-training#volcano#kubeflow

한국어English日本語中文

1. Kubernetes 기본 스케줄러의 한계

Kubernetes의 기본 스케줄러(kube-scheduler)는 Pod 단위의 스케줄링을 수행한다. 하나의 Pod가 요청하는 리소스(CPU, Memory, GPU)를 기반으로 적절한 Node에 배치하는 것이 핵심 역할이다. 이 구조는 웹 서비스, API 서버 같은 일반적인 워크로드에는 적합하지만, AI/ML 분산 학습과 같은 배치 워크로드에서는 근본적인 한계가 있다.

1.1 Gang Scheduling의 부재

분산 학습에서 가장 큰 문제는 Gang Scheduling의 부재다. 예를 들어, PyTorch DDP(Distributed Data Parallel) 학습에서 4개의 Worker Pod가 모두 동시에 시작되어야 학습이 진행된다. 기본 스케줄러는 Pod를 개별적으로 스케줄링하므로, 4개 중 3개만 스케줄링되고 1개는 리소스 부족으로 Pending 상태에 머물 수 있다. 이미 배치된 3개의 Pod는 GPU 리소스를 점유하고 있지만 학습은 시작되지 못하는 리소스 낭비 상태가 발생한다.

1.2 Deadlock 문제

더 심각한 것은 Deadlock이다. 두 개의 분산 학습 Job이 각각 4개의 GPU를 필요로 하는데, 클러스터에 총 6개의 GPU만 있다고 가정하자. 기본 스케줄러는 Job A에 3개, Job B에 3개의 GPU를 할당할 수 있다. 두 Job 모두 최소 4개가 필요하므로 어느 쪽도 학습을 시작하지 못한다. 이것이 바로 Gang Scheduling이 해결하는 "All or Nothing" 문제다.

1.3 Fair-share와 Queue 기반 리소스 관리 부재

기본 스케줄러에는 팀 간 리소스 공유를 위한 Queue 개념이 없다. AI 조직에서는 여러 연구 팀이 GPU 클러스터를 공유하는데, 팀별 Quota를 설정하고 유휴 리소스를 공유하는 Fair-share 정책을 기본 스케줄러만으로는 구현하기 어렵다.

이러한 한계를 해결하기 위해 Volcano, Kubeflow Training Operator, Kueue 같은 프로젝트가 등장했다.

2. Volcano Scheduler 공식 문서 분석

Volcano는 CNCF(Cloud Native Computing Foundation) 산하 프로젝트로, Kubernetes 위에서 고성능 배치 워크로드를 처리하기 위한 스케줄러이자 컨트롤러이다. AI/ML, BigData, HPC(High Performance Computing) 등의 워크로드를 위한 배치 스케줄링 시스템이다.

2.1 아키텍처

Volcano는 세 가지 핵심 컴포넌트로 구성된다.

Volcano Scheduler는 Kubernetes의 Scheduling Framework와 호환되며, predicates와 nodeorder 플러그인을 통해 기본 스케줄러의 PreFilter/Filter 및 Score 단계를 수행한다. 이 위에 Gang Scheduling, Fair-share, Binpack 등의 고급 스케줄링 전략을 플러그인 형태로 추가한다.

2.2 Queue CRD

Queue는 클러스터 리소스를 논리적으로 분할하여 멀티 테넌트 환경에서 리소스 격리를 제공하는 CRD이다.

apiVersion: scheduling.volcano.sh/v1beta1
kind: Queue
metadata:
  name: ml-research-queue
spec:
  weight: 4
  reclaimable: true
  capability:
    cpu: '64'
    memory: '256Gi'
    nvidia.com/gpu: '8'

2.3 PodGroup CRD

PodGroup은 서로 강하게 연관된 Pod들의 그룹을 정의하는 CRD로, Gang Scheduling의 핵심 단위이다.

apiVersion: scheduling.volcano.sh/v1beta1
kind: PodGroup
metadata:
  name: pytorch-training-pg
  namespace: default
spec:
  minMember: 4
  minResources:
    cpu: '16'
    memory: '64Gi'
    nvidia.com/gpu: '4'
  priorityClassName: high-priority
  queue: ml-research-queue

PodGroup의 Status Phase는 다음과 같이 전이한다: Pending (리소스 미충족) -> Inqueue (검증 완료, 노드 바인딩 대기) -> Running (minMember 이상 Pod 실행 중) -> Unknown (일부 Pod가 스케줄링 불가). VolcanoJob을 생성하면 PodGroup은 자동으로 생성된다.

3. Volcano Gang Scheduling과 Fair-share 스케줄링 정책

3.1 Gang Scheduling

Gang Scheduling은 Volcano의 가장 핵심적인 기능이다. "All or Nothing" 전략으로, Job을 구성하는 모든 Pod(또는 minMember에 지정된 최소 수의 Pod)가 동시에 스케줄링 가능할 때만 배치를 수행한다.

Gang Scheduling은 다음 시나리오에서 필수적이다:

Gang Plugin은 PodGroup의 minMember 필드를 확인하여, 클러스터에서 해당 수의 Pod를 동시에 스케줄링 가능한지 판단한다. 불가능하면 어떤 Pod도 스케줄링하지 않아 리소스 낭비와 Deadlock을 방지한다.

3.2 DRF (Dominant Resource Fairness) 스케줄링

DRF는 멀티 테넌트 환경에서 공정한 리소스 분배를 구현하는 알고리즘이다. 각 Job의 지배적 리소스(Dominant Resource) 를 기준으로 Fair-share를 계산한다.

예를 들어, Job A가 CPU 중심(CPU 8코어, Memory 4GB), Job B가 메모리 중심(CPU 2코어, Memory 32GB)이라면, 각 Job의 지배적 리소스 비율이 가장 낮은 Job을 우선 스케줄링하여 공정성을 확보한다. DRF Plugin은 다차원 리소스(CPU, Memory, GPU)를 동시에 고려하여 특정 리소스 유형에 편중되지 않는 공정한 분배를 수행한다.

3.3 Proportion Plugin

Proportion Plugin은 Queue별 리소스 비율을 관리한다. 각 Queue의 weight를 기준으로 전체 클러스터 리소스를 비례 배분하며, Queue의 capability 범위 내에서 리소스를 할당한다. 여러 팀이 GPU 클러스터를 공유할 때, 팀별 Queue를 생성하고 weight를 설정하여 리소스 배분 비율을 조정할 수 있다.

4. Volcano Job CRD와 Plugin 시스템

4.1 VolcanoJob CRD

VolcanoJob은 배치 워크로드를 정의하는 Volcano의 핵심 CRD다. 여러 Task(역할)를 하나의 Job으로 묶어 관리할 수 있다.

apiVersion: batch.volcano.sh/v1alpha1
kind: Job
metadata:
  name: pytorch-ddp-training
spec:
  minAvailable: 5
  schedulerName: volcano
  queue: ml-research-queue
  plugins:
    env: []
    svc: []
  policies:
    - event: PodEvicted
      action: RestartJob
  tasks:
    - replicas: 1
      name: master
      template:
        spec:
          containers:
            - image: pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
              name: master
              command: ['torchrun']
              args: ['--nproc_per_node=1', '--nnodes=5', '--node_rank=0', 'train.py']
              resources:
                limits:
                  nvidia.com/gpu: 1
          restartPolicy: OnFailure
    - replicas: 4
      name: worker
      template:
        spec:
          containers:
            - image: pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
              name: worker
              command: ['torchrun']
              args: ['--nproc_per_node=1', '--nnodes=5', 'train.py']
              resources:
                limits:
                  nvidia.com/gpu: 1
          restartPolicy: OnFailure

4.2 Plugin 시스템 상세

Volcano Scheduler는 플러그인 아키텍처를 채택하여 스케줄링 전략을 유연하게 확장할 수 있다. 주요 플러그인은 다음과 같다.

Plugin기능적합한 워크로드
GangAll or Nothing 스케줄링분산 학습, MPI, Spark
Binpack기존 노드를 최대한 채워서 배치소규모 Job, 리소스 효율 극대화
DRFDominant Resource 기반 Fair-share멀티 테넌트 배치 처리
ProportionQueue별 리소스 비율 관리팀별 리소스 격리
Priority우선순위 기반 스케줄링우선순위가 다른 다중 Job
Task-TopologyTask 간 Affinity/Anti-Affinity네트워크 최적화가 필요한 분산 학습
TDM시분할 리소스 공유Kubernetes + YARN 혼합 환경
SLA대기 시간 제한실시간 서비스 요구사항이 있는 배치
Numa-awareNUMA 토폴로지 인식 스케줄링CPU 캐시 친화성이 중요한 HPC
PredicatesGPU 리소스 기반 사전 필터링GPU 집약적 AI 워크로드
Nodeorder다차원 노드 스코어링복합 스케줄링 커스터마이징

Volcano Scheduler 설정 파일에서 actions와 tiers를 통해 플러그인 조합을 구성한다:

actions: 'enqueue, allocate, backfill'
tiers:
  - plugins:
      - name: priority
      - name: gang
      - name: conformance
  - plugins:
      - name: drf
      - name: predicates
      - name: proportion
      - name: nodeorder
      - name: binpack

5. Kubeflow Training Operator 공식 문서 분석

Kubeflow Training Operator는 Kubernetes 위에서 분산 AI 학습 Job을 관리하는 Operator이다. PyTorch, TensorFlow, MPI, XGBoost 등 다양한 프레임워크를 지원하며, 각 프레임워크에 맞는 CRD를 제공한다.

5.1 V1 (Legacy)과 V2 아키텍처

현재 Training Operator는 두 가지 버전이 병존한다.

V1 (Legacy): 프레임워크별 별도 CRD(PyTorchJob, TFJob, MPIJob 등)를 사용한다. 각 CRD는 해당 프레임워크의 분산 학습 패턴에 맞는 ReplicaSpec(Master, Worker, PS 등)을 제공한다.

V2 (Current): 단일 TrainJob CRD로 통합되었다. 프레임워크별 설정은 TrainingRuntime과 ClusterTrainingRuntime을 통해 관리된다. PyTorch, MLX, HuggingFace, DeepSpeed, JAX, XGBoost 등 더 넓은 범위의 프레임워크를 지원한다.

본 글에서는 실제 프로덕션에서 널리 사용되는 V1의 PyTorchJob을 중심으로 분석하되, V2의 방향성도 함께 살펴본다.

5.2 Training Operator의 역할

Training Operator는 다음을 자동으로 처리한다:

  1. 환경 변수 자동 설정: WORLD_SIZE, RANK, MASTER_ADDR, MASTER_PORT 등 분산 학습에 필요한 환경 변수를 Pod에 자동 주입한다.
  2. torchrun CLI 설정: PyTorch의 torchrun(이전 torch.distributed.launch)이 올바르게 동작하도록 환경을 구성한다.
  3. Job 상태 관리: Created -> Running -> Succeeded/Failed 등의 라이프사이클을 관리하며, 실패 시 재시작 정책을 적용한다.
  4. Service 생성: Pod 간 통신을 위한 Headless Service를 자동 생성한다.

6. PyTorchJob CRD 상세 분석 및 YAML 예제

PyTorchJob은 Kubernetes 위에서 PyTorch 분산 학습 Job을 실행하기 위한 CRD이다.

6.1 PyTorchJob 구조

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: pytorch-ddp-mnist
  namespace: ml-training
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        metadata:
          annotations:
            sidecar.istio.io/inject: 'false'
        spec:
          containers:
            - name: pytorch
              image: my-registry/pytorch-training:latest
              command:
                - 'torchrun'
                - '--nproc_per_node=2'
                - '--nnodes=3'
                - '--node_rank=0'
                - '--master_addr=$(MASTER_ADDR)'
                - '--master_port=$(MASTER_PORT)'
                - 'train.py'
                - '--epochs=100'
                - '--batch-size=256'
              env:
                - name: NCCL_DEBUG
                  value: 'INFO'
              resources:
                limits:
                  nvidia.com/gpu: 2
                  memory: '32Gi'
                  cpu: '8'
              volumeMounts:
                - name: training-data
                  mountPath: /data
                - name: checkpoints
                  mountPath: /checkpoints
          volumes:
            - name: training-data
              persistentVolumeClaim:
                claimName: training-data-pvc
            - name: checkpoints
              persistentVolumeClaim:
                claimName: checkpoints-pvc
    Worker:
      replicas: 2
      restartPolicy: OnFailure
      template:
        metadata:
          annotations:
            sidecar.istio.io/inject: 'false'
        spec:
          containers:
            - name: pytorch
              image: my-registry/pytorch-training:latest
              command:
                - 'torchrun'
                - '--nproc_per_node=2'
                - '--nnodes=3'
                - '--master_addr=$(MASTER_ADDR)'
                - '--master_port=$(MASTER_PORT)'
                - 'train.py'
                - '--epochs=100'
                - '--batch-size=256'
              env:
                - name: NCCL_DEBUG
                  value: 'INFO'
              resources:
                limits:
                  nvidia.com/gpu: 2
                  memory: '32Gi'
                  cpu: '8'
              volumeMounts:
                - name: training-data
                  mountPath: /data
                - name: checkpoints
                  mountPath: /checkpoints
          volumes:
            - name: training-data
              persistentVolumeClaim:
                claimName: training-data-pvc
            - name: checkpoints
              persistentVolumeClaim:
                claimName: checkpoints-pvc

6.2 핵심 필드 설명

6.3 Job 상태 모니터링

kubectl get pytorchjobs -n ml-training
kubectl describe pytorchjob pytorch-ddp-mnist -n ml-training
kubectl get pods -l training.kubeflow.org/job-name=pytorch-ddp-mnist -n ml-training

PyTorchJob의 상태는 conditions 필드에서 확인할 수 있으며, Created -> Running -> Succeeded 순으로 전이한다.

7. MPIJob, TFJob 지원

7.1 MPIJob

MPIJob은 MPI(Message Passing Interface) 기반 분산 학습을 위한 CRD이다. Horovod와 같은 allreduce 기반 프레임워크에서 주로 사용된다.

apiVersion: kubeflow.org/v2beta1
kind: MPIJob
metadata:
  name: horovod-training
spec:
  slotsPerWorker: 2
  mpiReplicaSpecs:
    Launcher:
      replicas: 1
      template:
        spec:
          containers:
            - name: mpi-launcher
              image: horovod/horovod:latest
              command:
                - mpirun
                - --allow-run-as-root
                - -np
                - '8'
                - -bind-to
                - none
                - -map-by
                - slot
                - python
                - train.py
              resources:
                limits:
                  cpu: '2'
                  memory: '4Gi'
    Worker:
      replicas: 4
      template:
        spec:
          containers:
            - name: mpi-worker
              image: horovod/horovod:latest
              resources:
                limits:
                  nvidia.com/gpu: 2
                  cpu: '8'
                  memory: '32Gi'

MPIJob의 특징:

7.2 TFJob

TFJob은 TensorFlow의 Parameter Server 분산 학습 전략을 위한 CRD이다.

apiVersion: kubeflow.org/v1
kind: TFJob
metadata:
  name: tf-ps-training
spec:
  tfReplicaSpecs:
    PS:
      replicas: 2
      template:
        spec:
          containers:
            - name: tensorflow
              image: tensorflow/tensorflow:2.14.0-gpu
              command: ['python', 'train.py']
              resources:
                limits:
                  cpu: '4'
                  memory: '16Gi'
    Worker:
      replicas: 4
      template:
        spec:
          containers:
            - name: tensorflow
              image: tensorflow/tensorflow:2.14.0-gpu
              command: ['python', 'train.py']
              resources:
                limits:
                  nvidia.com/gpu: 1
                  cpu: '8'
                  memory: '32Gi'

TFJob의 특징:

8. Multi-node Multi-GPU 분산 학습 설정 (PyTorch DDP on K8s)

8.1 PyTorch DDP의 동작 원리

PyTorch DDP(Distributed Data Parallel)는 데이터 병렬 처리 방식의 분산 학습 전략이다. 각 GPU에 모델 복제본을 배치하고, 미니 배치를 GPU 수만큼 분할하여 학습한 뒤, allreduce 연산으로 그래디언트를 동기화한다.

Kubernetes에서 Multi-node Multi-GPU DDP를 구성할 때 이해해야 할 핵심 개념:

8.2 학습 코드 예제

Training Operator를 사용할 때, 학습 코드 자체는 PyTorch의 네이티브 Distributed API를 사용한다:

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def main():
    # Training Operator가 환경 변수를 자동 설정하므로
    # torchrun이 init_process_group을 자동 처리한다.
    dist.init_process_group(backend="nccl")

    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    model = MyModel().cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])

    # DistributedSampler로 데이터를 GPU 수만큼 분할
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=64,
        sampler=train_sampler
    )

    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)
        for batch in train_loader:
            loss = model(batch)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

        # 체크포인트는 rank 0에서만 저장
        if dist.get_rank() == 0:
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.module.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
            }, f'/checkpoints/checkpoint_epoch_{epoch}.pt')

    dist.destroy_process_group()

8.3 NCCL 통신 최적화

Multi-node GPU 통신에서 NCCL(NVIDIA Collective Communication Library)의 성능은 학습 속도에 직접적인 영향을 미친다. Kubernetes 환경에서의 NCCL 최적화 팁:

env:
  - name: NCCL_DEBUG
    value: 'INFO' # 디버깅 시 NCCL 통신 로그 확인
  - name: NCCL_IB_DISABLE
    value: '0' # InfiniBand 사용 가능 시 활성화
  - name: NCCL_SOCKET_IFNAME
    value: 'eth0' # 통신에 사용할 네트워크 인터페이스 지정
  - name: NCCL_P2P_LEVEL
    value: 'NVL' # NVLink 사용 시 P2P 통신 레벨 설정

InfiniBand나 RoCE 같은 RDMA 네트워크가 있는 환경에서는 hostNetwork: true를 Pod에 설정하여 네트워크 성능을 극대화할 수도 있다.

9. PVC/PV로 학습 데이터 및 체크포인트 관리

분산 학습에서 스토리지 설계는 학습 데이터 접근, 체크포인트 저장, 모델 아티팩트 관리의 핵심이다.

9.1 학습 데이터용 PVC

학습 데이터는 모든 Worker에서 읽어야 하므로 ReadOnlyMany(ROX) 또는 ReadWriteMany(RWX) Access Mode가 필요하다.

apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: training-data-pvc
  namespace: ml-training
spec:
  accessModes:
    - ReadOnlyMany
  storageClassName: nfs-storage
  resources:
    requests:
      storage: 500Gi

9.2 체크포인트용 PVC

체크포인트는 Rank 0 프로세스만 저장하므로 ReadWriteOnce(RWO) 로도 충분하지만, Fault Tolerance를 고려하면 ReadWriteMany(RWX) 가 더 유연하다.

apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: checkpoints-pvc
  namespace: ml-training
spec:
  accessModes:
    - ReadWriteMany
  storageClassName: fast-ssd
  resources:
    requests:
      storage: 100Gi

9.3 스토리지 백엔드 선택

스토리지 백엔드Access Mode성능 특성적합한 용도
NFSRWX/ROX중간 대역폭, 높은 호환성학습 데이터 공유
CephFSRWX높은 대역폭, 분산 스토리지대규모 데이터셋
Lustre / GPFSRWX매우 높은 대역폭HPC 수준 I/O
Local SSDRWO최고 IOPS체크포인트, 로컬 캐시
S3 (CSI)RWX높은 대역폭, 무한 용량오브젝트 스토리지 기반 데이터셋

대규모 학습 데이터셋의 경우, NFS나 CephFS 위에 직접 접근하면 I/O 병목이 발생할 수 있다. 이 경우 학습 시작 전에 데이터를 Local SSD로 복사하는 Init Container 패턴을 사용하거나, S3 호환 오브젝트 스토리지를 CSI Driver를 통해 마운트하는 방법을 고려한다.

10. Kueue: 차세대 Job 큐잉 시스템

Kueue는 Kubernetes SIG(Special Interest Group) 산하에서 개발되는 클라우드 네이티브 Job 큐잉 시스템이다. Batch, HPC, AI/ML 워크로드를 위한 Job 레벨의 Admission Control을 제공한다.

10.1 Kueue vs Volcano

Kueue와 Volcano는 서로 다른 레이어에서 동작한다.

Kueue는 "Job이 시작되어야 하는가?"를 결정하고, Volcano(또는 kube-scheduler)는 "시작된 Job의 Pod를 어디에 배치할 것인가?"를 결정한다. 두 시스템은 상호 보완적으로 사용 가능하다.

10.2 Kueue의 핵심 개념

Kueue의 Admission 과정은 다음과 같다:

  1. 사용자가 Job을 생성하면 Kueue가 이를 Workload로 변환한다.
  2. Workload는 LocalQueue에 제출된다.
  3. LocalQueue는 ClusterQueue를 참조한다.
  4. ClusterQueue는 ResourceFlavor에 정의된 리소스 풀에서 Quota를 확인한다.
  5. Quota가 충분하면 Workload를 Admit하여 실제 Pod가 생성된다.

11. Kueue의 ResourceFlavor, ClusterQueue, LocalQueue

11.1 ResourceFlavor

ResourceFlavor는 클러스터에서 사용 가능한 리소스의 유형을 정의한다. 노드의 특성(가격, 아키텍처, 가용성 등)을 구분하는 데 사용된다.

apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: gpu-a100
spec:
  nodeLabels:
    cloud.google.com/gke-accelerator: nvidia-tesla-a100
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: gpu-t4-spot
spec:
  nodeLabels:
    cloud.google.com/gke-accelerator: nvidia-tesla-t4
    cloud.google.com/gke-provisioning: spot
  tolerations:
    - key: cloud.google.com/gke-spot
      operator: Equal
      value: 'true'
      effect: NoSchedule

위 예시에서 gpu-a100은 On-demand A100 GPU 노드를, gpu-t4-spot은 Spot T4 GPU 노드를 나타낸다. Toleration을 통해 Spot 노드에 Pod가 스케줄링될 수 있도록 설정한다.

11.2 ClusterQueue

ClusterQueue는 클러스터 범위의 리소스 풀을 정의하고, Quota와 Fair Sharing 규칙을 관리한다.

apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: ml-cluster-queue
spec:
  cohort: ml-team-cohort
  namespaceSelector: {}
  preemption:
    withinClusterQueue: LowerPriority
    reclaimWithinCohort: LowerPriority
    borrowWithinCohort:
      policy: LowerPriority
      maxPriorityThreshold: 100
  resourceGroups:
    - coveredResources: ['cpu', 'memory', 'nvidia.com/gpu']
      flavors:
        - name: gpu-a100
          resources:
            - name: 'cpu'
              nominalQuota: 64
            - name: 'memory'
              nominalQuota: 256Gi
            - name: 'nvidia.com/gpu'
              nominalQuota: 8
              borrowingLimit: 4
              lendingLimit: 2
        - name: gpu-t4-spot
          resources:
            - name: 'cpu'
              nominalQuota: 128
            - name: 'memory'
              nominalQuota: 512Gi
            - name: 'nvidia.com/gpu'
              nominalQuota: 16

핵심 필드 설명:

11.3 LocalQueue

LocalQueue는 네임스페이스 범위에서 사용자가 Job을 제출하는 진입점이다.

apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  name: ml-research-queue
  namespace: ml-training
spec:
  clusterQueue: ml-cluster-queue

사용자는 Job에 kueue.x-k8s.io/queue-name 레이블을 추가하여 LocalQueue에 제출한다:

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: my-training-job
  namespace: ml-training
  labels:
    kueue.x-k8s.io/queue-name: ml-research-queue
spec:
  # ... PyTorchJob spec

11.4 Kueue의 큐잉 전략

Kueue는 두 가지 큐잉 전략을 지원한다:

12. Spot/Preemptible Instance에서의 학습 (Fault Tolerance)

12.1 Spot Instance의 장점과 위험

클라우드 Spot Instance(AWS), Preemptible VM(GCP), Spot VM(Azure)은 On-demand 대비 60~90%까지 비용을 절감할 수 있다. AI 학습은 비용의 대부분이 GPU 컴퓨팅 비용이므로 Spot Instance의 활용은 매우 매력적이다.

하지만 Spot Instance는 클라우드 프로바이더에 의해 언제든 회수(Preemption)될 수 있으며, 보통 30초에서 2분 사이의 짧은 사전 통지만 제공된다. 체크포인트 없이 학습 중인 Job이 Preempt되면 모든 진행 상황을 잃게 된다.

12.2 Fault-tolerant 학습 전략

Checkpoint 기반 복구

가장 기본적인 전략은 주기적인 Checkpoint 저장이다. Rank 0 프로세스가 일정 주기(에폭 단위 또는 스텝 단위)로 모델 파라미터, 옵티마이저 상태, 학습 진행 상황을 영구 스토리지에 저장한다. Pod가 재시작되면 마지막 Checkpoint에서 학습을 재개한다.

# preStop Hook에서 긴급 체크포인트 저장
# Pod의 lifecycle.preStop에 이 스크립트를 호출하도록 설정한다.
import signal
import sys

def graceful_shutdown(signum, frame):
    """Spot Instance 회수 통지 수신 시 긴급 체크포인트 저장"""
    if dist.get_rank() == 0:
        torch.save({
            'epoch': current_epoch,
            'step': current_step,
            'model_state_dict': model.module.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
        }, '/checkpoints/emergency_checkpoint.pt')
    dist.barrier()
    sys.exit(0)

signal.signal(signal.SIGTERM, graceful_shutdown)

Kubernetes preStop Hook 설정

lifecycle:
  preStop:
    exec:
      command: ['/bin/sh', '-c', 'python save_checkpoint.py --emergency']

TorchElastic (Elastic Training)

PyTorch의 TorchElastic(현재는 torchrun에 통합됨)은 Spot Instance 환경에 최적화된 분산 학습 프레임워크이다.

주요 특징:

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: elastic-training
spec:
  elasticPolicy:
    rdzvBackend: etcd
    rdzvHost: etcd-service
    rdzvPort: 2379
    minReplicas: 2
    maxReplicas: 8
  pytorchReplicaSpecs:
    Worker:
      replicas: 4
      restartPolicy: OnFailure
      template:
        spec:
          nodeSelector:
            cloud.google.com/gke-provisioning: spot
          tolerations:
            - key: cloud.google.com/gke-spot
              operator: Equal
              value: 'true'
              effect: NoSchedule
          containers:
            - name: pytorch
              image: my-registry/elastic-training:latest
              resources:
                limits:
                  nvidia.com/gpu: 1

12.3 Hybrid 클러스터 전략

실제 프로덕션에서는 On-demand와 Spot 노드를 혼합한 Hybrid 전략을 사용한다:

Kueue의 Preemption 정책과 결합하면, Spot 노드가 회수될 때 Kueue가 자동으로 Workload를 다른 가용 리소스로 재배치할 수 있다.

13. References

본 글은 다음 공식 문서와 자료를 참고하여 작성되었다.

Volcano 공식 문서

Kubeflow Training Operator 공식 문서

Kueue 공식 문서

기타 참고 자료

원문 보기: https://www.youngju.dev/blog/kubernetes/ai_training_pipeline_k8s

댓글

아직 댓글이 없습니다.

로그인하면 댓글을 쓸 수 있습니다