Machine Learning for Load Balancing in the Linux Kernel

Paper site

論文內容摘要

Linux 核心 採用的 Completely Fair Scheduler 所採取的排程策略是每次挑選當前 virtual runtime 最小的 sched_entity 來執行，並且盡量平衡每個 CPU 的 utilization ，但這樣的作法並未考量到不同任務之間對於硬體資源的需求不同造成競爭帶來的負擔。本論文引入機器學習模型來做到以下幾件事

• 模擬原先 Linux 核心 的負載平衡
• 機器學習模型會將硬體資源需求納入考量
• 針對資源用量大與高度競爭的情境來訓練
• 利用 Reinforcement Learning 來加強負載平衡的表現

Linux 核心 對於每個 CPU 都配置一個 runqueue cfs_rq 來達到更快的 local access ，負載平衡就是在不同的 runqueue 之間進行任務的遷移。排程的單位雖然是 thread ，但並不直接針對 thread 進行排程，而是對 sched_entity 。 cfs_rq 利用 virtual runtime 作為 key ，利用紅黑樹對 sched_entity 排序。而且 cfs_rq 也擁有自己的 sched_entity 並且可以被其他 cfs_rq 排程。

CFS Load Balancing

CFS 會透過 timer interrupt handler 週期性的呼叫scheduler_tick() 。進行負載平衡，而負責處理負載平衡的程式會以 softirq 的方式運作，而是否進行 runqueue 間的負載平衡是透過 cache 和 NUMA locality 來判斷， CFS 會將不同的處理器切分到不同的 scheduling domain sched_domain ，而 sched_domain 當中的 CPU core 又會被分到不同的 sched_group ，負載平衡則是在不同的 groups 之間進行，來確保 group 的 workload 在不同 domain 之間都是平衡的。

什麼是 NUMA locality ?

scheduler_tick() 會呼叫 trigger_load_balance() 並且在該進行負載平衡時觸發 SCHED_SOFTIRQ 來執行 run_rebalance_domains() 。負載平衡的方式又分為兩種

• nohz_idle_balance : 檢查 idle cores ，它們當前的 periodic scheduler ticks 應當是被暫停以節省能源。
• normal load balance : 檢查 intra-domain load imbalance 。

不管是何種負載平衡，進行時都會將該 CPU 所在之 sched_domain 給掃過一遍，透過 load_balance() 確保每個 sched_domain 的 workload 是平衡的。不過為了避免重複的運算，通常只有第一個 idle CPU 或者在沒有 idle CPU 時的第一個 CPU 會進行平衡的運算。如果 domain 當中有非平衡的狀況發生，則會找出負擔最重的 sched_group 當中負擔最重的 core 。其中 load_balance() 會先呼叫 detach_tasks() 來走訪所有在 overload cfs_rq 當中的 runnable threads ，利用 can_migrate_task() 來辨識此任務能否被移動，接著透過 attach_tasks() 把先前 detach 的任務移動到目標的核上， can_migrate_task() 判斷某個執行緒是否能被 migrate 的方式是透過 cache 和 NUMA locality 資訊排判斷。

Imitation Learning in Linux

為了訓練可以輔助進行負載平衡的機器學習模型，作者把 can_migrate_task() 作為進入點並使用 supervised imitation learning ，來訓練一個 Multi-Layer Perceptron model 。透過 eBPF 在 can_migrate_task() 的進入與結束都加上 kprobe handlers 藉此搜集資料。

Data Collection

透過建立在 eBPF 之上的工具集 BCC 與 kprobe 做到動態追蹤，已進行資料搜集，在 can_migrate_task() 的進入與回傳附加上 kprobe handler ，蒐集的資料包含以下幾種

• Idleness of the target CPU core
• Source and destination NUMA node numbers
• Preferred NUMA node of the process
• Loads of the source and target CPU cores
• Lengths of the CFS runqueues of both cores
• Number of processes prefer to stay on the source core
• Time the process has been running on the source core
• Number of NUMA faults the process had on each node

本論文建立一個包含 500000 次 load-balancer 呼叫的資料集作為 training dataset ，利用 stress-ng 來模擬不同 workload level ，分別有 30, 60, 120 個 stressor threads ，分別對應到 50%, 100%, 200% average CPU load 。
從這九種 data fields 當中經過前處理得出 15 種 input features ，利用 one-hot encoded 來表示每個輸入欄位，而 class label 則是利用 can_migrate_task() 的回傳值，訓練時利用 Adam optimizer 。 Loss function 則選用 binary cross-entropy 。訓練完成後經過 10 次的 Monte Carlo cross-validation 來驗證模型表現，平均的 loss 是 0.0955 ，平均 evaluation accuracy 是 99.24% 。結果雖然好但要整合進入 kernel 需要改成 fixed-point implementation 。此處 word-length 採 32 bits ，若

$WL=QI+QF$ ，如何決定怎麼分配

$QI,QF$ ?
其中一個方法是先讓

$QF$ 的精度範圍能夠表示我們所需要的最小值，再來檢查

$QI$ 範圍是否足夠，對於

$QF$ 而言，它的 resolution 是

$\epsilon =\frac{1}{2^{QF}}$ ，所以對於 resulation 為

$\epsilon$ 的情況，我們的

$QF$ 至少要是

$QF=\lceil lo{g}_2(\frac{1}{\epsilon })\rceil$ ，在論文實驗中發現，

$\epsilon \le 0.0005$ 才能保證模型的品質，因此

$QF=\lceil lo{g}_2(\frac{1}{0.0005})\rceil =11$ ，是我們需要的小數位數，而整數部分剩下 21 bits ，範圍會是

$-2^{QI-1}\le \alpha \le 2^{QI-1}-1$ 也就是

$-1048576\le \alpha \le 1048575$ ，對於此模型而言是足夠的。

Results and Future work

量測系統的 maximum imbalance 也就是每個 CPU runqueue 上任務數量差距的最大值，結果可以發現有 ML model 的支援下， maximum imbalance 的分佈值會趨向較少的那側，相較原先的 kernel 有更好的表現。
另外量測 total runqueue length 的話，則 10~22 jobs 上是原先的 kernel 表現優於 ML model 支援的 kernel 。
至於跑數個 benchmarks 測試 runtime 後的結果， ML support kernel 在 blackscholes 與 dedup 上比原先 kernel 更快，其他 benchmarks 則無明顯區別。

Insight 當中我認為最值得關注的是資料的蒐集，由於 training dataset 會很大程度的影響到模型如何改變排程策略，如何蒐集更多樣的資料並且在不同的系統架構、狀態之下蒐集是減少 biased 很重要的因素。

Multilayer Perceptron

Introduction

又稱為 MLP ，代表現代 feedforward artificial neural network ，通常由 fully connection neurons 和 nonlinear activation function 組成，至少三層，在處理無法以線性方式分類的資料時很適合，和原本的 perceptron 不同的點在於原本的 perceptron 採用的 activation function 是 Heaviside step function 而 MLP 是採用 nonlinear activation function 。

Feed Forward Neural Network

Wiki

Introduction

ANN 的其中一種，它的資料流向是單向的，只會從 input nodes 流向 hidden nodes 最後留到 output nodes ，不包含任何 cycles 或 loop 。現在多數的 FNN 都是透過 backpropagation 的方式訓練。

Mathematical foundation

Activation function

• Hyperbolic tangent function
  
  $y(v_i)=tanh(v_i)$
  ranges from -1 to 1 ,
  $y_i$ 是第 i 個節點的輸出，
  $v_i$ 則是輸入節點的 weighted sum
• Logistic function
  
  $y(v_i)=(1+e^{-v_i}{)}^{-1}$
• Rectified Linear Unit (ReLU)

Learning

Learning 的方法是透過比較預期輸出和實際輸出的誤差量來判斷要如何改變 connection weights 。誤差表示方法為

$e_j(n)=d_j(n)-y_j(n)$ 其中

$d_j(n)$ 是第 n 個節點的預期輸出而

$y_j(n)$ 是實際輸出。
而對於第 n 個 data point 來說整體誤差是

$$\epsilon (n)=\frac{1}{2}\sum _j{e}_j^2(n)$$
我們的目標是透過改變節點權重來最小化上述算式的值，例如可以利用 gradient descent 來改變節點權重

$w_{ij}$

$$\mathrm{\Delta }{w}_{ij}(n)=-\eta \frac{\partial \epsilon (n)}{\partial v_j(n)}{y}_i(n)$$
其中

$y_i(n)$ 是前一個節點

$i$ 的輸出，而

$\eta$ 是 learning rate 。值得注意的是上述微分項的值會隨著

$v_j$ 的值變動，並且可以簡化為

$-\frac{\partial \epsilon (n)}{\partial v_j(n)}=e_j(n){\varphi }^{\prime }(v_j(n))$ ，

${\varphi }^{\prime }$ 是 activation function 的微分。 Hidden node 的 weights change 可以利用以下方式表示

$$-\frac{\partial \epsilon (n)}{\partial v_j(n)}={\varphi }^{\prime }(v_j(n))\sum _k-\frac{\partial \epsilon (n)}{\partial v_k(n)}{w}_{kj}(n)$$
其中

$k$ 節點代表 output layer ，所以計算 hidden layer weight 的改變就是計算 output layer weight change 和 activation function 的微分。