利用Scheduling Framework实现一个简单的gang调度器

写在前面，本项目gang-scheduler面向此项目中的cosheduler造轮子，如果只是想使用gang调度器，推荐阅读原项目；如果想了解如何自定义调度器，可以看看本项目，其改用ConfigMap而没有用CRD来管理PodGroup，也省略了一些优化，是原项目的简陋版，仅方便入门学习。

动机

在上一篇文章中，我们写了一个简单的MPI Operator，其可以根据MPIJob任务需求生成Pod来进行分布式训练。在使用过程中，我发现一个不是很好的现象：在资源不足的情况下，新建一个训练任务，虽然资源（例如GPU）不足以满足训练任务，但集群依旧会“尽量”（能分配一个Pod就分配一个Pod）分配资源生成Worker，使一部分Worker处于Running状态，等待Launcher的命令，而另一部分Worker则由于后续资源不足而处于Pending的状态，导致任务无法启动，而浪费了资源。

比如，假设我们需要5个GPU来进行训练任务，而集群中可用的GPU只有4个。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
apiVersion: batch.test.bdap.com/v1
kind: MPIJob
metadata:
  name: simple-train
  namespace: sw-mpi-operator
spec:
  numWorkers: 5
  launcherTemplate:
    spec:
      containers:
        - image: farawaya/horovod-torch-cpu
          name: horovod-master
  workerTemplate:
    spec:
      containers:
        - args:
            - sleep infinity
          command:
            - /bin/sh
            - -c
          image: farawaya/horovod-torch-cuda113
          name: horovod-worker
          resources:
            limits:
              nvidia.com/gpu: 1
      tolerations:
        - effect: NoSchedule
          key: gpu
          operator: Exists

根据MPI Operator中的流程，由于Workers未全部就绪，无法生成Launcher来启动任务

1
2
3
4
5
6
7
NAME                                      READY   STATUS    RESTARTS   AGE
mpi-controller-manager-6fdb6ffd94-sjlrr   1/1     Running   0          18d
simple-train-worker-0                     1/1     Running   0          3m12s
simple-train-worker-1                     1/1     Running   0          3m12s
simple-train-worker-2                     1/1     Running   0          3m12s
simple-train-worker-3                     1/1     Running   0          3m11s
simple-train-worker-4                     0/1     Pending   0          3m11s

可以查看Operator的日志进行验证：

1
1.6604734402540798e+09	INFO	controller.mpijob	workers not ready	{"reconciler group": "batch.test.bdap.com", "reconciler kind": "MPIJob", "name": "simple-train", "namespace": "sw-mpi-operator"}

并且Workers会一直占有GPU而等待，造成资源浪费。更坏的情况是，该任务在等待其他Pod完成后释放该资源，但不巧的是，占有该资源的Pod属于另外一个也处于等待状态的训练任务，那么就形成了死锁。

其形成原因是kubernetes的默认调度器以Pod为单位进行调度，如果我们以整个训练任务作为调度单位，要么给其中所有Pod都分配资源，要么就一个都不分配，“All or Nothing”，则可以解决这个问题，而gang调度正对应了该想法。

如何自定义Scheduler

我们先得知道如何去自定义kubernetes中的调度器。

一些弯路 我在网上搜索到了这篇自定义调度器的[文章](http://dockone.io/article/2434717)，其提到自定义调度器的几种方式：
  1. 修改默认调度器的源代码，加入自己的调度算法，然后重新编译和部署调度器，论文kcss和kubecg中的调度器研究基于此方案实现；
  2. 开发自己的调度器，和默认调度器同时运行在集群中；
  3. 基于Kubernetes Scheduler Extender机制，在扩展调度器中实现自定义算法，论文dynamic IO中的算法实现基于这种方案。

其文章中使用的Extender方式，通过给默认调度器注册webhook进行，当默认调度器进行到Filter（预选）、Score（优选）、Bind阶段的时候会通过http请求调用webhook来获取结果。

关于调度流程中的每个阶段介绍，可以参考K8S Scheduling Framework

尝试与思考之后，我认为仅用Kubernetes Scheduler Extender方法来实现gang调度比较困难，因为它能控制的模块有限，并且我们不想手动去写分配资源的过程，而是想让调度器预分配资源后再根据情况做决定，因此需要在Reserve资源之后的Permit阶段进行修改。并且http请求的速度也比较慢。

后面发现一篇文章中提到四种方式自定义调度器：

  • default-scheduler recoding: 直接在Kubernetes默认scheduler基础上进行添加，然后重新编译kube-scheduler
  • standalone: 实现一个与kube-scheduler平行的custom scheduler，单独或者和默认kube-scheduler一起运行在集群中
  • scheduler extender: 实现一个”scheduler extender”，kube-scheduler会调用它(http/https)作为默认调度算法(预选&优选&bind)的补充
  • scheduler framework: 实现scheduler framework plugins，重新编译kube-scheduler，类似于第一种方案，但是更加标准化，插件化

我认为最后scheduler framework的方式比较合适，我们只需要修改Permit阶段，检查每个Pod，根据其PodGroup对应的信息来判断就好。实现过程后面说，这里先说一下实现完后怎么使用。

  1. Pod如何找到调度器？Pod通过在spec中指定schedulerName来选择调度器
  2. 调度器如何注册到集群中？调度器本身是kube-scheduler的扩展版，而kube-scheduler启动时，会接收一个config文件路径作为参数，其内容为KubeSchedulerConfiguration类型的yaml文件，其指定了该调度器在集群中的schedulerName，而该名字需要与1中的名字一致。

一开始调试的时候，可以跑一个与kube-scheduler平行的gang-scheduler，那么即便存在问题也不会影响原本的集群。虽然集群中同时存在多个scheduler的时候，可能会有同一块资源被分配给不同Pod的情况，不过问题不大，等到该调度器稳定后，将其替换默认的调度器就可以了。

插件是write还是append？

btw，我们写插件的时候，不需要担心自己的某个阶段（如Filter）会覆盖原有的阶段，除了QueueSort阶段外，我们的阶段都是附加的，这一点可以通过看源代码验证，可以看到对于Taint的检测也是通过插件的方式来进行的。

如何实现gang调度

这部分我也没想到什么好思路，不过搜索到了阿里的一篇好文：进击的 Kubernetes 调度系统（二）：支持批任务的 Coscheduling/Gang scheduling，发现他们已经早在两年前就已经完成并开源了这份工作，源码。墙裂推荐去看一下，本文的调度器也是仿照它造轮子，不过也存在一些改动。

博客中记录的是第一个版本的简单做法，使用label即可实现，对应commit，源码中目前是第二版本，使用CRD PodGroup实现，我猜想有两个原因：1. 防止同一个PodGroup中的Pod写的label声明的minAvailable不一致的情况。2. 解耦，podgroup本身不应该由scheduler管理，而应该由单独的operator进行管理。那为什么不用ConfigMap呢？我猜想是需要记录下PodGroup的一些Status。不过我们的使用情况比较简单，我偷点懒，就不写PodGroup的Operator了，简单使用ConfigMap来记录PodGroup信息，也能把PodGroup的控制权交给用户，让用户来维护，不需要scheduler来考虑创建与回收，也没有引入CRD管理，无需controller。

Permit阶段

我们的目标是把带有PodGroup标记的Pod都拦截下来，看看集群能否满足PodGroup整体的资源需求，如果满足条件就放行。因为到达这一步的Pod都已经经过Reserve阶段，资源已经预留下来，因此只需看到达该阶段的Pod是否达到数量要求就行，一共就三步：

  1. 判断Pod的Label中是否标注了PodGroup，如果没有，说明是普通的Pod，不影响它，直接进入下一阶段
  2. 如果Label中写明了PodGroup，则记录该PodGroup对应需要的minAvailable
  3. 判断此时是否已经满足minAvailable，如果满足，则进入下一阶段，否则进入Wait状态

其中“判断此时是否已经满足minAvailable”也包括三步：

  1. 所有之前一定时间内的Wait状态的Pod都在一个队列中，我们查询该队列，给处于PodGroup的Pod计数。
  2. 查询当前集群，给属于该PodGroup的Pod计数，并且需要Pod处于Running状态下。
  3. 判断1和2的总和再加上当前Pod是否满足minAvailable
什么情况会出现该PodGroup已经有一些Pod在运行了，而一些Pod还在等待调度的队列中呢？ 当第一次分配时就满足minAvailable条件，已经跑起来之后，后续又新增属于该PodGroup中的Pod，在调度这些新增Pod的时候。

QueueSort阶段

仅修改Permit还不够。原因：Default Scheduler的队列并不能感知PodGroup的信息，如果有多个PodGroup，那么不同PodGroup中的Pod有可能在队列里面交错，那么进入Reserve阶段的Pod也是交错的，因此还是产生死锁。因此我们需要添加QueueSort Plugin。

我们的目标是让不同的PodGroup中的Pod不能交错。做法：查询两个Pod对应的PodGroup标签，看是否属于某个PodGroup

  1. 都不属于，则按照优先级，再按照create时间排序
  2. 都属于，则按照PodGroup的create时间排序，时间相同说明属于同一个PodGroup，则按照1进行
  3. 一个属于，一个不属于，则让不属于PodGroup的先行一步

如何使用gang-scheduler

安装与更加详细的例子，可以查看gang-scheduler

本文就拿之前动机部分的GPU举例，新增一个ConfigMap记录PodGroup信息：

1
2
3
4
5
6
7
8
apiVersion: v1
kind: ConfigMap
metadata:
  name: gpu-pg
  namespace: sw-mpi-operator
data:
  minAvailable: "5"
  scheduleTimeoutSeconds: "20"

修改Pod信息，在label中添加pod-group.scheduling.bdap.com/podgroup-configmap: gpu-pg，以及修改调度器schedulerName: gang-scheduler

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
apiVersion: batch.test.bdap.com/v1
kind: MPIJob
metadata:
  name: simple-train
  namespace: sw-mpi-operator
spec:
  numWorkers: 5
  launcherTemplate:
    spec:
      containers:
        - image: farawaya/horovod-torch-cpu
          name: horovod-master
  workerTemplate:
    metadata:
      labels:
        pod-group.scheduling.bdap.com/podgroup-configmap: gpu-pg
    spec:
      schedulerName: gang-scheduler
      containers:
        - args:
            - sleep infinity
          command:
            - /bin/sh
            - -c
          image: farawaya/horovod-torch-cuda113
          name: horovod-worker
          resources:
            limits:
              nvidia.com/gpu: 1
      tolerations:
        - effect: NoSchedule
          key: gpu
          operator: Exists

运行结果：

1
2
3
4
5
6
7
NAME                                      READY   STATUS    RESTARTS   AGE
mpi-controller-manager-6fdb6ffd94-np6pg   1/1     Running   0          79s
simple-train-worker-0                     0/1     Pending   0          11s
simple-train-worker-1                     0/1     Pending   0          11s
simple-train-worker-2                     0/1     Pending   0          11s
simple-train-worker-3                     0/1     Pending   0          11s
simple-train-worker-4                     0/1     Pending   0          11s

查看原因：

1
2
3
4
5
Events:
  Type     Reason            Age   From            Message
  ----     ------            ----  ----            -------
  Warning  FailedScheduling  55s   gang-scheduler  rejected due to timeout after waiting 20s at plugin sample
  Warning  FailedScheduling  33s   gang-scheduler  rejected due to timeout after waiting 20s at plugin sample

如果我们不创建PodGroup对应的ConfigMap就直接使用其作为label，也会出现提示：

1
2
3
4
5
Events:
  Type     Reason            Age   From            Message
  ----     ------            ----  ----            -------
  Warning  FailedScheduling  14s   gang-scheduler  running Permit plugin "sample": podgroup configmap not found, please create configmap first
  Warning  FailedScheduling  13s   gang-scheduler  running Permit plugin "sample": podgroup configmap not found, please create configmap first
利用Scheduling Framework实现一个简单的gang调度器
https://fffffaraway.github.io/2022/08/14/利用Scheduling-Framework实现一个简单的gang调度器/
Author
Song Wei
Posted on
August 14, 2022
Licensed under