python-detector-worker/core/processes/communication.py

"""
Inter-Process Communication (IPC) system for session processes.
Defines message types and protocols for main ↔ session communication.
"""

import time
from enum import Enum
from typing import Dict, Any, Optional, Union
from dataclasses import dataclass, field
import numpy as np


class MessageType(Enum):
    """Message types for IPC communication."""

    # Commands: Main → Session
    INITIALIZE = "initialize"
    PROCESS_FRAME = "process_frame"
    SET_SESSION_ID = "set_session_id"
    SHUTDOWN = "shutdown"
    HEALTH_CHECK = "health_check"

    # Responses: Session → Main
    INITIALIZED = "initialized"
    DETECTION_RESULT = "detection_result"
    SESSION_SET = "session_set"
    SHUTDOWN_COMPLETE = "shutdown_complete"
    HEALTH_RESPONSE = "health_response"
    ERROR = "error"


@dataclass
class IPCMessage:
    """Base class for all IPC messages."""
    type: MessageType
    session_id: str
    timestamp: float = field(default_factory=time.time)
    message_id: str = field(default_factory=lambda: str(int(time.time() * 1000000)))


@dataclass
class InitializeCommand(IPCMessage):
    """Initialize session process with configuration."""
    subscription_config: Dict[str, Any] = field(default_factory=dict)
    model_config: Dict[str, Any] = field(default_factory=dict)



@dataclass
class ProcessFrameCommand(IPCMessage):
    """Process a frame through the detection pipeline."""
    frame: Optional[np.ndarray] = None
    display_id: str = ""
    subscription_identifier: str = ""
    frame_timestamp: float = 0.0



@dataclass
class SetSessionIdCommand(IPCMessage):
    """Set the session ID for the current session."""
    backend_session_id: str = ""
    display_id: str = ""



@dataclass
class ShutdownCommand(IPCMessage):
    """Shutdown the session process gracefully."""



@dataclass
class HealthCheckCommand(IPCMessage):
    """Check health status of session process."""



@dataclass
class InitializedResponse(IPCMessage):
    """Response indicating successful initialization."""
    success: bool = False
    error_message: Optional[str] = None



@dataclass
class DetectionResultResponse(IPCMessage):
    """Detection results from session process."""
    detections: Dict[str, Any] = field(default_factory=dict)
    processing_time: float = 0.0
    phase: str = ""  # "detection" or "processing"



@dataclass
class SessionSetResponse(IPCMessage):
    """Response confirming session ID was set."""
    success: bool = False
    backend_session_id: str = ""



@dataclass
class ShutdownCompleteResponse(IPCMessage):
    """Response confirming graceful shutdown."""



@dataclass
class HealthResponse(IPCMessage):
    """Health status response."""
    status: str = "unknown"  # "healthy", "degraded", "unhealthy"
    memory_usage_mb: float = 0.0
    cpu_percent: float = 0.0
    gpu_memory_mb: Optional[float] = None
    uptime_seconds: float = 0.0
    processed_frames: int = 0



@dataclass
class ErrorResponse(IPCMessage):
    """Error message from session process."""
    error_type: str = ""
    error_message: str = ""
    traceback: Optional[str] = None



# Type aliases for message unions
CommandMessage = Union[
    InitializeCommand,
    ProcessFrameCommand,
    SetSessionIdCommand,
    ShutdownCommand,
    HealthCheckCommand
]

ResponseMessage = Union[
    InitializedResponse,
    DetectionResultResponse,
    SessionSetResponse,
    ShutdownCompleteResponse,
    HealthResponse,
    ErrorResponse
]

IPCMessageUnion = Union[CommandMessage, ResponseMessage]


class MessageSerializer:
    """Handles serialization/deserialization of IPC messages."""

    @staticmethod
    def serialize_message(message: IPCMessageUnion) -> Dict[str, Any]:
        """
        Serialize message to dictionary for queue transport.

        Args:
            message: Message to serialize

        Returns:
            Dictionary representation of message
        """
        result = {
            'type': message.type.value,
            'session_id': message.session_id,
            'timestamp': message.timestamp,
            'message_id': message.message_id,
        }

        # Add specific fields based on message type
        if isinstance(message, InitializeCommand):
            result.update({
                'subscription_config': message.subscription_config,
                'model_config': message.model_config
            })
        elif isinstance(message, ProcessFrameCommand):
            result.update({
                'frame': message.frame,
                'display_id': message.display_id,
                'subscription_identifier': message.subscription_identifier,
                'frame_timestamp': message.frame_timestamp
            })
        elif isinstance(message, SetSessionIdCommand):
            result.update({
                'backend_session_id': message.backend_session_id,
                'display_id': message.display_id
            })
        elif isinstance(message, InitializedResponse):
            result.update({
                'success': message.success,
                'error_message': message.error_message
            })
        elif isinstance(message, DetectionResultResponse):
            result.update({
                'detections': message.detections,
                'processing_time': message.processing_time,
                'phase': message.phase
            })
        elif isinstance(message, SessionSetResponse):
            result.update({
                'success': message.success,
                'backend_session_id': message.backend_session_id
            })
        elif isinstance(message, HealthResponse):
            result.update({
                'status': message.status,
                'memory_usage_mb': message.memory_usage_mb,
                'cpu_percent': message.cpu_percent,
                'gpu_memory_mb': message.gpu_memory_mb,
                'uptime_seconds': message.uptime_seconds,
                'processed_frames': message.processed_frames
            })
        elif isinstance(message, ErrorResponse):
            result.update({
                'error_type': message.error_type,
                'error_message': message.error_message,
                'traceback': message.traceback
            })

        return result

    @staticmethod
    def deserialize_message(data: Dict[str, Any]) -> IPCMessageUnion:
        """
        Deserialize dictionary back to message object.

        Args:
            data: Dictionary representation

        Returns:
            Deserialized message object
        """
        msg_type = MessageType(data['type'])
        session_id = data['session_id']
        timestamp = data['timestamp']
        message_id = data['message_id']

        base_kwargs = {
            'session_id': session_id,
            'timestamp': timestamp,
            'message_id': message_id
        }

        if msg_type == MessageType.INITIALIZE:
            return InitializeCommand(
                type=msg_type,
                subscription_config=data['subscription_config'],
                model_config=data['model_config'],
                **base_kwargs
            )
        elif msg_type == MessageType.PROCESS_FRAME:
            return ProcessFrameCommand(
                type=msg_type,
                frame=data['frame'],
                display_id=data['display_id'],
                subscription_identifier=data['subscription_identifier'],
                frame_timestamp=data['frame_timestamp'],
                **base_kwargs
            )
        elif msg_type == MessageType.SET_SESSION_ID:
            return SetSessionIdCommand(
                backend_session_id=data['backend_session_id'],
                display_id=data['display_id'],
                **base_kwargs
            )
        elif msg_type == MessageType.SHUTDOWN:
            return ShutdownCommand(**base_kwargs)
        elif msg_type == MessageType.HEALTH_CHECK:
            return HealthCheckCommand(**base_kwargs)
        elif msg_type == MessageType.INITIALIZED:
            return InitializedResponse(
                type=msg_type,
                success=data['success'],
                error_message=data.get('error_message'),
                **base_kwargs
            )
        elif msg_type == MessageType.DETECTION_RESULT:
            return DetectionResultResponse(
                type=msg_type,
                detections=data['detections'],
                processing_time=data['processing_time'],
                phase=data['phase'],
                **base_kwargs
            )
        elif msg_type == MessageType.SESSION_SET:
            return SessionSetResponse(
                type=msg_type,
                success=data['success'],
                backend_session_id=data['backend_session_id'],
                **base_kwargs
            )
        elif msg_type == MessageType.SHUTDOWN_COMPLETE:
            return ShutdownCompleteResponse(type=msg_type, **base_kwargs)
        elif msg_type == MessageType.HEALTH_RESPONSE:
            return HealthResponse(
                type=msg_type,
                status=data['status'],
                memory_usage_mb=data['memory_usage_mb'],
                cpu_percent=data['cpu_percent'],
                gpu_memory_mb=data.get('gpu_memory_mb'),
                uptime_seconds=data.get('uptime_seconds', 0.0),
                processed_frames=data.get('processed_frames', 0),
                **base_kwargs
            )
        elif msg_type == MessageType.ERROR:
            return ErrorResponse(
                type=msg_type,
                error_type=data['error_type'],
                error_message=data['error_message'],
                traceback=data.get('traceback'),
                **base_kwargs
            )
        else:
            raise ValueError(f"Unknown message type: {msg_type}")