こんにちは。PFN 大村(@everpeace)です。

この記事は、Kubernetes Advent Calendar 2021 の 11日目の記事です。

TL;DR

Preemption TolerationというPreemptionされる側の挙動を拡張できるScheduler Pluginをkuberenetes-sigs/scheduler-pluginで公開しました
このPluginを使うと↓みたいなPriorityClassが定義できます
優先度が低いのに(一定の範囲の優先度からは)PreemptionされないPriorityClass
- 商用ソフトウェアをつかった学習等でPreemption対応できない場合に便利
一定期間PreemptionされないPriorityClass
- 機械学習とかで1 Epochも回らずにPreemptionされ続けるのを防げて便利

Preemption Toleration Plugin

モチベーション

KubernetesではPod Priority and Preemptionの機能を使って、独自の PriorityClass を定義する¹ことで、Kuberntesクラスタの空き容量が足りない時、低優先度のPriority Classに属すPodをPreemption²して、高優先度を持つPriorityClassに属すPodをスケジュールする事ができます。

On-Premiseのようにクラスタ容量がElasticでないKubenetesクラスタや、クラウド上であっても予算等の何らかの制約によって容量上限が限られたKubernetesクラスタでは、PriorityClassを活用することで、優先度の高いPodに効率的に計算リソースを割り当てることが可能です。

Pod Priority and Preemptionに記述されている通り PriorityClass には PreemptionPolicy: Never | PreemptLowerPriority という設定がありますが、これはあまり柔軟とは言えません。

通常期待される挙動は PreemptLowerPriority で、文字通り優先度の低いPodをPreemptするPolicyです。もう一つ Never がありますが、これは Preemption を一切しないPriorityClassを作成できます。つまり、するかしないか？という極端な二択しか柔軟性がありません。しかも、これらの設定はすべてPreemptionする側の設定であり、Preemptionされる側の挙動は設定ができません。

機械学習やバッチジョブでは、Preemptionされたくない、Preemptionされてもいいけど条件が満たされてからにしたい、という要求が多く、これだと公式機能だけではカバーできません。そこで、PFNではPreemptionされる側(Victimと呼ばれたりします)の挙動を拡張できるようにするために、Preemption Toleration Pluginを開発しました。

これを作りたくなった具体的なユースケース

優先度は低いけどPreemptionされたくない

PFNのKubernetesクラスタでは主にhigh, lowという２つのPriorityClassを提供しています³。つまり、low PriorityClassをつかって処理を行う場合、クラスタが混雑してくると任意のタイミングでPodが停止する可能性があります。

通常、機械学習の典型的なワークロードでは、Preemptionで停止されても、停止時点から再学習可能にするために、Epochの切れ目等で学習途中のモデルの状態をPod外部のストレージに退避します(強化学習などでは、エージェントだけでなく環境も状態として退避する必要があったりもします)。

しかし、シミュレーションに用いる商用ソフトウェア等、処理の停止/再開に対応していないソフトウェアを使っている場合に、

優先度は低くて良い(空いてる時だけ動けばいい)けど、スケジュールされたらPreemptionされたくない

という要求が出てきます。

それって高優先度にすればPreemptionされないのでは？しかもPreemptionPolicy: Neverにすれば、lowをPreemptすることもないしほぼ所望の動作なのでは？と思われるかもしれません。そのとおりです。Preemptionの挙動としてはその設定で同じことが実現できます。が、その設定には２つの点でふさわしくない点があります

本質的には優先度は低い(空いているときだけ動けばいい)のに、クラスタ上では「Preemptionしない高優先度」という表現になってしまっている
優先度は空きリソースへのPodの割当順序(scheduler内部のqueue上の順序)にも影響するので、空きリソースがある時に高優先度として扱われてしまう

なので、やはり低優先度にもかかわらずPreemptionされたくないという要求が生まれるのです。

ただし、Clusterの動作に不可欠なPod向けにsystem-cluster-critical, system-node-criticalといった特殊(非常に高い優先度のPriorityClass)が定義されていたりするため、運用上低優先度なのに絶対Preemptionされないというのは危険な場合が多いです。つまり、言い換えると、

特定の優先度以上からしかPreemptされないPrioricyClassがほしい

という要求に言い換えることができます。具体的には、下記のような low-non-preemptiblePriorityClassが実現したいという要求です。

PriorityClass名	優先度	Preemptionされる時の挙動
system-critical	10000	最高優先度。Preemptされない。
high	9000	system-criticalにだけPreemptされる。
low-non-preemptible	8000	highからはPreemptionされない。 system-criticalからは即Preemptされる
low	8000	high以上から即Preemptされる。

一定期間はPreemptionされたくない

機械学習のワークロードでは、通常Epochというイテレーションの単位で進みます。Preemptionを想定してEpoch毎に状態を退避するようにプログラムされたPodであっても、困る場合があります。それは、1 Epochも完了せずにPreemptionされる場合です。

更に悪いことに、クラスタの容量が逼迫していて、high PriorityClassのPodの需要が多い場合、何度もPreemptionされてしまって、永遠に学習が進捗しないという自体が発生してしまいます。

そこで、Preemptionはされてもいいけど、せめて1 Epochは回るまではPreemptionを抑制してほしい。という要求が生まれます。

1 Epochが完了したことを外部から正確に検出することは大変むずかしい問題です。幸いなことに、機械学習のワークロードでは、学習データセットのサイズ、学習に利用する計算リソースから1 Epochの時間を見積もれることが多いため、

今回のPreemption Toleration Pluginでは「時間」を採用しました。具体的には下記のlow-non-premptible-15mのような　PriorityClassを定義したいという要求です。

PriorityClass名	優先度	Preemptionされる時の挙動
system-critical	10000	最高優先度。Preemptされない。
high	9000	system-criticalにだけPreemptされる。
low-non-premptible-15m	8000	highからのPreemptionは15分耐える。 system-criticalからは即Preemptされる
low	8000	high以上から即Preemptされる。

設計: PriorityClass上にPreemption Toleration Policyを定義する

Preemption Toleration Pluginでの拡張の設計に移りましょう。今回はPriorityClass のannotationに下記のように Preemption Toleration Policyを定義することにしました。

kind: PriorityClass
metadata:
  name: toleration-policy-sample
  annotation:
    # A: このPriorityClassが即Preemptされる事ができる最小の優先度 (デフォルト:value+1)
    preemption-toleration.scheduling.sigs.k8s.io/minimum-preemptable-priority: "10000"
    # value < p < A の範囲のPriorityからのPreemptionに耐えられる時間(デフォルト: ∞)
    preemption-toleration.scheduling.sigs.k8s.io/toleration-seconds: "3600"
# このPriorityClassの優先度
value: 8000

具体的なPremption Toleration Policyの例

このanntationを使うと上で上げた要求は下記のように実装できます

low-non-preemptible (優先度は低いけどPreemptionされたくない)

   kind: PriorityClass
   metadata:
     name: low-non-preemptible
     annotation:
      # 10000以上の優先度のPriorityClassからしかPreemptされない
      preemption-toleration.scheduling.sigs.k8s.io/minimum-preemptable-priority: "10000"
    value: 8000

low-non-preemptible-15m (15分間はPreemptionされたくない)

   kind: PriorityClass
   metadata:
     name: low-non-preemptible-15m
     annotation:
      # 10000以上の優先度のPriorityClassからしかPreemptされない
      preemption-toleration.scheduling.sigs.k8s.io/minimum-preemptable-priority: "10000"
      # 8000 < p < 10000の優先度のPriorityClassからのPreemptionには15分間耐える
      preemption-toleration.scheduling.sigs.k8s.io/toleration-seconds: "900"
    value: 8000

使い方

https://github.com/kubernetes-sigs/scheduler-plugins/blob/master/doc/install.md を参考に下記の設定を有効にしてscheduler-pluginをクラスタにデプロイしてください。

apiVersion: kubescheduler.config.k8s.io/v1beta2
kind: KubeSchedulerConfiguration
...
profiles:
- schedulerName: custom-scheduler
  plugins:
    postFilter:
      # DefaultPreemption プラグインを無効化して
      # PreemptionToleration　プラグインを有効にする
      enabled:
      - name: PreemptionToleration
      disabled:
      - name: DefaultPreemption

最後に

今回は、PFNのKubernetesクラスタで発生した独特のPreemptionにまつわるSchedule要求をScheduler Pluginとして、KEPの提案(kubernetes-sigs/scheduler-plugins#205)から実装(kubernetes-sigs/scheduler-plugins#224)まで行うという貴重な経験ができました。

是非試して頂いてkubernetes-sigs/scheduler-pluginsにrepoにどしどしフィードバックをいただけたらと思います。

Podのリソース要求量がImmutableである限り、PreemptionはKubernetesクラスタのリソース割当の柔軟性・効率性を担う機能だと考えていますが、Upstreamでは着々と"In-Place Update of Pod Resources" (KEP-1287, kubernetes/kubernetes#102884)の開発が進行しており、近い将来リリースが予定されています。この機能がリリースされたら、空き容量を生み出す手段としてPodを縮小するという手段が可能になるため、クラスタ逼迫時のリソース割当戦略も大きく変わるかもしれません。

We're Hiring

Preferred NetworksではKubernetesをはじめとするクラウドネイティブ技術を活用し大規模機械学習環境を一緒に作るエンジニアを募集中です！数千枚のGPUや独自チップで研究開発を加速させるべくツール開発からKubernetesを改良し資源配分の最適化まで行う、多くの技術が活用できるチャレンジングな環境です！興味のある方は是非連絡お待ちしています！！

その他にも数々のOpen Positionがありますので https://www.preferred.jp/ja/careers/ も御覧ください！！

条件付きでPreemptionを抑制(Tolerate)できるPreemption TolerationというKubernetes Scheduler Pluginを作った