production_deploy.py

#!/usr/bin/env python3
"""
Production TRM Deployment Example
Complete production-ready deployment with FastAPI, monitoring, and scaling
"""

import asyncio
import logging
import time
import uuid
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import List, Optional, Dict, Any

import torch
import uvicorn
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
import prometheus_client as prom

# Import our TRM model
from minimal_trm import MinimalTRM

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

# Prometheus metrics
REQUEST_COUNT = prom.Counter('trm_requests_total', 'Total requests', ['method', 'endpoint'])
REQUEST_DURATION = prom.Histogram('trm_request_duration_seconds', 'Request duration')
MODEL_INFERENCE_TIME = prom.Histogram('trm_inference_duration_seconds', 'Model inference time')
ACTIVE_CONNECTIONS = prom.Gauge('trm_active_connections', 'Active connections')
MODEL_MEMORY_USAGE = prom.Gauge('trm_model_memory_bytes', 'Model memory usage')

# Global model instance
model: Optional[MinimalTRM] = None


@dataclass
class ModelConfig:
    """Configuration for TRM model"""
    vocab_size: int = 10000
    d_model: int = 256
    num_recursive_steps: int = 4
    max_seq_len: int = 512
    model_path: Optional[str] = None


class GenerationRequest(BaseModel):
    """Request model for text generation"""
    input_text: str = Field(..., description="Input text for generation", min_length=1, max_length=1000)
    max_length: int = Field(50, description="Maximum generation length", ge=1, le=500)
    temperature: float = Field(1.0, description="Sampling temperature", ge=0.1, le=2.0)
    top_k: int = Field(50, description="Top-k sampling parameter", ge=1, le=1000)
    do_sample: bool = Field(True, description="Whether to use sampling")
    return_attention: bool = Field(False, description="Whether to return attention weights")


class GenerationResponse(BaseModel):
    """Response model for text generation"""
    generated_text: str = Field(..., description="Generated text")
    input_text: str = Field(..., description="Original input text")
    generation_time: float = Field(..., description="Time taken for generation (seconds)")
    model_info: Dict[str, Any] = Field(..., description="Model information")
    request_id: str = Field(..., description="Unique request identifier")


class HealthResponse(BaseModel):
    """Health check response"""
    status: str = Field(..., description="Service status")
    model_loaded: bool = Field(..., description="Whether model is loaded")
    memory_usage: float = Field(..., description="Memory usage (MB)")
    uptime: float = Field(..., description="Service uptime (seconds)")


class ErrorResponse(BaseModel):
    """Error response model"""
    error: str = Field(..., description="Error message")
    detail: Optional[str] = Field(None, description="Additional error details")
    request_id: Optional[str] = Field(None, description="Request identifier")


# Simple tokenizer for demonstration (in production, use proper tokenizer)
class SimpleTokenizer:
    """Simple character-level tokenizer for demonstration"""

    def __init__(self, vocab_size: int = 10000):
        self.vocab_size = vocab_size
        # Create simple vocabulary (in production, load from file)
        self.chars = list("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 .,!?:;'\n")
        self.char_to_id = {c: i for i, c in enumerate(self.chars)}
        self.id_to_char = {i: c for c, i in self.char_to_id.items()}

    def encode(self, text: str) -> List[int]:
        """Encode text to token IDs"""
        return [self.char_to_id.get(c, 0) for c in text[:512]]

    def decode(self, token_ids: List[int]) -> str:
        """Decode token IDs to text"""
        return ''.join([self.id_to_char.get(id, '<unk>') for id in token_ids if id < len(self.id_to_char)])


tokenizer = SimpleTokenizer()


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan manager"""
    global model

    # Startup
    logger.info("🚀 Starting TRM Production Server...")

    # Load model
    try:
        config = ModelConfig()
        model = MinimalTRM(
            vocab_size=config.vocab_size,
            d_model=config.d_model,
            num_recursive_steps=config.num_recursive_steps,
            max_seq_len=config.max_seq_len
        )

        # Load pre-trained weights if available
        if config.model_path:
            logger.info(f"Loading model from {config.model_path}")
            model.load_state_dict(torch.load(config.model_path, map_location='cpu'))

        model.eval()
        logger.info("✅ Model loaded successfully")

        # Log model info
        model_info = model.get_model_info()
        logger.info(f"📊 Model info: {model_info}")

        # Update memory usage metric
        if torch.cuda.is_available():
            MODEL_MEMORY_USAGE.set(torch.cuda.memory_allocated())

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

    yield

    # Shutdown
    logger.info("🛑 Shutting down TRM Production Server...")


# Create FastAPI app
app = FastAPI(
    title="TRM Production API",
    description="Production-ready Tiny Recursive Model deployment",
    version="1.0.0",
    lifespan=lifespan
)

# Add middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production, specify allowed origins
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
app.add_middleware(GZipMiddleware, minimum_size=1000)


@app.middleware("http")
async def add_process_time_header(request, call_next):
    """Add processing time header and metrics"""
    start_time = time.time()

    # Update active connections
    ACTIVE_CONNECTIONS.inc()

    try:
        response = await call_next(request)

        # Calculate processing time
        process_time = time.time() - start_time
        response.headers["X-Process-Time"] = str(process_time)

        # Update metrics
        REQUEST_DURATION.observe(process_time)
        REQUEST_COUNT.labels(method=request.method, endpoint=request.url.path).inc()

        return response
    finally:
        ACTIVE_CONNECTIONS.dec()


@app.get("/", response_model=dict)
async def root():
    """Root endpoint"""
    return {
        "message": "TRM Production API",
        "version": "1.0.0",
        "docs": "/docs",
        "health": "/health"
    }


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

    # Check model status
    model_loaded = model is not None

    # Get memory usage
    memory_usage = 0.0
    if torch.cuda.is_available():
        memory_usage = torch.cuda.memory_allocated() / (1024 * 1024)  # MB

    # Get uptime (simplified)
    uptime = time.time()

    status = "healthy" if model_loaded else "unhealthy"

    return HealthResponse(
        status=status,
        model_loaded=model_loaded,
        memory_usage=memory_usage,
        uptime=uptime
    )


@app.post("/generate", response_model=GenerationResponse)
async def generate_text(request: GenerationRequest, background_tasks: BackgroundTasks):
    """Generate text using TRM model"""
    global model

    request_id = str(uuid.uuid4())

    # Validate model is loaded
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        # Tokenize input
        input_ids = tokenizer.encode(request.input_text)
        if not input_ids:
            raise HTTPException(status_code=400, detail="Empty input after tokenization")

        # Convert to tensor
        input_tensor = torch.tensor([input_ids], dtype=torch.long)

        # Generate text
        start_time = time.time()

        with torch.no_grad():
            generated_ids = model.generate(
                input_tensor,
                max_length=request.max_length,
                temperature=request.temperature,
                do_sample=request.do_sample,
                top_k=request.top_k
            )

        generation_time = time.time() - start_time

        # Update metrics
        MODEL_INFERENCE_TIME.observe(generation_time)

        # Decode generated text
        generated_text = tokenizer.decode(generated_ids[0].tolist())

        # Get model info
        model_info = model.get_model_info()

        # Log generation
        logger.info(f"Generated text for request {request_id}: {len(generated_text)} chars in {generation_time:.3f}s")

        # Add background task for logging
        background_tasks.add_task(log_generation, request_id, request.input_text, generated_text, generation_time)

        return GenerationResponse(
            generated_text=generated_text,
            input_text=request.input_text,
            generation_time=generation_time,
            model_info=model_info,
            request_id=request_id
        )

    except Exception as e:
        logger.error(f"Generation error for request {request_id}: {e}")
        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")


@app.get("/model/info")
async def get_model_info():
    """Get model information"""
    global model

    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    return model.get_model_info()


@app.get("/metrics")
async def get_metrics():
    """Prometheus metrics endpoint"""
    from prometheus_client import generate_latest, CONTENT_TYPE_LATEST
    return JSONResponse(content=generate_latest().decode(), media_type=CONTENT_TYPE_LATEST)


async def log_generation(request_id: str, input_text: str, generated_text: str, generation_time: float):
    """Log generation details (could be sent to external logging service)"""
    logger.info(f"Generation log - ID: {request_id}, Input: {input_text[:50]}..., "
                f"Output: {generated_text[:50]}..., Time: {generation_time:.3f}s")


# Exception handlers
@app.exception_handler(HTTPException)
async def http_exception_handler(request, exc):
    """Handle HTTP exceptions"""
    return JSONResponse(
        status_code=exc.status_code,
        content=ErrorResponse(
            error=exc.detail,
            request_id=str(uuid.uuid4())
        ).dict()
    )


@app.exception_handler(Exception)
async def general_exception_handler(request, exc):
    """Handle general exceptions"""
    logger.error(f"Unhandled exception: {exc}")
    return JSONResponse(
        status_code=500,
        content=ErrorResponse(
            error="Internal server error",
            detail=str(exc),
            request_id=str(uuid.uuid4())
        ).dict()
    )


def main():
    """Main function to run the server"""
    import argparse

    parser = argparse.ArgumentParser(description="TRM Production Server")
    parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
    parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
    parser.add_argument("--workers", type=int, default=1, help="Number of workers")
    parser.add_argument("--reload", action="store_true", help="Enable auto-reload")
    parser.add_argument("--log-level", default="info", help="Log level")

    args = parser.parse_args()

    # Configure logging
    logging.basicConfig(level=getattr(logging, args.log_level.upper()))

    print("🚀 Starting TRM Production Server")
    print(f"📍 Host: {args.host}")
    print(f"🔌 Port: {args.port}")
    print(f"👥 Workers: {args.workers}")
    print(f"📚 Docs: http://{args.host}:{args.port}/docs")
    print(f"📊 Metrics: http://{args.host}:{args.port}/metrics")
    print("=" * 50)

    # Run server
    uvicorn.run(
        "production_deploy:app",
        host=args.host,
        port=args.port,
        workers=args.workers,
        reload=args.reload,
        log_level=args.log_level,
        access_log=True
    )


if __name__ == "__main__":
    main()