Skip to content
Published on

Kubernetes AI学習パイプライン：Volcano、Training Operator、Kueueの徹底分析

🇰🇷KO🇺🇸EN🇯🇵JA
Authors
  • Name
    Twitter

1. Kubernetesデフォルトスケジューラの限界

Kubernetesのデフォルトスケジューラ（kube-scheduler）は、Pod単位のスケジューリングを行う。各Podが要求するリソース（CPU、Memory、GPU）に基づいて適切なNodeに配置するのがその核心的な役割である。このアーキテクチャはWebサービスや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 デッドロック問題

さらに深刻なのはデッドロックである。2つの分散学習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は3つの核心コンポーネントで構成される。

  • Volcano Scheduler：kube-schedulerを置き換えるか並行して動作するスケジューラ。Pod単位ではなくJob/PodGroup単位のスケジューリングを行う。
  • Volcano Controller Manager：VolcanoJob、Queue、PodGroupなどのCRDライフサイクルを管理するコントローラ。
  • Volcano Webhook（Admission Controller）：CRD作成時のバリデーションとデフォルト値注入を担当するAdmission Webhook。

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'
  • weight：Queue間のリソース配分比率。weightが高いほどより多くのリソースが割り当てられる。
  • reclaimable：アイドルリソースを他のQueueに貸し出せるかどうか。trueに設定すると、このQueueのアイドルリソースを他のQueueが使用できる。
  • capability：Queueが使用可能な最大リソース上限。CPU、Memory、GPUなどを指定する。

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
  • minMember：PodGroup内で最低限実行されるべきPod数。クラスタリソースがこの数を満たせない場合、PodGroup内のどのPodもスケジューリングされない。
  • minResources：PodGroup実行に必要な最小リソース。利用可能なリソースがこれを満たさない場合、スケジューリングは保留される。
  • priorityClassName：PodGroupの優先度。スケジューラがQueue内のPodGroupをソートする際に使用する。
  • queue：PodGroupが所属するQueue名。QueueはOpen状態でなければならない。

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は以下のシナリオで必須である：

  • MPIベースの分散学習：すべてのWorkerが同時に起動しなければallreduce通信ができない。
  • PyTorch DDP：MasterとWorkerがすべて準備されてから学習が開始される。
  • Spark/Big Data：DriverとExecutorがすべて準備されてからJobが実行される。

Gang PluginはPodGroupのminMemberフィールドを確認し、クラスタ上でその数のPodを同時にスケジューリング可能かを判断する。不可能な場合はどのPodもスケジューリングせず、リソース浪費とデッドロックを防止する。

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（役割）を1つの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
  • minAvailable：Gang Schedulingの最小Pod数。PodGroupのminMemberと同じ役割を果たす。
  • schedulerName：volcanoを指定してVolcano Schedulerを使用する。
  • queue：Jobが投入されるQueue名。
  • plugins：有効化するVolcanoプラグインリスト。envは環境変数の自動注入、svcはサービスの自動生成を担当する。
  • policies：イベントベースのアクションポリシー。PodがEvictされた場合にJobを再起動するなどのポリシーを定義する。
  • tasks：Jobを構成するTaskリスト。各Taskは役割（master、worker）ごとにreplicasとPod Templateを定義する。

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には2つのバージョンが共存している。

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 主要フィールド説明

  • Master：常にreplicas: 1でなければならない。学習Jobの状態を代表し、モデルパラメータ集約の基準点となる。Training OperatorがこのPodのアドレスをMASTER_ADDRとして他のPodに注入する。
  • Worker：実際の学習を行うPod群。replicas数を調整して分散学習のスケールを設定する。
  • restartPolicy：OnFailureに設定すると、Pod失敗時に自動再起動される。
  • sidecar.istio.io/inject: "false"：Istioがインストールされたクラスタでで PyTorchJobが正しく動作するようサイドカー注入を無効化する。公式ドキュメントで明示的に推奨されている設定である。

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の特徴：

  • Launcher：mpirunコマンドを実行するPod。実際の学習は行わず、Workerにタスクを分配する。
  • Worker：実際の学習を行うPod。slotsPerWorkerでWorkerあたりのプロセス（GPU）数を指定する。
  • フレームワーク非依存のため、Horovod、PyTorch、TensorFlow、MXNetなどMPIをサポートするすべてのフレームワークで使用可能である。

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の特徴：

  • PS（Parameter Server）：モデルパラメータを保存し、Workerからグラディエントを受信してパラメータを更新する。
  • Worker：学習データの一部を処理してグラディエントを計算し、PSに送信する。
  • Training OperatorがTF_CONFIG環境変数を自動設定し、TensorFlowの分散戦略が正しく動作するようにする。

8. マルチノード・マルチGPU分散学習設定（PyTorch DDP on K8s）

8.1 PyTorch DDPの動作原理

PyTorch DDP（Distributed Data Parallel）はデータ並列方式の分散学習戦略である。各GPUにモデルのレプリカを配置し、ミニバッチをGPU数分に分割して学習した後、allreduce演算でグラディエントを同期する。

KubernetesでマルチノードマルチGPU DDPを構成する際に理解すべき核心概念：

  • WORLD_SIZE：全プロセス数。（ノード数）×（ノードあたりのGPU数）。例：3ノード、各2 GPU = WORLD_SIZE 6。
  • RANK：各プロセスの一意なID（0からWORLD_SIZE-1まで）。
  • LOCAL_RANK：ノード内でのプロセスID（0からノードあたりのGPU数-1まで）。
  • MASTER_ADDR：Masterノードのアドレス。Training Operatorが自動設定する。
  • MASTER_PORT：Masterノードの通信ポート。Training Operatorが自動設定する。

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通信最適化

マルチノード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ネットワークがある環境では、PodにhostNetwork: trueを設定してネットワーク性能を最大化することもできる。

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）でも十分だが、フォールトトレランスを考慮すると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ボトルネックが発生する可能性がある。その場合、学習開始前にInit ContainerパターンでデータをLocal SSDにコピーするか、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は異なるレイヤーで動作する。

  • Volcano：Schedulerレベルで動作する。PodをどのNodeに配置するかを決定する。kube-schedulerを置き換える。
  • Kueue：Admission Controlレベルで動作する。Jobをいつ開始（admit）できるかを決定する。kube-schedulerと連携して動作し、スケジューラを置き換えない。

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はオンデマンド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

主要フィールドの説明：

  • cohort：同じCohortに属するClusterQueue同士でアイドルQuotaを貸し借りできる。
  • nominalQuota：基本的に保証されるリソース量。
  • borrowingLimit：Cohort内の他のClusterQueueから借りられる最大リソース量。
  • lendingLimit：他のClusterQueueに貸し出せる最大リソース量。
  • preemption.withinClusterQueue：同一ClusterQueue内で優先度の低いWorkloadをプリエンプトできるかどうか。LowerPriorityに設定すると、高優先度のWorkloadが低優先度のWorkloadを中断できる。
  • preemption.reclaimWithinCohort：Cohort内の他のClusterQueueが借りたリソースを回収（プリエンプト）できるかどうか。
  • flavorFungibility：複数のResourceFlavorがある場合、BorrowingとPreemptionの優先順位を決定する。

11.3 LocalQueue

LocalQueueはNamespaceスコープでユーザーが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は2つのキューイング戦略をサポートする：

  • StrictFIFO：厳密な先入先出。キュー先頭のWorkloadがAdmitできない場合、後続のWorkloadもAdmitされない。
  • BestEffortFIFO：先頭のWorkloadがリソース不足でAdmitできなくても、後続のWorkloadがリソース要件が少ない場合は先にAdmitされることがある。

12. Spot/Preemptibleインスタンスでの学習（フォールトトレランス）

12.1 Spotインスタンスの利点とリスク

クラウドSpotインスタンス（AWS）、Preemptible VM（GCP）、Spot VM（Azure）は、オンデマンドと比較して60〜90%のコスト削減が可能である。AI学習ではコストの大部分がGPUコンピューティング費用であるため、Spotインスタンスの活用は非常に魅力的である。

ただし、Spotインスタンスはクラウドプロバイダーによっていつでも回収（Preemption）される可能性があり、通常30秒から2分の短い事前通知のみが提供される。チェックポイントなしで学習中のJobがPreemptされると、すべての進捗が失われる。

12.2 フォールトトレラントな学習戦略

チェックポイントベースの復旧

最も基本的な戦略は定期的なチェックポイント保存である。Rank 0プロセスが一定間隔（エポック単位またはステップ単位）でモデルパラメータ、オプティマイザの状態、学習進捗を永続ストレージに保存する。Podが再起動すると、最後のチェックポイントから学習を再開する。

# preStop Hookで緊急チェックポイントを保存
# Podのlifecycle.preStopにこのスクリプトを呼び出すよう設定する。
import signal
import sys

def graceful_shutdown(signum, frame):
    """Spotインスタンス回収通知受信時に緊急チェックポイントを保存"""
    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インスタンス環境に最適化された分散学習フレームワークである。

主な特徴：

  • 弾力的スケーリング：Workerが追加・削除されても学習が中断しない。
  • 自動再起動：Worker障害時に自動的に学習を再開する。
  • 部分実行：要求したGPU数の一部のみ確保されても学習を開始できる。
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 ハイブリッドクラスタ戦略

本番環境では、オンデマンドとSpotノードを混合したハイブリッド戦略を使用する：

  • オンデマンドノード：クラスタの核心コンポーネント（etcd、Kueue Controller、Training Operator）を配置する。Taintを設定して一般ワークロードが配置されないようにする。
  • Spotノード：実際の学習Worker Podを配置する。KueueのResourceFlavorでSpotノードを個別管理し、コスト効率の高いGPUを選択する。

KueueのPreemptionポリシーと組み合わせることで、Spotノードが回収された際にKueueが自動的にWorkloadを他の利用可能なリソースに再配置できる。

13. References

本記事は以下の公式ドキュメントおよび資料を参考に執筆した。

Volcano公式ドキュメント

Kubeflow Training Operator公式ドキュメント

Kueue公式ドキュメント

その他参考資料

Discuss on Twitter • View on GitHub