Update server.py

server.py

@@ -37,15 +37,10 @@ def verify_api_key(request):
 # ============================================================================
 # 🚀 PERFORMANCE OPTIMIZATION CONFIGURATION
 # ============================================================================
-# CPU Optimization
 MAX_WORKERS = min(4, (os.cpu_count() or 1) + 1)
 executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)
-
-# Memory Management
-MEMORY_THRESHOLD = 80  # Percentage
-MAX_CACHE_SIZE = 50   # Maximum number of cached results
-
-# Model Loading Optimization
+MEMORY_THRESHOLD = 80
+MAX_CACHE_SIZE = 50
 os.environ['OMP_NUM_THREADS'] = '2'
 os.environ['OPENBLAS_NUM_THREADS'] = '2'
 os.environ['MKL_NUM_THREADS'] = '2'
@@ -104,7 +99,6 @@ prediction_cache = {}
 cache_lock = threading.Lock()
 
 def cache_prediction(key, result):
-    """Cache prediction result with memory management"""
     with cache_lock:
         if len(prediction_cache) >= MAX_CACHE_SIZE:
             oldest_keys = list(prediction_cache.keys())[:MAX_CACHE_SIZE//2]
@@ -113,7 +107,6 @@ def cache_prediction(key, result):
         prediction_cache[key] = result
 
 def get_cached_prediction(key):
-    """Get cached prediction result"""
     with cache_lock:
         return prediction_cache.get(key)
 
@@ -139,7 +132,6 @@ CLASS_ALIASES = {
 # ============================================================================
 @lru_cache(maxsize=256)
 def normalize_text(text):
-    """Cached text normalization"""
     if not text: return ""
     text = text.lower().strip()
     text = re.sub(r'[^\w\s]', ' ', text)
@@ -148,7 +140,6 @@ def normalize_text(text):
 
 @lru_cache(maxsize=256)
 def find_class_match(input_text):
-    """Cached class matching"""
     if not input_text: return None
     normalized_input = normalize_text(input_text)
     for canonical_name, aliases in CLASS_ALIASES.items():
@@ -161,15 +152,13 @@ def find_class_match(input_text):
     return None
 
 # ============================================================================
-# 📁 MODEL LOADING & INITIALIZATION (OPTIMIZED FOR GUNICORN --PRELOAD)
+# 📁 MODEL LOADING & INITIALIZATION (STRATEGI: LAZY PER-WORKER)
 # ============================================================================
-logging.info("🔄 Initializing models at module level for Gunicorn --preload...")
-
-models = {}
+_worker_models = {}
+_model_lock = threading.Lock()
 model_class_maps = {}
 
 def load_yaml_classes(yaml_path):
-    """Load classes from YAML file"""
     try:
         with open(yaml_path, 'r', encoding='utf-8') as file:
             data = yaml.safe_load(file)
@@ -179,9 +168,6 @@ def load_yaml_classes(yaml_path):
     return {}
 
 try:
-    models['3x3'] = YOLO('best.onnx', task='classify')
-    models['4x4'] = YOLO('best4x4.onnx', task='segment')
-
     for model_type, yaml_file in [('3x3', 'data.yaml'), ('4x4', 'data4x4.yaml')]:
         class_map = {}
         yaml_classes = load_yaml_classes(yaml_file)
@@ -192,18 +178,48 @@ try:
             else:
                 class_map[class_name.lower()] = class_id
         model_class_maps[model_type] = class_map
-    
-    logging.info("✅ Models and class maps initialized successfully!")
-
+    logging.info("✅ Class maps loaded successfully!")
 except Exception as e:
-    logging.error(f"❌ FATAL: Failed to initialize models at module level: {e}")
+    logging.error(f"❌ FATAL: Failed to initialize class maps: {e}")
     raise
 
+
+def get_model(model_type: str):
+    """
+    Loads a model only once per worker process (lazy initialization).
+    This is the robust solution for multi-process servers like Gunicorn.
+    """
+    if model_type in _worker_models:
+        return _worker_models[model_type]
+    
+    with _model_lock:
+        if model_type in _worker_models:
+            return _worker_models[model_type]
+
+        logging.info(f"WORKER_INIT: Loading model '{model_type}' for worker PID: {os.getpid()}...")
+        
+        model_path, task_type = '', ''
+        if model_type == '3x3':
+            model_path, task_type = 'best.onnx', 'classify'
+        elif model_type == '4x4':
+            model_path, task_type = 'best4x4.onnx', 'segment'
+        else:
+            logging.error(f"Attempted to load unknown model type: {model_type}")
+            return None
+            
+        try:
+            model = YOLO(model_path, task=task_type)
+            _worker_models[model_type] = model
+            logging.info(f"WORKER_INIT: Model '{model_type}' loaded successfully for worker PID: {os.getpid()}.")
+            return model
+        except Exception as e:
+            logging.error(f"WORKER_INIT: Failed to load model '{model_path}' for worker PID: {os.getpid()}: {e}")
+            return None
+
 # ============================================================================
 # 🖼️ OPTIMIZED IMAGE PROCESSING
 # ============================================================================
 def decode_image_optimized(base64_string):
-    """Optimized image decoding"""
     try:
         image_data = base64.b64decode(base64_string.split(',')[1])
         image = Image.open(io.BytesIO(image_data)).convert("RGB")
@@ -213,7 +229,6 @@ def decode_image_optimized(base64_string):
         return None
 
 def divide_image_into_4x4_grid(image_cv2):
-    """Optimized grid division"""
     height, width = image_cv2.shape[:2]
     grid_height, grid_width = height // 4, width // 4
     grid_images, grid_coordinates = [], []
@@ -226,7 +241,6 @@ def divide_image_into_4x4_grid(image_cv2):
     return grid_images, grid_coordinates
 
 def is_object_in_grid_cell(mask_contour, grid_coords, min_coverage_percentage=MIN_COVERAGE_PERCENTAGE):
-    """Optimized object detection in grid cell"""
     x1, y1, x2, y2 = grid_coords
     grid_width, grid_height = x2 - x1, y2 - y1
     grid_area = grid_width * grid_height
@@ -246,7 +260,6 @@ def is_object_in_grid_cell(mask_contour, grid_coords, min_coverage_percentage=MI
 # 🔧 UTILITY FUNCTIONS
 # ============================================================================
 def get_target_class_index(input_title, model_type):
-    """Get target class index with caching"""
     model_classes = model_class_maps.get(model_type, {})
     if not input_title or not model_classes: return None
     canonical_name = find_class_match(input_title)
@@ -256,7 +269,6 @@ def get_target_class_index(input_title, model_type):
     return model_classes.get(normalized_input)
 
 def memory_cleanup():
-    """Perform memory cleanup"""
     gc.collect()
     current_memory = psutil.virtual_memory().percent
     if current_memory > MEMORY_THRESHOLD:
@@ -267,7 +279,6 @@ def memory_cleanup():
 # ============================================================================
 @app.before_request
 def check_api_key():
-    """Verify API key for all requests except health check"""
     if request.endpoint in ['health', 'stats']: return
     if not verify_api_key(request):
         return jsonify({"error": "Invalid or missing API key"}), 401
@@ -278,8 +289,10 @@ def check_api_key():
 @app.route('/health', methods=['GET'])
 def health():
     return jsonify({
-        "status": "healthy", "models_loaded": len(models),
-        "memory_usage": psutil.virtual_memory().percent, "cpu_usage": psutil.cpu_percent()
+        "status": "healthy",
+        "models_loaded_in_worker": len(_worker_models),
+        "memory_usage": psutil.virtual_memory().percent,
+        "cpu_usage": psutil.cpu_percent()
     })
 
 @app.route('/stats', methods=['GET'])
@@ -288,15 +301,15 @@ def stats():
 
 @app.route('/predict', methods=['POST'])
 def predict():
-    """Optimized 3x3 prediction endpoint"""
     import time
     start_time = time.time()
-    
     try:
         data = request.get_json(silent=True)
         if not data: return jsonify({"error": "Invalid request body"}), 400
-        model = models.get('3x3')
+        
+        model = get_model('3x3')
         if not model: return jsonify({"error": "3x3 model not loaded"}), 500
+        
         input_title = data.get('title', '')
         target_class_index = get_target_class_index(input_title, '3x3')
         if target_class_index is None:
@@ -314,12 +327,9 @@ def predict():
                 image = decode_image_optimized(item['base64'])
                 if image is None: return None
                 results = model(image, verbose=False)
-                
-                # Robust validation to prevent 'NoneType' error
                 if not results: return None
                 res = results[0]
                 if res.probs is None or res.probs.data is None: return None
-                
                 confidence = res.probs.data[target_class_index].item()
                 return {'index': item['index'], 'confidence': confidence, 'selected': confidence >= CONFIDENCE_THRESHOLD_3X3}
             except Exception as e:
@@ -338,20 +348,20 @@ def predict():
         if psutil.virtual_memory().percent > MEMORY_THRESHOLD: memory_cleanup()
         return jsonify(response)
     except Exception as e:
-        logging.error(f"Error in /predict: {e}")
+        logging.error(f"Error in /predict: {e}", exc_info=True)
         return jsonify({"error": "Internal server error"}), 500
 
 @app.route('/predict_4x4', methods=['POST'])
 def predict_4x4():
-    """Optimized 4x4 prediction endpoint"""
     import time
     start_time = time.time()
-    
     try:
         data = request.get_json(silent=True)
         if not data: return jsonify({"error": "Invalid request body"}), 400
-        model = models.get('4x4')
+        
+        model = get_model('4x4')
         if not model: return jsonify({"error": "4x4 model not loaded"}), 500
+        
         input_title = data.get('title', '')
         target_class_index = get_target_class_index(input_title, '4x4')
         if target_class_index is None:
@@ -367,7 +377,6 @@ def predict_4x4():
         image_pil = decode_image_optimized(data['image_b64'])
         if image_pil is None: return jsonify({"error": "Invalid image data"}), 400
         image_cv2 = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
-        grid_images, grid_coordinates = divide_image_into_4x4_grid(image_cv2)
         
         results = model(image_cv2, verbose=False)
         indices_to_click = []
@@ -390,7 +399,7 @@ def predict_4x4():
         if psutil.virtual_memory().percent > MEMORY_THRESHOLD: memory_cleanup()
         return jsonify(response)
     except Exception as e:
-        logging.error(f"Error in /predict_4x4: {e}")
+        logging.error(f"Error in /predict_4x4: {e}", exc_info=True)
         return jsonify({"error": "Internal server error"}), 500
 
 @app.route('/classes', methods=['GET'])