生产者消费者模式下实现多batch延时推理

需求分析

在实际推理过程中为了实现较高的吞吐量和较高的资源利用率，往往会使用多线程来收集多次请求，并组合形成多batch下的模型推理，一种常见的实现便是生产者和消费者模式，其需求如下：

1. 生产者收集提交的请求，消费者对请求进行消费，并将结果返回。
2. 资源是有限的，可运行的最大max_batch_size是有限的，队列中可等待的max_queue_num也是有限的。
3. 邻近的资源可进行等待，最长要等待timeout的时间
4. 资源封装完备，哪里使用就在哪里释放。

设计实现

1. 在设计实现上，首先是需要有线程安全的生产者消费者队列，通过promise和future完成返回消息。
2. 需要有设置各个参数的公开接口，以便于设置各个参数。
3. 需要有一把锁和延时机制，使用条件变量condition_variable的wait_for函数来等待后续参数的传入。
4. 资源的申请和释放都放置在消费者线程中，并使用promise和future来返回中间结果。
5. 使用抽象类实现接口和封装，使用多态进行实现。

具体实现

1. 头文件

这里声明了一个抽象类，关键要实现forward函数，create_infer外部函数实现多态调用。

#ifndef INFER_HPP
#define INFER_HPP
#include <memory>
#include <string>
#include <map>
#include <future>
namespace Infer {
using Tensor = std::string;
using RET_TYPE = std::map<std::string, Tensor>;
class InferInterface {
 public:
  // pic_tensor 表示输入给生产者的tensor, timeout 表示超时时间
  // wait_timeout 表示当队列满了的情况下，愿意最多等待 wait_timeout ms 来等待
  virtual std::shared_future<RET_TYPE> forward(const Tensor& pic_tensor, int wait_timeout = 10) = 0;
  virtual void set_timeout(const int timeout) = 0;
  virtual void set_max_queue_num(const int max_queue_num) = 0;
  virtual ~InferInterface() {}
  explicit InferInterface() {}
 protected:
  InferInterface(const InferInterface&) = delete;
  InferInterface(InferInterface&&) = delete;
  InferInterface& operator=(const InferInterface&) = delete;
};
std::shared_ptr<InferInterface> create_infer(const std::string& filepath, int max_bactch_size);
}
#endif // INFER_HPP

2. cpp实现类

InferImpl是InferInterface的具体实现，是一种推理引擎的多态封装。

这里有4个原子操作，分别对应着running、超时时间timeout、最大可推理的batch_size以及队列的最大长度max_queue_num。
有一个锁，用以等待后续的请求，并合为一个batch。
有两个条件变量，cond_var_对应的是消费者，用以wait_for不为空且系统在runnning状态，并且queue的请求数量大于batch_size时取出batch_size请求。cond_queue_overflow_对应的是生产者，需要等待当前队列中的数据不满时才能进行填充。

forward函数里的实现生产者，当队列不满的时候向队列中插入请求。
worker函数是消费者的具体实现，会等待当前队列中请求数大于要推理的请求时才会组成batch一同推理，并将结果通过promise返回得到最终结果。

#include "infer.hpp"
#include <cstdio>
#include <queue>

namespace Infer
{
    using std::atomic;
    using std::future;
    using std::map;
    using std::mutex;
    using std::promise;
    using std::queue;
    using std::string;
    using std::thread;
    class InferImpl : public InferInterface
    {
    protected:
        /* data */
        using RET_PROMISE_PTR = std::shared_ptr<promise<RET_TYPE>>;
        struct Job
        {
            /* data */
            Tensor data;
            RET_PROMISE_PTR ret_promise;
        };

        atomic<bool> running;
        atomic<int> timeout;
        atomic<int> max_batch_size;
        atomic<int> max_queue_num;

        mutex lock;
        thread thread_;
        queue<Job> queue_;
        string context;

        std::condition_variable cond_var_;
        std::condition_variable cond_queue_overflow_;

    public:
        void set_timeout(const int timeout)
        {
            this->timeout = timeout;
        }
        void set_max_queue_num(const int max_queue_num)
        {
            this->max_queue_num = max_queue_num;
        }
        bool set_batch_size(const int batch_size)
        {
            if (batch_size < 1)
                return false;
            this->max_batch_size = batch_size;
            return true;
        }

        std::shared_future<RET_TYPE> forward(const Tensor &pic_tensor, int wait_timeout = 10)
        {
            Job job;
            job.data = pic_tensor;
            job.ret_promise = RET_PROMISE_PTR(new promise<RET_TYPE>());
            {
                std::unique_lock<mutex> l(lock);
                if (queue_.size() >= max_queue_num)
                {
                    if (0 == wait_timeout)
                    {
                        throw std::runtime_error("exhausted resource");
                    }
                    cond_queue_overflow_.wait_for(l, std::chrono::milliseconds(wait_timeout), [&]()
                                                 { return queue.size() < max_queue_num; });
                }

                if (queue_.size() >= max_queue_num)
                {
                    throw std::runtime_error("exhausted resource");
                }

                queue_.push(job);
            }
            cond_var_.notify_one();
            return job.ret_promise->get_future();
        }
        explicit InferImpl()
        {
            running = false;
            max_batch_size = 1;
            max_queue_num = 5;
            timeout = 0;
        }
        ~InferImpl() override {
            running = false;
            cond_var_.notify_one();
            if (thread_.joinable())
            {
                thread_.join();
            }
        }

        bool load_model(const string &filepath)
        {
            promise<bool> init_promise;
            thread_ = thread(&InferImpl::worker, this, filepath, std::ref(init_promise));
            running = true;
            return init_promise.get_future().get();
        }

    protected:
        void worker(const string &filepath, promise<bool> &init_promise)
        {
            // 加载模型
            context = filepath;
            if (context.empty())
            {
                init_promise.set_value(false);
                return;
            }

            init_promise.set_value(true);
            std::vector<Job> jobs;
            int batch_id = 0;

            while (running)
            {
                {
                    std::unique_lock<mutex> l(lock);
                    cond_var_.wait(l, [&](){
                        // true 则退出等待
                        return !queue_.empty() || !running; });
                    if (!running)
                        break;
                    if (0 != timeout)
                    {
                        cond_var_.wait(l, [&]()
                                       { return queue_.size() >= max_batch_size; });
                    }
                    for (int i = 0; !queue_.empty() && i < max_batch_size; i++)
                    {
                        jobs.emplace_back(queue_.front());
                        queue_.pop();
                    }
                }

                // 此处假装inference，得到运行结果

                int sz = jobs.size();
                for (auto& job:jobs){
                    auto bbox = job.data + "_result";
                    RET_TYPE handle_result;
                    handle_result["bbox"] = bbox;
                    job.ret_promise->set_value(handle_result);
                }

                jobs.clear();
                std::this_thread::sleep_for(std::chrono::milliseconds(1000 * sz));
                printf("batch id: %d job size: %d \n", batch_id, sz);
                cond_queue_overflow_.notify_one();
                ++batch_id;
            }

            context.clear();
            puts("context_ has cleared");
            puts("Workder done!");      
        }
    };

    std::shared_ptr<InferInterface> create_infer(const string& file_path, int max_batch_size){
        auto infer_ptr = new InferImpl();
        infer_ptr->set_batch_size(max_batch_size);

        if (! infer_ptr->load_model(file_path)){
            delete infer_ptr;
            return nullptr;
        }
        return std::shared_ptr<InferInterface>(infer_ptr);
    }
}

3. main文件

在main文件中，创建了24个请求，每5个请求将组成一个batch进行推理。
这里有个小问题是，当最后一个请求小于5个时会一直等待，因此应当设置超时机制对剩余的请求进行一并处理。

#include <iostream>
#include <vector>
#include "infer.hpp"

using std::string;
using std::errc;

using RET_FUTURE = std::shared_future<Infer::RET_TYPE>;

int main(int argc, char const *argv[])
{
    /* code */
    string file_path = ""; 
    std::shared_ptr<Infer::InferInterface>infer_ptr = Infer::create_infer(file_path, 5);
    if (infer_ptr == nullptr){
        printf("create infer engine error\n");
        return -1;
    }
    // 每个实例推理愿意等待的时间
    infer_ptr->set_timeout(5);
    infer_ptr->set_max_queue_num(10);
    printf("create infer engine success!\n");

    std::vector<RET_FUTURE> shared_ptrs;
    char buffer[100];
    for (int i=0; i<24; ++i){
        sprintf(buffer, "%d.tensor", i);

        // 在队列满的时候，生产者愿意等1000ms
        shared_ptrs.push_back(infer_ptr->forward(buffer, 1000));
    }

    for (auto &shard_ptr:shared_ptrs){
        shard_ptr.get();
    }

    return 0;
}

本栏目推荐文章
• spring与设计模式之三代理模式
• 深入理解spring框架：剖析多线程模式下数据库连接
• 代理模式（Proxy）
• AntDesignBlazor示例——暗黑模式
• 单例模式（Singleton）的6种实现
• 设计模式之中介者模式
• freeswitch: esl inbound模式下外呼拨号
• 第二章 Spring Boot 整合 Kafka消息队列 生产者
• C++ 单例模式以及内存管理
• vim模式用法及命令

生产者 消费者 模式 batch生产者 消费者 模式batch 生产者 简易 消费者 模式 生产者 消费者 语言 模式 生产者 消费者 模式notify 产者 消费者 模式golang 生产者 线程 消费者 模式 编程技巧 生产者 消费者 模式 生产者 消费者 模式 生产者 模型 消费者 生产者 队列 消费者