Building Custom Model Context Protocol Servers for Real‑Time Data Retrieval Systems

syntax = "proto3";

package realtime.inference;

// -------------------------------------------------------------------
// Versioning
// -------------------------------------------------------------------
enum ProtocolVersion {
  V1 = 0;
  V2 = 1;
}

// -------------------------------------------------------------------
// Core request/response messages
// -------------------------------------------------------------------
message InferenceRequest {
  // Protocol version – enables backward compatibility
  ProtocolVersion version = 1;

  // Identifier of the model context (e.g., "sentiment_v2")
  string context_id = 2;

  // Arbitrary binary payload – can be JSON, base64, or raw tensors
  bytes payload = 3;

  // Optional per‑request metadata (e.g., user ID, timestamp)
  map<string, string> metadata = 4;

  // For streaming: indicate if more chunks will follow
  bool end_of_stream = 5;
}

message InferenceResponse {
  // Mirrors request version
  ProtocolVersion version = 1;

  // Unique request identifier for correlation
  string request_id = 2;

  // Inference result – structured as a JSON string for flexibility
  string result_json = 3;

  // Server‑side latency in milliseconds
  int32 latency_ms = 4;

  // Optional error field (populated only on failure)
  string error_message = 5;
}

// -------------------------------------------------------------------
// Service definition
// -------------------------------------------------------------------
service InferenceService {
  // Unary RPC for single request/response
  rpc Predict(InferenceRequest) returns (InferenceResponse);

  // Bidirectional streaming for continuous inference
  rpc PredictStream(stream InferenceRequest) returns (stream InferenceResponse);
}

mkdir realtime-inference-server
cd realtime-inference-server
python -m venv .venv
source .venv/bin/activate
pip install grpcio grpcio-tools torch torchvision

python -m grpc_tools.protoc -I./protos --python_out=./server --grpc_python_out=./server protos/inference.proto

import struct

def encode_token_ids(token_ids: list[int]) -> bytes:
    # 32‑bit little‑endian integers packed consecutively
    return struct.pack(f'<{len(token_ids)}I', *token_ids)

def decode_token_ids(payload: bytes) -> list[int]:
    count = len(payload) // 4
    return list(struct.unpack(f'<{count}I', payload))

# server/context_registry.py
import torch
from pathlib import Path
from typing import Dict

class ModelContext:
    def __init__(self, model_path: Path, device: str = "cpu"):
        self.model_path = model_path
        self.device = device
        self.model = torch.jit.load(str(model_path), map_location=device)
        self.model.eval()

    def infer(self, token_ids: list[int]) -> dict:
        # Simple example: a text classification model expecting [batch, seq_len]
        input_tensor = torch.tensor([token_ids], dtype=torch.long, device=self.device)
        with torch.no_grad():
            logits = self.model(input_tensor)
        probs = torch.softmax(logits, dim=-1).cpu().numpy().tolist()[0]
        return {"probabilities": probs}

class ContextRegistry:
    """Thread‑safe registry for multiple model contexts."""
    _registry: Dict[str, ModelContext] = {}
    _lock = threading.RLock()

    @classmethod
    def add_context(cls, context_id: str, model_path: Path, device: str = "cpu"):
        with cls._lock:
            cls._registry[context_id] = ModelContext(model_path, device)

    @classmethod
    def get_context(cls, context_id: str) -> ModelContext:
        with cls._lock:
            ctx = cls._registry.get(context_id)
            if not ctx:
                raise KeyError(f"Context '{context_id}' not found")
            return ctx

    @classmethod
    def remove_context(cls, context_id: str):
        with cls._lock:
            cls._registry.pop(context_id, None)

# server/inference_service.py
import time
import json
import grpc
from concurrent import futures
from .context_registry import ContextRegistry
from . import inference_pb2_grpc, inference_pb2

class InferenceServicer(inference_pb2_grpc.InferenceServiceServicer):
    def Predict(self, request, context):
        start = time.time()
        try:
            model_ctx = ContextRegistry.get_context(request.context_id)
            token_ids = decode_token_ids(request.payload)
            result = model_ctx.infer(token_ids)
            latency_ms = int((time.time() - start) * 1000)

            response = inference_pb2.InferenceResponse(
                version=request.version,
                request_id=str(uuid.uuid4()),
                result_json=json.dumps(result),
                latency_ms=latency_ms,
                error_message=""
            )
            return response
        except Exception as e:
            context.set_details(str(e))
            context.set_code(grpc.StatusCode.INTERNAL)
            return inference_pb2.InferenceResponse(
                version=request.version,
                request_id=str(uuid.uuid4()),
                result_json="",
                latency_ms=0,
                error_message=str(e)
            )

    def PredictStream(self, request_iterator, context):
        for request in request_iterator:
            yield self.Predict(request, context)

# server/main.py
import signal
import sys
import grpc
from concurrent import futures
from .inference_service import InferenceServicer
from . import inference_pb2_grpc
from .context_registry import ContextRegistry
from pathlib import Path

def serve():
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=8))
    inference_pb2_grpc.add_InferenceServiceServicer_to_server(InferenceServicer(), server)

    # Bind to all interfaces, port 50051 (TLS can be added later)
    server.add_insecure_port('[::]:50051')
    server.start()
    print("🚀 Inference server listening on port 50051")

    # Load default contexts
    ContextRegistry.add_context(
        "sentiment_v1",
        Path("../models/sentiment_v1.pt"),
        device="cuda:0" if torch.cuda.is_available() else "cpu"
    )

    # Graceful shutdown handling
    def handle_sigterm(*_):
        print("\n🛑 Received termination signal – shutting down gracefully...")
        server.stop(grace=5)  # 5‑second grace period for ongoing RPCs
        sys.exit(0)

    signal.signal(signal.SIGINT, handle_sigterm)
    signal.signal(signal.SIGTERM, handle_sigterm)

    # Keep the main thread alive
    try:
        while True:
            time.sleep(86400)
    except KeyboardInterrupt:
        handle_sigterm()

if __name__ == "__main__":
    serve()

# server/monitoring.py
from prometheus_client import Counter, Histogram, start_http_server

REQUEST_COUNT = Counter('inference_requests_total', 'Total inference requests', ['method', 'status'])
LATENCY_HIST = Histogram('inference_latency_seconds', 'Inference latency', ['method'])

def monitoring_interceptor(method):
    def wrapper(request, context):
        start = time.time()
        try:
            response = method(request, context)
            REQUEST_COUNT.labels(method=method.__name__, status='OK').inc()
            return response
        finally:
            elapsed = time.time() - start
            LATENCY_HIST.labels(method=method.__name__).observe(elapsed)
    return wrapper

class BatchingQueue:
    def __init__(self, max_batch_size=32, max_wait_ms=5):
        self.max_batch_size = max_batch_size
        self.max_wait_ms = max_wait_ms
        self.queue = []
        self.lock = threading.Lock()
        self.cv = threading.Condition(self.lock)

    def add(self, request, future):
        with self.lock:
            self.queue.append((request, future))
            if len(self.queue) >= self.max_batch_size:
                self.cv.notify()
            else:
                # Wake up after timeout to process smaller batch
                self.cv.wait(timeout=self.max_wait_ms / 1000)

    def process_batch(self):
        while True:
            with self.lock:
                while not self.queue:
                    self.cv.wait()
                batch = self.queue[:self.max_batch_size]
                self.queue = self.queue[self.max_batch_size:]

            # Extract payloads, run a single model forward pass, then
            # set results on each future.