ml-service-integration

ML Service Integration Patterns

🤖 GPU-ACCELERATED ML SERVICE PATTERNS (From ml_inference_fastapi_app)

Proven patterns for building high-performance ML inference services with GPU optimization.

⚡ GPU RESOURCE MANAGEMENT:

1. GPU Lock Pattern for Exclusive Access

# ✅ MANDATORY for GPU services - prevents OOM crashes
import asyncio
import torch

# Global GPU resource lock
gpu_lock = asyncio.Lock()

async def gpu_inference_operation(data):
    """Ensure only one operation uses GPU at a time."""
    async with gpu_lock:
        try:
            # GPU computation here
            with torch.inference_mode():
                with torch.amp.autocast("cuda", enabled=(device.type == "cuda")):
                    result = model(data)
            return result
        finally:
            # CRITICAL: Always cleanup GPU memory
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

2. GPU Memory Probing Pattern

# ✅ Dynamic batch size calculation based on available GPU memory
def probe_safe_batch_size(model, input_shape):
    """Determine safe batch size to prevent OOM errors."""
    if device.type != "cuda":
        return 1
    
    try:
        # Create dummy input
        dummy_input = torch.zeros(*input_shape).to(device)
        if device.type == "cuda":
            dummy_input = dummy_input.half()
        
        # Measure memory usage for single item
        torch.cuda.empty_cache()
        mem_before = torch.cuda.memory_allocated()
        
        with torch.inference_mode():
            model(dummy_input)
        
        mem_after = torch.cuda.memory_allocated()
        torch.cuda.empty_cache()
        
        per_item_memory = mem_after - mem_before
        free_memory, _ = torch.cuda.mem_get_info()
        
        # Use 80% of free memory as safety margin
        safe_batch_size = int((free_memory * 0.8) // per_item_memory)
        return max(1, safe_batch_size)
        
    except Exception as e:
        logger.error(f"GPU probing failed: {e}")
        return 1