refactor: done phase 1
This commit is contained in:
ziesorx
parent
f7c464be21
commit
cbbed3d933
13 changed files with 1084 additions and 891 deletions
219
core/communication/state.py
Normal file
219
core/communication/state.py
Normal file
|
|
@ -0,0 +1,219 @@
|
|||
"""
|
||||
Worker state management for system metrics and subscription tracking.
|
||||
"""
|
||||
import logging
|
||||
import psutil
|
||||
import threading
|
||||
from typing import Dict, Set, Optional, List
|
||||
from dataclasses import dataclass, field
|
||||
from .models import CameraConnection, SubscriptionObject
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Try to import torch for GPU monitoring
|
||||
try:
|
||||
import torch
|
||||
TORCH_AVAILABLE = True
|
||||
except ImportError:
|
||||
TORCH_AVAILABLE = False
|
||||
logger.warning("PyTorch not available, GPU metrics will not be collected")
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkerState:
|
||||
"""Central state management for the detector worker."""
|
||||
|
||||
# Active subscriptions indexed by subscription identifier
|
||||
subscriptions: Dict[str, SubscriptionObject] = field(default_factory=dict)
|
||||
|
||||
# Session ID mapping: display_identifier -> session_id
|
||||
session_ids: Dict[str, int] = field(default_factory=dict)
|
||||
|
||||
# Progression stage mapping: display_identifier -> stage
|
||||
progression_stages: Dict[str, str] = field(default_factory=dict)
|
||||
|
||||
# Active camera connections for state reporting
|
||||
camera_connections: List[CameraConnection] = field(default_factory=list)
|
||||
|
||||
# Thread lock for state synchronization
|
||||
_lock: threading.RLock = field(default_factory=threading.RLock)
|
||||
|
||||
def set_subscriptions(self, new_subscriptions: List[SubscriptionObject]) -> None:
|
||||
"""
|
||||
Update active subscriptions with declarative list from backend.
|
||||
|
||||
Args:
|
||||
new_subscriptions: Complete list of desired subscriptions
|
||||
"""
|
||||
with self._lock:
|
||||
# Convert to dict for easy lookup
|
||||
new_sub_dict = {sub.subscriptionIdentifier: sub for sub in new_subscriptions}
|
||||
|
||||
# Log changes for debugging
|
||||
current_ids = set(self.subscriptions.keys())
|
||||
new_ids = set(new_sub_dict.keys())
|
||||
|
||||
added = new_ids - current_ids
|
||||
removed = current_ids - new_ids
|
||||
updated = current_ids & new_ids
|
||||
|
||||
if added:
|
||||
logger.info(f"Adding subscriptions: {added}")
|
||||
if removed:
|
||||
logger.info(f"Removing subscriptions: {removed}")
|
||||
if updated:
|
||||
logger.info(f"Updating subscriptions: {updated}")
|
||||
|
||||
# Replace entire subscription dict
|
||||
self.subscriptions = new_sub_dict
|
||||
|
||||
# Update camera connections for state reporting
|
||||
self._update_camera_connections()
|
||||
|
||||
def get_subscription(self, subscription_identifier: str) -> Optional[SubscriptionObject]:
|
||||
"""Get subscription by identifier."""
|
||||
with self._lock:
|
||||
return self.subscriptions.get(subscription_identifier)
|
||||
|
||||
def get_all_subscriptions(self) -> List[SubscriptionObject]:
|
||||
"""Get all active subscriptions."""
|
||||
with self._lock:
|
||||
return list(self.subscriptions.values())
|
||||
|
||||
def set_session_id(self, display_identifier: str, session_id: Optional[int]) -> None:
|
||||
"""
|
||||
Set or clear session ID for a display.
|
||||
|
||||
Args:
|
||||
display_identifier: Display identifier
|
||||
session_id: Session ID to set, or None to clear
|
||||
"""
|
||||
with self._lock:
|
||||
if session_id is None:
|
||||
self.session_ids.pop(display_identifier, None)
|
||||
logger.info(f"Cleared session ID for display {display_identifier}")
|
||||
else:
|
||||
self.session_ids[display_identifier] = session_id
|
||||
logger.info(f"Set session ID {session_id} for display {display_identifier}")
|
||||
|
||||
def get_session_id(self, display_identifier: str) -> Optional[int]:
|
||||
"""Get session ID for display identifier."""
|
||||
with self._lock:
|
||||
return self.session_ids.get(display_identifier)
|
||||
|
||||
def get_session_id_for_subscription(self, subscription_identifier: str) -> Optional[int]:
|
||||
"""Get session ID for subscription by extracting display identifier."""
|
||||
from .messages import extract_display_identifier
|
||||
|
||||
display_id = extract_display_identifier(subscription_identifier)
|
||||
if display_id:
|
||||
return self.get_session_id(display_id)
|
||||
return None
|
||||
|
||||
def set_progression_stage(self, display_identifier: str, stage: Optional[str]) -> None:
|
||||
"""
|
||||
Set or clear progression stage for a display.
|
||||
|
||||
Args:
|
||||
display_identifier: Display identifier
|
||||
stage: Progression stage to set, or None to clear
|
||||
"""
|
||||
with self._lock:
|
||||
if stage is None:
|
||||
self.progression_stages.pop(display_identifier, None)
|
||||
logger.info(f"Cleared progression stage for display {display_identifier}")
|
||||
else:
|
||||
self.progression_stages[display_identifier] = stage
|
||||
logger.info(f"Set progression stage '{stage}' for display {display_identifier}")
|
||||
|
||||
def get_progression_stage(self, display_identifier: str) -> Optional[str]:
|
||||
"""Get progression stage for display identifier."""
|
||||
with self._lock:
|
||||
return self.progression_stages.get(display_identifier)
|
||||
|
||||
def _update_camera_connections(self) -> None:
|
||||
"""Update camera connections list for state reporting."""
|
||||
connections = []
|
||||
|
||||
for sub in self.subscriptions.values():
|
||||
connection = CameraConnection(
|
||||
subscriptionIdentifier=sub.subscriptionIdentifier,
|
||||
modelId=sub.modelId,
|
||||
modelName=sub.modelName,
|
||||
online=True, # TODO: Add actual online status tracking
|
||||
cropX1=sub.cropX1,
|
||||
cropY1=sub.cropY1,
|
||||
cropX2=sub.cropX2,
|
||||
cropY2=sub.cropY2
|
||||
)
|
||||
connections.append(connection)
|
||||
|
||||
self.camera_connections = connections
|
||||
|
||||
def get_camera_connections(self) -> List[CameraConnection]:
|
||||
"""Get current camera connections for state reporting."""
|
||||
with self._lock:
|
||||
return self.camera_connections.copy()
|
||||
|
||||
|
||||
class SystemMetrics:
|
||||
"""System metrics collection for state reporting."""
|
||||
|
||||
@staticmethod
|
||||
def get_cpu_usage() -> float:
|
||||
"""Get current CPU usage percentage."""
|
||||
try:
|
||||
return psutil.cpu_percent(interval=0.1)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get CPU usage: {e}")
|
||||
return 0.0
|
||||
|
||||
@staticmethod
|
||||
def get_memory_usage() -> float:
|
||||
"""Get current memory usage percentage."""
|
||||
try:
|
||||
return psutil.virtual_memory().percent
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get memory usage: {e}")
|
||||
return 0.0
|
||||
|
||||
@staticmethod
|
||||
def get_gpu_usage() -> Optional[float]:
|
||||
"""Get current GPU usage percentage."""
|
||||
if not TORCH_AVAILABLE:
|
||||
return None
|
||||
|
||||
try:
|
||||
if torch.cuda.is_available():
|
||||
# PyTorch doesn't provide direct GPU utilization
|
||||
# This is a placeholder - real implementation might use nvidia-ml-py
|
||||
if hasattr(torch.cuda, 'utilization'):
|
||||
return torch.cuda.utilization()
|
||||
else:
|
||||
# Fallback: estimate based on memory usage
|
||||
allocated = torch.cuda.memory_allocated()
|
||||
reserved = torch.cuda.memory_reserved()
|
||||
if reserved > 0:
|
||||
return (allocated / reserved) * 100
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get GPU usage: {e}")
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_gpu_memory_usage() -> Optional[float]:
|
||||
"""Get current GPU memory usage in MB."""
|
||||
if not TORCH_AVAILABLE:
|
||||
return None
|
||||
|
||||
try:
|
||||
if torch.cuda.is_available():
|
||||
return torch.cuda.memory_reserved() / (1024 ** 2) # Convert to MB
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get GPU memory usage: {e}")
|
||||
return None
|
||||
|
||||
|
||||
# Global worker state instance
|
||||
worker_state = WorkerState()
|
||||
Loading…
Add table
Add a link
Reference in a new issue