工作上的原因，最近学习了一下 Run:ai，也考察了一下它的替代技术。在这里记录一下。
Run:ai 功能特性
Nvidia 当年收购 Run:ai，还是很有前瞻性的。并且它并没有阻拦开源版本的继续演进，只不过，商业分支提供了更强大的能力。
Run:ai 是一个企业级的 AI 基础设施编排和管理平台，核心能力就是最大化 GPU 的资源利用和调度能力，在它的主页上，列出了四大特性：
- AI-Native Workload Orchestration
- Dynamic GPU Allocation
- Policy-Driven Governance
- Open Architecture
抛开那些商业宣传和拥抱开源等等方面，我们只讨论技术层面，我觉得拥有核心竞争力的 feature 主要是：
- Dynamic Quotas and Pooling，因为原生 Kubernetes 的资源分配是静态的，有一个简单的排队机制。但是 Run:ai 的 scheduling 支持一个复杂得多的规则系统，比如可以保证每一个 team 或者 project 的基础配合，但是也可以超额借用，而对于资源竞争的情况，还允许定义 preemption 的方式，把低优先级的任务踢掉。
- Factional GPU and Hard Isolation，单个 GPU 可以被切分成多个虚拟的 GPU，这个非常实用；而底层硬隔离是商业版本的卖点——Kai Scheduler 是 Run:ai 调度逻辑的开源版本，但是它只能做到 Kubernetes 层面上资源的软隔离，而 Run:ai 深入到底层 virtualization，能够做到物理级别的 CUDA 调用拦截，无论是 GPU 还是 memory，都可以设置硬 quota。
- Advanced AI Scheduling，比如它支持需要多个 Pods 的任务，同时获得多个 Pod，或者等待，避免获取部分 Pods 是饿死情况出现；它支持 Bin Packing 的装箱算法和碎片整理，把碎片化的 GPU 能力拼合成一个大的；还有 Topology Awareness，把频繁通信的分布式训练任务放到物理距离更近的 GPU 上，这也是和硬件层面更紧密的结合。
Run:ai 的架构
Run:ai 的官方文档上这张组件图是最好的 overview 的材料：
从远处看，Run:ai 分成 control plane 和 cluster 这样两个部分，一般说来，一个 control plane，一个或多个 cluster。这让我想起了前段时间学习过的 Flyte 的架构，类似的部署方式。值得一提的是，出于安全等原因的考虑，没有从 control plane 主动去连接和发送请求到 cluster 的情况，只有从 cluster 发起去读取 control plane 的情况。
- Control Plane 是指挥和回应用户的，提供包括 resource management，用户侧 RBAC，处理工作提交和监控分析 cluster 的功能；
- Cluster 是真正做事的，负责 scheduling 和 cluster 内 workload 的管理。
除了官方文档，DeepWiki 上面也有很有用的材料可以学习，比如下面这张组件图：
在 Control Plane，scheduling 等 metadata 是存放在 PostgreSQL 里面的；而每个 cluster 则是扩展 Kubernetes，数据存放在 Etcd。
部署方式上，其中的 Control Plane 可以部署在用户自己这里，甚至包括 docker registry（Air-Gapped 模式），也可以使用 Nvidia 提供的远程服务（SaaS 模式）。很多公司出于安全考虑都是使用本地部署的。
接着，从和 Kubernetes 结合的机制上去理解 Run:ai：
- 当有一个 job 部署到已经集成 Run:ai 的 Kubernetes cluster 上的时候，Mutating Admission Webhook 会介入，部署好的 MutatingWebhookConfiguration 会被触发，强制改写资源的 schedulerName 为 runai-scheduler，从而把它从 Kubernetes 的默认调度器接管过来。
- 如果是 factional GPU 的任务，Webhook 还会往容器配置中挂载特定的系统目录，并注入环境变量，让容器启动的时候，把 Run:ai 的 CUDA 拦截动态链接库塞进 Python 环境中。
在运行的时候，我们能够从 job 资源的 labels 和 annotations 上面看到被 Run:ai 劫持后的标记。
KAI Scheduler
Run:ai 是收费的，KAI Scheduler 是一定程度上 Run:ai 调度模块的平替，它是 Kubernetes 原生的调度器。功能实现的层面不太一样，总体来说 Run:ai 要更强大一些。
- 对于 Dynamic Quotas and Pooling 部分，Run:ai 支持多集群的算力池化控制和调度，但是 KAI Scheduler 自身只支持单一集群内部。
- 对于 Factional GPU and Hard Isolation 部分，GPU 切分没问题，但是 Hard Isolation 只有 Run:ai 支持。KAI 只能在 Kubernetes 层面做逻辑控制，尤其对于 multi-tenancy 的情况，如果某个容器超额占用，会影响到同一张显卡上的其它 tenant。
- 对于 Advanced AI Scheduling 部分，Run:ai 的功能覆盖面更广，比如任务的全生命周期覆盖等等。
说说其中最核心的软隔离和硬隔离的问题。如果使用 KAI Scheduler，如果说用户提交一个 PyTorch 任务，上面申请了一定量的显存，但是如果 PyTorch 任务的实际代码中，用户申请更大的显存，KAI Scheduler 根本不会去阻止。这就会引发 Noisy Neighbor 的问题。这个问题不只是显存层面，在 GPU 层面也有，KAI Scheduler 的调度逻辑能把任务安排到相应的显卡上，但是没法强制设置某一张显卡部分使用时的百分比上限。
对此，如果要寻找解决方法：
- 如果是新的企业级 GPU 显卡，可以使用 Nvidia 的 MIG（Multi-Instance GPU）显卡硬件支持的技术，一张显卡可以切割成最多 7 张小显卡（并非任意分隔单元）；
- 或者使用 Nvidia 的 MPS（Multi-Process Service）模式，这是把时间切片改成基于 CUDA 进程的资源切片，显存和 GPU 都可以使用这种基于进程的模式。
- 再有就是纯软件工程上来缓解这个问题，使用一个 Kubernetes Mutating Webhook，给用户的容器外面包一层（比如可以做 Python 运行库或者 C/C++动态链接库级别的拦截），来控制资源获取操作。这种方法存在被恶意绕过的可能。
在 KAI Scheduler 的 GitHub 页面上，有一个很有帮助的 presentation 列表和一个 key features 列表。
文章未经特殊标明皆为本人原创，未经许可不得用于任何商业用途，转载请保持完整性并注明来源链接 《四火的唠叨》