matrice-inference 0.1.2__py3-none-any.whl

matrice_inference/server/stream/inference_worker.py

@@ -0,0 +1,163 @@
+import asyncio
+import json
+import time
+import logging
+import threading
+import queue
+from typing import Any, Dict
+from concurrent.futures import ThreadPoolExecutor
+
+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):
+        self.worker_id = worker_id
+        self.inference_queue = inference_queue
+        self.postproc_queue = postproc_queue
+        self.inference_executor = inference_executor
+        self.message_timeout = message_timeout
+        self.inference_timeout = inference_timeout
+        self.inference_interface = inference_interface
+        self.running = False
+        self.logger = logging.getLogger(f"{__name__}.inference.{worker_id}")
+    
+    def start(self):
+        """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.start()
+        return thread
+    
+    def stop(self):
+        """Stop the inference worker."""
+        self.running = False
+    
+    def _run(self):
+        """Main inference dispatcher loop."""
+        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")
+    
+    def _process_inference_task(self, priority: int, task_data: Dict[str, Any]):
+        """Process a single inference task."""
+        try:
+            message = task_data["message"]
+            
+            # Submit to thread pool for async execution
+            start_time = time.time()
+            future = self.inference_executor.submit(self._run_inference, task_data)
+            result = future.result(timeout=self.inference_timeout)
+            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
+                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 _run_inference(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
+        """Run inference in thread pool."""
+        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", {})
+            extra_params = task_data.get("extra_params", {})
+            
+            # Ensure extra_params is a dictionary
+            if not isinstance(extra_params, dict):
+                logging.warning(f"extra_params is not a dict in inference worker, converting from {type(extra_params)}: {extra_params}")
+                if isinstance(extra_params, list):
+                    # Convert list to dict if possible
+                    if len(extra_params) == 0:
+                        extra_params = {}
+                    elif all(isinstance(item, dict) for item in extra_params):
+                        # Merge all dictionaries in the list
+                        merged_params = {}
+                        for item in extra_params:
+                            merged_params.update(item)
+                        extra_params = merged_params
+                    else:
+                        extra_params = {}
+                else:
+                    extra_params = {}
+            
+            if self.inference_interface is None:
+                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
+            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
+                )
+            )
+            
+            return {
+                "model_result": model_result,
+                "metadata": metadata,
+                "success": True,
+                "error": None
+            }
+            
+        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)
+            }
+

matrice_inference/server/stream/post_processing_worker.py

@@ -0,0 +1,230 @@
+import asyncio
+import logging
+import threading
+import queue
+import time
+from typing import Any, Dict
+from concurrent.futures import ThreadPoolExecutor
+
+
+class PostProcessingWorker:
+    """Handles post-processing using threading."""
+    
+    def __init__(self, worker_id: int, postproc_queue: queue.PriorityQueue,
+                 output_queue: queue.PriorityQueue, postprocessing_executor: ThreadPoolExecutor,
+                 message_timeout: float, inference_timeout: float, post_processor=None,
+                 frame_cache=None):
+        self.worker_id = worker_id
+        self.postproc_queue = postproc_queue
+        self.output_queue = output_queue
+        self.postprocessing_executor = postprocessing_executor
+        self.message_timeout = message_timeout
+        self.inference_timeout = inference_timeout
+        self.post_processor = post_processor
+        self.frame_cache = frame_cache
+        self.running = False
+        self.logger = logging.getLogger(f"{__name__}.postproc.{worker_id}")
+    
+    def start(self):
+        """Start the post-processing worker in a separate thread."""
+        self.running = True
+        thread = threading.Thread(target=self._run, name=f"PostProcWorker-{self.worker_id}", daemon=False)
+        thread.start()
+        return thread
+    
+    def stop(self):
+        """Stop the post-processing worker."""
+        self.running = False
+    
+    def _run(self):
+        """Main post-processing dispatcher loop."""
+        self.logger.info(f"Started post-processing worker {self.worker_id}")
+        
+        while self.running:
+            try:
+                # Get task from post-processing queue
+                try:
+                    priority, timestamp, task_data = self.postproc_queue.get(timeout=self.message_timeout)
+                except queue.Empty:
+                    continue
+                
+                # Process post-processing task
+                self._process_postproc_task(priority, task_data)
+                
+            except Exception as e:
+                self.logger.error(f"Post-processing worker error: {e}")
+        
+        self.logger.info(f"Post-processing worker {self.worker_id} stopped")
+    
+    def _process_postproc_task(self, priority: int, task_data: Dict[str, Any]):
+        """Process a single post-processing task."""
+        try:
+            # Submit to thread pool for async execution
+            future = self.postprocessing_executor.submit(self._run_post_processing, task_data)
+            result = future.result(timeout=self.inference_timeout)
+            
+            if result["success"]:
+                # Cache disabled: preserving content in output and not pushing to Redis
+                # try:
+                #     orig_input = task_data.get("input_stream", {}) or {}
+                #     content = orig_input.get("content")
+                #     frame_id_for_cache = task_data.get("frame_id") or orig_input.get("frame_id")
+                #     if content and frame_id_for_cache and self.frame_cache:
+                #         if isinstance(content, bytes):
+                #             import base64
+                #             try:
+                #                 content = base64.b64encode(content).decode("ascii")
+                #             except Exception:
+                #                 content = None
+                #         if isinstance(content, str):
+                #             self.frame_cache.put(frame_id_for_cache, content)
+                # except Exception:
+                #     pass
+
+                # Create final output message
+                # Prepare input_stream for output: ensure frame_id is present and strip bulky content
+                safe_input_stream = {}
+                try:
+                    if isinstance(task_data.get("input_stream"), dict):
+                        safe_input_stream = dict(task_data["input_stream"])  # shallow copy
+                        # Ensure frame_id propagation
+                        if "frame_id" not in safe_input_stream and "frame_id" in task_data:
+                            safe_input_stream["frame_id"] = task_data["frame_id"]
+                        # Do not strip content; keep as-is in output
+                        # if "content" in safe_input_stream:
+                        #     safe_input_stream["content"] = ""
+                except Exception:
+                    safe_input_stream = task_data.get("input_stream", {})
+
+                # Determine frame_id for top-level convenience
+                frame_id = task_data.get("frame_id")
+                if not frame_id and isinstance(safe_input_stream, dict):
+                    frame_id = safe_input_stream.get("frame_id")
+
+                output_data = {
+                    "camera_id": task_data["original_message"].camera_id,
+                    "message_key": task_data["original_message"].message_key,
+                    "timestamp": task_data["original_message"].timestamp.isoformat(),
+                    "frame_id": frame_id,
+                    "model_result": task_data["model_result"],
+                    "input_stream": safe_input_stream,
+                    "post_processing_result": result["post_processing_result"],
+                    "processing_time_sec": task_data["processing_time"],
+                    "metadata": task_data.get("metadata", {})
+                }
+                
+                # Add to output queue
+                output_task = {
+                    "camera_id": task_data["original_message"].camera_id,
+                    "message_key": task_data["original_message"].message_key,
+                    "data": output_data,
+                }
+                # Add to output queue with timestamp as tie-breaker
+                self.output_queue.put((priority, time.time(), output_task))
+            else:
+                self.logger.error(f"Post-processing failed: {result['error']}")
+                
+        except Exception as e:
+            self.logger.error(f"Post-processing task error: {e}")
+
+    def _run_post_processing(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
+        """Run post-processing in thread pool."""
+        try:
+            if self.post_processor is None:
+                raise ValueError("Post processor not initialized")
+            
+            # Extract task data
+            model_result = task_data["model_result"]
+            input_stream_data = task_data.get("input_stream", {})
+            input_content = input_stream_data.get("content")
+            
+            # Handle base64 encoded content
+            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)}")
+                    input_content = None
+            
+            stream_key = task_data.get("stream_key")
+            stream_info = input_stream_data.get("stream_info", {})
+            camera_config = task_data.get("camera_config", {})
+            
+            # 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 post-processing
+            result = loop.run_until_complete(
+                self.post_processor.process(
+                    data=model_result,
+                    input_bytes=input_content if isinstance(input_content, bytes) else None,
+                    stream_key=stream_key,
+                    stream_info=stream_info
+                )
+            )
+            
+            # For face recognition use case, return empty raw results
+            processed_raw_results = []
+            try:
+                if hasattr(result, 'usecase') and not result.usecase == 'face_recognition':
+                    processed_raw_results = model_result
+            except Exception as e:
+                logging.warning(f"Failed to get processed raw results: {str(e)}")
+            
+            # Extract agg_summary from result data if available
+            agg_summary = {}
+            try:
+                if hasattr(result, 'data') and isinstance(result.data, dict):
+                    agg_summary = result.data.get("agg_summary", {})
+            except Exception as e:
+                logging.warning(f"Failed to get agg summary: {str(e)}")
+
+            # Format result similar to InferenceInterface
+            if result.is_success():
+                post_processing_result = {
+                    "status": "success",
+                    "processing_time": result.processing_time,
+                    "usecase": getattr(result, 'usecase', ''),
+                    "category": getattr(result, 'category', ''),
+                    "summary": getattr(result, 'summary', ''),
+                    "insights": getattr(result, 'insights', []),
+                    "metrics": getattr(result, 'metrics', {}),
+                    "predictions": getattr(result, 'predictions', []),
+                    "agg_summary": agg_summary,
+                    "raw_results": processed_raw_results,
+                    "stream_key": stream_key
+                }
+            else:
+                post_processing_result = {
+                    "status": "post_processing_failed",
+                    "error": result.error_message,
+                    "error_type": getattr(result, 'error_type', 'ProcessingError'),
+                    "processing_time": result.processing_time,
+                    "stream_key": stream_key,
+                    "agg_summary": agg_summary,
+                    "raw_results": model_result
+                }
+            
+            return {
+                "post_processing_result": post_processing_result,
+                "success": True,
+                "error": None
+            }
+            
+        except Exception as e:
+            logging.error(f"Post-processing worker error: {str(e)}", exc_info=True)
+            return {
+                "post_processing_result": {
+                    "status": "post_processing_failed",
+                    "error": str(e),
+                    "error_type": type(e).__name__,
+                    "stream_key": task_data.get("stream_key")
+                },
+                "success": False,
+                "error": str(e)
+            }

matrice_inference/server/stream/producer_worker.py

@@ -0,0 +1,147 @@
+import asyncio
+import json
+import time
+import logging
+import threading
+import queue
+from typing import Any, Dict
+from matrice_common.stream.matrice_stream import MatriceStream
+from matrice_inference.server.stream.utils import CameraConfig
+
+class ProducerWorker:
+    """Handles message production to streams."""
+    
+    def __init__(self, worker_id: int, output_queue: queue.PriorityQueue,
+                 camera_configs: Dict[str, CameraConfig], message_timeout: float):
+        self.worker_id = worker_id
+        self.output_queue = output_queue
+        self.camera_configs = camera_configs
+        self.message_timeout = message_timeout
+        self.running = False
+        self.producer_streams = {}  # Will be created in worker thread's event loop
+        self.logger = logging.getLogger(f"{__name__}.producer.{worker_id}")
+    
+    def start(self):
+        """Start the producer worker in a separate thread."""
+        self.running = True
+        thread = threading.Thread(target=self._run, name=f"ProducerWorker-{self.worker_id}", daemon=False)
+        thread.start()
+        return thread
+    
+    def stop(self):
+        """Stop the producer worker."""
+        self.running = False
+    
+    def _run(self):
+        """Main producer loop."""
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        
+        self.logger.info(f"Started producer worker {self.worker_id}")
+        
+        try:
+            # Initialize streams for all cameras in this event loop
+            loop.run_until_complete(self._initialize_streams())
+            
+            while self.running:
+                try:
+                    # Get task from output queue
+                    try:
+                        priority, timestamp, task_data = self.output_queue.get(timeout=self.message_timeout)
+                    except queue.Empty:
+                        continue
+                    
+                    # Send message to stream
+                    loop.run_until_complete(self._send_message_safely(task_data))
+                    
+                except Exception as e:
+                    self.logger.error(f"Producer error: {e}")
+                    time.sleep(0.1)
+        
+        finally:
+            # Clean up streams
+            for stream in self.producer_streams.values():
+                try:
+                    loop.run_until_complete(stream.async_close())
+                except Exception as e:
+                    self.logger.error(f"Error closing producer stream: {e}")
+            loop.close()
+            self.logger.info(f"Producer worker {self.worker_id} stopped")
+
+    async def _initialize_streams(self):
+        """Initialize producer streams for all cameras in the current event loop."""
+        try:
+            from matrice_common.stream.matrice_stream import MatriceStream, StreamType
+            
+            for camera_id, camera_config in self.camera_configs.items():
+                try:
+                    stream_config = camera_config.stream_config
+                    output_topic = camera_config.output_topic
+                    
+                    # Determine stream type
+                    stream_type = StreamType.KAFKA if stream_config.get("stream_type", "kafka").lower() == "kafka" else StreamType.REDIS
+                    
+                    # Create stream configuration
+                    if stream_type == StreamType.KAFKA:
+                        stream_params = {
+                            "bootstrap_servers": stream_config.get("bootstrap_servers", "localhost:9092"),
+                            "sasl_username": stream_config.get("sasl_username", "matrice-sdk-user"),
+                            "sasl_password": stream_config.get("sasl_password", "matrice-sdk-password"),
+                            "sasl_mechanism": stream_config.get("sasl_mechanism", "SCRAM-SHA-256"),
+                            "security_protocol": stream_config.get("security_protocol", "SASL_PLAINTEXT"),
+                        }
+                    else:  # Redis
+                        stream_params = {
+                            "host": stream_config.get("host", "localhost"),
+                            "port": stream_config.get("port", 6379),
+                            "password": stream_config.get("password"),
+                            "username": stream_config.get("username"),
+                            "db": stream_config.get("db", 0),
+                        }
+                    
+                    # Create and setup producer stream
+                    producer_stream = MatriceStream(stream_type, **stream_params)
+                    await producer_stream.async_setup(output_topic)
+                    self.producer_streams[camera_id] = producer_stream
+                    
+                    self.logger.info(f"Initialized {stream_type.value} producer stream for camera {camera_id} in worker {self.worker_id}")
+                    
+                except Exception as e:
+                    self.logger.error(f"Failed to initialize producer stream for camera {camera_id}: {e}")
+                    continue
+                    
+        except Exception as e:
+            self.logger.error(f"Failed to initialize producer streams: {e}")
+            raise
+    
+    async def _send_message_safely(self, task_data: Dict[str, Any]):
+        """Send message to the appropriate stream safely."""
+        try:
+            camera_id = task_data["camera_id"]
+            message_key = task_data["message_key"]
+            data = task_data["data"]
+            
+            # Check if camera and stream still exist
+            if camera_id not in self.producer_streams or camera_id not in self.camera_configs:
+                self.logger.warning(f"Camera {camera_id} not found in producer streams or configs")
+                return
+            
+            camera_config = self.camera_configs[camera_id]
+            if not camera_config.enabled:
+                self.logger.debug(f"Camera {camera_id} is disabled, skipping message")
+                return
+            
+            # Get producer stream for camera
+            producer_stream = self.producer_streams[camera_id]
+            output_topic = camera_config.output_topic
+            
+            # Send message to stream
+            await producer_stream.async_add_message(
+                output_topic,
+                json.dumps(data),
+                key=message_key
+            )
+            
+        except Exception as e:
+            self.logger.error(f"Error sending message: {e}")
+