读论文：Towards a Machine Learning-Assisted Kernel with LAKE

本文最后更新于：2024年11月21日

  • 发表会议：ASPLOS 2023
  • 作者：Henrique Fingler, Isha Tarte, Hangchen Yu, Ariel Szekely, Bodun Hu, Aditya Akella, Christopher J. Rossbach
  • 单位：The University of Texas at Austin, Meta, MIT

概述：

  • 作者提出了内核辅助决策框架LAKE，使得内核能访问GPU加速 ML 算法

问题

内核中比如内存管理、进程调度等子系统目前依赖手工调整的启发式方法来提供合理的平均性能，ML 可以更好地权衡这些 tradeoff。

例子：LinnOS@OSDI2020

问题：SSD复杂的内部活动，如垃圾回收、负载均衡等，和用户的请求冲突时会产生严重的尾延迟。

观察：存储系统可能存在一定冗余（RAID），而多个不同的SSD的内部行为同时和用户请求产生冲突的概率非常低。 所以可以将一个请求发给一个SSD后，若等待请求完成的时间超过了阈值，则重发请求到另一个可用的SSD。但是对于比较快的SSD来说，“等待慢SSD处理的时间超过阈值”这件事还是很慢。

解法：LinnOS学习SSD的特征，预测将要变慢的SSD而及时将请求重发到快SSD

看起来一系列信息都和SSD快或慢有关：读写请求？请求的块内偏移？长期的写入历史？然而，作者发现这些请求都对提高精确度没有明显帮助。首先，由于当前SSD常有内置写缓存，写之后的读延迟常常没有明显提高，更为常见的其实是数据从缓存”冲“（flush）入SSD后，读延迟会更高。其次，一组I/O请求会通过条带均匀地写入各个通道或者芯片，它们写入同一个芯片的概率很低，所以块内偏移这个特征其实并不重要。最后，GC或者flush通常发生时间短，短期写入历史就可以预测。

因此，可以使用SSD当前I/O队列长度来预测SSD快或者慢：一个直观的感受是，当I/O队列较长时，SSD处理通常比较慢。但是这样并不能体现SSD的内部活动的发生，因此额外增加了历史四条请求进入SSD时的队列长度和完成请求的时间。若某个请求进入SSD时队列短而完成请求的时间长，意味着SSD内部行为可能和用户请求冲突了。

如何最小化预测错误的影响？

作者分析发现，若将一个快的SSD预测为慢的从而错误地重发了，将带来微秒级延迟，而若将一个慢的SSD预测为快，将带来毫秒级延迟，比第一种情况严重许多，所以作者在训练时对第二种情况施加了更加严重的惩罚以减少它们的发生。此外，还补充了hedged request以减少预测失败的损失。

问题描述：How to best support ML in the OS?

挑战1：内核无法用GPU

  • 现状：ML is compute expensive
    • 部分决策存在低时延需求
    • 无法使用比较大的模型（PSS@ASPLOS'23只用单层感知机）
  • 想法：使用 GPU/TPU 可降低 ML 性能开销
    • 内核空间无法利用现有的硬件加速器（如英伟达的GPU驱动闭源并且仅针对用户空间开发）
    • 硬件加速需要均摊数据传输的成本
    • 用户空间和内核空间对GPU/TPU存在资源争用

挑战2：训练和推理需要跨模块收集数据

为了收集来自于CPU、内存、IO等模块的信息，如果没有任何基础设施，可能需要手动维护数据结构和控制临界区，但是这样的实现很复杂，加入新的代码可能带来死锁等一系列问题。

1
2
3
4
5
6
7
void submit_bio() { 
    ts_tart = now(); 
    new_fv = allocate_fv(); 
    new_fv_size = ...; 
    new_fv.io_size = new_fv_size; 
    store_fv(); 
}

方案

内核如何用GPU：提供内核态加速器调用接口，将请求转发到用户空间处理

LAKE 中主要有三个组件：

  • lakelib：在内核中提供一组与用户空间相同的 API，如cuXXX
  • lakeD：实现 API 的用户端守护进程，负责接受内核态请求，再调用用户空间的接口。
  • lakeShm：为了避免数据拷贝，LAKE使用共享内存的方式在用户空间和内核之间共享数据

解决争用：动态调节使用GPU/CPU

LAKE采取了如下两条策略

  • GPU只有在大规模输入的情况下才有优势 ----> batch大于一定值(>32)才使用GPU推理
  • 避免用户空间产生资源争抢 -----> GPU使用率低于一定值的时候才使用GPU推理

LAKE支持使用eBPF自定义策略：在GPU占用率过高或者batch过小的时候会使用CPU推理。

提供收集feature的内核公共基础设施

LAKE实现了一套用于跨模块收集数据的异步API以及与之对应的无锁数据结构。主要功能包括：

  • 注册特征
  • 记录特征
  • 加载和更新模型

效果

  • 2 * 16-core Intel Xeon Gold 6226R CPUs
  • 376 GiB DDR4 RAM
  • 2 * NVIDIA A100 GPUs
  • 3 * Samsung 980 Pro 1TB (PCIe 4.0) NVMes
  • Ubuntu 22.04 with our modified Linux kernel based on version 6.0

作者测试了一个端到端的IO预测问题，使用的模型为发表在OSDI2020上的LinnOS，baseline是没有重定向的IO

NN指neural networks。默认大小为 [256,2]，+1代表加一层256个神经元的网络

总结

  • 优点
    • 在内核中使用GPU加速机器学习推理
    • 实现了一套跨模块收集数据的基础设施
  • 不足
    • 没有体现更大模型的必要性
    • 没有明显的端到端的性能优势
    • 跨层优化：如果用户态和内核态同时需要用GPU推理，本文会牺牲内核态，但未必是全局最优。