production_server.py

#!/usr/bin/env python3
"""
Production TTS Server with Micro-Batching

High-concurrency, low-latency TTS server implementing:
- Token-level micro-batching across sessions
- Pre-allocated KV cache for N concurrent sessions
- Streaming audio with TTFB optimization
- Fair scheduling with priority for new sessions

Target: 16-32 concurrent sessions on RTX 4090 with <150ms TTFB
"""

import os
import sys
import time
import asyncio
import logging
import argparse
import struct
from pathlib import Path
from typing import Optional, AsyncGenerator
from contextlib import asynccontextmanager

import numpy as np
import yaml
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel, Field

# Add src to path
sys.path.insert(0, str(Path(__file__).parent))

from src.scheduler import (
    BatchInferenceEngine,
    SchedulerConfig,
    SessionConfig,
    TTSSession,
)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


# ============================================================================
# Configuration
# ============================================================================

def load_config(path: str) -> dict:
    """Load configuration from YAML file."""
    with open(path, 'r') as f:
        return yaml.safe_load(f)


# ============================================================================
# Mock engines for development (replace with real engines in production)
# ============================================================================

class MockLLMEngine:
    """Mock LLM engine for development/testing."""
    
    def __init__(self):
        from transformers import AutoTokenizer
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(
                "meta-llama/Llama-3.2-3B-Instruct",
                trust_remote_code=True,
            )
        except Exception as e:
            logger.warning(f"Could not load tokenizer: {e}")
            self.tokenizer = None
    
    def format_prompt(self, text: str, voice: str = "tara") -> str:
        """Format the TTS prompt."""
        return f"<|audio|>{voice}: {text}<|eoa|>"
    
    async def generate_tokens(self, prompt: str, max_tokens: int = 100):
        """Generate tokens (mock)."""
        for i in range(max_tokens):
            # Simulate audio tokens
            token = 128266 + (i % 4096)
            yield token
            await asyncio.sleep(0.005)  # ~200 tok/s simulation


class MockSNACDecoder:
    """Mock SNAC decoder for development/testing."""
    
    AUDIO_TOKEN_OFFSET = 128266
    TOKENS_PER_FRAME = 7
    SAMPLE_RATE = 24000
    
    def decode(self, tokens: list) -> np.ndarray:
        """Decode tokens to audio (mock - generates silence with click)."""
        if not tokens:
            return np.array([], dtype=np.float32)
        
        # ~10.67ms per token frame
        samples_per_frame = int(0.01067 * self.SAMPLE_RATE)
        total_samples = (len(tokens) // self.TOKENS_PER_FRAME) * samples_per_frame
        
        if total_samples == 0:
            return np.array([], dtype=np.float32)
        
        # Generate low-amplitude white noise as placeholder
        audio = np.random.randn(total_samples).astype(np.float32) * 0.01
        
        return audio


# ============================================================================
# API Models
# ============================================================================

class TTSRequest(BaseModel):
    """TTS request body."""
    text: str = Field(..., min_length=1, max_length=10000)
    voice: str = Field(default="tara")
    response_format: str = Field(default="wav")
    speed: float = Field(default=1.0, ge=0.5, le=2.0)
    
    # Advanced options
    temperature: float = Field(default=0.6, ge=0.0, le=1.5)
    top_p: float = Field(default=0.95, ge=0.0, le=1.0)
    max_tokens: int = Field(default=4096, ge=1, le=8192)


class HealthResponse(BaseModel):
    """Health check response."""
    status: str
    gpu_available: bool
    engine_stats: dict


class StatsResponse(BaseModel):
    """Server statistics response."""
    scheduler: dict
    kv_cache: dict
    total_batches: int
    total_tokens: int
    total_audio_chunks: int


# ============================================================================
# Audio Utilities
# ============================================================================

def create_wav_header(sample_rate: int = 24000, bits_per_sample: int = 16) -> bytes:
    """Create WAV header for streaming (unknown length)."""
    return struct.pack(
        '<4sI4s4sIHHIIHH4sI',
        b'RIFF',
        0xFFFFFFFF,  # Placeholder for unknown length
        b'WAVE',
        b'fmt ',
        16,  # Subchunk1Size
        1,   # AudioFormat (PCM)
        1,   # NumChannels
        sample_rate,
        sample_rate * bits_per_sample // 8,  # ByteRate
        bits_per_sample // 8,  # BlockAlign
        bits_per_sample,
        b'data',
        0xFFFFFFFF,  # Placeholder for unknown length
    )


# ============================================================================
# Global state
# ============================================================================

engine: Optional[BatchInferenceEngine] = None


# ============================================================================
# FastAPI Application
# ============================================================================

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan handler."""
    global engine
    
    logger.info("Starting production TTS server...")
    
    # Load configuration
    config_path = os.environ.get("CONFIG_PATH", "config/model_config.yaml")
    
    try:
        config = load_config(config_path)
    except Exception as e:
        logger.warning(f"Could not load config: {e}, using defaults")
        config = {}
    
    # Create scheduler config
    server_config = config.get("server", {})
    scheduler_config = SchedulerConfig(
        max_concurrent_sessions=server_config.get("max_concurrent_sessions", 32),
        max_batch_size=server_config.get("max_batch_size", 16),
        ttfb_priority_tokens=server_config.get("ttfb_priority_tokens", 20),
        flush_first_n_tokens=7,   # SNAC needs 7 tokens for first frame
        flush_every_n_tokens=14,  # ~160ms audio chunks
        max_new_tokens_per_call=4096,
        kv_cache_slots=32,
        max_context_len=4096,
    )
    
    # Initialize engines (mock for development)
    logger.info("Initializing engines...")
    llm = MockLLMEngine()
    snac = MockSNACDecoder()
    
    # Create batch inference engine
    engine = BatchInferenceEngine(
        llm_engine=llm,
        snac_decoder=snac,
        config=scheduler_config,
    )
    
    # Start the engine
    await engine.start()
    
    logger.info("Production TTS server ready")
    logger.info(f"  Max concurrent sessions: {scheduler_config.max_concurrent_sessions}")
    logger.info(f"  Max batch size: {scheduler_config.max_batch_size}")
    
    yield
    
    # Shutdown
    logger.info("Shutting down...")
    if engine:
        await engine.stop()


app = FastAPI(
    title="Orpheus TTS Production Server",
    description="High-concurrency, low-latency TTS with micro-batching",
    version="2.0.0",
    lifespan=lifespan,
)


# ============================================================================
# Endpoints
# ============================================================================

@app.get("/health", response_model=HealthResponse)
async def health():
    """Health check endpoint."""
    import torch
    
    stats = engine.get_stats() if engine else {}
    
    return HealthResponse(
        status="healthy" if engine else "starting",
        gpu_available=torch.cuda.is_available(),
        engine_stats=stats,
    )


@app.get("/v1/stats", response_model=StatsResponse)
async def get_stats():
    """Get detailed server statistics."""
    if not engine:
        raise HTTPException(status_code=503, detail="Engine not ready")
    
    stats = engine.get_stats()
    return StatsResponse(**stats)


@app.get("/v1/voices")
async def list_voices():
    """List available voices."""
    return {
        "voices": [
            {"id": "tara", "name": "Tara", "description": "Default female voice"},
            {"id": "leah", "name": "Leah", "description": "Female voice"},
            {"id": "jess", "name": "Jess", "description": "Female voice"},
            {"id": "leo", "name": "Leo", "description": "Male voice"},
            {"id": "dan", "name": "Dan", "description": "Male voice"},
            {"id": "mia", "name": "Mia", "description": "Female voice"},
            {"id": "zac", "name": "Zac", "description": "Male voice"},
            {"id": "zoe", "name": "Zoe", "description": "Female voice"},
        ]
    }


async def stream_audio(session: TTSSession) -> AsyncGenerator[bytes, None]:
    """Stream audio from a session."""
    # Send WAV header first
    yield create_wav_header()
    
    bytes_sent = 0
    start_time = time.perf_counter()
    
    try:
        while True:
            try:
                # Wait for audio chunk with timeout
                chunk = await asyncio.wait_for(
                    session.audio_queue.get(),
                    timeout=30.0
                )
                
                if chunk is None:
                    # End of stream
                    break
                
                bytes_sent += len(chunk)
                yield chunk
                
            except asyncio.TimeoutError:
                logger.warning(f"Session {session.session_id[:8]} timed out")
                break
                
    finally:
        elapsed = time.perf_counter() - start_time
        logger.info(
            f"Session {session.session_id[:8]} streamed "
            f"{bytes_sent/1024:.1f}KB in {elapsed:.2f}s"
        )


@app.post("/v1/audio/speech")
async def create_speech(request: TTSRequest):
    """
    Generate speech from text with streaming response.
    
    Compatible with OpenAI TTS API.
    """
    if not engine:
        raise HTTPException(status_code=503, detail="Engine not ready")
    
    try:
        # Create session config
        session_config = SessionConfig(
            temperature=request.temperature,
            top_p=request.top_p,
            max_new_tokens=request.max_tokens,
        )
        
        # Submit request
        session = await engine.submit_request(
            text=request.text,
            voice=request.voice,
            config=session_config,
        )
        
        # Return streaming response
        return StreamingResponse(
            stream_audio(session),
            media_type="audio/wav",
            headers={
                "X-Session-ID": session.session_id,
                "Transfer-Encoding": "chunked",
            }
        )
        
    except RuntimeError as e:
        raise HTTPException(status_code=429, detail=str(e))
    except Exception as e:
        logger.error(f"Speech generation error: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.websocket("/v1/audio/stream")
async def websocket_stream(websocket: WebSocket):
    """
    WebSocket endpoint for real-time streaming TTS.
    
    Protocol:
    - Client sends: {"text": "...", "voice": "...", ...}
    - Server sends: binary audio chunks
    - Server sends: {"done": true, "metrics": {...}} when complete
    """
    await websocket.accept()
    
    try:
        while True:
            # Receive request
            data = await websocket.receive_json()
            
            text = data.get("text", "")
            voice = data.get("voice", "tara")
            
            if not text:
                await websocket.send_json({"error": "No text provided"})
                continue
            
            try:
                # Submit request
                session = await engine.submit_request(text=text, voice=voice)
                
                # Stream audio chunks
                while True:
                    try:
                        chunk = await asyncio.wait_for(
                            session.audio_queue.get(),
                            timeout=30.0
                        )
                        
                        if chunk is None:
                            # Send completion message
                            metrics = session.get_metrics()
                            await websocket.send_json({
                                "done": True,
                                "metrics": metrics
                            })
                            break
                        
                        # Send binary audio
                        await websocket.send_bytes(chunk)
                        
                    except asyncio.TimeoutError:
                        await websocket.send_json({"error": "Timeout"})
                        break
                        
            except RuntimeError as e:
                await websocket.send_json({"error": str(e)})
                
    except WebSocketDisconnect:
        logger.info("WebSocket client disconnected")
    except Exception as e:
        logger.error(f"WebSocket error: {e}")


# ============================================================================
# gRPC-style endpoint for backend-to-backend
# ============================================================================

@app.post("/v1/audio/generate")
async def generate_audio(request: TTSRequest):
    """
    Generate complete audio (non-streaming).
    
    For batch processing or backend-to-backend calls.
    """
    if not engine:
        raise HTTPException(status_code=503, detail="Engine not ready")
    
    try:
        session = await engine.submit_request(
            text=request.text,
            voice=request.voice,
        )
        
        # Collect all audio
        chunks = []
        while True:
            try:
                chunk = await asyncio.wait_for(
                    session.audio_queue.get(),
                    timeout=60.0
                )
                
                if chunk is None:
                    break
                chunks.append(chunk)
                
            except asyncio.TimeoutError:
                break
        
        # Combine and return
        audio_data = b''.join(chunks)
        wav_header = create_wav_header()
        
        return StreamingResponse(
            iter([wav_header + audio_data]),
            media_type="audio/wav",
            headers={
                "X-Session-ID": session.session_id,
                "X-TTFB-Ms": str(int(session.get_ttfb_ms())),
                "X-Total-Tokens": str(session.total_tokens),
            }
        )
        
    except RuntimeError as e:
        raise HTTPException(status_code=429, detail=str(e))


# ============================================================================
# Main
# ============================================================================

def main():
    parser = argparse.ArgumentParser(description="Production TTS 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("--config", default="config/model_config.yaml", help="Config file")
    parser.add_argument("--workers", type=int, default=1, help="Number of workers")
    parser.add_argument("--reload", action="store_true", help="Enable auto-reload")
    
    args = parser.parse_args()
    
    # Set config path
    os.environ["CONFIG_PATH"] = args.config
    
    import uvicorn
    
    uvicorn.run(
        "production_server:app",
        host=args.host,
        port=args.port,
        workers=args.workers,
        reload=args.reload,
        log_level="info",
    )


if __name__ == "__main__":
    main()