Add model management endpoints and database fetch functionality

api/prediction.py

@@ -1,5 +1,5 @@
 import torch
-from fastapi import APIRouter, HTTPException
+from fastapi import APIRouter, HTTPException, Depends
 from fastapi.responses import JSONResponse
 from PIL import Image
 from io import BytesIO
@@ -12,11 +12,12 @@ from schemas.requests import BatchPredictRequest
 from models.loader import get_model
 from steps.preprocess import process_image
 from config.settings import IMAGE_SIZE, NUM_THREADS
+from api.router import verify_api_key 
 
 logger = logging.getLogger(__name__)
 router = APIRouter()
 
-@router.post("/batch_predict")
+@router.post("/batch_predict", dependencies=[Depends(verify_api_key)])
 async def batch_predict(request: BatchPredictRequest):
     """Endpoint pour prédire à partir de plusieurs images."""
     model_name = request.modelName

api/router.py

@@ -3,7 +3,10 @@ import logging
 import os
 
 from api import prediction
-from config.settings import API_KEY
+from config.settings import API_KEY, MANAGEMENT_API_KEY
+from db.models import fetch_model_by_id
+from models.loader import model_pipelines, _load_single_model_pipeline, get_model
+from models.schemas import PredictionRequest, PredictionResponse
 
 logger = logging.getLogger(__name__)
 
@@ -34,5 +37,113 @@ async def verify_api_key(request: Request, call_next):
     response = await call_next(request)
     return response
 
+# Dépendance pour la Sécurité de l'API de Gestion
+async def verify_management_api_key(x_api_key: str = Header(...)):
+    """Vérifie si la clé API fournie correspond à celle configurée."""
+    if not MANAGEMENT_API_KEY:
+        logger.warning("MANAGEMENT_API_KEY is not set. Management endpoints are unsecured!")
+        # Décider si on bloque ou on autorise sans clé définie
+        # Pour la sécurité, il vaut mieux bloquer par défaut
+        raise HTTPException(status_code=500, detail="Management API key not configured on server.")
+    if x_api_key != MANAGEMENT_API_KEY:
+        logger.warning(f"Invalid or missing API key attempt for management endpoint.")
+        raise HTTPException(status_code=401, detail="Invalid or missing API Key")
+    return True # Clé valide
+
 # Inclure les routes des autres modules
 router.include_router(prediction.router, tags=["Prediction"])
+
+# Nouvel Endpoint de Gestion
+@router.post(
+    "/manage/load_model/{model_db_id}", 
+    summary="Load a specific model into memory", 
+    dependencies=[Depends(verify_management_api_key)] # Sécurise l'endpoint
+)
+async def load_single_model(model_db_id: Any): # L'ID peut être int ou str
+    """Charge un modèle spécifique en mémoire en utilisant son ID de la base de données."""
+    logger.info(f"Received request to load model with DB ID: {model_db_id}")
+
+    # 1. Vérifier si le modèle est déjà chargé
+    if model_db_id in model_pipelines:
+        logger.info(f"Model ID {model_db_id} is already loaded.")
+        return {"status": "success", "message": f"Model ID {model_db_id} is already loaded."}
+
+    # 2. Récupérer les informations du modèle depuis la DB
+    try:
+        model_data = await fetch_model_by_id(model_db_id)
+        if not model_data:
+            logger.error(f"Model ID {model_db_id} not found in database.")
+            raise HTTPException(status_code=404, detail=f"Model ID {model_db_id} not found in database.")
+    except Exception as e:
+        logger.exception(f"Database error fetching model ID {model_db_id}: {e}")
+        raise HTTPException(status_code=500, detail=f"Database error checking model ID {model_db_id}.")
+
+    # 3. Charger le modèle
+    try:
+        logger.info(f"Attempting to load model ID {model_db_id} ('{model_data.get('name', 'N/A')}') into memory...")
+        pipeline = await _load_single_model_pipeline(model_data)
+        
+        # 4. Ajouter au dictionnaire des modèles chargés
+        model_pipelines[model_db_id] = pipeline
+        logger.info(f"Successfully loaded and added model ID {model_db_id} to running pipelines.")
+        return {"status": "success", "message": f"Model ID {model_db_id} loaded successfully."}
+        
+    except Exception as e:
+        logger.exception(f"Failed to load model ID {model_db_id}: {e}")
+        # Ne pas laisser un pipeline potentiellement corrompu dans le dictionnaire
+        if model_db_id in model_pipelines:
+            del model_pipelines[model_db_id]
+        raise HTTPException(status_code=500, detail=f"Failed to load model ID {model_db_id}. Check server logs for details.")
+
+@router.post(
+    "/manage/update_model/{model_db_id}",
+    summary="Reload/Update a specific model already in memory",
+    dependencies=[Depends(verify_management_api_key)] # Sécurise l'endpoint
+)
+async def update_single_model(model_db_id: Any):
+    """Retélécharge et met à jour un modèle spécifique qui est déjà chargé en mémoire."""
+    logger.info(f"Received request to update model with DB ID: {model_db_id}")
+
+    # 1. Vérifier si le modèle est actuellement chargé
+    if model_db_id not in model_pipelines:
+        logger.error(f"Attempted to update model ID {model_db_id}, but it is not loaded.")
+        raise HTTPException(
+            status_code=404, 
+            detail=f"Model ID {model_db_id} is not currently loaded. Use load_model first."
+        )
+
+    # 2. Récupérer les informations du modèle depuis la DB (pour s'assurer qu'elles sont à jour si besoin)
+    try:
+        model_data = await fetch_model_by_id(model_db_id)
+        if not model_data:
+            # Ceci indiquerait une incohérence si le modèle est dans model_pipelines mais pas dans la DB
+            logger.error(f"Inconsistency: Model ID {model_db_id} loaded but not found in database during update.")
+            raise HTTPException(status_code=500, detail=f"Inconsistency: Model ID {model_db_id} not found in database.")
+    except Exception as e:
+        logger.exception(f"Database error fetching model ID {model_db_id} during update: {e}")
+        raise HTTPException(status_code=500, detail=f"Database error checking model ID {model_db_id} for update.")
+
+    # 3. Recharger le modèle
+    try:
+        logger.info(f"Attempting to reload model ID {model_db_id} ('{model_data.get('name', 'N/A')}') from source...")
+        # Supprimer l'ancien modèle de la mémoire avant de charger le nouveau pour libérer des ressources GPU/CPU si possible
+        # Attention : ceci pourrait causer une brève indisponibilité du modèle pendant le rechargement.
+        # Une stratégie alternative serait de charger le nouveau d'abord, puis de remplacer.
+        if model_db_id in model_pipelines:
+            del model_pipelines[model_db_id] 
+            # Potentiellement forcer le nettoyage de la mémoire GPU ici si nécessaire (torch.cuda.empty_cache() - à utiliser avec prudence)
+            logger.debug(f"Removed old instance of model ID {model_db_id} from memory before update.")
+            
+        pipeline = await _load_single_model_pipeline(model_data)
+        
+        # 4. Mettre à jour le dictionnaire avec le nouveau pipeline
+        model_pipelines[model_db_id] = pipeline
+        logger.info(f"Successfully updated model ID {model_db_id} in running pipelines.")
+        return {"status": "success", "message": f"Model ID {model_db_id} updated successfully."}
+        
+    except Exception as e:
+        logger.exception(f"Failed to reload model ID {model_db_id}: {e}")
+        # Si le rechargement échoue, l'ancien modèle a déjà été supprimé.
+        # Il faut soit tenter de recharger l'ancien, soit le laisser déchargé.
+        # Pour l'instant, on le laisse déchargé et on signale l'erreur.
+        raise HTTPException(status_code=500, detail=f"Failed to reload model ID {model_db_id}. Model is now unloaded. Check server logs.")

config/settings.py

@@ -1,6 +1,10 @@
 import os
 import logging
 import torch
+from dotenv import load_dotenv
+
+# Charger les variables d'environnement depuis le fichier .env
+load_dotenv()
 
 # Configuration de base des logs
 logging.basicConfig(level=logging.INFO)
@@ -22,6 +26,11 @@ HF_TOKEN = get_env_or_fail("api_read")
 RESOURCE_GROUP = get_env_or_fail("RESOURCE_GROUP")
 DATABASE_URL = get_env_or_fail("DATABASE_URL")
 
+# Configuration Gestion
+MANAGEMENT_API_KEY = os.getenv("MANAGEMENT_API_KEY")
+if not MANAGEMENT_API_KEY:
+    print("Warning: MANAGEMENT_API_KEY environment variable is not set. Management endpoints will be inaccessible.")
+
 # Log des paramètres importants (sans détails sensibles)
 logger.info(f"RESOURCE_GROUP set to: {RESOURCE_GROUP}")
 

db/models.py

@@ -46,3 +46,50 @@ async def fetch_models_for_group(resource_group: str) -> List[Dict[str, Any]]:
         if conn and not conn.is_closed():
             await conn.close()
             logger.debug("Database connection closed")
+
+
+async def fetch_model_by_id(model_id: str) -> Dict[str, Any] | None:
+    """Récupérer les détails d'un modèle spécifique par son ID de base de données.
+
+    Args:
+        model_id: L'ID du modèle dans la base de données (peut être int ou str selon le schéma).
+
+    Returns:
+        Un dictionnaire contenant les informations du modèle si trouvé, sinon None.
+
+    Raises:
+        Exception: Si une erreur se produit lors de la connexion ou de la requête.
+    """
+    conn = None
+    try:
+        conn = await asyncpg.connect(DATABASE_URL)
+        logger.debug(f"Successfully connected to database to fetch model ID: {model_id}")
+
+        # Récupérer le modèle spécifique par son ID
+        query = """
+            SELECT 
+                model_id, 
+                name, 
+                display_name, 
+                hf_repo_id, 
+                hf_subfolder, 
+                hf_filename
+            FROM models
+            WHERE model_id = $1
+        """
+        row = await conn.fetchrow(query, model_id)
+        
+        if row:
+            logger.info(f"Found model with ID '{model_id}': {row['name']}")
+            return dict(row)
+        else:
+            logger.warning(f"No model found with ID '{model_id}'")
+            return None
+            
+    except Exception as e:
+        logger.error(f"Database error fetching model ID {model_id}: {e}", exc_info=True)
+        raise
+    finally:
+        if conn and not conn.is_closed():
+            await conn.close()
+            logger.debug(f"Database connection closed after fetching model ID: {model_id}")

models/loader.py

@@ -1,7 +1,7 @@
-import torch
 import logging
-from typing import Dict, List, Any
+import torch
 from huggingface_hub import hf_hub_download
+from typing import List, Dict, Any
 
 from config.settings import DEVICE, HF_TOKEN, NUM_THREADS
 from architecture.resnet import ResNet
@@ -12,62 +12,109 @@ logger = logging.getLogger(__name__)
 torch.set_num_threads(NUM_THREADS)
 
 # Instance de base pour le modèle ResNet
-base_model = ResNet("resnet152", num_output_neurons=2).to(DEVICE)
+# Note: Peut-être pas nécessaire de l'instancier ici si chaque chargement en crée une nouvelle
+# base_model = ResNet("resnet152", num_output_neurons=2).to(DEVICE)
 
-# Dictionnaire global pour stocker les modèles chargés
-model_pipelines = {}
+# Dictionnaire global pour stocker les modèles chargés (pipelines)
+# Clé: ID du modèle (provenant de la DB), Valeur: Pipeline/Modèle chargé
+model_pipelines: Dict[Any, Any] = {}
 
-async def load_models(models_data: List[Dict[str, Any]]) -> Dict[str, Any]:
-    """Charger les modèles depuis Hugging Face à partir des données de la base de données.
+async def _load_single_model_pipeline(model_data: Dict[str, Any]) -> Any:
+    """Charge un seul pipeline de modèle à partir de ses données.
     
     Args:
-        models_data: Liste de dictionnaires contenant les informations des modèles
+        model_data: Dictionnaire contenant les informations du modèle (hf_repo_id, etc.).
         
     Returns:
-        Dictionnaire des modèles chargés
+        Le pipeline/modèle chargé.
+        
+    Raises:
+        Exception: Si le chargement échoue.
+    """
+    model_id = model_data['model_id'] # Utiliser l'ID de la DB comme clé
+    model_name = model_data['hf_filename']
+    repo_id = model_data['hf_repo_id']
+    subfolder = model_data['hf_subfolder']
+    
+    logger.info(f"Loading model ID {model_id}: {model_name} (repo: {repo_id}, subfolder: {subfolder})")
+    
+    try:
+        model_weight_path = hf_hub_download(
+            repo_id=repo_id,
+            subfolder=subfolder,
+            filename=model_name,
+            token=HF_TOKEN, # Assurez-vous que HF_TOKEN est géré correctement
+        )
+        
+        logger.debug(f"Model weights downloaded to: {model_weight_path}")
+        
+        # Créer une nouvelle instance de modèle ResNet pour ce chargement spécifique
+        # Assurez-vous que ResNet et ses arguments sont corrects
+        model = ResNet("resnet152", num_output_neurons=2).to(DEVICE) 
+        
+        # Charger les poids
+        # Attention: la méthode de chargement dépend du format des poids (state_dict, etc.)
+        state_dict = torch.load(model_weight_path, map_location=DEVICE)
+        
+        # Gérer les cas où les poids sont dans une sous-clé (ex: 'state_dict', 'model')
+        if isinstance(state_dict, dict) and 'state_dict' in state_dict:
+            state_dict = state_dict['state_dict']
+        elif isinstance(state_dict, dict) and 'model' in state_dict: # Autre cas commun
+            state_dict = state_dict['model']
+            
+        # Adapter les clés si nécessaire (ex: supprimer le préfixe 'module.' de DataParallel/DDP)
+        state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
+
+        model.load_state_dict(state_dict)
+        model.eval() # Mettre le modèle en mode évaluation
+        
+        logger.info(f"Successfully loaded model ID {model_id}: {model_name}")
+        return model # Retourner le modèle chargé (ou un pipeline si vous en créez un)
+        
+    except Exception as e:
+        logger.error(f"Failed to load model ID {model_id} ({model_name}): {e}", exc_info=True)
+        raise # Propage l'exception pour que l'appelant puisse la gérer
+
+
+async def load_models(models_data: List[Dict[str, Any]]) -> None:
+    """Charger les modèles depuis Hugging Face et les stocker dans model_pipelines.
+    
+    Args:
+        models_data: Liste de dictionnaires contenant les informations des modèles.
         
     Raises:
-        RuntimeError: Si aucun modèle n'est trouvé ou ne peut être chargé
+        RuntimeError: Si aucun modèle n'est trouvé.
     """
-    logger.info(f"Attempting to load {len(models_data)} models...")
+    logger.info(f"Attempting to load {len(models_data)} models into memory...")
     
     if not models_data:
-        error_msg = "No models found. API cannot start without models."
+        error_msg = "No models data provided. Cannot load models."
         logger.error(error_msg)
-        raise RuntimeError(error_msg)
+        # On ne lève plus d'erreur ici, on logge juste. L'API démarrera sans modèles.
+        return
     
     loaded_count = 0
+    failed_models = []
     for model_data in models_data:
+        model_id = model_data.get('model_id', 'N/A') # Assurez-vous que model_id est présent
         try:
-            model_name = model_data['hf_filename']
-            logger.info(f"Loading model: {model_name} (repo: {model_data['hf_repo_id']}, subfolder: {model_data['hf_subfolder']})")
-            
-            model_weight = hf_hub_download(
-                repo_id=model_data['hf_repo_id'],
-                subfolder=model_data['hf_subfolder'],
-                filename=model_name,
-                token=HF_TOKEN,
-            )
-            
-            # Créer une nouvelle instance pour chaque modèle pour tenir ses poids spécifiques
-            model = base_model.__class__("resnet152", num_output_neurons=2).to(DEVICE)
-            model.load_state_dict(
-                torch.load(model_weight, weights_only=True, map_location=DEVICE)
-            )
-            model.eval()
-            model_pipelines[model_name] = model
+            # Utilise la nouvelle fonction pour charger un seul modèle
+            pipeline = await _load_single_model_pipeline(model_data)
+            # Stocke le pipeline chargé dans le dictionnaire global
+            model_pipelines[model_id] = pipeline
             loaded_count += 1
         except Exception as e:
-            logger.error(f"Error loading model {model_data.get('hf_filename', 'N/A')}: {e}", exc_info=True)
-    
-    logger.info(f"Model loading finished. Successfully loaded {loaded_count}/{len(models_data)} models.")
-    
-    if loaded_count == 0:
-        error_msg = "Failed to load any models. API cannot start without models."
-        logger.error(error_msg)
-        raise RuntimeError(error_msg)
-    
-    return model_pipelines
+            # Log l'échec mais continue avec les autres modèles
+            logger.error(f"Failed to load model ID {model_id}: {e}")
+            failed_models.append(model_data.get('name', f'ID {model_id}'))
+
+    logger.info(f"Finished loading models. Successfully loaded: {loaded_count}/{len(models_data)}")
+    if failed_models:
+        logger.warning(f"Failed to load the following models: {', '.join(failed_models)}")
+
+    # Pas besoin de retourner les pipelines, ils sont dans le dictionnaire global
+    # return model_pipelines # Ancienne logique
+
 
 def get_model(model_name: str):
     """Récupérer un modèle chargé par son nom.