tuoheng_alg_move_20240715/service/Dispatcher.py

# -*- coding: utf-8 -*-
import time
import GPUtil
from util import YmlUtils, FileUtils, LogUtils
from loguru import logger
from multiprocessing import Queue
from concurrency.IntelligentRecognitionProcess import OnlineIntelligentRecognitionProcess, OfflineIntelligentRecognitionProcess
from concurrency.MessagePollingThread import OfflineMessagePollingThread, OnlineMessagePollingThread
from util import GPUtils
'''
   分发服务
'''


class DispatcherService():

    # 初始化
    def __init__(self):
        # 获取DSP环境所需要的配置
        self.content = YmlUtils.getConfigs()
        # 初始化日志
        LogUtils.init_log(self.content)
        # 检查视频保存地址，不存在创建文件夹,迁移初始化
        FileUtils.create_dir_not_exist(self.content["video"]["file_path"])
        # 记录当前正在执行的实时流分析任务
        self.onlineProcesses = {}
        # 记录当前正在执行的离线视频分析任务
        self.offlineProcesses = {}
        # 记录当前正在执行的图片分析任务
        # self.photoProcesses = {}
        self.onlineMpt = None
        self.offlineMpt = None

    # 服务调用启动方法
    def start_service(self):
        # 解决进程启动模型的环境问题，具体问题百度一下
        # torch.multiprocessing.set_start_method('spawn')
        # 启动实时，离线kafka消息拉取线程
        self.Kafka_message_listening()
        # 循环消息处理
        while True:
            time.sleep(1)
            # 检查任务进程运行情况，去除活动的任务
            self.check_process_task()
            gpu_ids = GPUtils.get_gpu_ids(self.content)
            if gpu_ids is not None and len(gpu_ids) > 0:
                ################## 消息驱动实时流分析进程执行 ##################
                onlineMsg = self.onlineMpt.poll()
                if onlineMsg is not None and len(onlineMsg) > 0:
                    for k, v in onlineMsg.items():
                        for m in v:
                            try:
                                msg = m.value
                                # 校验kafka消息
                                check_result = self.check_online_msg(msg)
                                if not check_result:
                                    raise Exception("实时任务消息格式非法")
                                if 'start' == msg.get("command"):
                                    logger.info("开始实时分析")
                                    self.startOnlineProcess(msg, self.content, gpu_ids)
                                elif 'stop' == msg.get("command"):
                                    self.stopOnlineProcess(msg)
                                else:
                                    pass
                            except Exception as e:
                                logger.error("实时消息监听异常：")
                                logger.exception(e)

                ################## 消息驱动离线视频分析进程执行 ##################
                offlineMsg = self.offlineMpt.poll()
                if offlineMsg is not None and len(offlineMsg) > 0:
                    for k, v in offlineMsg.items():
                        for m in v:
                            try:
                                msg = m.value
                                # 校验kafka消息
                                check_result = self.check_offline_msg(msg)
                                if not check_result:
                                    raise Exception("离线任务消息格式非法")

                                if 'start' == msg.get("command"):
                                    logger.info("开始离线分析")
                                    self.startOfflineProcess(msg, self.content, gpu_ids)
                                elif 'stop' == msg.get("command"):
                                    self.stopOfflineProcess(msg)
                                else:
                                    pass
                            except Exception as e:
                                logger.error("离线消息监听异常：")
                                logger.exception(e)
            else:
                logger.info("当前可用gpu数量: {}", gpu_ids)
                GPUtil.showUtilization()

    def startOnlineProcess(self, msg, content, gpu_ids):
        # 相同的requestId不在执行
        if self.onlineProcesses.get(msg.get("request_id"), None) is not None:
            logger.info("重复任务，请稍后再试！requestId:{}", msg.get("request_id"))
            raise Exception("重复任务，请稍后再试！requestId:{}".format(msg.get("request_id")))
        # 反馈队列
        fbQueue = Queue()
        # 图片队列
        imageQueue = Queue()
        # 创建在线识别进程并启动
        oirp = OnlineIntelligentRecognitionProcess(fbQueue, content, msg, imageQueue, gpu_ids)
        oirp.start()
        # 记录请求与进程映射
        self.onlineProcesses[msg.get("request_id")] = oirp

    def stopOnlineProcess(self, msg):
        ps = self.onlineProcesses.get(msg.get("request_id"), None)
        if ps is None:
            logger.info("未查询到该任务，无法停止任务！requestId:{}", msg.get("request_id"))
            raise Exception("未查询到该任务，无法停止任务！requestId:{}".format(msg.get("request_id")))
        ps.sendEvent({'command': 'stop'})

    # 检查实时、离线进程任务运行情况，去除不活动的任务
    def check_process_task(self):
        for requestId in list(self.onlineProcesses.keys()):
            if not self.onlineProcesses[requestId].is_alive():
                del self.onlineProcesses[requestId]
        for requestId in list(self.offlineProcesses.keys()):
            if not self.offlineProcesses[requestId].is_alive():
                del self.offlineProcesses[requestId]

    # 开启离线进程
    def startOfflineProcess(self, msg, content, gpu_ids):
        # 相同的requestId不在执行
        if self.offlineProcesses.get(msg.get("request_id"), None) is not None:
            logger.info("重复任务，请稍后再试！requestId:{}", msg.get("request_id"))
            raise Exception("重复任务，请稍后再试！requestId:{}".format(msg.get("request_id")))
        # 反馈队列
        fbQueue = Queue()
        # 图片队列
        imageQueue = Queue()
        # 创建在线识别进程并启动
        ofirp = OfflineIntelligentRecognitionProcess(fbQueue, content, msg, imageQueue, gpu_ids)
        ofirp.start()
        self.offlineProcesses[msg.get("request_id")] = ofirp

    def stopOfflineProcess(self, msg):
        ps = self.offlineProcesses.get(msg.get("request_id"), None)
        if ps is None:
            logger.info("未查询到该任务，无法停止任务！requestId:{}", msg.get("request_id"))
            raise Exception("未查询到该任务，无法停止任务！requestId:{}".format(msg.get("request_id")))
        ps.sendEvent({'command': 'stop'})

    # 校验实时kafka消息
    def check_online_msg(self, msg):
        requestId = msg.get("request_id")
        command = msg.get("command")
        models = msg.get("models")
        pull_url = msg.get("pull_url")
        push_url = msg.get("push_url")
        results_base_dir = msg.get("results_base_dir")
        if command is None:
            return False
        if requestId is None:
            return False
        if command == "start" and models is None:
            return False
        if models is not None:
            for model in models:
                if model.get("code") is None:
                    return False
                if model.get("categories") is None:
                    return False
        if command == "start" and pull_url is None:
            return False
        if command == "start" and push_url is None:
            return False
        if command == "start" and results_base_dir is None:
            return False
        return True

    # 校验实时kafka消息
    def check_offline_msg(self, msg):
        requestId = msg.get("request_id")
        models = msg.get("models")
        command = msg.get("command")
        original_url = msg.get("original_url")
        original_type = msg.get("original_type")
        results_base_dir = msg.get("results_base_dir")
        if command is None:
            return False
        if requestId is None:
            return False
        if command == 'start' and models is None:
            return False
        if models is not None:
            for model in models:
                if model.get("code") is None:
                    return False
                if model.get("categories") is None:
                    return False
        if command == 'start' and original_url is None:
            return False
        if command == 'start' and original_type is None:
            return False
        if command == 'start' and results_base_dir is None:
            return False
        return True

    # 实时、离线kafka消息监听
    def Kafka_message_listening(self):
        # 实时流分析消息拉取
        self.onlineMpt = OnlineMessagePollingThread('online_thread', {'content': self.content,
                                                                      'topics': [self.content["kafka"]["topic"][
                                                                                     "dsp-alg-online-tasks-topic"]]})

        # 离线视频分析消息拉取
        self.offlineMpt = OfflineMessagePollingThread('offline_thread', {'content': self.content,
                                                                         'topics': [self.content["kafka"]["topic"][
                                                                                        "dsp-alg-offline-tasks-topic"]]})
        # 开启监听线程
        self.onlineMpt.start()
        self.offlineMpt.start()