Designing Production-Grade ML Systems

Static vs Dynamic Training & Serving

Production ML systems must make two fundamental architecture decisions: how to train (static vs dynamic) and how to serve (static vs dynamic). These choices cascade through every other design decision.

Causes, Detection, and Prevention

Training-serving skew is one of the most insidious bugs in production ML. It occurs when the data or feature transformations used during training differ from those used during serving, causing the model to receive inputs it was not trained on. The model may appear healthy in offline evaluation but perform poorly in production.

Common causes of training-serving skew:

• Feature transformation code duplication: training uses Python/Spark for preprocessing, serving reimplements the same logic in Java/C++ (or vice versa). Even tiny differences (rounding, null handling) cause skew.
• Stale statistics: normalization uses training-time mean/std, but data distribution has shifted. Serving data gets normalized with wrong parameters.
• Data leakage during training: features computed using information from the future (e.g., averages that include the label period). Model learns shortcuts that are unavailable at serving time.
• Different data sources: training reads from a data warehouse (cleaned, batch-processed), serving reads from a real-time stream (raw, uncleaned).
• Feature store inconsistency: training features are point-in-time correct, but serving features use current values, creating temporal leakage.

Prevention strategies:

Use a shared preprocessing layer. The single most effective prevention is to embed preprocessing into the model itself using tf.Transform (TFX) or Keras preprocessing layers. This guarantees identical transformations during training and serving because the same code path executes in both environments.

Use a Feature Store (Vertex AI Feature Store). Feature stores provide a single source of truth for feature values, ensuring that training and serving use the same feature computation logic and the same feature values.

# Prevent skew: embed preprocessing in the model
preprocessing = tf.keras.Sequential([
    tf.keras.layers.Normalization(axis=-1),   # Learned mean/std
    tf.keras.layers.CategoryEncoding(num_tokens=100),
])
preprocessing.adapt(training_data)  # Compute statistics from training data

# Full model with preprocessing baked in
model = tf.keras.Sequential([
    preprocessing,               # Same transforms at train + serve
    tf.keras.layers.Dense(128, activation="relu"),
    tf.keras.layers.Dense(1)
])
# SavedModel includes preprocessing — no separate serving code needed

The seminal Google paper "Hidden Technical Debt in Machine Learning Systems" (2015) argues that ML systems have all the maintenance problems of traditional software plus a set of ML-specific issues. The actual model code is often a tiny fraction of the total system — the surrounding infrastructure (data collection, feature extraction, configuration, monitoring) is far larger and harder to maintain.

Data dependencies are the ML equivalent of code dependencies, but harder to track. If a feature comes from an upstream system and that system changes its data format, units, or semantics, your model silently receives different inputs. There is no compiler to catch this. Key antipatterns to avoid:

• Unstable data dependencies: features from systems you do not own that may change without notice. Mitigate with data contracts and schema validation (TFDV).
• Underutilized features: features added during experimentation that provide marginal value but increase complexity. Regularly audit feature importance and remove low-value features.
• Feedback loops: model predictions influence future training data (e.g., recommendation models show what the model predicts users want, creating a self-reinforcing loop). Detect with holdout experiments and randomized exploration.
• Pipeline jungles: tangled data preparation code that grows organically. Refactor into clean, tested, versioned pipeline stages with TFX or Vertex AI Pipelines.

Concept Drift & Data Drift

Production models degrade over time because the world changes. There are two distinct forms of change that must be monitored and addressed:

Data drift (also called covariate shift) occurs when the statistical distribution of input features changes, even though the underlying relationship between features and labels remains the same. Example: a model trained on summer traffic patterns receives winter data with different driving behaviors. Detection: compare feature distributions between training data and recent serving data using statistical tests (KS test, PSI, chi-squared).

Concept drift occurs when the relationship between inputs and outputs changes. The same features now predict different outcomes. Example: a fraud model trained on pre-pandemic transaction patterns fails post-pandemic because legitimate purchasing behavior has changed. Detection: monitor prediction performance metrics (accuracy, AUC) over time using labeled serving data.

Mitigation strategies:

• Scheduled retraining: retrain on recent data at fixed intervals (daily, weekly). Simple but effective for gradual drift.
• Triggered retraining: automatically retrain when monitoring detects significant drift (TFDV schema violations, performance below threshold).
• Windowed training: train only on recent data (e.g., last 30 days) to keep the model adapted to current patterns. Risk: lose rare events.
• Continuous evaluation: compare model predictions against actual outcomes and alert on degradation. Vertex AI Model Monitoring provides this out of the box.

Distributed Training: Data vs Model Parallelism

Data parallelism is the most common approach: the model is replicated on each device, each replica processes a different mini-batch, and gradients are averaged across replicas via all-reduce. Total throughput scales linearly with device count (ideally). This is what MirroredStrategy and MultiWorkerMirroredStrategy implement.

Model parallelism splits the model itself across devices. Layer 1-24 on GPU 0, layers 25-48 on GPU 1, etc. Required when the model is too large to fit on a single device (e.g., GPT-3 at 175B parameters requires ~350GB in float16 — no single GPU has this). More complex to implement, lower GPU utilization due to pipeline bubbles. Pipeline parallelism (e.g., GPipe) mitigates this by splitting mini-batches into micro-batches.

TPU vs GPU: When to Use Each

Google Cloud offers both GPUs (NVIDIA A100, V100, T4, L4) and TPUs (v2, v3, v4, v5e) for training and inference. The choice depends on your workload characteristics.

TPU-specific considerations:

• TPUs use bfloat16 natively for matrix multiplications. This is different from IEEE float16 — bfloat16 has the same exponent range as float32 but fewer mantissa bits, making it more numerically stable for training.
• Batch size should be a multiple of 8 (one per TPU core) and ideally a multiple of 128 for optimal memory layout.
• Avoid Python-side data processing in tf.data map functions. TPUs cannot execute Python operations efficiently. Use pure TF ops.
• Use tf.data pipelines (not NumPy arrays or generators). TPUs require efficient data streaming from GCS.

Mixed Precision Training

Mixed precision training uses both float16 (or bfloat16) and float32 during training. Most computations (forward pass, backward pass) use float16 for speed and memory savings, while critical operations (loss computation, gradient accumulation, weight updates) use float32 for numerical stability.

Benefits: 2x memory reduction (enabling larger batch sizes or models), 2-3x training speedup on modern GPUs (V100, A100 with Tensor Cores), with minimal accuracy loss.

# Enable mixed precision globally
tf.keras.mixed_precision.set_global_policy("mixed_float16")

# Build model — layers automatically use float16 for compute
model = tf.keras.Sequential([
    tf.keras.layers.Dense(256, activation="relu"),
    tf.keras.layers.Dense(128, activation="relu"),
    # Final layer should output float32 for numerical stability
    tf.keras.layers.Dense(10, dtype="float32")
])

# Optimizer with loss scaling to prevent float16 underflow
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
# In TF 2.x, loss scaling is handled automatically by Keras

model.compile(optimizer=optimizer, loss="sparse_categorical_crossentropy")

Not all inference can happen in the cloud. Latency requirements, connectivity limitations, privacy constraints, or cost considerations may require running models on-device (edge) or in a hybrid cloud + on-prem architecture.

TensorFlow Lite is the primary framework for on-device inference. The workflow is: train a full model in TensorFlow → convert to TFLite format → optimize (quantize/prune) → deploy to device.

# Convert Keras model to TF Lite
converter = tf.lite.TFLiteConverter.from_keras_model(model)

# Post-training quantization — reduces model size by ~4x
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# Optional: full integer quantization for maximum speed
converter.target_spec.supported_types = [tf.int8]
def representative_dataset():
    for data in calibration_data.batch(1).take(100):
        yield [data]
converter.representative_dataset = representative_dataset

tflite_model = converter.convert()

# Save the quantized model
with open("model.tflite", "wb") as f:
    f.write(tflite_model)

Hybrid architectures combine cloud-based training with edge inference, or split inference between edge and cloud. Common patterns include: edge device runs a lightweight model for initial filtering, sends uncertain cases to a larger cloud model for final prediction. This reduces latency and bandwidth while maintaining accuracy on difficult cases.

Production ML systems fail in ways that are subtler and harder to debug than traditional software. The model may still return predictions — they are just wrong. Key troubleshooting areas:

Data Quality Issues

Symptoms: sudden accuracy drop, unexpected prediction distributions, NaN losses during retraining. Root causes: upstream schema changes, missing values where none existed before, duplicated records, label corruption, stale feature store. Tools: TensorFlow Data Validation (TFDV) for schema monitoring, Great Expectations for data testing, Vertex AI Model Monitoring for serving-time data validation.

Model Performance Degradation

Symptoms: gradual accuracy decline over days/weeks, A/B test metrics dropping. Root causes: concept drift (world has changed), data drift (input distribution shifted), training-serving skew introduced by a code change. Diagnosis: compare current serving data distribution with training data distribution, check feature importance shifts, examine recent code changes to preprocessing.

Latency & Throughput Problems

Symptoms: P99 latency exceeding SLO, prediction timeouts, queueing. Root causes: model too large for available hardware, inefficient preprocessing at serving time, cold starts on autoscaled endpoints, network bottlenecks fetching features. Solutions: model optimization (quantization, pruning, distillation), GPU/TPU acceleration, precompute features, horizontal scaling, batch prediction for non-latency-sensitive use cases.

Memory & Compute Optimization

Symptoms: OOM errors during training or serving, GPU memory fragmentation, slow batch processing. Solutions: mixed precision training (halves memory), gradient accumulation (simulate large batches), model sharding, reduce batch size, optimize tf.data pipeline memory usage (.cache(filename) for disk-based caching).

# Batch prediction pipeline pattern
import pandas as pd

# Load batch data
df = pd.read_gbq("SELECT * FROM my_dataset.scoring_table")

# Predict in batches (memory-efficient)
BATCH_SIZE = 10000
results = []
for i in range(0, len(df), BATCH_SIZE):
    batch = df.iloc[i:i+BATCH_SIZE]
    features = batch[feature_columns].values
    preds = model.predict(features, batch_size=256)
    results.append(preds)

# Write predictions back
df["prediction"] = np.concatenate(results)
df.to_gbq("my_dataset.predictions", if_exists="replace")

# Vertex AI Batch Prediction (managed)
# batch_prediction_job = model.batch_predict(
#     job_display_name="nightly-scoring",
#     gcs_source="gs://bucket/input.jsonl",
#     gcs_destination_prefix="gs://bucket/output/",
#     machine_type="n1-standard-4",
#     accelerator_type="NVIDIA_TESLA_T4",
#     accelerator_count=1,
# )

This section summarizes the key concepts from Course 08 that are most likely to appear on the GCP Professional Machine Learning Engineer exam, mapping to Section 3 (Scaling prototypes), Section 4 (Serving & scaling models), and Section 6 (ML solutions monitoring).

Training Strategy Selection

• Data parallelism (MirroredStrategy): model fits on one GPU, need faster training. Most common.
• Model parallelism: model too large for single GPU (>10B parameters). Split across devices.
• Mixed precision: 2x memory reduction, 2-3x speedup on A100/V100. Always enable on modern hardware.

TPU vs GPU Scenarios

• TPU: large-batch training, TensorFlow/JAX, matrix-heavy models, cost-efficient at scale.
• GPU: PyTorch, custom CUDA ops, small/medium models, diverse workloads, inference serving.
• TPU batch size must be a multiple of 8. Use bfloat16. Avoid Python-side preprocessing.

Debugging Production Issues

• Training-serving skew: use tf.Transform or Keras preprocessing layers. Single transformation code path.
• Data drift: feature distribution changed. Retrain on recent data. Monitor with TFDV/Vertex AI Model Monitoring.
• Concept drift: real-world relationship changed. Monitor prediction metrics. Retrain with new labeled data.
• Latency issues: quantize model, use GPU serving, cache features, batch non-critical predictions.

How to Explain This in 2 Minutes

Likely Interview Questions

Model Answers

System Design Scenario

Common Mistakes

• Confusing data drift with concept drift — Data drift means input distributions changed; concept drift means the target relationship changed. The distinction matters because data drift can be detected by monitoring features alone (no labels needed), while concept drift requires monitoring prediction quality (which needs labels). Treating all model degradation as "the model needs more data" misses the root cause and leads to wasted retraining cycles.
• Deploying models without A/B testing — Replacing a production model wholesale because it has better offline metrics is a common and costly mistake. Offline evaluation does not capture all production dynamics (traffic patterns, feature distributions, business metrics). Always A/B test in production. Even a 1% traffic canary for 24 hours can catch catastrophic issues that offline evaluation misses.
• Using float32 for the output layer with mixed precision training — Mixed precision (float16/bfloat16) dramatically speeds up training and reduces memory, but the output layer must remain float32 to maintain numerical stability in the loss computation. Forgetting this causes NaN losses or degraded accuracy. In Keras, use tf.keras.mixed_precision.set_global_policy('mixed_float16') and explicitly set dtype='float32' on the final Dense layer.