kaggle_llm_engine.py

"""
Kaggle-Optimized LLM Engine - 4-bit Qwen3.5-35B with multi-GPU sharding
Production-ready for T4x2 with aggressive VRAM optimizations
"""

import torch
import json
import logging
import os
from typing import Dict, List, Optional, Union
from dataclasses import dataclass

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class KaggleConfig:
    """Kaggle-specific configuration"""
    mock_mode: bool = False  # Toggle between Mock and Real LLM
    model_name: str = "Qwen/Qwen2.5-35B-Instruct"
    max_memory_per_gpu: float = 15.5  # GB per T4 (leave some margin)
    use_flash_attention: bool = True
    use_torch_compile: bool = False  # Disable for compatibility
    inference_mode: bool = True
    gradient_checkpointing: bool = False

class KaggleLLMEngine:
    """
    Production-ready LLM engine for Kaggle T4x2
    """
    
    def __init__(self, config: KaggleConfig):
        self.config = config
        self.model = None
        self.tokenizer = None
        self.device_map = "auto"
        
        # Mock LLM for testing
        if config.mock_mode:
            from agent.mock_llm import get_mock_llm
            self.mock_llm = get_mock_llm()
            logger.info("Using Mock LLM for testing")
        else:
            self.mock_llm = None
            self._setup_quantization_config()
            self._setup_generation_config()
    
    def _setup_quantization_config(self):
        """Setup 4-bit quantization with aggressive memory optimization"""
        try:
            from transformers import BitsAndBytesConfig
            
            self.quantization_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
            )
            logger.info("4-bit quantization config loaded")
        except ImportError:
            logger.error("bitsandbytes not available")
            self.quantization_config = None
    
    def _setup_generation_config(self):
        """Setup generation config for speed and memory efficiency"""
        self.generation_config = {
            "max_new_tokens": 256,  # Aggressive token limit
            "temperature": 0.1,
            "top_p": 0.95,
            "do_sample": False,
            "pad_token_id": 151643,
            "eos_token_id": 151645,
            "use_cache": True,
            "return_dict_in_generate": True,
        }
    
    def load_model(self) -> bool:
        """
        Load model with multi-GPU sharding and memory optimizations
        """
        if self.config.mock_mode:
            logger.info("Mock mode enabled - skipping model loading")
            return True
        
        try:
            logger.info("Loading tokenizer...")
            from transformers import AutoTokenizer
            
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.config.model_name,
                trust_remote_code=True,
                padding_side="left"
            )
            
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token
            
            logger.info("Loading model with multi-GPU sharding...")
            self._load_model_with_sharding()
            
            logger.info("Applying inference optimizations...")
            self._apply_inference_optimizations()
            
            logger.info("Model loaded successfully")
            self._log_memory_usage()
            return True
            
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            return False
    
    def _load_model_with_sharding(self):
        """Load model with aggressive memory management"""
        from transformers import AutoModelForCausalLM
        import gc
        
        # Clear cache before loading
        gc.collect()
        torch.cuda.empty_cache()
        
        # Load with device_map="auto" for automatic sharding
        self.model = AutoModelForCausalLM.from_pretrained(
            self.config.model_name,
            quantization_config=self.quantization_config,
            device_map=self.device_map,
            trust_remote_code=True,
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True,
            # Memory optimizations
            use_cache=True,
            attn_implementation="flash_attention_2" if self.config.use_flash_attention else "sdpa",
        )
        
        # Configure model for inference
        self.model.config.use_cache = True
        if self.config.gradient_checkpointing:
            self.model.gradient_checkpointing_disable()
    
    def _apply_inference_optimizations(self):
        """Apply inference-specific optimizations"""
        if self.config.inference_mode:
            self.model = torch.inference_mode(self.model)
        
        # Disable gradient computation
        for param in self.model.parameters():
            param.requires_grad = False
        
        # Enable memory-efficient attention if available
        if hasattr(self.model, 'enable_attention_slicing'):
            self.model.enable_attention_slicing()
        
        # Enable CPU offloading for very large models if needed
        if hasattr(self.model, 'enable_cpu_offload'):
            # Only enable if we're running out of GPU memory
            pass
    
    def _log_memory_usage(self):
        """Log detailed memory usage"""
        if torch.cuda.is_available():
            for i in range(torch.cuda.device_count()):
                allocated = torch.cuda.memory_allocated(i) / 1024**3
                cached = torch.cuda.memory_reserved(i) / 1024**3
                max_allocated = torch.cuda.max_memory_allocated(i) / 1024**3
                
                logger.info(f"GPU {i}: {allocated:.2f}GB allocated, {cached:.2f}GB cached, {max_allocated:.2f}GB max")
    
    def generate_response(self, prompt: str, max_new_tokens: Optional[int] = None) -> str:
        """
        Generate response with optimized inference
        """
        if self.config.mock_mode:
            # Use Mock LLM
            response = self.mock_llm.generate_response(prompt, max_new_tokens)
            return response
        
        if not self.model or not self.tokenizer:
            raise RuntimeError("Model not loaded")
        
        try:
            # Override generation config if provided
            gen_config = self.generation_config.copy()
            if max_new_tokens:
                gen_config["max_new_tokens"] = max_new_tokens
            
            # Tokenize with memory efficiency
            inputs = self.tokenizer(
                prompt,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=2048,  # Conservative context window
            ).to(self.model.device)
            
            # Generate with memory management
            with torch.inference_mode():
                outputs = self.model.generate(
                    **inputs,
                    **gen_config
                )
            
            # Decode response
            response = self.tokenizer.decode(
                outputs[0][inputs['input_ids'].shape[1]:],
                skip_special_tokens=True
            )
            
            # Cleanup
            del inputs, outputs
            torch.cuda.empty_cache()
            
            return response.strip()
            
        except Exception as e:
            logger.error(f"Generation failed: {e}")
            return ""
    
    def generate_json_response(self, prompt: str) -> Dict:
        """
        Generate JSON response for ReAct agent
        """
        if self.config.mock_mode:
            return self.mock_llm.generate_json_response(prompt)
        
        # Add JSON formatting instruction
        json_prompt = f"""{prompt}

Respond in strict JSON format with these keys:
- "thought": Your reasoning about the current situation
- "action": The tool to execute (e.g., "get_node_status", "execute_command")
- "action_input": The parameters for the action (JSON object)

Example:
{{"thought": "I need to check the node status", "action": "get_node_status", "action_input": {{"node_id": "Node-01"}}}}

Your response:"""
        
        response = self.generate_response(json_prompt, max_new_tokens=256)
        
        try:
            # Extract JSON from response
            import re
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                return json.loads(json_match.group())
            else:
                return {
                    "thought": response[:100],
                    "action": "error",
                    "action_input": {"error": "Failed to parse JSON response"}
                }
        except Exception as e:
            logger.error(f"JSON parsing failed: {e}")
            return {
                "thought": "JSON parsing failed",
                "action": "error",
                "action_input": {"error": str(e)}
            }
    
    def get_memory_stats(self) -> Dict:
        """Get detailed memory statistics"""
        stats = {}
        
        if torch.cuda.is_available():
            for i in range(torch.cuda.device_count()):
                stats[f"gpu_{i}"] = {
                    "allocated_gb": torch.cuda.memory_allocated(i) / 1024**3,
                    "cached_gb": torch.cuda.memory_reserved(i) / 1024**3,
                    "max_allocated_gb": torch.cuda.max_memory_allocated(i) / 1024**3,
                    "utilization_percent": (torch.cuda.memory_allocated(i) / (self.config.max_memory_per_gpu * 1024**3)) * 100
                }
        
        return stats
    
    def cleanup(self):
        """Clean up model and tokenizer from memory"""
        if self.model:
            del self.model
        if self.tokenizer:
            del self.tokenizer
        
        torch.cuda.empty_cache()
        import gc
        gc.collect()
        
        logger.info("Model and tokenizer cleaned up from memory")
    
    def health_check(self) -> Dict:
        """Perform health check on the LLM engine"""
        if self.config.mock_mode:
            return {
                "status": "healthy",
                "mode": "mock",
                "model": self.mock_llm.model_name
            }
        
        try:
            # Test generation
            test_prompt = "Test prompt"
            response = self.generate_response(test_prompt, max_new_tokens=10)
            
            # Get memory stats
            memory_stats = self.get_memory_stats()
            
            return {
                "status": "healthy",
                "mode": "production",
                "model": self.config.model_name,
                "test_response": response[:50],
                "memory_stats": memory_stats
            }
            
        except Exception as e:
            return {
                "status": "unhealthy",
                "error": str(e)
            }


# Factory function
def create_kaggle_llm_engine(mock_mode: bool = False) -> KaggleLLMEngine:
    """Create LLM engine with Kaggle optimizations"""
    config = KaggleConfig(mock_mode=mock_mode)
    return KaggleLLMEngine(config)


if __name__ == "__main__":
    # Test the Kaggle engine
    print("Testing Kaggle LLM Engine...")
    
    # Test with Mock mode first
    engine = create_kaggle_llm_engine(mock_mode=True)
    
    if engine.load_model():
        print("Mock engine loaded successfully")
        
        # Test generation
        response = engine.generate_response("Test prompt", max_new_tokens=50)
        print(f"Mock response: {response}")
        
        # Test JSON generation
        json_response = engine.generate_json_response("I need to check node status")
        print(f"JSON response: {json_response}")
        
        # Health check
        health = engine.health_check()
        print(f"Health check: {health}")
        
        engine.cleanup()
    
    print("Kaggle engine test completed")