Model Serving Patterns

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import torch
import asyncio
import time
from typing import List, Optional, Dict, Any
import logging
from contextlib import asynccontextmanager

# Request/Response models
class InferenceRequest(BaseModel):
    text: str = Field(..., description="Input text for inference")
    max_length: Optional[int] = Field(100, description="Maximum output length")
    temperature: Optional[float] = Field(0.7, description="Sampling temperature")
    top_p: Optional[float] = Field(0.9, description="Top-p sampling")
    
class InferenceResponse(BaseModel):
    generated_text: str
    inference_time: float
    model_version: str
    request_id: str

class BatchInferenceRequest(BaseModel):
    requests: List[InferenceRequest]
    batch_id: Optional[str] = None

class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    gpu_available: bool
    memory_usage: Dict[str, float]

# Global model instance
model_manager = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    global model_manager
    model_manager = ModelManager()
    await model_manager.load_model()
    yield
    # Shutdown
    await model_manager.cleanup()

# Initialize FastAPI app
app = FastAPI(
    title="Production ML Model API",
    description="High-performance model serving with FastAPI",
    version="1.0.0",
    lifespan=lifespan
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class ModelManager:
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.device = None
        self.batch_processor = None
        self.model_version = "1.0.0"
        
    async def load_model(self):
        """Load model with optimizations"""
        try:
            # Check GPU availability
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            
            # Load model and tokenizer
            from transformers import AutoModelForCausalLM, AutoTokenizer
            
            model_name = "gpt2"  # Replace with your model
            
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            self.model = AutoModelForCausalLM.from_pretrained(
                model_name,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                device_map="auto" if torch.cuda.is_available() else None
            )
            
            # Optimize model
            if torch.cuda.is_available():
                self.model = torch.compile(self.model)
            
            self.model.eval()
            
            # Initialize batch processor
            self.batch_processor = BatchProcessor(self.model, self.tokenizer, self.device)
            
            logging.info(f"Model loaded successfully on {self.device}")
            
        except Exception as e:
            logging.error(f"Failed to load model: {e}")
            raise
    
    async def generate_text(self, 
                          text: str, 
                          max_length: int = 100,
                          temperature: float = 0.7,
                          top_p: float = 0.9) -> Dict[str, Any]:
        """Generate text with the model"""
        
        start_time = time.time()
        
        try:
            # Tokenize input
            inputs = self.tokenizer.encode(text, return_tensors="pt").to(self.device)
            
            # Generate
            with torch.inference_mode():
                outputs = self.model.generate(
                    inputs,
                    max_new_tokens=max_length,
                    temperature=temperature,
                    top_p=top_p,
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id
                )
            
            # Decode output
            generated_text = self.tokenizer.decode(
                outputs[0][inputs.shape[1]:], 
                skip_special_tokens=True
            )
            
            inference_time = time.time() - start_time
            
            return {
                "generated_text": generated_text,
                "inference_time": inference_time,
                "model_version": self.model_version
            }
            
        except Exception as e:
            logging.error(f"Inference failed: {e}")
            raise HTTPException(status_code=500, detail=str(e))
    
    async def batch_generate(self, requests: List[InferenceRequest]) -> List[Dict[str, Any]]:
        """Process batch of requests efficiently"""
        
        return await self.batch_processor.process_batch(requests)
    
    def get_health_status(self) -> HealthResponse:
        """Get system health status"""
        
        gpu_available = torch.cuda.is_available()
        model_loaded = self.model is not None
        
        memory_usage = {}
        if gpu_available:
            memory_usage = {
                "gpu_allocated_gb": torch.cuda.memory_allocated() / 1e9,
                "gpu_reserved_gb": torch.cuda.memory_reserved() / 1e9,
            }
        
        status = "healthy" if (model_loaded and (gpu_available or not torch.cuda.is_available())) else "unhealthy"
        
        return HealthResponse(
            status=status,
            model_loaded=model_loaded,
            gpu_available=gpu_available,
            memory_usage=memory_usage
        )
    
    async def cleanup(self):
        """Clean up resources"""
        if self.model:
            del self.model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

class BatchProcessor:
    """Efficient batch processing for multiple requests"""
    
    def __init__(self, model, tokenizer, device):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.max_batch_size = 8  # Adjust based on GPU memory
        
    async def process_batch(self, requests: List[InferenceRequest]) -> List[Dict[str, Any]]:
        """Process batch of requests"""
        
        batch_size = min(len(requests), self.max_batch_size)
        results = []
        
        # Process in chunks
        for i in range(0, len(requests), batch_size):
            batch = requests[i:i + batch_size]
            batch_results = await self._process_batch_chunk(batch)
            results.extend(batch_results)
        
        return results
    
    async def _process_batch_chunk(self, requests: List[InferenceRequest]) -> List[Dict[str, Any]]:
        """Process a single batch chunk"""
        
        start_time = time.time()
        
        try:
            # Prepare batch inputs
            texts = [req.text for req in requests]
            max_lengths = [req.max_length for req in requests]
            
            # Tokenize batch
            batch_inputs = self.tokenizer(
                texts,
                padding=True,
                truncation=True,
                return_tensors="pt",
                max_length=512
            ).to(self.device)
            
            # Generate for batch
            with torch.inference_mode():
                batch_outputs = self.model.generate(
                    **batch_inputs,
                    max_new_tokens=max(max_lengths),
                    temperature=requests[0].temperature,  # Use first request's params
                    top_p=requests[0].top_p,
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id
                )
            
            # Process outputs
            results = []
            for i, (request, output) in enumerate(zip(requests, batch_outputs)):
                # Extract new tokens
                input_length = batch_inputs.input_ids[i].shape[0]
                generated_tokens = output[input_length:]
                
                generated_text = self.tokenizer.decode(
                    generated_tokens, 
                    skip_special_tokens=True
                )
                
                results.append({
                    "generated_text": generated_text,
                    "inference_time": (time.time() - start_time) / len(requests),
                    "model_version": "1.0.0"
                })
            
            return results
            
        except Exception as e:
            logging.error(f"Batch processing failed: {e}")
            # Return error for all requests in batch
            return [{"error": str(e)} for _ in requests]

# API Endpoints
@app.post("/v1/generate", response_model=InferenceResponse)
async def generate_text(request: InferenceRequest, background_tasks: BackgroundTasks):
    """Generate text from input"""
    
    request_id = f"req_{int(time.time() * 1000)}"
    
    try:
        result = await model_manager.generate_text(
            text=request.text,
            max_length=request.max_length,
            temperature=request.temperature,
            top_p=request.top_p
        )
        
        # Log request in background
        background_tasks.add_task(log_request, request_id, request.text, result)
        
        return InferenceResponse(
            generated_text=result["generated_text"],
            inference_time=result["inference_time"],
            model_version=result["model_version"],
            request_id=request_id
        )
        
    except Exception as e:
        logging.error(f"Generation failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/v1/batch-generate")
async def batch_generate_text(request: BatchInferenceRequest):
    """Generate text for batch of requests"""
    
    try:
        results = await model_manager.batch_generate(request.requests)
        
        return {
            "results": results,
            "batch_id": request.batch_id,
            "processed_count": len(results)
        }
        
    except Exception as e:
        logging.error(f"Batch generation failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint"""
    return model_manager.get_health_status()

@app.get("/metrics")
async def get_metrics():
    """Prometheus metrics endpoint"""
    # Return metrics in Prometheus format
    return {"message": "Implement Prometheus metrics here"}

def log_request(request_id: str, input_text: str, result: Dict):
    """Log request for monitoring"""
    logging.info(f"Request {request_id}: {len(input_text)} chars -> {result['inference_time']:.3f}s")

# Run with: uvicorn app:app --host 0.0.0.0 --port 8000 --workers 1