LLM Inference Serving: Architecture, Routing & Auto-Scaling › Putting It Together: A Complete Serving Pipeline

import asyncio
import random
import time
from collections import deque
from dataclasses import dataclass, field
from typing import AsyncIterator, List, Optional


@dataclass
class InferenceRequest:
    """A request to the inference server."""

    request_id: str
    prompt: str
    max_tokens: int
    priority: int = 0
    arrival_time: float = field(default_factory=time.time)


@dataclass
class InferenceResponse:
    """Response from inference, supports streaming."""

    request_id: str
    tokens: List[str]
    finished: bool
    ttft_ms: Optional[float] = None
    total_ms: Optional[float] = None


class MockModel:
    """Simulates model inference with realistic timing."""

    def __init__(self, tokens_per_second: float = 50):
        self.tokens_per_second = tokens_per_second

    async def generate(
        self, prompt: str, max_tokens: int
    ) -> AsyncIterator[str]:
        """Simulate token-by-token generation."""
        # Simulate prefill time based on prompt length
        prompt_tokens = len(prompt.split())
        prefill_time = prompt_tokens * 0.0001  # 0.1ms per token
        await asyncio.sleep(prefill_time)

        # Simulate decode, yielding tokens
        words = [
            "The",
            "quick",
            "brown",
            "fox",
            "jumps",
            "over",
            "the",
            "lazy",
            "dog",
            ".",
            "This",
            "is",
            "a",
            "test",
            ".",
        ]

        for i in range(min(max_tokens, 50)):
            await asyncio.sleep(1.0 / self.tokens_per_second)
            yield random.choice(words)


class InferenceServer:
    """Simple inference server with request handling."""

    def __init__(self, model: MockModel, max_concurrent: int = 10):
        self.model = model
        self.max_concurrent = max_concurrent
        self.active_requests = 0
        self.request_queue: deque = deque()
        self.metrics = MetricsCollector()

    async def handle_request(
        self, request: InferenceRequest
    ) -> AsyncIterator[InferenceResponse]:
        """Process a request and yield streaming responses."""

        # Wait if at capacity
        while self.active_requests >= self.max_concurrent:
            await asyncio.sleep(0.01)

        self.active_requests += 1
        tokens = []
        first_token_time = None

        try:
            async for token in self.model.generate(
                request.prompt, request.max_tokens
            ):
                if first_token_time is None:
                    first_token_time = time.time()

                tokens.append(token)
                yield InferenceResponse(
                    request_id=request.request_id,
                    tokens=tokens.copy(),
                    finished=False,
                    ttft_ms=(first_token_time - request.arrival_time) * 1000,
                )

            # Final response
            completion_time = time.time()
            ttft_ms = (first_token_time - request.arrival_time) * 1000
            total_ms = (completion_time - request.arrival_time) * 1000

            self.metrics.record_request(
                latency_ms=total_ms,
                ttft_ms=ttft_ms,
                prompt_tokens=len(request.prompt.split()),
                generated_tokens=len(tokens),
            )

            yield InferenceResponse(
                request_id=request.request_id,
                tokens=tokens,
                finished=True,
                ttft_ms=ttft_ms,
                total_ms=total_ms,
            )

        finally:
            self.active_requests -= 1

async def simulate_traffic(
    server: InferenceServer,
    num_requests: int = 20,
    arrival_rate: float = 5.0,  # requests per second
):
    """Simulate concurrent traffic to the server."""

    async def process_request(request: InferenceRequest):
        final_response = None
        async for response in server.handle_request(request):
            if response.finished:
                final_response = response
        return final_response

    # Create requests with Poisson arrivals
    tasks = []
    for i in range(num_requests):
        request = InferenceRequest(
            request_id=f"req-{i:03d}",
            prompt=f"Request number {i}: "
            + " ".join(["word"] * random.randint(10, 100)),
            max_tokens=random.randint(20, 50),
        )
        task = asyncio.create_task(process_request(request))
        tasks.append(task)

        # Random inter-arrival time
        await asyncio.sleep(random.expovariate(arrival_rate))

    # Wait for all requests to complete
    responses = await asyncio.gather(*tasks)
    return responses

# Run the simulation
model = MockModel(tokens_per_second=50)
server = InferenceServer(model, max_concurrent=5)

# Run async simulation
responses = await simulate_traffic(server, num_requests=30, arrival_rate=3.0)

ttfts = [r.ttft_ms for r in responses if r.ttft_ms]
totals = [r.total_ms for r in responses if r.total_ms]

# Server metrics
server_metrics = server.metrics.get_current_metrics()