model_manager.py

"""
Model manager module with strict typing using newtype pattern.
"""

import torch
import asyncio
from typing import Dict, List, Tuple, Any, Set, cast
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.configuration_utils import GenerationConfig
import gc
from loguru import logger
from dataclasses import dataclass
from huggingface_hub import login

from config import ModelConfig, MODELS_CONFIG, HF_TOKEN
from custom_types import (
    ModelId, ModelMemoryUsage, Timestamp, GPUMemoryUsageWrapper,
    create_model_id, create_gpu_memory_usage, ModelStatus,
    HuggingFaceModel, HuggingFaceTokenizer, HuggingFaceGenerationConfig, TorchDevice,
    create_huggingface_model, create_huggingface_tokenizer, create_huggingface_generation_config,
    create_torch_tensor, create_torch_device
)


@dataclass(frozen=True)
class ModelInfo:
    """Information about a loaded model with strict typing."""
    model: HuggingFaceModel
    tokenizer: HuggingFaceTokenizer
    memory_used: ModelMemoryUsage
    last_used: Timestamp
    generation_config: HuggingFaceGenerationConfig
    status: ModelStatus


class ModelManager:
    """Manages multiple models with memory optimization and multiplexing."""

    models: Dict[ModelId, ModelInfo]
    max_memory_usage: GPUMemoryUsageWrapper
    device: TorchDevice
    lock: asyncio.Lock

    def __init__(self, max_memory_usage: float = 0.9):
        """
        Initialize the model manager.

        Args:
            max_memory_usage: Maximum fraction of GPU memory to use (0.9 = 90%)
        """
        self.models = {}
        self.max_memory_usage = create_gpu_memory_usage(max_memory_usage)
        self.device = create_torch_device(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
        self.lock = asyncio.Lock()

        # Login to HuggingFace
        if HF_TOKEN:
            login(token=str(HF_TOKEN))
            logger.info("Logged into HuggingFace Hub")
        else:
            logger.warning("No HuggingFace token provided - private models may not be accessible")

    async def initialize(self) -> None:
        """Initialize the model manager and load default models."""
        logger.info("Initializing Model Manager...")

        # Check GPU availability
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA is not available. This service requires GPU.")

        logger.info(f"Using device: {self.device}")
        logger.info(f"Available GPU memory: {self.get_total_gpu_memory():.2f} GB")

        # Load Mistral 7B as the default model
        await self.load_model("mistralai/Mistral-7B-Instruct-v0.3")

    async def load_model(self, model_id: str) -> None:
        """Load a model into memory.

        Avoids holding the manager lock during heavy I/O or GPU work to prevent
        deadlocks and to allow concurrent requests to progress.
        """
        # Validate and normalize model id
        try:
            model_id_wrapper = create_model_id(model_id)
            model_id_typed = model_id_wrapper.unwrap()
        except ValueError as e:
            logger.error(f"Invalid model ID: {model_id}, error: {e}")
            raise

        # Fast-path: if already loaded, return early
        async with self.lock:
            if model_id_typed in self.models:
                logger.info(f"Model {model_id_typed} already loaded")
                return

        logger.info(f"Loading model: {model_id_typed}")

        # If memory is tight, free up space (this will acquire the lock internally)
        current_usage = self.get_gpu_memory_usage()
        if current_usage.unwrap() > self.max_memory_usage.unwrap():
            await self._free_memory()

        # Heavy work happens outside the lock
        try:
            # Get model config
            config = MODELS_CONFIG.get(model_id_typed, ModelConfig())

            # Load tokenizer
            tokenizer_raw: Any = AutoTokenizer.from_pretrained(
                model_id_typed,
                trust_remote_code=True,
                token=str(HF_TOKEN) if HF_TOKEN else None
            )

            # Set padding token if not set
            if tokenizer_raw.pad_token is None:
                tokenizer_raw.pad_token = tokenizer_raw.eos_token

            # Load model
            model_raw: Any = AutoModelForCausalLM.from_pretrained(
                model_id_typed,
                device_map="auto",
                torch_dtype=torch.float16,  # Use fp16 for memory efficiency
                trust_remote_code=True,
                token=str(HF_TOKEN) if HF_TOKEN else None,
                low_cpu_mem_usage=True,
                **config.model_kwargs
            )

            # Create generation config
            generation_config_raw: Any = GenerationConfig(
                eos_token_id=tokenizer_raw.eos_token_id,
                pad_token_id=tokenizer_raw.pad_token_id,
                **config.generation_defaults
            )

            # Wrap external library objects
            model = create_huggingface_model(model_raw)
            tokenizer = create_huggingface_tokenizer(tokenizer_raw)
            generation_config = create_huggingface_generation_config(generation_config_raw)

            # Calculate memory usage
            memory_used = self._calculate_model_memory(model)

        except Exception as e:
            logger.error(f"Failed to load model {model_id_typed}: {str(e)}")
            raise

        # Store the model info under the lock; handle races where another task
        # finished loading the same model while we were working
        async with self.lock:
            if model_id_typed in self.models:
                # Another coroutine loaded it; dispose of the duplicate we created
                if 'model_raw' in locals():
                    del model_raw
                if 'tokenizer_raw' in locals():
                    del tokenizer_raw
                _ = gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                logger.info(f"Model {model_id_typed} was loaded concurrently; discarded duplicate")
                return

            self.models[model_id_typed] = ModelInfo(
                model=model,
                tokenizer=tokenizer,
                memory_used=memory_used,
                last_used=Timestamp(asyncio.get_event_loop().time()),
                generation_config=generation_config,
                status="loaded"
            )

        logger.info(f"Successfully loaded {model_id_typed} ({memory_used:.2f} GB)")

    async def unload_model(self, model_id: str) -> None:
        """Unload a model from memory."""
        async with self.lock:
            try:
                model_id_wrapper = create_model_id(model_id)
                model_id_typed = model_id_wrapper.unwrap()
            except ValueError as e:
                logger.error(f"Invalid model ID: {model_id}, error: {e}")
                return

            if model_id_typed not in self.models:
                logger.warning(f"Model {model_id_typed} not loaded")
                return

            logger.info(f"Unloading model: {model_id_typed}")

            # Remove from models dict
            model_info = self.models.pop(model_id_typed)

            # Clear references to allow garbage collection
            model = model_info.model.unwrap()
            tokenizer = model_info.tokenizer.unwrap()

            # Delete references
            del model
            del tokenizer
            del model_info

            # Force garbage collection
            _ = gc.collect()
            torch.cuda.empty_cache()

            logger.info(f"Successfully unloaded {model_id_typed}")

    async def generate_response(
        self,
        model_id: str,
        prompt: str,
        max_new_tokens: int = 80,
        temperature: float = 0.6,
        top_p: float = 0.85,
        repetition_penalty: float = 1.25,
        no_repeat_ngram_size: int = 3,
        do_sample: bool = True
    ) -> Tuple[str, int]:
        """Generate a response using the specified model.

        Ensures the model is loaded without holding the manager lock, and
        only locks around brief metadata updates. This avoids re-entrant lock
        deadlocks and improves concurrency.
        """
        # Validate and wrap inputs
        try:
            model_id_wrapper = create_model_id(model_id)
            model_id_typed = model_id_wrapper.unwrap()
            from custom_types import (
                create_max_new_tokens,
                create_temperature,
                create_top_p,
                create_repetition_penalty,
                create_no_repeat_ngram_size,
            )
            max_tokens_wrapper = create_max_new_tokens(max_new_tokens)
            temp_wrapper = create_temperature(temperature)
            top_p_wrapper = create_top_p(top_p)
            rep_penalty_wrapper = create_repetition_penalty(repetition_penalty)
            ngram_size_wrapper = create_no_repeat_ngram_size(no_repeat_ngram_size)
        except ValueError as e:
            logger.error(f"Invalid generation parameters: {e}")
            raise

        # Ensure the model is loaded (this acquires the internal lock itself)
        async with self.lock:
            is_loaded = model_id_typed in self.models
        if not is_loaded:
            await self.load_model(model_id)

        # Snapshot the model info and update last_used under the lock
        async with self.lock:
            model_info = self.models[model_id_typed]
            self.models[model_id_typed] = ModelInfo(
                model=model_info.model,
                tokenizer=model_info.tokenizer,
                memory_used=model_info.memory_used,
                last_used=Timestamp(asyncio.get_event_loop().time()),
                generation_config=model_info.generation_config,
                status=model_info.status,
            )

        # Heavy work outside the lock
        try:
            # Get model config for context length
            config = MODELS_CONFIG.get(model_id_typed, ModelConfig())
            context_length = int(config.context_length)

            inputs_raw = model_info.tokenizer.unwrap()(
                prompt,
                return_tensors="pt",
                truncation=True,
                max_length=context_length,
            ).to(self.device.unwrap())
            inputs = create_torch_tensor(inputs_raw)

            # Generate response
            with torch.no_grad():
                outputs_raw = model_info.model.unwrap().generate(
                    **inputs.unwrap(),
                    max_new_tokens=max_tokens_wrapper.unwrap(),
                    temperature=temp_wrapper.unwrap(),
                    top_p=top_p_wrapper.unwrap(),
                    repetition_penalty=rep_penalty_wrapper.unwrap(),
                    no_repeat_ngram_size=ngram_size_wrapper.unwrap(),
                    do_sample=do_sample,
                    eos_token_id=model_info.tokenizer.unwrap().eos_token_id,
                    pad_token_id=model_info.tokenizer.unwrap().pad_token_id,
                    use_cache=True,
                )
                outputs = create_torch_tensor(outputs_raw)

            # Decode response (raw, no post-processing)
            input_length: int = inputs.unwrap().input_ids.shape[1]
            generated_tokens = create_torch_tensor(outputs.unwrap()[0][input_length:])
            response: str = model_info.tokenizer.unwrap().decode(
                generated_tokens.unwrap(), skip_special_tokens=True
            )

            return response, len(generated_tokens.unwrap())

        except Exception as e:
            logger.error(f"Error generating response with {model_id_typed}: {str(e)}")
            raise

    def get_available_models(self) -> List[str]:
        """Get list of available models."""
        # Include explicitly configured models
        configured_models: List[str] = cast(List[str], list(MODELS_CONFIG.keys()))

        # Add additional whitelisted models
        whitelisted_models: List[str] = [
            "mistralai/Mistral-7B-Instruct-v0.3"
        ]

        available_models: List[str] = configured_models + whitelisted_models

        # Remove duplicates while preserving order
        seen: Set[str] = set()
        unique_models: List[str] = []

        for model in available_models:
            typed_model = cast(str, model)
            if typed_model not in seen:
                seen.add(typed_model)
                unique_models.append(typed_model)

        return unique_models

    def get_loaded_models(self) -> List[str]:
        """Get list of currently loaded models."""
        return [str(model_id) for model_id in self.models.keys()]

    def get_gpu_memory_usage(self) -> GPUMemoryUsageWrapper:
        """Get current GPU memory usage as a fraction."""
        if not torch.cuda.is_available():
            return create_gpu_memory_usage(0.0)

        allocated: float = torch.cuda.memory_allocated() / 1024**3  # Convert to GB
        total: float = float(torch.cuda.get_device_properties(0).total_memory / 1024**3)  # Convert to GB

        usage: float = allocated / total if total > 0 else 0.0
        return create_gpu_memory_usage(usage)

    def get_total_gpu_memory(self) -> float:
        """Get total GPU memory in GB."""
        if not torch.cuda.is_available():
            return 0.0

        return float(torch.cuda.get_device_properties(0).total_memory / 1024**3)

    def _calculate_model_memory(self, model: HuggingFaceModel) -> ModelMemoryUsage:
        """Calculate approximate memory usage of a model in GB."""
        num_params: int = sum(p.numel() for p in model.unwrap().parameters())
        # Approximate: 2 bytes per parameter for fp16
        memory_gb: float = (num_params * 2) / 1024**3
        return ModelMemoryUsage(memory_gb)

    async def _free_memory(self) -> None:
        """Free memory by unloading least recently used models."""
        if not self.models:
            return

        # Sort models by last used time
        sorted_models = sorted(
            self.models.items(),
            key=lambda x: x[1].last_used
        )

        # Unload least recently used model
        oldest_model_id = sorted_models[0][0]
        logger.info(f"Freeing memory by unloading {oldest_model_id}")
        await self.unload_model(str(oldest_model_id))

    async def cleanup(self) -> None:
        """Clean up all models and resources."""
        logger.info("Cleaning up model manager...")

        # Unload all models
        for model_id in list(self.models.keys()):
            await self.unload_model(str(model_id))

        # Final cleanup
        _ = gc.collect()
        torch.cuda.empty_cache()

        logger.info("Model manager cleanup complete")