matrice-inference 0.1.2__py3-none-any.whl → 0.1.22__py3-none-any.whl

matrice_inference/server/stream/frame_cache.py

@@ -1,50 +1,80 @@
 import logging
-import threading
 import queue
-
+import threading
+from typing import Optional
 
 try:
     import redis  # type: ignore
-except Exception:  # pragma: no cover
+except ImportError:  # pragma: no cover
     redis = None  # type: ignore
 
 
 class RedisFrameCache:
-    """Non-blocking Redis cache for frames keyed by frame_id.
+    """Non-blocking Redis cache for frames with optimized resource management.
 
     Stores base64 string content under key 'stream:frames:{frame_id}' with field 'frame'.
     Each insert sets or refreshes the TTL.
     """
 
+    DEFAULT_TTL_SECONDS = 300
+    DEFAULT_MAX_QUEUE = 10000
+    DEFAULT_WORKER_THREADS = 2
+    DEFAULT_CONNECT_TIMEOUT = 2.0
+    DEFAULT_SOCKET_TIMEOUT = 0.5
+    DEFAULT_HEALTH_CHECK_INTERVAL = 30
+    DEFAULT_PREFIX = "stream:frames:"
+
     def __init__(
         self,
         host: str = "localhost",
         port: int = 6379,
         db: int = 0,
-        password: str = None,
-        username: str = None,
-        ttl_seconds: int = 300,
-        prefix: str = "stream:frames:",
-        max_queue: int = 10000,
-        worker_threads: int = 2,
-        connect_timeout: float = 2.0,
-        socket_timeout: float = 0.5,
+        password: Optional[str] = None,
+        username: Optional[str] = None,
+        ttl_seconds: int = DEFAULT_TTL_SECONDS,
+        prefix: str = DEFAULT_PREFIX,
+        max_queue: int = DEFAULT_MAX_QUEUE,
+        worker_threads: int = DEFAULT_WORKER_THREADS,
+        connect_timeout: float = DEFAULT_CONNECT_TIMEOUT,
+        socket_timeout: float = DEFAULT_SOCKET_TIMEOUT,
     ) -> None:
-        self.logger = logging.getLogger(__name__ + ".frame_cache")
-        self.ttl_seconds = int(ttl_seconds)
+        self.logger = logging.getLogger(f"{__name__}.frame_cache")
+        self.ttl_seconds = max(1, int(ttl_seconds))
         self.prefix = prefix
-        self.queue: "queue.Queue" = queue.Queue(maxsize=max_queue)
-        self.threads = []
         self.running = False
-        self._client = None
         self._worker_threads = max(1, int(worker_threads))
 
+        self.queue: queue.Queue = queue.Queue(maxsize=max_queue)
+        self.threads: list = []
+        self._client: Optional[redis.Redis] = None
+
+        if not self._is_redis_available():
+            return
+
+        self._client = self._create_redis_client(
+            host, port, db, password, username, connect_timeout, socket_timeout
+        )
+
+    def _is_redis_available(self) -> bool:
+        """Check if Redis package is available."""
         if redis is None:
             self.logger.warning("redis package not installed; frame caching disabled")
-            return
+            return False
+        return True
 
+    def _create_redis_client(
+        self,
+        host: str,
+        port: int,
+        db: int,
+        password: Optional[str],
+        username: Optional[str],
+        connect_timeout: float,
+        socket_timeout: float
+    ) -> Optional[redis.Redis]:
+        """Create Redis client with proper error handling."""
         try:
-            self._client = redis.Redis(
+            return redis.Redis(
                 host=host,
                 port=port,
                 db=db,
@@ -52,76 +82,141 @@ class RedisFrameCache:
                 username=username,
                 socket_connect_timeout=connect_timeout,
                 socket_timeout=socket_timeout,
-                health_check_interval=30,
+                health_check_interval=self.DEFAULT_HEALTH_CHECK_INTERVAL,
                 retry_on_timeout=True,
-                decode_responses=True,  # store strings directly
+                decode_responses=True,
             )
         except Exception as e:
-            self.logger.warning("Failed to init Redis client: %s", e)
-            self._client = None
+            self.logger.warning(f"Failed to initialize Redis client: {e}")
+            return None
 
     def start(self) -> None:
+        """Start the frame cache with worker threads."""
         if not self._client or self.running:
             return
+
         self.running = True
+        self._start_worker_threads()
+
+    def _start_worker_threads(self) -> None:
+        """Start worker threads for processing cache operations."""
         for i in range(self._worker_threads):
-            t = threading.Thread(target=self._worker, name=f"FrameCache-{i}", daemon=True)
-            t.start()
-            self.threads.append(t)
+            thread = threading.Thread(
+                target=self._worker,
+                name=f"FrameCache-{i}",
+                daemon=True
+            )
+            thread.start()
+            self.threads.append(thread)
 
     def stop(self) -> None:
+        """Stop the frame cache and cleanup resources."""
         if not self.running:
             return
+
         self.running = False
+        self._stop_worker_threads()
+        self.threads.clear()
+
+    def _stop_worker_threads(self) -> None:
+        """Stop all worker threads gracefully."""
+        # Signal threads to stop
         for _ in self.threads:
             try:
                 self.queue.put_nowait(None)
-            except Exception:
+            except queue.Full:
                 pass
-        for t in self.threads:
+
+        # Wait for threads to finish
+        for thread in self.threads:
             try:
-                t.join(timeout=2.0)
-            except Exception:
-                pass
-        self.threads.clear()
+                thread.join(timeout=2.0)
+            except Exception as e:
+                self.logger.warning(f"Error joining thread {thread.name}: {e}")
 
     def put(self, frame_id: str, base64_content: str) -> None:
         """Enqueue a cache write for the given frame.
 
-        - frame_id: unique identifier
-        - base64_content: base64-encoded image string
+        Args:
+            frame_id: unique identifier for the frame
+            base64_content: base64-encoded image string
         """
-        if not self._client or not self.running:
+        if not self._is_cache_ready():
             return
-        if not frame_id or not base64_content:
+
+        if not self._validate_input(frame_id, base64_content):
             return
+
         try:
             key = f"{self.prefix}{frame_id}"
             self.queue.put_nowait((key, base64_content))
         except queue.Full:
-            # Drop silently; never block pipeline
-            self.logger.debug("Frame cache queue full; dropping frame_id=%s", frame_id)
+            self._handle_queue_full(frame_id)
+
+    def _is_cache_ready(self) -> bool:
+        """Check if cache is ready for operations."""
+        return bool(self._client and self.running)
+
+    def _validate_input(self, frame_id: str, base64_content: str) -> bool:
+        """Validate input parameters."""
+        if not frame_id:
+            self.logger.warning("Empty frame_id provided")
+            return False
+        if not base64_content:
+            self.logger.warning("Empty base64_content provided")
+            return False
+        return True
+
+    def _handle_queue_full(self, frame_id: str) -> None:
+        """Handle queue full condition."""
+        self.logger.debug(f"Frame cache queue full; dropping frame_id={frame_id}")
 
     def _worker(self) -> None:
+        """Worker thread for processing cache operations."""
         while self.running:
-            try:
-                item = self.queue.get(timeout=0.5)
-            except queue.Empty:
-                continue
+            item = self._get_work_item()
             if item is None:
+                continue
+            if self._is_stop_signal(item):
                 break
+
+            self._process_cache_item(item)
+
+    def _get_work_item(self) -> Optional[tuple]:
+        """Get work item from queue with timeout."""
+        try:
+            return self.queue.get(timeout=0.5)
+        except queue.Empty:
+            return None
+
+    def _is_stop_signal(self, item: tuple) -> bool:
+        """Check if item is a stop signal."""
+        return item is None
+
+    def _process_cache_item(self, item: tuple) -> None:
+        """Process a single cache item."""
+        try:
             key, base64_content = item
-            try:
-                # Store base64 string in a Redis hash field 'frame', then set TTL
-                # Mimics the Go backend behavior
-                self._client.hset(key, "frame", base64_content)
-                self._client.expire(key, self.ttl_seconds)
-            except Exception as e:
-                self.logger.debug("Failed to cache frame %s: %s", key, e)
-            finally:
-                try:
-                    self.queue.task_done()
-                except Exception:
-                    pass
+            self._store_frame_data(key, base64_content)
+        except Exception as e:
+            self.logger.debug(f"Failed to process cache item: {e}")
+        finally:
+            self._mark_task_done()
+
+    def _store_frame_data(self, key: str, base64_content: str) -> None:
+        """Store frame data in Redis with TTL."""
+        try:
+            # Store base64 string in Redis hash field 'frame', then set TTL
+            self._client.hset(key, "frame", base64_content)
+            self._client.expire(key, self.ttl_seconds)
+        except Exception as e:
+            self.logger.debug(f"Failed to cache frame {key}: {e}")
+
+    def _mark_task_done(self) -> None:
+        """Mark queue task as done."""
+        try:
+            self.queue.task_done()
+        except Exception:
+            pass
 
 

matrice_inference/server/stream/inference_worker.py

@@ -1,18 +1,25 @@
 import asyncio
-import json
-import time
+import base64
 import logging
-import threading
 import queue
-from typing import Any, Dict
+import threading
+import time
 from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, Optional
 
 class InferenceWorker:
-    """Handles inference processing using threading."""
-    
-    def __init__(self, worker_id: int, inference_queue: queue.PriorityQueue,
-                 postproc_queue: queue.PriorityQueue, inference_executor: ThreadPoolExecutor,
-                 message_timeout: float, inference_timeout: float, inference_interface=None):
+    """Handles inference processing with clean resource management and error handling."""
+
+    def __init__(
+        self,
+        worker_id: int,
+        inference_queue: queue.PriorityQueue,
+        postproc_queue: queue.PriorityQueue,
+        inference_executor: ThreadPoolExecutor,
+        message_timeout: float,
+        inference_timeout: float,
+        inference_interface: Optional[Any] = None
+    ):
         self.worker_id = worker_id
         self.inference_queue = inference_queue
         self.postproc_queue = postproc_queue
@@ -23,10 +30,14 @@ class InferenceWorker:
         self.running = False
         self.logger = logging.getLogger(f"{__name__}.inference.{worker_id}")
     
-    def start(self):
+    def start(self) -> threading.Thread:
         """Start the inference worker in a separate thread."""
         self.running = True
-        thread = threading.Thread(target=self._run, name=f"InferenceWorker-{self.worker_id}", daemon=False)
+        thread = threading.Thread(
+            target=self._run,
+            name=f"InferenceWorker-{self.worker_id}",
+            daemon=False
+        )
         thread.start()
         return thread
     
@@ -34,72 +45,84 @@ class InferenceWorker:
         """Stop the inference worker."""
         self.running = False
     
-    def _run(self):
-        """Main inference dispatcher loop."""
+    def _run(self) -> None:
+        """Main inference dispatcher loop with proper error handling."""
         self.logger.info(f"Started inference worker {self.worker_id}")
-        
-        while self.running:
-            try:
-                # Get task from inference queue
-                try:
-                    priority, timestamp, task_data = self.inference_queue.get(timeout=self.message_timeout)
-                except queue.Empty:
-                    continue
-                
-                # Process inference task
-                self._process_inference_task(priority, task_data)
-                
-            except Exception as e:
-                self.logger.error(f"Inference worker error: {e}")
-        
-        self.logger.info(f"Inference worker {self.worker_id} stopped")
+
+        try:
+            while self.running:
+                task = self._get_task_from_queue()
+                if task:
+                    self._process_inference_task(*task)
+        except Exception as e:
+            self.logger.error(f"Fatal error in inference worker: {e}")
+        finally:
+            self.logger.info(f"Inference worker {self.worker_id} stopped")
+
+    def _get_task_from_queue(self) -> Optional[tuple]:
+        """Get task from inference queue with timeout handling."""
+        try:
+            return self.inference_queue.get(timeout=self.message_timeout)
+        except queue.Empty:
+            return None
+        except Exception as e:
+            self.logger.error(f"Error getting task from queue: {e}")
+            return None
     
-    def _process_inference_task(self, priority: int, task_data: Dict[str, Any]):
-        """Process a single inference task."""
+    def _process_inference_task(self, priority: int, timestamp: float, task_data: Dict[str, Any]) -> None:
+        """Process a single inference task with proper error handling."""
         try:
-            message = task_data["message"]
-            
-            # Submit to thread pool for async execution
+            if not self._validate_task_data(task_data):
+                return
+
             start_time = time.time()
-            future = self.inference_executor.submit(self._run_inference, task_data)
-            result = future.result(timeout=self.inference_timeout)
+            result = self._execute_inference(task_data)
             processing_time = time.time() - start_time
-            
+
             if result["success"]:
-                # Create post-processing task
-                postproc_task = {
-                    "original_message": message,
-                    "model_result": result["model_result"],
-                    "metadata": result["metadata"],
-                    "processing_time": processing_time,
-                    "input_stream": task_data["input_stream"],
-                    "stream_key": task_data["stream_key"],
-                    "camera_config": task_data["camera_config"]
-                }
-                
-                # Add to post-processing queue with timestamp as tie-breaker
+                postproc_task = self._create_postprocessing_task(
+                    task_data, result, processing_time
+                )
                 self.postproc_queue.put((priority, time.time(), postproc_task))
             else:
                 self.logger.error(f"Inference failed: {result['error']}")
-                
+
         except Exception as e:
             self.logger.error(f"Inference task error: {e}")
 
+    def _validate_task_data(self, task_data: Dict[str, Any]) -> bool:
+        """Validate that task data contains required fields."""
+        required_fields = ["message", "input_stream", "stream_key", "camera_config"]
+        for field in required_fields:
+            if field not in task_data:
+                self.logger.error(f"Missing required field '{field}' in task data")
+                return False
+        return True
+
+    def _execute_inference(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
+        """Execute inference task in thread pool."""
+        future = self.inference_executor.submit(self._run_inference, task_data)
+        return future.result(timeout=self.inference_timeout)
+
+    def _create_postprocessing_task(
+        self, task_data: Dict[str, Any], result: Dict[str, Any], processing_time: float
+    ) -> Dict[str, Any]:
+        """Create post-processing task from inference result."""
+        return {
+            "original_message": task_data["message"],
+            "model_result": result["model_result"],
+            "metadata": result["metadata"],
+            "processing_time": processing_time,
+            "input_stream": task_data["input_stream"],
+            "stream_key": task_data["stream_key"],
+            "camera_config": task_data["camera_config"]
+        }
+
     def _run_inference(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
-        """Run inference in thread pool."""
+        """Run inference in thread pool with proper error handling and cleanup."""
         try:
             # Extract task data - handle camera streamer format
             input_stream_data = task_data.get("input_stream", {})
-            input_content = input_stream_data.get("content")
-            
-            # Handle base64 encoded content from camera streamer
-            if input_content and isinstance(input_content, str):
-                import base64
-                try:
-                    input_content = base64.b64decode(input_content)
-                except Exception as e:
-                    logging.warning(f"Failed to decode base64 input: {str(e)}")
-            
             stream_key = task_data.get("stream_key")
             stream_info = input_stream_data.get("stream_info", {})
             camera_info = input_stream_data.get("camera_info", {})
@@ -123,41 +146,107 @@ class InferenceWorker:
                 else:
                     extra_params = {}
             
-            if self.inference_interface is None:
+            if not self.inference_interface:
                 raise ValueError("Inference interface not initialized")
-            
-            # Create event loop for this thread if it doesn't exist
-            try:
-                loop = asyncio.get_event_loop()
-            except RuntimeError:
-                loop = asyncio.new_event_loop()
-                asyncio.set_event_loop(loop)
-            
-            # Perform inference
+
+            inference_params = self._extract_inference_params(task_data)
+            loop = self._get_or_create_event_loop()
+
             model_result, metadata = loop.run_until_complete(
-                self.inference_interface.inference(
-                    input=input_content,
-                    extra_params=extra_params,
-                    apply_post_processing=False,  # Inference only
-                    stream_key=stream_key,
-                    stream_info=stream_info,
-                    camera_info=camera_info
-                )
+                self.inference_interface.inference(**inference_params)
             )
-            
-            return {
-                "model_result": model_result,
-                "metadata": metadata,
-                "success": True,
-                "error": None
-            }
-            
+
+            return self._create_success_result(model_result, metadata)
+
         except Exception as e:
-            logging.error(f"Inference worker error: {str(e)}", exc_info=True)
-            return {
-                "model_result": None,
-                "metadata": None,
-                "success": False,
-                "error": str(e)
-            }
+            self.logger.error(f"Inference execution error: {e}", exc_info=True)
+            return self._create_error_result(str(e))
+
+    def _extract_inference_params(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
+        """Extract and validate inference parameters from task data."""
+        input_stream_data = task_data.get("input_stream", {})
+        # Prefer decoded bytes if provided by upstream stages
+        input_bytes = task_data.get("decoded_input_bytes")
+        if not isinstance(input_bytes, (bytes, bytearray)):
+            content = input_stream_data.get("content")
+            if isinstance(content, str) and content:
+                try:
+                    input_bytes = base64.b64decode(content)
+                except Exception as e:
+                    self.logger.warning(f"Failed to decode base64 content for inference: {e}")
+                    input_bytes = None
+            elif isinstance(content, (bytes, bytearray)):
+                input_bytes = content
+            else:
+                input_bytes = None
+
+        extra_params = self._normalize_extra_params(task_data.get("extra_params", {}))
+
+        return {
+            "input": input_bytes,
+            "extra_params": extra_params,
+            "apply_post_processing": False,
+            "stream_key": task_data.get("stream_key"),
+            "stream_info": input_stream_data.get("stream_info", {}),
+            "camera_info": input_stream_data.get("camera_info", {})
+        }
+
+    def _decode_input_content(self, content: Any) -> Any:
+        """Decode base64 content if it's a string."""
+        if content and isinstance(content, str):
+            try:
+                return base64.b64decode(content)
+            except Exception as e:
+                self.logger.warning(f"Failed to decode base64 input: {e}")
+        return content
+
+    def _normalize_extra_params(self, extra_params: Any) -> Dict[str, Any]:
+        """Normalize extra_params to ensure it's a dictionary."""
+        if isinstance(extra_params, dict):
+            return extra_params
+        elif isinstance(extra_params, list):
+            return self._merge_list_params(extra_params)
+        else:
+            self.logger.warning(f"Invalid extra_params type {type(extra_params)}, using empty dict")
+            return {}
+
+    def _merge_list_params(self, params_list: list) -> Dict[str, Any]:
+        """Merge list of dictionaries into single dictionary."""
+        if not params_list:
+            return {}
+
+        if all(isinstance(item, dict) for item in params_list):
+            merged = {}
+            for item in params_list:
+                merged.update(item)
+            return merged
+
+        return {}
+
+    def _get_or_create_event_loop(self) -> asyncio.AbstractEventLoop:
+        """Get existing event loop or create a new one for this thread."""
+        try:
+            return asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            return loop
+
+    def _create_success_result(self, model_result: Any, metadata: Any) -> Dict[str, Any]:
+        """Create successful inference result."""
+        return {
+            "model_result": model_result,
+            "metadata": metadata,
+            "success": True,
+            "error": None
+        }
+
+    def _create_error_result(self, error_message: str) -> Dict[str, Any]:
+        """Create error inference result."""
+        return {
+            "model_result": None,
+            "metadata": None,
+            "success": False,
+            "error": error_message
+        }