September 17, 2021

Decima:用强化学习解决调度问题

用强化学习解决调度问题

特别感谢:RL for Scheduling,如何用强化学习解决调度问题? - 知乎 (zhihu.com)

Decima (mit.edu)

一、解决的问题

这篇文章为了解决多个有向无环图(DAG)任务在多个Executor(处理器)运行时的调度问题。这个任务是NP难的(For example, scheduling jobs with dependency structure is inherently NP hard)。对于这个问题论文作者使用RL和GNN来解决它。

任务描述

  • 每一节点表示一个计算阶段(computation stage),每一个阶段都包含可以并行运算的任务
  • 每一条边表示数据依赖(Data dependencies),父节点跑完了才能跑子节点

img

交互过程:

  • 调度器接受任务,并将挂起的计算阶段映射到可用服务器上运行。
  • 在这个过程中不断有新的、随机的、任务请求出现,也即待分配任务池和可用服务器池均是动态变化的
  • Agent的可观测量:所有可用作业的信息和服务器在每个调度事件上的状态。(the information of all available jobs and the status of the servers at every scheduling event.)
  • 目标:最小化任务的平均完成时间

img

论文贡献:

  • 针对具有依赖关系的任务,使用RL训练workload-specific调度算法:Decima
  • 使用可伸缩的GNN来表示调度策略,可以处理任意大形状和大小的DAG状任务
  • 任务是随机到达的情况下,可以在线学习

二、强化学习模型:

  • Action:(node, 服务器数量)
  • 当服务器池中有空闲服务器时,选择一个node、分配相应的节点数量,重复这一过程直到所有的服务器均在忙。这种设计相比其他设计可以有效减少动作空间的大小和动作序列的长度,降低强化学习的难度。作者称这是唯一work的设计。

img

  • state:DAG信息,包括任务数量、平均任务时间,可用服务器数量、待决策的任务数量等

GNN(DAG信息聚合)

其中x表示节点的特征,小w表示所有的子节点。我们可以看到,节点的打分由子节点打分之和、节点特征决定。

img

这里使用两个非线性函数的原因是,两层的非线性能让Decima表达更加广泛多变的聚合函数,例如 $f \sim \log (\cdot / n),g \sim \exp (n \times \cdot)$时,Decima能表示下图的关键路径

img

1

上式为Per-node embedddings表达式,体现了子节点信息向当前节点汇聚的过程。与Per-node embedddings类似,作者为每个Job和所有的任务分别设计了Per-job embeddings

2

以及global embeddings

3

GNN(节点信息聚合)的全过程可由下图表示:

img

因此,每个层级的信息聚合由两个非线性变换(f和g)表示; 整个GNN由六个非线性变换构成。

决策网络设计

调度过程: 当某个job完成或者有新job到达时,系统开始调度。

动作(action) ⟨ $v,l_i$ ⟩有两个维度, 决策网络的输出包括两部分, $v$ 是选择哪个任务节点( $P_{node} $), $ l_i$ 是为这个任务多少计算资源(Parallelism limit on job)。调度包括以下三阶段:

  • 当任务 $i$ 使用的服务器比 $l_i$ 少时, Decima为节点 $v$ 分配executors至 $l_i$ 。
  • 上一步完成后,若有多余的executors, Decima 继续运行agent得到stage( $v$ ) 和parallelism limit( $l_i$ ).
  • 上述步骤一直运行,直至没有空闲的executors和未处理的stage( $v$ ).

网络结构整体如下图所示,包括GNN部分和决策网络部分,GNN负责聚合DAG信息,决策网络负责输出stage( $v$ ) 和 parallelism limit( $l_i$ ):

img

决策网络运行过程:

1、得到stage(v)

  • 用非线性网络 $q(\cdot)$ ,得到 $q_{v}^{i} \triangleq q\left(\mathbf{e}_{v}^{i}, \mathbf{y}^{i}, \mathbf{z}\right) $

  • 使用softmax层,选择调度的stage(v)

  • $$ P(\text { node }=v)=\frac{\exp \left(q_{v}^{i}\right)}{\sum_{u \in \mathcal{A}_{t}} \exp \left(q_{u}^{j(u)}\right)} $$

2、为任务节点分配服务器数量parallelism limit($l_i$

计算过程与得到stage(v)这一步类似

训练模型与实验

系统整体结构如下图所示:

img

模型的训练采用策略梯度法(PG):

img

img

作者训练时采用以下公式更新策略梯度:

$$ \theta \leftarrow \theta+\alpha \sum_{k=1}^{T} \nabla_{\theta} \log \pi_{\theta}\left(s_{k}, a_{k}\right)\left(\sum_{k^{\prime}=k}^{T} r_{k^{\prime}}-b_{k}\right) $$

其中阿尔法为学习率, $b_k$ 为baseline,添加baseline的原因可见

李宏毅深度强化学习(国语)课程(2018)

的33分28秒。

img

训练初期使用较短的任务序列,逐渐加长序列。

  • 这是因为在训练初期,agent很智障,你给他个长序列他也不知道咋整,导致任务堆积,他学习效率也很低。所以作者就想,先让他处理简单点的,再逐渐增加难度,效率会高一些。
  • 这是一种课程学习(curriculum learning)的模式(确实令人想起了中学时代……)

与(朴素的)传统方法比较

  • 朴素的先进任务先做(demo用时225)、最短任务优先(demo用时135)会导致资源闲置等待。朴素的公平分配(demo用时120)虽然会使服务器没有闲置,但是会造成大量的IO开销从而拖慢速度。而Decima(demo用时98):
  • 有选择地将服务器集中在某些小任务上,同时仍然不留下任何资源空闲。
  • 为每个作业分配适当数量的服务器,这样就不会产生不必要的IO工作。
  • (因为我也不知道传统一般怎么做大规模调度,它对比的这两个看上去就很简单,所以我擅自加了“朴素”这个形容词
  • (这一块ppt里有动画演示,感兴趣的可以去找talk或者下载pptx来播放下看看,可以直观感受一下

参考文献:

RL for Scheduling,如何用强化学习解决调度问题? - 知乎 (zhihu.com)

Decima (mit.edu)

李宏毅深度强化学习(国语)课程(2018)

版权声明