File size: 3,663 Bytes
1424118 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | # ta_ml/api/admin.py
"""
API endpoints for administrative tasks.
"""
import logging
from fastapi import APIRouter, Depends, HTTPException, Request
from api.dependencies import get_model_manager
from models.api import ReloadResponse
from services.model_manager import ModelManager
from services.prediction_service import PredictionService
from services.tokenization_service import TokenizationService
from services.training_service import TrainingService
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/admin", tags=["Admin"])
def _reinitalize_ml_services(request: Request, model_manager: ModelManager):
"""
Re-initializes all services that depend on the ModelManager's components.
This must be called after a model reload to ensure services use the new models.
"""
logger.info("Re-initializing ML-dependent services...")
# Prediction Service
request.app.state.prediction_service = PredictionService(
tokenizer=model_manager.tokenizer,
base_model=model_manager.base_model,
classifier=model_manager.classifier,
label_encoder=model_manager.label_encoder,
device=model_manager.device,
)
# Tokenization Service
request.app.state.tokenization_service = TokenizationService(
tokenizer=model_manager.tokenizer)
# Training Service
request.app.state.training_service = TrainingService(
settings=request.app.state.settings,
tokenizer=model_manager.tokenizer,
base_model=model_manager.base_model,
device=model_manager.device,
)
logger.info("ML-dependent services re-initialized successfully.")
@router.post(
"/reload-models",
response_model=ReloadResponse,
summary="Hot-reload all machine learning models",
description=(
"Triggers a hot-reload of all models and tokenizers from disk or the Hub. "
"This is a powerful administrative action that allows for updating the deployed "
"models (e.g., after a new training run) without restarting the entire API service. "
"Upon successful reload, all dependent services (Prediction, Tokenization, etc.) "
"are automatically re-initialized to use the new model instances."
),
responses={
200: {
"description": "The models were reloaded and services were re-initialized successfully.",
"content": {
"application/json": {
"example": {"status": "Models reloaded successfully", "detail": None}
}
},
},
500: {
"description": "The model reload process failed. Check application logs for details.",
"content": {
"application/json": {
"example": {"status": "Error reloading models", "detail": "Specific error message here"}
}
},
},
},
)
def reload_models(
request: Request,
model_manager: ModelManager = Depends(get_model_manager),
):
"""
Triggers a hot-reload of all models and tokenizers from disk or the Hub.
This allows updating the model artifacts without restarting the API service.
After the models are reloaded, all dependent services are re-initialized
to use the new model instances.
"""
try:
model_manager.reload_models()
_reinitalize_ml_services(request, model_manager)
return ReloadResponse(status="Models reloaded successfully")
except Exception as e:
logger.exception("Error during model reload.", exc_info=e)
raise HTTPException(
status_code=500, detail=f"Error reloading models: {e}"
)
|