Cost Optimization & Infrastructure — AI Tutorials

The Cost Structure of ML Systems

ML infrastructure cost has two distinct components:

Component	Cost Driver	Optimization Lever
Training	GPU compute hours, storage	Spot instances, efficient data loading, mixed precision
Serving	GPU/CPU hours, memory, network	Auto-scaling, model optimization, caching, right-sizing

For most production systems, serving costs dominate because training happens occasionally but serving is 24/7.

Cost-Per-Prediction

The single most useful cost metric:

$\text{cost per prediction} = \frac{\text{monthly infrastructure cost}}{\text{monthly prediction count}}$

This metric:

Normalizes cost across different traffic volumes
Lets you compare model serving configurations
Drives prioritization: reduce cost per prediction or increase predictions per dollar

Training Cost Optimization

Spot / Preemptible Instances

Cloud providers offer unused compute capacity at 60–90% discounts as "spot" (AWS), "preemptible" (GCP), or "spot" (Azure) instances. The trade-off: they can be reclaimed with 2 minutes' notice.

Mitigation strategies:

Checkpoint frequently: save model state every N steps; resume from checkpoint on reclaim
Use managed training jobs: AWS SageMaker, Vertex AI, and Azure ML handle interruptions automatically
Mixed instance pools: combine spot and on-demand instances — fall back to on-demand if spot unavailable

Fault-Tolerant Training with Checkpointingpython

import torch
import os
from pathlib import Path

CHECKPOINT_DIR = Path("checkpoints/")
CHECKPOINT_DIR.mkdir(exist_ok=True)

def save_checkpoint(model, optimizer, epoch, step, loss, path: Path):
    torch.save({
        "epoch": epoch,
        "step": step,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss,
    }, path)

def load_checkpoint(model, optimizer, path: Path):
    if not path.exists():
        return 0, 0  # start fresh
    checkpoint = torch.load(path)
    model.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    return checkpoint["epoch"], checkpoint["step"]

# Training loop with checkpointing
checkpoint_path = CHECKPOINT_DIR / "latest.pt"
start_epoch, start_step = load_checkpoint(model, optimizer, checkpoint_path)

for epoch in range(start_epoch, NUM_EPOCHS):
    for step, batch in enumerate(dataloader):
        if epoch == start_epoch and step < start_step:
            continue  # skip already-processed steps

        loss = train_step(model, batch)

        # Checkpoint every 500 steps
        if step % 500 == 0:
            save_checkpoint(model, optimizer, epoch, step, loss, checkpoint_path)

AWS SageMaker Spot Training Jobyaml

# sagemaker_training.py
import sagemaker
from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    entry_point="train.py",
    source_dir="src/",
    role="arn:aws:iam::123456789:role/SageMakerRole",
    instance_type="ml.p3.2xlarge",
    instance_count=1,
    framework_version="2.1",
    py_version="py310",

    # Spot instance configuration
    use_spot_instances=True,
    max_run=3600 * 8,           # 8 hour max runtime
    max_wait=3600 * 24,         # wait up to 24h for spot capacity
    checkpoint_s3_uri="s3://my-bucket/checkpoints/",

    hyperparameters={
        "epochs": 50,
        "learning-rate": 0.001,
    },
)

estimator.fit({"training": "s3://my-bucket/data/train/"})
print(f"Training cost: ${estimator.training_job_analytics().training_job_billing_in_seconds / 3600 * 3.06:.2f}")

Serving Infrastructure: Auto-Scaling

Model serving must handle variable traffic efficiently — expensive GPUs should not idle at 5% utilization at night.

Kubernetes Horizontal Pod Autoscaler (HPA)

HPA scales Kubernetes deployments based on CPU, memory, or custom metrics.

Kubernetes HPA for Model Servingyaml

# deployment.yaml
apiVersion: apps/v1
kind: Deployment
metadata:
  name: fraud-detector
spec:
  replicas: 2
  selector:
    matchLabels:
      app: fraud-detector
  template:
    spec:
      containers:
        - name: fraud-detector
          image: my-registry/fraud-detector:2.1.0
          resources:
            requests:
              cpu: "1"
              memory: "2Gi"
            limits:
              cpu: "2"
              memory: "4Gi"

---
# hpa.yaml — scale on RPS using KEDA
apiVersion: keda.sh/v1alpha1
kind: ScaledObject
metadata:
  name: fraud-detector-scaler
spec:
  scaleTargetRef:
    name: fraud-detector
  minReplicaCount: 1    # scale to zero when idle
  maxReplicaCount: 20
  cooldownPeriod: 60
  triggers:
    - type: prometheus
      metadata:
        serverAddress: http://prometheus:9090
        metricName: http_requests_per_second
        query: |
          sum(rate(http_requests_total{app="fraud-detector"}[1m]))
        threshold: "100"   # 100 RPS per pod target

Right-Sizing: Matching Instance to Model

Right-sizing means choosing the smallest instance that meets your latency SLA. Common mistakes:

Over-provisioning GPU memory for CPU-bound models
Using GPU for small sklearn/ONNX models that run faster on CPU
Ignoring memory: a model that fits in GPU VRAM runs 10× faster than one that swaps

Profiling Inference Cost and Latencypython

import time
import numpy as np
import psutil
import GPUtil

def profile_inference(model, test_batch: np.ndarray, n_warmup: int = 10, n_runs: int = 100):
    """Profile latency, throughput, and resource usage for a serving configuration."""
    # Warmup
    for _ in range(n_warmup):
        model.predict(test_batch)

    # Measure
    latencies = []
    cpu_usage = []

    for _ in range(n_runs):
        cpu_before = psutil.cpu_percent(interval=None)
        start = time.perf_counter()
        _ = model.predict(test_batch)
        elapsed = time.perf_counter() - start
        latencies.append(elapsed * 1000)  # ms
        cpu_usage.append(psutil.cpu_percent(interval=None))

    batch_size = len(test_batch)
    latencies = np.array(latencies)

    return {
        "batch_size": batch_size,
        "p50_ms": np.percentile(latencies, 50),
        "p95_ms": np.percentile(latencies, 95),
        "p99_ms": np.percentile(latencies, 99),
        "throughput_rps": batch_size / (np.mean(latencies) / 1000),
        "avg_cpu_pct": np.mean(cpu_usage),
    }

# Profile different batch sizes to find optimal
for batch_size in [1, 8, 16, 32, 64]:
    batch = np.random.randn(batch_size, 128).astype(np.float32)
    profile = profile_inference(ort_session, batch)
    cost_per_k = (1 / profile["throughput_rps"]) * 1000 * COST_PER_SECOND
    print(f"Batch={batch_size}: p99={profile['p99_ms']:.1f}ms, "
          f"RPS={profile['throughput_rps']:.0f}, "
          f"$/1k predictions={cost_per_k:.4f}")

Cost Reduction Techniques Summary

Technique	Typical Savings	Complexity
Spot instances (training)	60–90%	Low (with checkpointing)
Model quantization (INT8)	30–50% serving cost	Low
Dynamic batching	50–80% GPU idle time	Medium
Scale-to-zero (KEDA)	100% during off-hours	Medium
Knowledge distillation	50–90% model size	High
Shared model server (Triton)	40–60%	Medium
Request caching	30–70% (for skewed inputs)	Low

Cost Dashboard Metrics to Track

Cost per 1k predictions (by model, by environment)
GPU utilization % (idle time is wasted money)
Instance uptime (off-hours scaling efficiency)
Storage cost (model artifacts, training data, logs)

Knowledge check

What is the primary trade-off when using spot/preemptible instances for model training?

Summary

Cost-per-prediction is the key unit economics metric for ML systems
Spot instances reduce training compute costs by 60–90% with checkpointing for fault tolerance
Auto-scaling with Kubernetes HPA/KEDA eliminates idle serving capacity, especially during off-hours
Right-sizing means choosing the smallest instance that meets your latency SLA — don't use a GPU for a small sklearn model
Quantization and dynamic batching are the highest-ROI serving optimizations
Build a cost dashboard tracking cost-per-prediction, GPU utilization, and instance uptime

Final chapter: ML Platform Architecture — how all of these components fit together into an end-to-end platform.

← MLOps & AI Engineering