sched_ext(5): Integrate Machine Learning into scx_rusty Load balancer

scx_rusty overview

scx_rusty 身為 scx 專案目前收錄的一個排程器，將 CPU 分到不同 domain 當中，在 bpf program 當中實作主要的排程邏輯，而特別的是透過一個使用者層級的 load balancer 來幫助進行負載平衡，定時進行 load balance 運算並將資料透過 BPF map 更新， bpf program 進行 select_cpu 或 enqueue, dispatch 等行為時可透過讀取 BPF map 當中的資訊而得知是否改進行負載平衡。

Userspace part

使用者層級的程式透過 Rust 程式語言實作，主要負責兩個部分，第一個是提供 higher frequency tuning decision 。找出當前負載不沉重的 CPU 並將他們標記起來，使得他們可以從過載的 CPU 當中進行 task pull 。

第二個部分是運行 lower frequency load balancing ，透過比較 domain 之間的 load averages 來決定是否進行 load balancing 。如果 load difference 夠大則會檢查最多 1024 個 task 來判斷哪個 task 應該被移動（這也是我們可以優化之處）。

在使用者層級的操作成本低，而且負載平衡成本也不像過往高，進行的頻率也低，不過 work-conservation 依舊透過 tuning 和 greedy execution 來進行。

目前 scx_rusty 假設所有 domain 之間的距離都一樣，而且 processing power 也都一樣 。

Load Balancer

負載平衡的流程大致如下

1. 利用 dom_ctx_map_elem bpf map 當中的 duty cycle buckets 來計算 domain load average ，利用 scx_utils::LoadCalculator 來計算 load 並利用 scx_utils::LoadLedger 來判斷 load distribution 。
2. 拆分 NumaNode 和 Domain 的 hierarchy ，詳細可見 scheds/rust/scx_rusty/src/load_balance.rs
3. 從 LB root 開始尋遍每個 NUMA node ，從 push nodes 也就是負載過高的 node 將任務遷移到 pull nodes 當中。此處進行 task migration 的代價會比 node 當中的 domain 進行 task migration 的代價高。因此這種 migration 只有當 imbalance 情況非常嚴重才會發生。
4. 若 NUMA nodes 之間的負載平衡結束後，接下來會針對每個 NUMA nodes 當中的 domain 進行負載平衡，此處 task migration 的成本低。

BPF part

負責進行在每個 domain 間進行 round robin ，當一個 task 初次進入系統當中時（ rusty_prep_enable ），它們會透過 round robin 順序進行其中一個 domain 。
排程最主要在 rusty_select_cpu 當中進行，當一個 task 狀態變為 runnable 時會被呼叫， lb_data bpf map 是用來判斷此 task 是否需要被遷移到新的 domain ，這個 bpf map 的 producer 就是 userspace part 的 load balancer 。若不需要的話 task 排程的 priority order 如下

1. 若 task 被同步喚醒而系統當中有其他 idle CPU ，則放到當前的 core 當中
2. 若前一個 core 為 idle ，則放到它上面
3. 若 task 有被指派到特定的 core 上則放到該 core 上
4. domain 當中任意的 idle CPU

若上述四項都不滿足，該 task 會被放到對應 domain 的 dispatch queue 當中。

rusty_dispatch 會嘗試從對應 domain 的 dispatch queue 當中 consume 一個 task 來運行。若沒有找到 task ，則會從其他 dispatch queue 當中進行 greedy load stealing 。

基本上 load balancing 完全交給使用者層級處理， BPF part 負責將 task weight, domain mask 和 current domain 生產進入 task_data bpf map 當中，並透過 lb_data bpf map 當中的資訊判斷是否進行 load balance 。

整合機器學習進入 Load balancer

Main idea

在使用者層級的 load balancer 實作當中，有一個函式稱為 try_find_move_task ，用來找到 push_dom 當中要遷移到 pull_dom 當中的 task ，若找到的話會更新對應 lb_data bpf map 當中的資料，並且回傳移動的 load 。

該函式當中尋找可以被遷移的 task 的實作當中，若可以適當的加入機器學習模型的 inference 來判斷哪個 task 應該被遷移，或許可以帶來更好的效果。

Training

透過 MLLB 當中的程式碼為基礎，以 tensorflow 訓練一個模擬 EEVDF/CFS 的 load balancing 機制的模型。也就是模仿 can_migrate_task() 的行為。

原本論文當中透過掛載 kprobe 在 can_migrate_task() 上蒐集資料的方式可以採納，但拿來訓練的 feature 必須改變，因為在 scx_rusty 當中並不具有 struct lb_env 這個結構體，用來進行負載平衡的 lb_data, task_data 成員須擴充，但無法完全滿足 struct lb_env 的所有成員。

當前 scx_rusty 完全只依據 task load 來判定是否進行 load balance 和 task migration 。我們依照原論文方法重新蒐集資料，也就是分別在 stress-ng --cpu 30 -l 100 --timeout 120s, stress-ng --cpu 60 -l 100 --timeout 120s, stress-ng --cpu 120 -l 100 --timeout 120s 情況下蒐集系統資料，之後只利用四項 feature 進行訓練，分別是 cpu_idle, cpu_not_idle, src_load, dst_load （應該要再加上 NUMA 資訊，但該程式嘗試取得 NUMA 相關資料時會出錯，尚待解決）。

把訓練好的模型檔案搬移至 scx_rusty 資料夾底下，準備進行 inference 。

Inference

Rust 剛好有能整合 tensorflow 的套件，安裝該套件後實作 inference 應能變得很方便。
首先要在 scx_rusty 的 Cargo.toml 檔案當中新增套件

[dependencies]
tensorflow = "0.21.0"

之後擴充 struct task_ctx 使我們可以從 kernel space 取得要拿來 inference 的資料 (位在 scheds/rust/scx_rusty/src/bpf/initf.h)

+	/* Machine Learning inference needed data */
+	unsigned long src_dom_load;
+	unsigned long dst_dom_load;
+	int cpu_idle;
+	int cpu_not_idle;
+	int cpu;
+};

另外在 main.bpf.c 當中修改 init 。

+		.cpu = bpf_get_smp_processor_id(),
+		.cpu_idle = 1,
+		.cpu_not_idle = 0,

之後在 scheds/rust/scx_rusty/src/load_balance.rs 當中定義 tensorflow 的 struct 和相關 method

extern crate tensorflow;

use tensorflow::Graph;
use tensorflow::Session;
use tensorflow::SessionOptions;
use tensorflow::SessionRunArgs;
use tensorflow::Tensor;
use std::error::Error;

struct TensorFlowModel {
    graph: Graph,
    session: Session,
}

impl TensorFlowModel {
    fn new(model_dir: &str) -> Result<Self, Box<dyn Error>> {
        let mut graph = Graph::new();
        let bundle = tensorflow::SavedModelBundle::load(
            &SessionOptions::new(),
            &["serve"],
            &mut graph,
            model_dir,
        )?;
        let session = bundle.session;
        
        Ok(TensorFlowModel { graph, session })
    }

    fn predict(&self, input_data: Vec<f64>) -> Result<bool, Box<dyn Error>> {
        let input_tensor = Tensor::new(&[input_data.len() as u64]).with_values(&input_data)?;

        let input_op = self.graph.operation_by_name_required("serving_default_input")?;
        let output_op = self.graph.operation_by_name_required("StatefulPartitionedCall")?;

        let mut args = SessionRunArgs::new();
        args.add_feed(&input_op, 0, &input_tensor);
        let output_token = args.request_fetch(&output_op, 0);

        self.session.run(&mut args)?;

        let output_tensor: Tensor<f64> = args.fetch(output_token).unwrap();
        let output_value = output_tensor[0];
        Ok(output_value == 1.0)
    }
}

可以看到 inference 函式就是 predict 。
之後實作 inference call back 函式

    fn migrate_inference(&self, cpu: &i32, cpu_idle: &i32, cpu_not_idle: &i32, src_dom_load: &f64, dst_dom_load: &f64) -> bool {
        let mut input_vec: Vec<f64> = Vec::with_capacity(5 as usize);
        input_vec.push(f64::from(*cpu));
        input_vec.push(f64::from(*cpu_idle));
        input_vec.push(f64::from(*cpu_not_idle));
        input_vec.push(*src_dom_load);
        input_vec.push(*dst_dom_load);

        self.inference_model.predict(input_vec).unwrap()
    }

呼叫的方式很簡單，在 try_find_move_task 的最後檢查挑選出的 task 是否適合遷移

+        let migrate_or_not = self.migrate_inference(&task.cpu, &task.cpu_idle, &task.cpu_not_idle, &push_load, &pull_load);
+        task.migrated.set(migrate_or_not);

之後進行測試，首先進行同時編譯 linux kernel 與進行 fio 測試情況下 task migration 次數

$ sudo perf stat -e sched:sched_migrate_task -a sleep 120

 Performance counter stats for 'system wide':

            6,1749      sched:sched_migrate_task

     120.200994178 seconds time elapsed

可以看到遷移次數相對原本 scx_rusty 有顯著的下降，再來測試編譯一次 linux kernel 的時間

$ time sudo make -j $(nproc)
...

real    4m23.921s
user    0m0.046s
sys     0m0.094s

居然還比 scx_rusty 編譯的更快！

目前模型的訓練與 inference 都還有很大的改進空間，但已經可以顯示原本 scx_rusty 過多的 task migration 是不必要的，不只對於 load imbalance 沒有過多幫助，整體效能也有提升的空間，當然已經比原本的 EEVDF 好上許多。