Spaces:

kota101
/

ta_ml

Sleeping

ta_ml / api /admin.py

muhammadjalaludin

Implement service layers for backup, model management, file handling, tokenization, and training

1424118 over 1 year ago

3.66 kB

# 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}"
        )

	# 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}"
	)