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