跳转至

对单机多卡AI模型推理场景下计算资源分配问题的思考

一 写在前面

出差将近一个月，角色从评委到学员，这是当前这份工作带给我的魔幻经历（捂脸笑）。出差的过程中感染了奥密克戎，幸运的是没有发高烧，流了五天清鼻涕和鼻塞，更多的是心理上的压力，是对于未知的恐惧和家人的担忧，不过也已经慢慢化解了。个人觉得更多的要养成长期防疫下的习惯，戴好N95少去人多地方，及时消毒。好长时间没吃水饺了，回来后第一天一个饺子一口蒜，一顿吃了36个，真的太香了，顿时感叹年轻真好。

不说废话了，关于单机多卡AI模型推理场景下计算资源分配问题，是业务抽象出来的问题，本文将针对这类共性问题的不同情况进行讨论，思考不同的解决方案下的优劣。

本文涉及到的完整代码在：

https://github.com/thb1314/mutual_share_memory_with_priority_queue

二 抽象问题

首先对”单机多卡AI模型推理场景下计算资源分配"这一句话进行拆解

单机多卡: 在一个主板上，存在多张计算卡（gpu、npu等），多张计算卡之间可以通信，计算卡与主控芯片（cpu）之间可以通信

AI模型推理: 人工智能模型的推理过程，这里特指神经网络模型的推理过程（即大量的矩阵乘法）

计算资源: 计算资源指的是多张计算卡共同提供的算力，也就是其提供的计算能力和存储能力共同组成的东西。

分配：毫无疑问指的是上述计算资源的划分问题，归根结底就是当有一个模型推理任务创建时应该采用上述单机多卡中采用哪张卡的问题。

拓展：

说明白了单机多卡上的问题，那么多机多卡就显而易见了，无非就是一个主板换成了多个主板，先解决单机多卡的问题后，再解决多机之间的通讯问题即可。

三 解决方案

3.1 算法层面

对于上述问题，实际上可以等价为如下问题：

已知：

  • N张卡，单张卡的存储容量为M，计算能力为P
  • X个任务：每个任务需要消耗存储容量\(M_i\)和计算能力\(P_i\)，消耗时间\(T_i\)（\(i\in[1,X]\)）
  • 目标：构造一个任务表，每次可以跑1个或多个任务，使得总体代价函数最小。代价函数比如可以是推理时间

讨论：

  1. 当X的值已知时，即将X个任务分批次安排到N张卡，每次安排的任务个数不限制（受存储容量和计算能力限制），安排多少次不限制，目标是总体代价函数最小。比如总体代价函数为总体消耗时间，那么此类最优化问题或者说是分配问题。
  2. 当X的值未知时，此时X应该是无穷大的情况，类似于用户点击或者请求的事件，每隔一段时间就会有一个任务，等待时间未知。

针对第一个问题我暂时没有找到比较好的解决方案，不过此类问题是我假象的，现实中X的值更多是未知的，比如在一个APP中调用端侧推理人物，用户一共调用了多少次什么时候调用都是未知的。

下面说一下第二个问题的解决方案：

考虑到任务的时效性，实际上局部最优方案是针对每次的安排的任务都使得当前代价函数计算的最小。有人或许会考虑局部最优并不代表全局最优，是这样的，但是下一个任务的到来时间是未知的，这也就意味着等待时间是没办法保证的，假设可以将当前任务和下一个任务合并，那么随之而来增加的等待时间统计进去后又该如何取舍呢？

我们总是希望我们做的选择是全局最优解，可由于等待时间的存在我们没办法去判断是否是最优解，只能让其发生然后再去观察，可那样对于已经发生对的结果毫无意义，人生或许亦是如此，我们总是从一个局部最优解选择到下一个局部最优解。

当然我们也允许等待时间的存在，设置一个最大等待时间，超时后才不等，比如我在生产者消费者模式在多batch推理下的应用(延时队列) 中的实现里面就存在一个最大等待时间，如果不超过就放入当前的推理batch中，如果超过了就不等了，放到下个batch中。

无论等待还是不等待都需要根据场景实际测试。

3.2 通讯层面

针对单机多卡场景，承担推理请求的有可能是进程、线程、协程或者是多个虚拟机、docker容器等，本文将其定义为“运行体”，根据是否在同一个操作系统中，我们可以将其分为两类，即是否可以共用一个操作系统内核。

这里需要说明的是，不同内存的方法自然也可以用于同一个内核，只不过同一个内核当中有更适合的办法。

同理，同一个内核当中，进程之间的通信方式自然也可以用于线程和协程（线程和协程不存在通信问题，需要解决的是同步问题），只不过后者有自己更好的同步方式。

3.2.1 同一个内核

在同一个内核的有同一个操作系统内部的进程、线程、协程，下面分开进行讨论。

进程之间通信：

  • 非IPC（Inter-Process Communication）通信
    • 管道：半双工通信，一条管道只能一个进程写，一个进程读，且不可以用时读写。
      • 匿名管道：适用于有亲缘关系的进程之间的通信（父子进程、兄弟进程）
      • 命名管道：适用于独立进程之间通信
    • 文件锁：主要利用IO操作可以做成互斥访问这一特性
    • 互斥锁：进程之间的互斥锁存在于内核当中，主要看操作系统是否支持
  • IPC通信：
    • 消息队列：将消息体以队列形式进行封装，可以实现边发送边接收。
    • 共享内存：共享内存的机制是不同进程拿出一块虚拟内存空间，映射到相同的物理内存空间
    • 信号：Windows和Linux中都有，用于响应各种事件，也可以用于自定义事件比如进程同步
    • 信号灯：实现共享资源的互斥访问与进程间的同步操作

Note:

在Linux/Unix系统中，IPC通信有两套API：

  • System-V IPC 接口：来自较早的Unix操作系统，是Unix众多分支中的一员，他有很多经典的用例，例如 ”SysV 初始化脚本“ (/etc/init.d)
  • Posix IPC 接口：则是来自IEEE所开发的一簇标准，基于Unix的经验和实践所实现的一堆调用服务接口，致力于在多种操作系统之间源代码级别移植。

Posix IPC和System-V IPC都是应用于系统级的接口，不仅适用于多进程间通信，也适用于多线程间通信。Posix相对于System-V可以说是比较新的标准，语法相对简单。

线程之间同步方法：

首先需要说明的是，同一个进程内部的线程间全局资源是共享的，不存在通信问题，更多的是对共享资源的保护问题。

  • 互斥量 mutex (windows/linux都有)
  • 条件变量 condition_variable（可以看做计数为1的信号量， linux特有）
  • 信号量 semaphore (linux/windows都有)
  • 临界区 critical section (windows特有)
  • 事件 event

协程之间同步方法：

协程主要是针对异步IO来做的，需要同步的情况不多，但是也是存在的，比如在python协程中，采用asyncio中的Lock和Event都是可以做到同步的，这一点与线程大同小异，不过需要时刻牢记的是，协程中的代码一定在cpu的单个核心上，且在同一个线程中，只是切分时间分段运行而已。

3.2.2 不同内核

不同内核之间，包括（不同docker容器，不同虚拟机，不同物理机）网络通信自然是首选，谈到网络自然离不开Socket，通过Socket我们可以创建机器之间的tcp和udp连接，针对共享资源我们可以保存在某一台主机或者集群中的数据库或者redis中，通过数据库的事务操作或者redis中的同步锁机制来保证资源的互斥访问。

四 具体实现案例

场景：单机多卡同一个内核，每个推理任务消耗的显存和计算资源相同，有多少个推理任务未知，“运行体”是独立进程，比如apache中的prefork模式，一次性开多个进程来实现并发处理，尽管每个进程是兄弟关系，但是对其采用管道方式是费力不讨好的，所以本文给出一种共享内存+进程间互斥锁的方式来实现共享资源的互斥访问。

作者目的是为了给一个python web服务来提供接口，web服务通常运行在linux机器上，所以作者采用System-V IPC接口来操作共享内存。采用C++调用linux应用层API，封装成C++ class来管理共享内存的读写与释放，然后采用pybind11给python提供接口。下面需要理清以下几个问题：

  1. 共享内存什么时候创建什么时候打开已存在？ 如果已存在就选择打开，如果不存在就创建，这里要做一个判断（创建报错就说明已经存在）

  2. 共享内存什么时候释放？ 使用的进程数小于等于1，且当前进程中的共享内存管理对象仅剩下一个的时候就释放。所以在共享内存里面开辟一小块区间用于记录当前进程中的对象个数，并在拷贝构造中维护该个数，在析构函数中加以判断。

  3. 如何保证计算资源的均匀分配呢？ 这里的计算资源可以视为卡的id号，采用优先队列的形式，先从优先队列中取出，再将需要判断的值优先级降低再放回队列中用于下次调用

  4. 如何保证对资源的互斥访问呢？ 可以将mutex创建在共享内存中，生命周期跟共享内存一样，创建共享内存的时候创建mutex

  5. 如何对共享内存进行初始化？ 如果共享内存已经存在就不需要初始化，不存在需要创建的时候初始化来保证仅初始化一次。

  6. 如何适配多种数据？ 采用STL模板编程思想，创建一个优先队列模板类，将数据类型和数据类型比较函数做成模板中的typename用于适配自定义数据

  7. 如何对python用户开放接口？ 采用pybind11封装一个带lock的共享内存优先队列类，针对特定场景将c++的一些的接口暴露给python

类之间关系

1
2
3
MutualSharedMemory: 用于对带有加锁功能的共享内存的管理（也可以继续抽象一个仅针对共享内存的管理类，这里简化）
MutualShmemWithPriorityQueue: 继承MutualSharedMemory，在上述功能基础上在共享内存内部维护一个优先队列
MutualShmemPriorityIndexedQueue: 继承MutualShmemWithPriorityQueue，可以看做MutualShmemWithPriorityQueue的一种模板特化与扩展，索引堆管理类，也是pybind11中的适配类型，其大部分api都对pybind11暴露。

代码这里就不贴出了，没有基础知识的读者可以看一下linux中共享内存的操作，然后再去查看源代码。

源码地址：

https://github.com/thb1314/mutual_share_memory_with_priority_queue

五 总结

2023-01-08回到西安，此篇文章即是我路上的碎碎念，也是出差这段时间对该问题的思考的一个总结。一周过去，本文的相关代码和文章终于写完了，也算是对自己有所交代。

六 参考文献

  1. https://blog.csdn.net/woyimibayi/article/details/80096275
最后更新: March 21, 2024
创建日期: March 21, 2024

评论