matrice-streaming 0.1.14__py3-none-any.whl → 0.1.65__py3-none-any.whl

matrice_streaming/streaming_gateway/camera_streamer/nvdec_worker_manager.py

@@ -0,0 +1,412 @@
+"""NVDEC Worker Manager for StreamingGateway integration.
+
+This module provides a simplified manager for the NVDEC hardware decoding backend.
+Unlike other backends, NVDEC uses static camera configuration at startup and outputs
+to CUDA IPC ring buffers (NV12 format) for zero-copy GPU inference pipelines.
+"""
+
+import logging
+import multiprocessing as mp
+import time
+from typing import Dict, List, Optional, Any
+
+from .nvdec import (
+    nvdec_pool_process,
+    StreamConfig,
+    CUPY_AVAILABLE,
+    PYNVCODEC_AVAILABLE,
+    RING_BUFFER_AVAILABLE,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def is_nvdec_available() -> bool:
+    """Check if NVDEC backend is available.
+
+    Requires:
+        - CuPy with CUDA support
+        - PyNvVideoCodec for NVDEC hardware decode
+        - cuda_shm_ring_buffer module for CUDA IPC
+    """
+    return CUPY_AVAILABLE and PYNVCODEC_AVAILABLE and RING_BUFFER_AVAILABLE
+
+
+def get_available_gpu_count() -> int:
+    """Detect the number of available CUDA GPUs.
+
+    Returns:
+        Number of available GPUs, or 1 if detection fails.
+    """
+    if not CUPY_AVAILABLE:
+        return 1
+
+    try:
+        import cupy as cp
+        return cp.cuda.runtime.getDeviceCount()
+    except Exception as e:
+        logger.warning(f"Failed to detect GPU count: {e}, defaulting to 1")
+        return 1
+
+
+class NVDECWorkerManager:
+    """Manager for NVDEC worker processes - static camera configuration.
+
+    This manager wraps the existing nvdec_pool_process function to integrate
+    with StreamingGateway. Key differences from other worker managers:
+
+    - Static camera configuration (no dynamic add/remove)
+    - Outputs to CUDA IPC ring buffers (not Redis/Kafka)
+    - NV12 format output (50% smaller than RGB)
+    - One worker process per GPU
+    """
+
+    def __init__(
+        self,
+        camera_configs: List[Dict[str, Any]],
+        stream_config: Dict[str, Any],  # Unused but kept for interface consistency
+        gpu_id: int = 0,
+        num_gpus: int = 0,  # 0 = auto-detect all available GPUs
+        nvdec_pool_size: int = 8,
+        nvdec_burst_size: int = 4,
+        frame_width: int = 640,
+        frame_height: int = 640,
+        num_slots: int = 32,
+        target_fps: int = 0,  # 0 = use per-camera FPS from config
+        duration_sec: float = 0,  # 0 = infinite
+    ):
+        """Initialize NVDEC Worker Manager.
+
+        Args:
+            camera_configs: List of camera configuration dicts with keys:
+                - camera_id or stream_key: Unique identifier (used for ring buffer naming)
+                - source: Video file path or RTSP URL
+                - width: Optional frame width (default: frame_width)
+                - height: Optional frame height (default: frame_height)
+                - fps: FPS limit for this camera (used by default)
+            stream_config: Stream configuration (unused, for interface consistency)
+            gpu_id: Primary GPU device ID (starting GPU for round-robin assignment)
+            num_gpus: Number of GPUs to use (0 = auto-detect all available GPUs)
+            nvdec_pool_size: Number of NVDEC decoders per GPU
+            nvdec_burst_size: Frames per stream before rotating to next
+            frame_width: Default output frame width (used if camera config doesn't specify)
+            frame_height: Default output frame height (used if camera config doesn't specify)
+            num_slots: Ring buffer slots per camera
+            target_fps: Global FPS override (0 = use per-camera FPS from config)
+            duration_sec: Duration to run (0 = infinite until stop)
+        """
+        if not is_nvdec_available():
+            raise RuntimeError(
+                "NVDEC not available. Requires CuPy, PyNvVideoCodec, and cuda_shm_ring_buffer"
+            )
+
+        self.camera_configs = camera_configs
+        self.stream_config = stream_config
+        self.gpu_id = gpu_id
+
+        # Auto-detect GPUs if num_gpus is 0
+        if num_gpus <= 0:
+            detected_gpus = get_available_gpu_count()
+            self.num_gpus = min(detected_gpus, 8)  # Max 8 GPUs
+            logger.info(f"Auto-detected {detected_gpus} GPU(s), using {self.num_gpus}")
+        else:
+            self.num_gpus = min(num_gpus, 8)  # Max 8 GPUs
+        self.nvdec_pool_size = nvdec_pool_size
+        self.nvdec_burst_size = nvdec_burst_size
+        self.frame_width = frame_width
+        self.frame_height = frame_height
+        self.num_slots = num_slots
+        self.target_fps = target_fps
+        self.duration_sec = duration_sec if duration_sec > 0 else float('inf')
+
+        self._workers: List[mp.Process] = []
+        self._stop_event: Optional[mp.Event] = None
+        self._result_queue: Optional[mp.Queue] = None
+        self._shared_frame_count: Optional[mp.Value] = None
+        self._gpu_frame_counts: Dict[int, mp.Value] = {}  # Per-GPU counters
+        self._start_time: Optional[float] = None
+        self._is_running = False
+
+        # Convert camera configs to StreamConfig objects and assign to GPUs
+        self._stream_configs: List[StreamConfig] = []
+        self._gpu_camera_assignments: Dict[int, List[StreamConfig]] = {
+            i: [] for i in range(self.num_gpus)
+        }
+        self._camera_to_gpu: Dict[str, int] = {}
+
+        self._prepare_camera_configs()
+
+        logger.info(
+            f"NVDECWorkerManager initialized: {len(camera_configs)} cameras, "
+            f"{self.num_gpus} GPU(s), pool_size={nvdec_pool_size}"
+        )
+
+    def _prepare_camera_configs(self):
+        """Convert dict configs to StreamConfig and distribute across GPUs.
+
+        Ring buffers are named using camera_id for SHM identification.
+        Per-camera FPS from config is used by default (target_fps=0 means use config FPS).
+        """
+        for i, config in enumerate(self.camera_configs):
+            # Extract camera ID (support both camera_id and stream_key)
+            # This ID is used for naming the CUDA IPC ring buffer
+            camera_id = config.get('camera_id') or config.get('stream_key') or f"cam_{i:04d}"
+
+            # Extract video source
+            source = config.get('source') or config.get('video_path')
+            if not source:
+                logger.warning(f"Camera {camera_id} has no source, skipping")
+                continue
+
+            # Extract dimensions (use per-camera config or fallback to defaults)
+            width = config.get('width') or self.frame_width
+            height = config.get('height') or self.frame_height
+
+            # Determine FPS: use global override if set, otherwise per-camera FPS from config
+            if self.target_fps > 0:
+                # Global FPS override is set
+                fps = self.target_fps
+            else:
+                # Use per-camera FPS from config (default streaming FPS)
+                fps = config.get('fps', 10)  # Default to 10 FPS if not specified
+
+            # Assign to GPU (round-robin starting from gpu_id)
+            gpu_id = (self.gpu_id + i) % self.num_gpus
+
+            stream_config = StreamConfig(
+                camera_id=camera_id,
+                video_path=source,
+                width=width,
+                height=height,
+                target_fps=fps,
+                gpu_id=gpu_id,
+            )
+
+            self._stream_configs.append(stream_config)
+            self._gpu_camera_assignments[gpu_id].append(stream_config)
+            self._camera_to_gpu[camera_id] = gpu_id
+
+            logger.debug(f"Camera {camera_id}: source={source}, {width}x{height}@{fps}fps, GPU{gpu_id}")
+
+    def start(self) -> None:
+        """Start NVDEC worker processes (one per GPU)."""
+        if self._is_running:
+            logger.warning("NVDECWorkerManager is already running")
+            return
+
+        if not self._stream_configs:
+            logger.warning("No cameras configured, nothing to start")
+            return
+
+        ctx = mp.get_context("spawn")
+        self._stop_event = ctx.Event()
+        self._result_queue = ctx.Queue()
+        self._shared_frame_count = ctx.Value('L', 0)  # Global counter (all GPUs)
+        self._start_time = time.perf_counter()
+
+        # Create per-GPU frame counters
+        self._gpu_frame_counts = {}
+        for gpu_id in range(self.num_gpus):
+            if self._gpu_camera_assignments[gpu_id]:  # Only if GPU has cameras
+                self._gpu_frame_counts[gpu_id] = ctx.Value('L', 0)
+
+        total_num_streams = len(self._stream_configs)
+        total_num_gpus = len([g for g in range(self.num_gpus) if self._gpu_camera_assignments[g]])
+
+        logger.info(f"Starting NVDEC: {total_num_streams} cameras across {total_num_gpus} GPUs")
+
+        # Start one process per GPU that has cameras
+        for gpu_id in range(self.num_gpus):
+            gpu_cameras = self._gpu_camera_assignments[gpu_id]
+            if not gpu_cameras:
+                continue
+
+            p = ctx.Process(
+                target=nvdec_pool_process,
+                args=(
+                    gpu_id,                     # process_id
+                    gpu_cameras,                # camera_configs (List[StreamConfig])
+                    self.nvdec_pool_size,       # pool_size
+                    self.duration_sec,          # duration_sec
+                    self._result_queue,         # result_queue
+                    self._stop_event,           # stop_event
+                    self.nvdec_burst_size,      # burst_size
+                    self.num_slots,             # num_slots
+                    self.target_fps,            # target_fps
+                    self._shared_frame_count,   # shared_frame_count (global)
+                    self._gpu_frame_counts,     # gpu_frame_counts (per-GPU dict)
+                    total_num_streams,          # total_num_streams
+                    total_num_gpus,             # total_num_gpus
+                ),
+                name=f"NVDECWorker-GPU{gpu_id}",
+                daemon=False,
+            )
+            p.start()
+            self._workers.append(p)
+            logger.info(f"Started NVDEC worker on GPU {gpu_id} with {len(gpu_cameras)} cameras")
+
+        self._is_running = True
+        logger.info(f"NVDECWorkerManager started: {len(self._workers)} workers")
+
+    def stop(self, timeout: float = 15.0) -> None:
+        """Stop all worker processes.
+
+        Args:
+            timeout: Maximum time to wait for each worker to stop gracefully
+        """
+        if not self._is_running:
+            logger.warning("NVDECWorkerManager is not running")
+            return
+
+        logger.info("Stopping NVDECWorkerManager...")
+
+        # Signal workers to stop
+        if self._stop_event:
+            self._stop_event.set()
+
+        # Wait for workers to finish
+        for p in self._workers:
+            p.join(timeout=timeout)
+            if p.is_alive():
+                logger.warning(f"Worker {p.name} did not stop gracefully, terminating")
+                p.terminate()
+                p.join(timeout=2.0)
+
+        self._workers.clear()
+        self._is_running = False
+        logger.info("NVDECWorkerManager stopped")
+
+    def get_worker_statistics(self) -> Dict[str, Any]:
+        """Return statistics from workers.
+
+        Returns:
+            Dict with keys:
+                - num_workers: Number of worker processes
+                - running_workers: Number of currently running workers
+                - total_cameras: Total cameras across all workers
+                - gpu_assignments: Cameras per GPU
+                - total_frames: Total frames processed (from shared counter)
+                - elapsed_sec: Time since start
+                - aggregate_fps: Overall FPS
+                - per_stream_fps: Average FPS per camera
+                - backend: 'nvdec'
+                - gpu_results: Per-GPU results from result queue
+        """
+        stats = {
+            'backend': 'nvdec',
+            'num_workers': len(self._workers),
+            'running_workers': sum(1 for p in self._workers if p.is_alive()),
+            'total_cameras': len(self._stream_configs),
+            'gpu_assignments': {
+                gpu_id: len(cameras)
+                for gpu_id, cameras in self._gpu_camera_assignments.items()
+            },
+            'nvdec_config': {
+                'gpu_id': self.gpu_id,
+                'num_gpus': self.num_gpus,
+                'pool_size': self.nvdec_pool_size,
+                'burst_size': self.nvdec_burst_size,
+                'frame_size': f"{self.frame_width}x{self.frame_height}",
+                'num_slots': self.num_slots,
+                'target_fps': self.target_fps,
+            },
+        }
+
+        # Add frame count and FPS
+        if self._shared_frame_count:
+            total_frames = self._shared_frame_count.value
+            stats['total_frames'] = total_frames
+
+            if self._start_time:
+                elapsed = time.perf_counter() - self._start_time
+                stats['elapsed_sec'] = elapsed
+                stats['aggregate_fps'] = total_frames / elapsed if elapsed > 0 else 0
+                stats['per_stream_fps'] = (
+                    stats['aggregate_fps'] / len(self._stream_configs)
+                    if self._stream_configs else 0
+                )
+
+        # Add per-GPU frame counts and FPS
+        if self._gpu_frame_counts and self._start_time:
+            elapsed = time.perf_counter() - self._start_time
+            gpu_stats = {}
+            for gpu_id, counter in self._gpu_frame_counts.items():
+                gpu_frames = counter.value
+                num_cams = len(self._gpu_camera_assignments.get(gpu_id, []))
+                gpu_fps = gpu_frames / elapsed if elapsed > 0 else 0
+                gpu_per_cam = gpu_fps / num_cams if num_cams > 0 else 0
+                gpu_stats[f'GPU{gpu_id}'] = {
+                    'frames': gpu_frames,
+                    'cameras': num_cams,
+                    'fps': gpu_fps,
+                    'fps_per_cam': gpu_per_cam,
+                }
+            stats['per_gpu_stats'] = gpu_stats
+
+        # Collect any available results from queue (non-blocking)
+        gpu_results = []
+        if self._result_queue:
+            while True:
+                try:
+                    result = self._result_queue.get_nowait()
+                    gpu_results.append(result)
+                except:
+                    break
+        stats['gpu_results'] = gpu_results
+
+        return stats
+
+    def get_camera_assignments(self) -> Dict[str, int]:
+        """Return mapping of camera_id to GPU ID.
+
+        Returns:
+            Dict mapping camera_id -> gpu_id
+        """
+        return self._camera_to_gpu.copy()
+
+    def add_camera(self, camera_config: Dict[str, Any]) -> bool:
+        """Not supported - NVDEC uses static camera configuration.
+
+        Raises:
+            NotImplementedError: Always raised
+        """
+        raise NotImplementedError(
+            "NVDEC backend uses static camera configuration. "
+            "Cameras must be configured at initialization."
+        )
+
+    def remove_camera(self, stream_key: str) -> bool:
+        """Not supported - NVDEC uses static camera configuration.
+
+        Raises:
+            NotImplementedError: Always raised
+        """
+        raise NotImplementedError(
+            "NVDEC backend uses static camera configuration. "
+            "Cameras cannot be removed at runtime."
+        )
+
+    def update_camera(self, camera_config: Dict[str, Any]) -> bool:
+        """Not supported - NVDEC uses static camera configuration.
+
+        Raises:
+            NotImplementedError: Always raised
+        """
+        raise NotImplementedError(
+            "NVDEC backend uses static camera configuration. "
+            "Cameras cannot be updated at runtime."
+        )
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the manager is currently running."""
+        return self._is_running
+
+    def __enter__(self):
+        """Context manager entry."""
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.stop()