video_monitor/app/analysis/inference_pool.py

"""共享推理子进程池 — 多路摄像头复用同一组 YOLO 引擎，避免每路重复加载模型。

阶段2架构：
- 主进程启动 N 个 InferenceWorker 子进程（默认 1，可配置 analysisInferenceWorkers）
- 子进程内按 algorithm_id 缓存引擎实例
- 请求/响应通过 multiprocessing.Queue 传递；帧以 JPEG 压缩跨进程（体积可控）
"""
import json
import logging
import multiprocessing as mp
import os
import queue
import threading
import time
import uuid

logger = logging.getLogger("analysis.inference_pool")

_POOL = None
_POOL_LOCK = threading.Lock()


def _resolve_model_path(model_file):
    from app.analysis.worker_pool import resolve_model_path
    return resolve_model_path(model_file)


def _inference_worker_loop(req_queue, resp_queue, worker_id):
    """子进程推理循环（不 import Django ORM）"""
    from app.utils.Logger import LOG_FORMAT
    if not logging.getLogger().handlers:
        logging.basicConfig(level=logging.INFO, format=LOG_FORMAT)
    engines = {}
    log = logging.getLogger("analysis.inference_worker.%s" % worker_id)
    log.info("推理 worker 启动")

    while True:
        try:
            msg = req_queue.get(timeout=1.0)
        except queue.Empty:
            continue
        if msg is None or msg.get("cmd") == "stop":
            break
        req_id = msg.get("req_id")
        try:
            algo = msg.get("algorithm") or {}
            jpeg = msg.get("jpeg")
            if not jpeg:
                resp_queue.put({"req_id": req_id, "ok": False, "error": "no frame"})
                continue
            import cv2
            import numpy as np
            arr = np.frombuffer(jpeg, dtype=np.uint8)
            frame = cv2.imdecode(arr, cv2.IMREAD_COLOR)
            if frame is None:
                resp_queue.put({"req_id": req_id, "ok": False, "error": "decode failed"})
                continue

            algo_id = algo.get("id", 0)
            key = (
                algo_id,
                algo.get("inference_engine", ""),
                algo.get("model_file", ""),
                float(algo.get("conf_threshold", 0.4)),
                float(algo.get("iou_threshold", 0.5)),
                int(algo.get("input_width", 640)),
                int(algo.get("input_height", 640)),
                algo.get("task_type", "detect"),
                algo.get("device", "cpu"),
                tuple(sorted(algo.get("target_labels") or [])),
            )
            eng = engines.get(key)
            if eng is None:
                from app.analysis.engines.factory import EngineFactory
                labels = algo.get("labels", [])
                if isinstance(labels, str):
                    try:
                        labels = json.loads(labels)
                    except Exception:
                        labels = []
                eng = EngineFactory.create(
                    algo.get("inference_engine", "yolo_pytorch"),
                    model_file=_resolve_model_path(algo.get("model_file", "")),
                    labels=labels,
                    input_size=(int(algo.get("input_width", 640)), int(algo.get("input_height", 640))),
                    conf_threshold=float(algo.get("conf_threshold", 0.4)),
                    iou_threshold=float(algo.get("iou_threshold", 0.5)),
                    algorithm_type=algo.get("algorithm_type", "yolo8"),
                    task_type=algo.get("task_type", "detect"),
                    device=algo.get("device", "cpu"),
                    target_labels=algo.get("target_labels") or [],
                )
                if not eng.load():
                    resp_queue.put({"req_id": req_id, "ok": False, "error": "engine load failed"})
                    continue
                engines[key] = eng

            results = eng.detect(frame)
            for r in results:
                r["algorithm_id"] = algo_id
                r["algorithm_name"] = algo.get("name", "")
            resp_queue.put({"req_id": req_id, "ok": True, "detections": results})
        except Exception as e:
            log.exception("推理异常: %s", e)
            resp_queue.put({"req_id": req_id, "ok": False, "error": str(e)})

    for eng in engines.values():
        try:
            if hasattr(eng, "unload"):
                eng.unload()
        except Exception:
            pass
    log.info("推理 worker 退出")


class InferenceProcessPool(object):
    """主进程侧推理池客户端"""

    def __init__(self, num_workers=1):
        self.num_workers = max(1, int(num_workers))
        self._req_queue = mp.Queue(maxsize=64)
        self._resp_queue = mp.Queue(maxsize=64)
        self._workers = []
        self._pending = {}
        self._lock = threading.Lock()
        self._running = False
        self._drain_thread = None
        self._timeout_count = 0
        self._last_timeout_ts = 0.0

    def start(self):
        if self._running:
            alive = sum(1 for p in self._workers if p.is_alive())
            if alive > 0:
                return
            logger.warning("InferenceProcessPool 标记运行中但 worker 已死，重新启动")
            self.stop()
        ctx = mp.get_context("spawn")
        self._req_queue = ctx.Queue(maxsize=64)
        self._resp_queue = ctx.Queue(maxsize=64)
        for i in range(self.num_workers):
            p = ctx.Process(
                target=_inference_worker_loop,
                args=(self._req_queue, self._resp_queue, i),
                name="inference-worker-%d" % i,
                daemon=True,
            )
            p.start()
            self._workers.append(p)
        self._running = True
        self._drain_thread = threading.Thread(target=self._drain_responses, name="inference-drain", daemon=True)
        self._drain_thread.start()
        logger.info("InferenceProcessPool 已启动 workers=%d", self.num_workers)

    def _drain_responses(self):
        while self._running:
            try:
                resp = self._resp_queue.get(timeout=0.5)
            except queue.Empty:
                continue
            req_id = resp.get("req_id")
            with self._lock:
                evt = self._pending.pop(req_id, None)
            if evt:
                evt["resp"] = resp
                evt["event"].set()

    def detect(self, frame, algorithm, timeout=30.0):
        """同步推理：将 frame + algorithm dict 发往子进程，等待结果"""
        if not self._running:
            self.start()
        try:
            import cv2
            ok, buf = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
            if not ok:
                return []
            jpeg = buf.tobytes()
        except Exception as e:
            logger.warning("帧编码失败: %s", e)
            return []
        return self.detect_jpeg(jpeg, algorithm, timeout=timeout)

    def _ensure_workers_alive(self):
        if not self._workers:
            if self._running:
                self.stop()
            self.start()
            return
        alive = sum(1 for p in self._workers if p.is_alive())
        if alive == 0:
            logger.warning("推理 worker 已全部退出，正在重启 InferenceProcessPool")
            self.stop()
            self.start()

    def detect_jpeg(self, jpeg, algorithm, timeout=30.0, raise_errors=False):
        """同步推理：直接传已编码的 JPEG bytes，跳过主进程 imencode/imdecode，
        消除 _inference_forwarder_loop 中的双重编解码与主进程 GIL 占用。"""
        self._ensure_workers_alive()
        if not self._running:
            self.start()
        if not jpeg:
            if raise_errors:
                raise ValueError("no frame")
            return []
        req_id = str(uuid.uuid4())
        evt = {"event": threading.Event(), "resp": None}
        with self._lock:
            self._pending[req_id] = evt

        algo_payload = algorithm if isinstance(algorithm, dict) else _algorithm_to_dict(algorithm)
        try:
            self._req_queue.put({"req_id": req_id, "algorithm": algo_payload, "jpeg": jpeg}, timeout=2.0)
        except Exception as e:
            with self._lock:
                self._pending.pop(req_id, None)
            logger.warning("推理请求入队失败: %s", e)
            if raise_errors:
                raise RuntimeError("推理请求入队失败: %s" % e) from e
            return []

        if not evt["event"].wait(timeout=timeout):
            with self._lock:
                self._pending.pop(req_id, None)
                self._timeout_count += 1
                self._last_timeout_ts = time.time()
            logger.warning("推理超时 req_id=%s (累计 %d)", req_id, self._timeout_count)
            if raise_errors:
                raise TimeoutError("模型推理超时（%ss）" % timeout)
            return []
        resp = evt.get("resp") or {}
        if not resp.get("ok"):
            logger.warning("推理失败: %s", resp.get("error"))
            if raise_errors:
                raise RuntimeError(resp.get("error") or "模型推理失败")
            return []
        return resp.get("detections") or []

    def status(self):
        self._ensure_workers_alive()
        alive = sum(1 for p in self._workers if p.is_alive())
        with self._lock:
            tc = self._timeout_count
            lts = self._last_timeout_ts
        degraded = alive > 0 and tc >= 3 and (time.time() - lts) < 120
        return {
            "workers": self.num_workers,
            "workers_alive": alive,
            "running": self._running,
            "timeout_count": tc,
            "inference_degraded": degraded,
        }

    def instance_count(self):
        """共享推理 worker 进程数（模型在 worker 内按需加载，非 worker_pool 计数）。"""
        return sum(1 for p in self._workers if p.is_alive())

    def stop(self):
        self._running = False
        for _ in self._workers:
            try:
                self._req_queue.put({"cmd": "stop"}, timeout=1.0)
            except Exception:
                pass
        for p in self._workers:
            try:
                p.join(timeout=3)
                if p.is_alive():
                    p.terminate()
            except Exception:
                pass
        self._workers = []
        logger.info("InferenceProcessPool 已停止")


def _algorithm_to_dict(a):
    labels = a.labels
    if isinstance(labels, str):
        try:
            labels = json.loads(labels)
        except Exception:
            labels = []
    return {
        "id": a.id,
        "name": a.name,
        "inference_engine": a.inference_engine,
        "model_file": a.model_file,
        "labels": labels,
        "input_width": a.input_width,
        "input_height": a.input_height,
        "conf_threshold": a.conf_threshold,
        "iou_threshold": a.iou_threshold,
        "algorithm_type": a.algorithm_type,
        "task_type": getattr(a, "task_type", "detect"),
        "device": getattr(a, "device", "cpu"),
    }


def get_inference_pool(num_workers=None):
    global _POOL
    with _POOL_LOCK:
        if _POOL is None:
            n = num_workers
            if n is None:
                try:
                    from app.utils.GlobalUtils import g_config
                    n = int(getattr(g_config, "analysisInferenceWorkers", 1))
                except Exception:
                    n = 1
            _POOL = InferenceProcessPool(num_workers=n)
            _POOL.start()
        return _POOL


def shutdown_inference_pool():
    global _POOL
    with _POOL_LOCK:
        if _POOL:
            _POOL.stop()
            _POOL = None