security-ai-edge/core/tensorrt_engine.py

"""
TensorRT推理引擎模块
实现引擎加载、显存优化、异步推理、性能监控
"""

import logging
import threading
import time
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

try:
    import tensorrt as trt
    TRT_AVAILABLE = True
except ImportError:
    TRT_AVAILABLE = False
    trt = None

from config.settings import get_settings, InferenceConfig
from utils.logger import get_logger

logger = logging.getLogger(__name__)


class TensorRTEngine:
    """TensorRT引擎管理类
    
    实现engine文件加载、显存管理、异步推理
    """
    
    def __init__(self, config: Optional[InferenceConfig] = None):
        """
        初始化TensorRT引擎
        
        Args:
            config: 推理配置
        """
        if not TRT_AVAILABLE:
            raise RuntimeError("TensorRT未安装，请先安装tensorrt库")
        
        if config is None:
            settings = get_settings()
            config = settings.inference
        
        self.config = config
        self._engine = None
        self._context = None
        self._input_binding = None
        self._output_bindings = []
        self._stream = None
        
        self._logger = get_logger("tensorrt")
        self._lock = threading.Lock()
        
        self._memory_pool: Dict[str, np.ndarray] = {}
        
        self._performance_stats = {
            "inference_count": 0,
            "total_inference_time_ms": 0.0,
            "avg_inference_time_ms": 0.0,
            "throughput_fps": 0.0,
            "last_inference_time_ms": 0.0,
        }
        
        self._logger.info(
            f"TensorRT引擎初始化配置: "
            f"模型={config.model_path}, "
            f"输入尺寸={config.input_width}x{config.input_height}, "
            f"Batch={config.batch_size}, "
            f"FP16={config.fp16_mode}"
        )
    
    def load_engine(self, engine_path: Optional[str] = None) -> bool:
        """
        加载TensorRT engine文件
        
        Args:
            engine_path: engine文件路径
        
        Returns:
            是否加载成功
        """
        if engine_path is None:
            engine_path = self.config.model_path
        
        with self._lock:
            try:
                if self._context is not None:
                    self._release_resources()
                
                TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
                
                with open(engine_path, "rb") as f:
                    runtime = trt.Runtime(TRT_LOGGER)
                    self._engine = runtime.deserialize_cuda_engine(f.read())
                
                self._context = self._engine.create_execution_context()
                
                self._setup_bindings()
                
                self._allocate_memory_pool()
                
                self._logger.log_connection_event(
                    "load", "TensorRT", engine_path, True
                )
                self._logger.info(f"TensorRT引擎加载成功: {engine_path}")
                
                return True
                
            except Exception as e:
                self._logger.error(f"TensorRT引擎加载失败: {e}")
                return False
    
    def _setup_bindings(self):
        """设置输入输出绑定"""
        self._input_binding = None
        self._output_bindings = []
        
        for i in range(self._engine.num_bindings):
            binding_name = self._engine.get_binding_name(i)
            binding_shape = self._engine.get_binding_shape(i)
            binding_dtype = self._engine.get_binding_dtype(i)
            
            if self._engine.binding_is_input(i):
                self._input_binding = {
                    "name": binding_name,
                    "shape": binding_shape,
                    "dtype": binding_dtype,
                    "index": i,
                }
            else:
                self._output_bindings.append({
                    "name": binding_name,
                    "shape": binding_shape,
                    "dtype": binding_dtype,
                    "index": i,
                })
        
        self._logger.debug(
            f"输入绑定: {self._input_binding}, "
            f"输出绑定: {len(self._output_bindings)}"
        )
    
    def _allocate_memory_pool(self):
        """分配显存池"""
        self._memory_pool.clear()
        
        if self._input_binding:
            shape = self._input_binding["shape"]
            shape = tuple(max(1, s) if s < 0 else s for s in shape)
            dtype = self._get_numpy_dtype(self._input_binding["dtype"])
            self._memory_pool["input"] = np.zeros(shape, dtype=dtype)
        
        for output in self._output_bindings:
            shape = output["shape"]
            shape = tuple(max(1, s) if s < 0 else s for s in shape)
            dtype = self._get_numpy_dtype(output["dtype"])
            self._memory_pool[output["name"]] = np.zeros(shape, dtype=dtype)
    
    def _get_numpy_dtype(self, trt_dtype) -> np.dtype:
        """转换TensorRT数据类型到numpy"""
        if trt_dtype == trt.float16:
            return np.float16
        elif trt_dtype == trt.float32:
            return np.float32
        elif trt_dtype == trt.int32:
            return np.int32
        elif trt_dtype == trt.int8:
            return np.int8
        else:
            return np.float32
    
    def _allocate_device_memory(self, batch_size: int) -> Tuple[np.ndarray, List[np.ndarray]]:
        """
        分配设备显存
        
        Returns:
            tuple: (输入数据, 输出数据列表)
        """
        input_shape = list(self._input_binding["shape"])
        input_shape[0] = batch_size
        
        input_data = np.zeros(input_shape, dtype=np.float16 if self.config.fp16_mode else np.float32)
        
        output_data_list = []
        for output in self._output_bindings:
            output_shape = list(output["shape"])
            output_shape[0] = batch_size
            output_data = np.zeros(output_shape, dtype=self._get_numpy_dtype(output["dtype"]))
            output_data_list.append(output_data)
        
        return input_data, output_data_list
    
    def set_input_shape(self, batch_size: int, height: int, width: int):
        """
        动态设置输入形状
        
        Args:
            batch_size: 批次大小
            height: 输入高度
            width: 输入宽度
        """
        if self._context is None:
            raise RuntimeError("引擎未加载")
        
        self._context.set_input_shape(
            self._input_binding["name"],
            [batch_size, 3, height, width]
        )
        
        self._logger.debug(f"输入形状已设置为: [{batch_size}, 3, {height}, {width}]")
    
    def infer(
        self,
        input_data: np.ndarray,
        async_mode: bool = False
    ) -> Tuple[List[np.ndarray], float]:
        """
        执行推理
        
        Args:
            input_data: 输入数据 (NCHW格式)
            async_mode: 是否使用异步模式
        
        Returns:
            tuple: (输出列表, 推理耗时ms)
        """
        if self._engine is None or self._context is None:
            raise RuntimeError("引擎未加载")
        
        start_time = time.perf_counter()
        
        batch_size = input_data.shape[0]
        
        input_data = input_data.astype(np.float16 if self.config.fp16_mode else np.float32)
        
        self._context.set_input_shape(
            self._input_binding["name"],
            input_data.shape
        )
        
        input_tensor = input_data
        output_tensors = []
        
        for output in self._output_bindings:
            output_shape = list(output["shape"])
            output_shape[0] = batch_size
            output_tensor = np.zeros(output_shape, dtype=self._get_numpy_dtype(output["dtype"]))
            output_tensors.append(output_tensor)
        
        bindings = [input_tensor] + output_tensors
        
        self._context.execute_v2(bindings=bindings)
        
        inference_time_ms = (time.perf_counter() - start_time) * 1000
        
        self._update_performance_stats(inference_time_ms, batch_size)
        
        return output_tensors, inference_time_ms
    
    def infer_async(self, input_data: np.ndarray) -> Tuple[List[np.ndarray], float]:
        """
        执行异步推理
        
        Args:
            input_data: 输入数据
        
        Returns:
            tuple: (输出列表, 推理耗时ms)
        """
        return self.infer(input_data, async_mode=True)
    
    def infer_batch(
        self,
        batch_data: np.ndarray,
        batch_size: int
    ) -> Tuple[List[np.ndarray], float]:
        """
        推理批次数据
        
        Args:
            batch_data: 批次数据
            batch_size: 实际批次大小
        
        Returns:
            tuple: (输出列表, 推理耗时ms)
        """
        if batch_data.shape[0] != batch_size:
            batch_data = batch_data[:batch_size]
        
        return self.infer(batch_data)
    
    def _update_performance_stats(self, inference_time_ms: float, batch_size: int):
        """更新性能统计"""
        stats = self._performance_stats
        
        stats["inference_count"] += 1
        stats["total_inference_time_ms"] += inference_time_ms
        stats["last_inference_time_ms"] = inference_time_ms
        stats["avg_inference_time_ms"] = (
            stats["total_inference_time_ms"] / stats["inference_count"]
        )
        
        total_time = stats["total_inference_time_ms"] / 1000.0
        if total_time > 0:
            stats["throughput_fps"] = stats["inference_count"] / total_time
    
    def get_performance_stats(self) -> Dict[str, Any]:
        """获取性能统计"""
        stats = dict(self._performance_stats)
        stats["memory_usage"] = self.get_memory_usage()
        return stats
    
    def get_memory_usage(self) -> Dict[str, float]:
        """获取显存使用情况"""
        try:
            import pynvml
            pynvml.nvmlInit()
            handle = pynvml.nvmlDeviceGetHandleByIndex(self.config.device_id)
            mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            pynvml.nvmlShutdown()
            
            return {
                "total_mb": mem_info.total / (1024 ** 2),
                "used_mb": mem_info.used / (1024 ** 2),
                "free_mb": mem_info.free / (1024 ** 2),
            }
        except Exception as e:
            self._logger.warning(f"获取显存信息失败: {e}")
            return {"total_mb": 0, "used_mb": 0, "free_mb": 0}
    
    def _release_resources(self):
        """释放资源（Python TensorRT 由 GC 管理，无需 destroy）"""
        if self._stream:
            try:
                self._stream.synchronize()
            except Exception:
                pass
            self._stream = None
        
        if self._context:
            self._context = None
        
        if self._engine:
            self._engine = None
        
        self._memory_pool.clear()
    
    def release(self):
        """释放引擎资源"""
        with self._lock:
            self._release_resources()
            self._logger.info("TensorRT引擎资源已释放")
    
    def __del__(self):
        """析构函数"""
        self.release()


class EngineManager:
    """引擎管理器类
    
    管理多个TensorRT引擎实例
    """
    
    def __init__(self):
        self._engines: Dict[str, TensorRTEngine] = {}
        self._logger = get_logger("tensorrt")
        self._lock = threading.Lock()
    
    def load_engine(
        self,
        engine_id: str,
        engine_path: str,
        config: Optional[InferenceConfig] = None
    ) -> bool:
        """
        加载引擎
        
        Args:
            engine_id: 引擎标识
            engine_path: engine文件路径
            config: 推理配置
        
        Returns:
            是否加载成功
        """
        with self._lock:
            if engine_id in self._engines:
                self._engines[engine_id].release()
            
            engine = TensorRTEngine(config)
            
            if engine.load_engine(engine_path):
                self._engines[engine_id] = engine
                self._logger.info(f"引擎已加载: {engine_id}")
                return True
            else:
                return False
    
    def get_engine(self, engine_id: str) -> Optional[TensorRTEngine]:
        """获取引擎"""
        return self._engines.get(engine_id)
    
    def release_engine(self, engine_id: str):
        """释放引擎"""
        with self._lock:
            if engine_id in self._engines:
                self._engines[engine_id].release()
                del self._engines[engine_id]
                self._logger.info(f"引擎已释放: {engine_id}")
    
    def release_all(self):
        """释放所有引擎"""
        with self._lock:
            for engine_id in list(self._engines.keys()):
                self._engines[engine_id].release()
            self._engines.clear()
            self._logger.info("所有引擎已释放")
    
    def get_all_stats(self) -> Dict[str, Any]:
        """获取所有引擎统计"""
        with self._lock:
            return {
                engine_id: engine.get_performance_stats()
                for engine_id, engine in self._engines.items()
            }


def create_tensorrt_engine(
    engine_path: str,
    config: Optional[InferenceConfig] = None
) -> TensorRTEngine:
    """
    创建TensorRT引擎的便捷函数
    
    Args:
        engine_path: engine文件路径
        config: 推理配置
    
    Returns:
        TensorRTEngine实例
    """
    engine = TensorRTEngine(config)
    if engine.load_engine(engine_path):
        return engine
    else:
        raise RuntimeError(f"无法加载TensorRT引擎: {engine_path}")