242 lines
No EOL
9 KiB
Python
242 lines
No EOL
9 KiB
Python
"""
|
|
Shared Model Registry for Memory Optimization
|
|
|
|
This module implements a global shared model registry to prevent duplicate model loading
|
|
in memory when multiple cameras use the same model. This significantly reduces RAM and
|
|
GPU VRAM usage by ensuring only one instance of each unique model is loaded.
|
|
|
|
Key Features:
|
|
- Thread-safe model loading and access
|
|
- Reference counting for proper cleanup
|
|
- Automatic model lifecycle management
|
|
- Maintains compatibility with existing pipeline system
|
|
"""
|
|
|
|
import os
|
|
import threading
|
|
import logging
|
|
from typing import Dict, Any, Optional, Set
|
|
import torch
|
|
from ultralytics import YOLO
|
|
|
|
# Create a logger for this module
|
|
logger = logging.getLogger("detector_worker.model_registry")
|
|
|
|
class ModelRegistry:
|
|
"""
|
|
Singleton class for managing shared YOLO models across multiple cameras.
|
|
|
|
This registry ensures that each unique model is loaded only once in memory,
|
|
dramatically reducing RAM and GPU VRAM usage when multiple cameras use the
|
|
same model.
|
|
"""
|
|
|
|
_instance = None
|
|
_lock = threading.Lock()
|
|
|
|
def __new__(cls):
|
|
if cls._instance is None:
|
|
with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = super(ModelRegistry, cls).__new__(cls)
|
|
cls._instance._initialized = False
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
if self._initialized:
|
|
return
|
|
|
|
self._initialized = True
|
|
|
|
# Thread-safe storage for loaded models
|
|
self._models: Dict[str, YOLO] = {} # modelId -> YOLO model instance
|
|
self._model_files: Dict[str, str] = {} # modelId -> file path
|
|
self._reference_counts: Dict[str, int] = {} # modelId -> reference count
|
|
self._model_lock = threading.RLock() # Reentrant lock for nested calls
|
|
|
|
logger.info("🏭 Shared Model Registry initialized - ready for memory-optimized model loading")
|
|
|
|
def get_model(self, model_id: str, model_file_path: str) -> YOLO:
|
|
"""
|
|
Get or load a YOLO model. Returns shared instance if already loaded.
|
|
|
|
Args:
|
|
model_id: Unique identifier for the model
|
|
model_file_path: Path to the model file
|
|
|
|
Returns:
|
|
YOLO model instance (shared across all callers)
|
|
"""
|
|
with self._model_lock:
|
|
if model_id in self._models:
|
|
# Model already loaded - increment reference count and return
|
|
self._reference_counts[model_id] += 1
|
|
logger.info(f"📖 Model '{model_id}' reused (ref_count: {self._reference_counts[model_id]}) - SAVED MEMORY!")
|
|
return self._models[model_id]
|
|
|
|
# Model not loaded yet - load it
|
|
logger.info(f"🔄 Loading NEW model '{model_id}' from {model_file_path}")
|
|
|
|
if not os.path.exists(model_file_path):
|
|
raise FileNotFoundError(f"Model file {model_file_path} not found")
|
|
|
|
try:
|
|
# Load the YOLO model
|
|
model = YOLO(model_file_path)
|
|
|
|
# Move to GPU if available
|
|
if torch.cuda.is_available():
|
|
logger.info(f"🚀 CUDA available. Moving model '{model_id}' to GPU VRAM")
|
|
model.to("cuda")
|
|
else:
|
|
logger.info(f"💻 CUDA not available. Using CPU for model '{model_id}'")
|
|
|
|
# Store in registry
|
|
self._models[model_id] = model
|
|
self._model_files[model_id] = model_file_path
|
|
self._reference_counts[model_id] = 1
|
|
|
|
logger.info(f"✅ Model '{model_id}' loaded and registered (ref_count: 1)")
|
|
self._log_registry_status()
|
|
|
|
return model
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ Failed to load model '{model_id}' from {model_file_path}: {e}")
|
|
raise
|
|
|
|
def release_model(self, model_id: str) -> None:
|
|
"""
|
|
Release a reference to a model. If reference count reaches zero,
|
|
the model may be unloaded to free memory.
|
|
|
|
Args:
|
|
model_id: Unique identifier for the model to release
|
|
"""
|
|
with self._model_lock:
|
|
if model_id not in self._reference_counts:
|
|
logger.warning(f"⚠️ Attempted to release unknown model '{model_id}'")
|
|
return
|
|
|
|
self._reference_counts[model_id] -= 1
|
|
logger.info(f"📉 Model '{model_id}' reference count decreased to {self._reference_counts[model_id]}")
|
|
|
|
# For now, keep models in memory even when ref count reaches 0
|
|
# This prevents reload overhead if the same model is needed again soon
|
|
# In the future, we could implement LRU eviction policy
|
|
# if self._reference_counts[model_id] <= 0:
|
|
# logger.info(f"💤 Model '{model_id}' has 0 references but keeping in memory for reuse")
|
|
# Optionally: self._unload_model(model_id)
|
|
|
|
def _unload_model(self, model_id: str) -> None:
|
|
"""
|
|
Internal method to unload a model from memory.
|
|
Currently not used to prevent reload overhead.
|
|
"""
|
|
with self._model_lock:
|
|
if model_id in self._models:
|
|
logger.info(f"🗑️ Unloading model '{model_id}' from memory")
|
|
|
|
# Clear GPU memory if model was on GPU
|
|
model = self._models[model_id]
|
|
if hasattr(model, 'model') and hasattr(model.model, 'cuda'):
|
|
try:
|
|
# Move model to CPU before deletion to free GPU memory
|
|
model.to('cpu')
|
|
except Exception as e:
|
|
logger.warning(f"⚠️ Failed to move model '{model_id}' to CPU: {e}")
|
|
|
|
# Remove from registry
|
|
del self._models[model_id]
|
|
del self._model_files[model_id]
|
|
del self._reference_counts[model_id]
|
|
|
|
# Force garbage collection
|
|
import gc
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
logger.info(f"✅ Model '{model_id}' unloaded and memory freed")
|
|
self._log_registry_status()
|
|
|
|
def get_registry_status(self) -> Dict[str, Any]:
|
|
"""
|
|
Get current status of the model registry.
|
|
|
|
Returns:
|
|
Dictionary with registry statistics
|
|
"""
|
|
with self._model_lock:
|
|
return {
|
|
"total_models": len(self._models),
|
|
"models": {
|
|
model_id: {
|
|
"file_path": self._model_files[model_id],
|
|
"reference_count": self._reference_counts[model_id]
|
|
}
|
|
for model_id in self._models
|
|
},
|
|
"total_references": sum(self._reference_counts.values())
|
|
}
|
|
|
|
def _log_registry_status(self) -> None:
|
|
"""Log current registry status for debugging."""
|
|
status = self.get_registry_status()
|
|
logger.info(f"📊 Model Registry Status: {status['total_models']} unique models, {status['total_references']} total references")
|
|
for model_id, info in status['models'].items():
|
|
logger.debug(f" 📋 '{model_id}': refs={info['reference_count']}, file={os.path.basename(info['file_path'])}")
|
|
|
|
def cleanup_all(self) -> None:
|
|
"""
|
|
Clean up all models from the registry. Used during shutdown.
|
|
"""
|
|
with self._model_lock:
|
|
model_ids = list(self._models.keys())
|
|
logger.info(f"🧹 Cleaning up {len(model_ids)} models from registry")
|
|
|
|
for model_id in model_ids:
|
|
self._unload_model(model_id)
|
|
|
|
logger.info("✅ Model registry cleanup complete")
|
|
|
|
|
|
# Global singleton instance
|
|
_registry = ModelRegistry()
|
|
|
|
def get_shared_model(model_id: str, model_file_path: str) -> YOLO:
|
|
"""
|
|
Convenience function to get a shared model instance.
|
|
|
|
Args:
|
|
model_id: Unique identifier for the model
|
|
model_file_path: Path to the model file
|
|
|
|
Returns:
|
|
YOLO model instance (shared across all callers)
|
|
"""
|
|
return _registry.get_model(model_id, model_file_path)
|
|
|
|
def release_shared_model(model_id: str) -> None:
|
|
"""
|
|
Convenience function to release a shared model reference.
|
|
|
|
Args:
|
|
model_id: Unique identifier for the model to release
|
|
"""
|
|
_registry.release_model(model_id)
|
|
|
|
def get_registry_status() -> Dict[str, Any]:
|
|
"""
|
|
Convenience function to get registry status.
|
|
|
|
Returns:
|
|
Dictionary with registry statistics
|
|
"""
|
|
return _registry.get_registry_status()
|
|
|
|
def cleanup_registry() -> None:
|
|
"""
|
|
Convenience function to cleanup the entire registry.
|
|
"""
|
|
_registry.cleanup_all() |