Production monitoring for PyTorch models

import torch
import torch.nn as nn
import time
import numpy as np
from dataclasses import dataclass
from typing import Dict, List
import warnings

warnings.filterwarnings('ignore')

@dataclass
class InferenceMetrics:
    latency_ms: float
    peak_memory_mb: float
    input_mean: float
    input_std: float
    output_class_dist: Dict[int, float]
    prediction_confidence: float

class ProductionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, 3)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

class MonitoredInference:
    def __init__(self, model: nn.Module, device: str = 'cpu', baseline_input_mean: float = 0.0, baseline_input_std: float = 1.0):
        self.model = model.to(device)
        self.device = device
        self.baseline_input_mean = baseline_input_mean
        self.baseline_input_std = baseline_input_std
        self.metrics_history: List[InferenceMetrics] = []
    
    def compute_kl_divergence(self, current_mean: float, current_std: float) -> float:
        if current_std <= 0:
            return float('inf')
        p_mean, p_std = self.baseline_input_mean, self.baseline_input_std
        q_mean, q_std = current_mean, current_std
        kl = torch.log(torch.tensor(q_std / p_std)) + (p_std**2 + (p_mean - q_mean)**2) / (2 * q_std**2) - 0.5
        return kl.item()
    
    def predict(self, x: torch.Tensor) -> tuple[torch.Tensor, InferenceMetrics]:
        self.model.eval()
        torch.cuda.reset_peak_memory_stats(self.device) if 'cuda' in self.device else None
        
        start_time = time.perf_counter()
        
        with torch.no_grad():
            x_input = x.to(self.device)
            logits = self.model(x_input)
            predictions = torch.softmax(logits, dim=1)
        
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        
        if 'cuda' in self.device:
            peak_memory_mb = torch.cuda.max_memory_allocated(self.device) / (1024 ** 2)
        else:
            peak_memory_mb = 0.0
        
        input_mean = x.mean().item()
        input_std = x.std().item()
        
        pred_classes = torch.argmax(predictions, dim=1)
        class_counts = torch.bincount(pred_classes, minlength=3)
        class_dist = {int(i): (class_counts[i].item() / len(pred_classes)) for i in range(3)}
        
        max_confidence = predictions.max().item()
        
        metrics = InferenceMetrics(
            latency_ms=elapsed_ms,
            peak_memory_mb=peak_memory_mb,
            input_mean=input_mean,
            input_std=input_std,
            output_class_dist=class_dist,
            prediction_confidence=max_confidence
        )
        
        self.metrics_history.append(metrics)
        
        kl_div = self.compute_kl_divergence(input_mean, input_std)
        
        return predictions, metrics, kl_div
    
    def check_anomalies(self, metrics: InferenceMetrics, latency_threshold_ms: float = 100.0, confidence_threshold: float = 0.5) -> Dict[str, bool]:
        anomalies = {
            'high_latency': metrics.latency_ms > latency_threshold_ms,
            'low_confidence': metrics.prediction_confidence < confidence_threshold,
            'high_memory': metrics.peak_memory_mb > 500.0,
            'imbalanced_output': max(metrics.output_class_dist.values()) > 0.9
        }
        return anomalies
    
    def summarize_metrics(self) -> Dict:
        if not self.metrics_history:
            return {}
        latencies = [m.latency_ms for m in self.metrics_history]
        confidences = [m.prediction_confidence for m in self.metrics_history]
        return {
            'avg_latency_ms': np.mean(latencies),
            'p95_latency_ms': np.percentile(latencies, 95),
            'p99_latency_ms': np.percentile(latencies, 99),
            'avg_confidence': np.mean(confidences),
            'min_confidence': np.min(confidences),
            'inference_count': len(self.metrics_history)
        }

model = ProductionModel()
monitor = MonitoredInference(model, device='cpu', baseline_input_mean=0.0, baseline_input_std=1.0)

X_train = torch.randn(100, 10)
monitor.baseline_input_mean = X_train.mean().item()
monitor.baseline_input_std = X_train.std().item()

X_normal = torch.randn(5, 10)
for i, batch in enumerate(X_normal):
    pred, metrics, kl = monitor.predict(batch.unsqueeze(0))
    anomalies = monitor.check_anomalies(metrics)
    print(f"Batch {i}: latency={metrics.latency_ms:.2f}ms, confidence={metrics.prediction_confidence:.3f}, kl_div={kl:.4f}, anomalies={anomalies}")

X_drift = torch.randn(5, 10) + 5.0
print("\n--- DISTRIBUTION SHIFT DETECTED ---")
for i, batch in enumerate(X_drift):
    pred, metrics, kl = monitor.predict(batch.unsqueeze(0))
    anomalies = monitor.check_anomalies(metrics)
    print(f"Drift Batch {i}: latency={metrics.latency_ms:.2f}ms, confidence={metrics.prediction_confidence:.3f}, kl_div={kl:.4f}, anomalies={anomalies}")

summary = monitor.summarize_metrics()
print(f"\nMetrics Summary: {summary}")