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

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

この記事は、[Kubernetes Advent Calendar 2021](https://qiita.com/advent-calendar/2021/kubernetes) の 11日目の記事です。

# TL;DR

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

# Preemption Toleration Plugin

## モチベーション

Kubernetesでは[Pod Priority and Preemption](https://kubernetes.io/docs/concepts/scheduling-eviction/pod-priority-preemption/)の機能を使って、独自の `PriorityClass` を定義する[^1]ことで、Kuberntesクラスタの空き容量が足りない時、低優先度のPriority Classに属すPodをPreemption[^5]して、高優先度を持つPriorityClassに属すPodをスケジュールする事ができます。

[^1]: Kubernetesではデフォルトで`system-cluster-critical`, `system-node-critical` という２つのpriority classが定義されています
[^5]: 高優先度のPodをスケジュールするために低優先度のPodを停止する処理をPreemptionと呼びます

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

[Pod Priority and Preemption](https://kubernetes.io/docs/concepts/scheduling-eviction/pod-priority-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を提供しています[^2]。つまり、`low` PriorityClassをつかって処理を行う場合、クラスタが混雑してくると任意のタイミングでPodが停止する可能性があります。

[^2]: 実際はより多くのPriorityClassを運用していますが、ここでは`high, low`の２つに注目して説明します。

通常、機械学習の典型的なワークロードでは、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-preemptible`PriorityClassが実現したいという要求です。

| PriorityClass名 | 優先度 | Preemptionされる時の挙動 |
| ------------- | ------------- | ------------- |
| system-critical | 10000 | 最高優先度。Preemptされない。|
| high  | 9000  | system-criticalにだけPreemptされる。|
| low-non-preemptible  | 8000  | highからはPreemptionされない。<br/>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分耐える。<br/>system-criticalからは即Preemptされる|
| low  | 8000  | high以上から即Preemptされる。|


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

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

```yaml
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されたくない)

   ```yaml
   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されたくない)

   ```yaml
   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をクラスタにデプロイしてください。

```yaml
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](https://github.com/kubernetes-sigs/scheduler-plugins/pull/205))から実装([kubernetes-sigs/scheduler-plugins#224](https://github.com/kubernetes-sigs/scheduler-plugins/pull/224))まで行うという貴重な経験ができました。

是非試して頂いて[kubernetes-sigs/scheduler-plugins](https://github.com/kubernetes-sigs/scheduler-plugins/)にrepoにどしどしフィードバックをいただけたらと思います。

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

# We're Hiring

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

- [機械学習プラットフォームエンジニア (Infrastructure) - エンジニアリング - 株式会社Preferred Networks](https://apply.workable.com/preferred-networks/j/150BE9CDDC/)
- [大規模計算基盤エンジニア (Infrastructure) - エンジニアリング - 株式会社Preferred Networks](https://apply.workable.com/preferred-networks/j/6CDF8CA1A8/)
- [大規模計算基盤リサーチャー (Infrastructure) - リサーチ - 株式会社Preferred Networks](https://apply.workable.com/preferred-networks/j/D85B7B005E/)

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

# Links

- KEP: https://github.com/kubernetes-sigs/scheduler-plugins/tree/master/kep/205-preemption-toleration 
- Code: https://github.com/kubernetes-sigs/scheduler-plugins/tree/master/pkg/preemptiontoleration