ACK Fleet での Gang スケジューリングによる PyTorchJob の実行 - Container Service for Kubernetes

PyTorch 分散トレーニングジョブでは、ジョブ内のすべての Pod を同時に実行する必要があります。一部のみが起動すると、リソースが無駄になり、ジョブが無限に停止（一部サービス停止）する可能性があります。ガンスケジューリングは「全または無し（all-or-nothing）」の保証を実施します。すなわち、ジョブ内のすべての Pod が同時にスケジュールされるか、まったくスケジュールされないかのいずれかです。これにより、マルチ GPU・マルチマシンでのトレーニングシナリオにおけるリソースデッドロックが防止されます。

本トピックでは、ACK Fleet インスタンス上で Kube Queue を構成して PyTorchJob をキューに登録する方法、およびすべての Pod をアトミックに同一のメンバークラスター上に配置するためのガンスケジューリングの適用方法について説明します。

仕組み

Fleet インスタンスは、以下の 2 つのコンポーネントを用いて、メンバークラスター間で PyTorchJob のスケジューリングを調整します。

Kube Queue：ジョブキューを管理し、弾力的なクォータ制限を適用します。メンバークラスターに十分なリソースが利用可能になるまで、ジョブを一時的に保持します。
ACK Scheduler：Fleet インスタンスが Pod をメンバークラスターに配信する際に、ガンスケジューリングのセマンティクスを適用します。これにより、すべてのレプリカ（Master および Worker）がアトミックに配置されます。

スケジューリングのフローは以下のとおりです。

PropagationPolicy を指定して PyTorchJob を Fleet インスタンスに送信します。この Policy には customSchedulingType: Gang を設定します。
キュー管理が有効になっている場合（suspension.scheduling: true）、ジョブは Kube Queue に入り、クォータスロットが利用可能になるまで待機します。
Fleet インスタンスは、メンバークラスター全体の利用可能なリソースを評価し、ターゲットクラスターを選択します。
ACK Scheduler は、選択されたクラスター上ですべての Pod をアトミックに配置し、ガンセマンティクスを維持します。
Fleet インスタンスはジョブをモニターし、ステータスを同期します。

前提条件

開始する前に、以下の条件を満たしていることを確認してください。

メンバークラスターに Cloud-native AI スイートがインストールされていること — Arena コンポーネントのみをデプロイしてください。
RAM（Resource Access Management）ユーザーに AliyunAdcpFullAccess RAM ポリシーがアタッチされていること。詳細については、「RAM ユーザーへの権限付与」をご参照ください。
AMC コマンドラインツールがインストールされていること。詳細については、「AMC の使用」をご参照ください。
（任意）Fleet インスタンスがメンバークラスターとのスケジューリング整合性を保証する場合、リソース予約を有効化します。リソース予約には Kubernetes 1.28 以降および ACK Scheduler 6.8.0 以降が必要です。

（任意）リソース予約の有効化

リソース予約を有効化しない場合、Fleet インスタンスはメンバークラスター内のすべてのノードの残りリソースを合計して、利用可能な容量を推定します。一方、リソース予約を有効化すると、Fleet インスタンスはコミット前にターゲットクラスター上で実際の容量を確保するため、Fleet レベルのスケジューリング判断がメンバークラスターの結果と一致します。

ACK コンソールにログインし、左側のナビゲーションウィンドウで [クラスター] をクリックします。
クラスター名をクリックします。左側のナビゲーションウィンドウで [アドオン] をクリックします。
[アドオン] ページで、Kube Scheduler を見つけ、[構成] をクリックします。
[Kube Scheduler パラメーター] ダイアログボックスで、[enableReservation] を true に設定し、[OK] をクリックします。

スケジューリングモードの選択

次の 2 つのモードが利用可能です。

モード	使用タイミング	主な構成
ガンスケジューリングのみ	キュー管理を行わず、Pod をアトミックに配置したい場合	PropagationPolicy に `customSchedulingType: Gang` を設定します。
ガンスケジューリング＋キュー管理	多数のジョブが限られたリソースを競合しており、クォータ制限付きの順序立てられたキュー処理が必要な場合	`customSchedulingType: Gang` および `suspension.scheduling: true`

キュー管理が必要な場合は「手順 1」に従ってください。ガンスケジューリングのみが必要な場合は、「手順 2」に進んでください。

手順 1（任意）：Kube Queue を使用したジョブキューの設定

ElasticQuotaTree を使用して、クォータ制限を定義し、名前空間間で同時に実行できるジョブ数を制御します。

Fleet インスタンスに ElasticQuotaTree を送信します。以下の例では、default 名前空間向けに、同時実行ジョブ数を 1 つに制限し、最大 CPU 数を 10,000、最大メモリ量を 10,000 GiB、最大 GPU 数を 10,000 とするクォータを構成しています。

apiVersion: scheduling.sigs.k8s.io/v1beta1
kind: ElasticQuotaTree
metadata:
  name: elasticquotatree  # ElasticQuotaTree は単一のみサポートされます。
  namespace: kube-system   # kube-system 名前空間内に作成する必要があります。
spec:
  root:
    name: root
    max:
      cpu: 999900
      memory: 400000Gi
      kube-queue/max-jobs: 10000000000
      nvidia.com/gpu: 100000
    min:
      cpu: 999900
      memory: 400000Gi
      kube-queue/max-jobs: 10000000000
      nvidia.com/gpu: 100000
    children:
    - name: child-2
      max:
        kube-queue/max-jobs: 1  # 同時にデキュー可能なジョブは 1 つのみです。
        cpu: 10000
        nvidia.com/gpu: 10000
        memory: 10000Gi
      namespaces:
        - default

Kube Queue が対応するキューを作成したことを確認します。

kubectl get queue -n kube-queue

期待される出力：

NAME                 AGE
root-child-2-v5zxz   15d
root-kdzw7           15d

手順 2：マルチクラスター向け PyTorchJob の送信

PropagationPolicy の送信

PropagationPolicy は、Fleet インスタンスに対して、PyTorchJob をメンバークラスター間でどのように分散するか、および適用するスケジューリングモードを指示します。

ガンスケジューリングのみ

キュー処理なしでアトミックな Pod 配置を有効化するには、customSchedulingType: Gang を設定します。

apiVersion: policy.one.alibabacloud.com/v1alpha1
kind: PropagationPolicy
metadata:
  name: example-policy
  namespace: default
spec:
  propagateDeps: true
  failover:
    application:
      decisionConditions:
        tolerationSeconds: 30
      purgeMode: Immediately
  placement:
    replicaScheduling:
      replicaSchedulingType: Divided
      customSchedulingType: Gang
  resourceSelectors:
    - apiVersion: kubeflow.org/v1
      kind: PyTorchJob

キュー管理付きガンスケジューリング

suspension.scheduling: true を追加することで、Fleet インスタンスはクォータスロットが利用可能になるまでジョブを Kube Queue に保持し、その後すべての Pod をアトミックに配置します。

apiVersion: policy.one.alibabacloud.com/v1alpha1
kind: PropagationPolicy
metadata:
  name: example-policy
  namespace: default
spec:
  suspension:
    scheduling: true
  propagateDeps: true
  failover:
    application:
      decisionConditions:
        tolerationSeconds: 30
      purgeMode: Immediately
  placement:
    replicaScheduling:
      replicaSchedulingType: Divided
      customSchedulingType: Gang
  resourceSelectors:
    - apiVersion: kubeflow.org/v1
      kind: PyTorchJob

PyTorchJob の送信

以下の PyTorchJob を Fleet インスタンスに送信します。このジョブは 1 つの Master Pod と 2 つの Worker Pod を定義します。

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  labels:
    app: pytorchjob
  name: pytorch-test
  namespace: default
spec:
  cleanPodPolicy: None
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: Never
      template:
        metadata:
          labels:
            app: pytorchjob
          name: pytorch-test
        spec:
          schedulerName: default-scheduler
          containers:
          - command:
            - sh
            - -c
            - sleep 1h
            env:
            - name: NVIDIA_VISIBLE_DEVICES
              value: void
            - name: gpus
              value: "0"
            - name: workers
              value: "8"
            image: registry-cn-hangzhou.ack.aliyuncs.com/acs/nginx
            imagePullPolicy: Always
            name: pytorch
            resources:
              limits:
                cpu: "3"
              requests:
                cpu: "10m"
            volumeMounts:
            - mountPath: /dev/shm
              name: dshm
            workingDir: /root
          volumes:
          - emptyDir:
              medium: Memory
              sizeLimit: 2Gi
            name: dshm
    Worker:
      replicas: 2
      restartPolicy: OnFailure
      template:
        metadata:
          labels:
            app: pytorchjob
          name: pytorch-test
        spec:
          containers:
          - command:
            - bash
            - -c
            - |
              echo "$WORKER_INDEX"
              sleep 1h
            env:
            - name: WORKER_INDEX
              valueFrom:
                fieldRef:
                  fieldPath: metadata.labels['pytorch-replica-index']
            - name: NVIDIA_VISIBLE_DEVICES
              value: void
            - name: gpus
              value: "0"
            - name: workers
              value: "8"
            image: registry-cn-hangzhou.ack.aliyuncs.com/acs/nginx
            imagePullPolicy: Always
            name: pytorch
            resources:
              limits:
                cpu: "2"
              requests:
                cpu: "2"
                memory: "2Gi"
            volumeMounts:
            - mountPath: /dev/shm
              name: dshm
            workingDir: /root
          volumes:
          - emptyDir:
              medium: Memory
              sizeLimit: 2Gi
            name: dshm

手順 3：ジョブステータスの確認

Fleet インスタンス上で以下のコマンドを実行し、ジョブが正しくスケジュールされ、すべての Pod が実行中であることを確認します。

Fleet インスタンス上の PyTorchJob の状態を確認します。
```
kubectl get pytorchjob
```
期待される出力：
```
NAME           STATE     AGE
pytorch-test   Created   3m44s
```
ジョブがどのメンバークラスターにスケジュールされたかを確認します。
```
kubectl describe pytorchjob pytorch-test
```
イベントから ScheduleBindingSucceed を探します。結果フィールドにはターゲットクラスターとレプリカ数が表示されます。
```
Normal   ScheduleBindingSucceed  4m59s   default-scheduler   Binding has been scheduled successfully. Result: {cfxxxxxx:0,[{master 1} {worker 2}]}
```
cfxxxxxx は、すべての Pod が実行されるメンバークラスターの ID です。
メンバークラスター上でジョブが実行中であることを確認します。
```
kubectl amc get pytorchjob -M
```
期待される出力：
```
NAME           CLUSTER    STATE     AGE     ADOPTION
pytorch-test   cfxxxxxx   Running   6m23s   Y
```
ADOPTION: Y は、Fleet インスタンスがこのジョブのスケジューリングを引き継いだことを意味します。
すべての Pod が実行中であることを確認します。
```
kubectl amc get pod -M
```
期待される出力：
```
NAME                    CLUSTER    READY   STATUS    RESTARTS   AGE
pytorch-test-master-0   cfxxxxxx   1/1     Running   0          7m16s
pytorch-test-worker-0   cfxxxxxx   1/1     Running   0          7m16s
pytorch-test-worker-1   cfxxxxxx   1/1     Running   0          7m16s
```
3 つの Pod（1 つの Master および 2 つの Worker）が同一クラスター上で実行中であり、ガンスケジューリングによってアトミックに配置されたことが確認できます。
メンバークラスター上の完全な PyTorchJob YAML を確認するには、以下のコマンドを実行します。
```
kubectl amc get pytorchjob pytorch-test -m ${member clusterid} -oyaml
```