AI System Scalability Questions

Covers designing and operating machine learning systems to handle growth in data volume, model complexity, and traffic. Topics include distributed training strategies such as data parallelism, model parallelism, and pipeline parallelism; coordination and orchestration approaches like parameter servers, gradient aggregation, and framework tools such as PyTorch distributed, Horovod, and TensorFlow strategies; data pipeline and I O considerations including sharding, efficient formats, preprocessing bottlenecks, streaming and batch ingestion; serving and inference scaling including model sharding, batching for throughput, autoscaling, request routing, caching, and latency versus throughput tradeoffs. Also includes monitoring, profiling, checkpointing and recovery, reproducibility, cost and resource optimization, and common bottleneck analysis across network, storage, CPU preprocessing, and accelerator utilization.

HardSystem Design

47 practiced

Architect a multi-region ML training and serving platform for a company serving users across North America, Europe, and APAC. Requirements: model training can be centralized but serving must be regional with <100ms latency; model updates are frequent (daily); regulatory constraints require that raw user data never leaves origin region. Describe data replication strategy, model artifact distribution, and how you ensure consistent feature computation across regions.

Sample Answer

Requirements clarification:- Training may be centralized but raw user data must never leave its origin region.- Serving must be regional with <100ms latency.- Model updates daily — automation and fast artifact distribution required.- Strong feature consistency across regions for model correctness and reproducibility.

High-level approach:- Keep raw data and online feature stores in-region.- Compute region-local feature views (deterministic transforms) and export only privacy-safe, aggregated, or encrypted artifacts for central training — OR use federated learning for sensitive cases.- Central training consumes either (a) aggregated/hashed feature tables from each region, or (b) region-specific model updates (gradients/weights) via secure aggregation; produces global model artifacts.- Distribute model artifacts to regional model stores/edge caches and swap traffic via phased rollout.

Architecture components:1. Ingest & local storage (per-region): Raw events -> regional data lake (S3/GCS/Azure Blob in-region). Access controls + encryption-at-rest.2. Local feature computation (compute-at-write/ETL): Run nightly/streaming Spark/Flink pipelines to compute deterministic feature tables (feature store). Feature tables are the canonical source for both online serving and for export.3. Privacy-preserving export: - Preferred: Export aggregated summaries, de-identified feature vectors, or differentially-private statistics to central training bucket (cross-region transfer allowed since no raw PII leaves). - Alternative (stronger privacy): Federated learning — local training or gradient computation per-region; send encrypted gradients to central aggregator that performs secure aggregation and global model update.4. Central training cluster: Pulls aggregated data / aggregated gradients from regions; trains model (daily), evaluates with holdouts, signs & versions artifacts.5. Model artifact distribution: Central CI/CD publishes versioned model artifacts to a global registry and pushes to per-region model stores (regional object storage/CDN edge). Use atomic swap (blue/green) and traffic canary.6. Regional serving: Low-latency inference layer in each region uses: - Local online feature store (Redis/RocksDB) populated from the same feature pipelines. - Local model artifact from regional store loaded into model servers (TF Serving / TorchServe / ONNX runtime) behind autoscaling.7. Feature store consistency: Enforce single source of truth: - Use identical feature transformation code (versioned libraries) and schema definitions stored in Git. - Feature store (offline + online) with versioning (feature version, transform version, timestamp). - Run CI tests: unit tests for transforms, integration tests comparing offline vs online feature outputs (shadow runs). - Maintain a metadata catalog (feature-versions, lineage, checksums). Before serving a new model, run data drift and offline-online parity checks for each region.

Data replication strategy (respecting regulation):- Do NOT replicate raw user data across regions.- Replicate processed feature tables (non-PII) and aggregated statistics to central training only when compliant.- Use region-local backups and cross-region replication for encrypted metadata/catalog only (no payload).- Use encryption in transit (TLS) and at rest; for federated gradients use secure aggregation (homomorphic encryption or secure multiparty compute if required).

Model artifact distribution details:- Central registry with immutable artifacts and metadata (model id, feature-version). Push to regional object stores using signed URLs or a pull-based agent in-region.- Use content-addressable storage + CDN for efficient distribution. Deploy via orchestration (Kubernetes rolling update) with canary and automatic rollback on error.- Ensure artifact signatures to prevent tampering.

Ensuring consistent feature computation:- Same versioned transform code deployed in each region; run transform tests daily.- Maintain deterministic pipelines (idempotent, fixed random seeds where used).- Use feature-store SDK so both offline training and online serving call same transformation code or use precomputed feature tables written by the same pipeline.- Monitor offline-online feature drift with alerts; run checksum comparisons and small-sample parity tests after every pipeline run.

Operational concerns and trade-offs:- Centralized training + aggregated exports is simpler but requires rigorous de-identification and may lose some utility; federated learning preserves privacy but increases complexity and training time.- To meet <100ms latency, colocate feature caches and model executors in-region and use efficient in-memory stores.- Automate testing (parity, data quality), monitoring (latency, feature drift), and compliance logs (audit trail of dataset transfers).

Metrics & SLOs:- Serving latency P95 <100ms per region; model update success rate 99%; parity failure rate <0.1%; end-to-end deployment time <1 hour.

This design balances regulatory constraints, low-latency serving, consistent features, and daily model updates by keeping raw data local, using versioned feature engineering, protecting transfers via aggregation or federated updates, and automating artifact distribution and verification into regional serving clusters.

EasyTechnical

28 practiced

Explain what gradient aggregation is and compare a parameter-server approach with an all-reduce approach (e.g., NCCL/torch.distributed). Focus on communication patterns, scalability, fault tolerance, and where the bottlenecks typically appear.

Sample Answer

Gradient aggregation is the process of combining gradients computed on multiple workers (GPUs/nodes) during distributed training to produce a single update for model parameters. It ensures all workers apply consistent parameter updates.

Comparison: parameter-server (PS) vs all-reduce (e.g., NCCL/torch.distributed)

Communication patterns:- Parameter-server: workers push gradients to one or more centralized servers; servers aggregate and send updated parameters back. Traffic is many-to-one and one-to-many.- All-reduce: workers participate in decentralized collective operations (ring/all-reduce tree) to aggregate gradients in a peer-to-peer fashion; data is exchanged in phases so each node ends up with the sum.

Scalability:- PS: scales poorly as model size or worker count grows if servers become hot spots; can mitigate by sharding parameter servers, but cross-rack latency limits scale.- All-reduce: scales better for homogeneous GPU clusters using RDMA/NIC-accelerated collectives (NCCL). Bandwidth is used more evenly; ring all-reduce achieves near-optimal bandwidth utilization as nodes increase.

Fault tolerance:- PS: easier to handle by replicating or checkpointing parameter servers; a failed server can be recovered or its shard reallocated, but complexity increases.- All-reduce: sensitive to node failures—collectives assume all participants; a single failure often aborts the operation. Workarounds require reconfiguration or using fault-tolerant collective libraries.

Bottlenecks:- PS: parameter servers become CPU/network bottlenecks (egress/incoming bandwidth), higher latency for large models, and serialization overhead.- All-reduce: bottlenecks are network bandwidth and NIC capabilities; poorly balanced network (cross-rack) increases latency. With optimized libraries (NCCL) and high-bandwidth interconnects (InfiniBand), compute and GPU-memory copy can become dominant.

When to use:- PS: heterogeneous environments, elastic clusters, or when asynchronous updates or sparse gradients are common (parameter sharding helps).- All-reduce: dense models on homogeneous GPU clusters with fast interconnects for synchronous SGD—preferred for throughput and scaling.

Practical notes for a data-engineer:- Choose all-reduce for stable GPU clusters with RDMA/NCCL; monitor network and GPU utilization.- Use PS when supporting elastic worker pools, multi-tenant clusters, or sparse update patterns; ensure server redundancy and sharding to avoid hotspots.

HardTechnical

37 practiced

Training stragglers: in a distributed training job, a few workers repeatedly finish epochs much slower causing overall job slowdown. Propose detection methods, dynamic load-balancing or re-sharding techniques, and specific implementations using Spark or PyTorch to mitigate the straggler problem.

Sample Answer

Clarify assumptions:- Synchronous SGD across N workers with data sharded per worker; stragglers repeatedly finish epochs slower (I/O, skewed shard size, hardware heterogeneity, GC).- Goal: detect stragglers quickly and mitigate without losing statistical correctness.

Detection methods:- Per-epoch/step latency metrics: track time per batch and per epoch per worker; compute z-score / percentile and flag persistent > threshold (e.g., >95th percentile for 3 consecutive epochs).- Progress heartbeats with logical step counters; detect lagging steps.- Resource telemetry: CPU, GPU utilization, disk IO, network stats correlated with latency.

Dynamic mitigation strategies:1. Adaptive re-sharding (data redistribution) - Move portions of hot shards from slow workers to faster ones; keep sampling deterministic with epoch-level shuffle seeds to preserve randomness.2. Load-aware batching - Reduce batch-size for stragglers (less compute per step) or give them smaller micro-batches and aggregate gradients appropriately.3. Backup / speculative workers - Launch redundant workers for the slowest partitions; use earliest gradient and drop the late one.4. Elastic training - Allow worker set to change: remove persistently slow worker and redistribute its shard (reshard) or add temporary workers.

Spark implementation (data-engineer friendly):- Detect via Spark metrics (executor metrics, task durations). Use Driver to compute per-executor task-time percentiles.- Re-shard using mapPartitions + repartitionByRange or custom Partitioner:

python

# identify slow executors -> compute new partition map and redistribute
rdd = rdd.map(lambda x: (assign_new_partition(x, partition_map), x)).partitionBy(new_num_partitions).values()

- Use coalesce(repartition) to rebalance without full shuffle if hardware constraints permit.- For streaming pipelines, use dynamic allocation + blacklist slow nodes in Spark’s scheduler to avoid scheduling new tasks there.

PyTorch implementation:- Detection: instrument each rank with timestamps and report to rank 0; compute moving average step times and flag stragglers.- Re-shard dataset dynamically: use DistributedSampler with epoch-aware indices. When a rank is removed/added, compute new world_size and regenerate samplers:

python

from torch.utils.data.distributed import DistributedSampler
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
loader = DataLoader(dataset, batch_size=..., sampler=sampler)

- Elastic approach: use torch.distributed.elastic (or TorchElastic/Kubernetes) to allow workers to join/leave; implement checkpointing to continue after reshard.- Speculative gradients: allow asynchronous gradient commit for bounded staleness (trade-off: possible slower convergence). Alternatively use gradient averaging with weighting to account for smaller batches.

Trade-offs and best practices:- Re-sharding incurs shuffle cost; trigger only after persistent detection window.- Speculative duplicates waste resources but are low-latency to implement.- Elastic removal simplifies scheduling but must preserve randomness and reproducibility—use deterministic seeds tied to epoch and global partition map.- Monitor and log resource signals to distinguish true hardware problems vs transient spikes.

Outcome metrics to track:- Epoch wall-clock time, variance across ranks, validation convergence (to ensure mitigation doesn’t hurt model quality), resource utilization.

MediumTechnical

34 practiced

You are rolling out a new model to production and your evaluation pipeline shows a small but significant regression in a critical metric. Describe a safe deployment and rollback strategy that minimizes customer impact. Include A/B or canary options, monitoring signals to watch, and automated rollback triggers.

Sample Answer

Situation: A new ML model shows a small but significant regression in a critical metric during evaluation. As a data engineer, my goal is to deploy safely, limit user impact, and enable fast rollback.

Safe deployment strategy- Staged rollout (canary + gradual traffic ramp): 1. Canary: route 0.5–1% traffic to model variant behind feature flag; run for N hours/days. 2. Incremental ramp: 1% → 5% → 20% → 50% → 100% only if signals remain healthy. 3. A/B experiment (parallel): for metrics requiring statistical comparison, run randomized A/B for enough users to reach significance while keeping exposure low.- Data-path parity: ensure canary uses identical preprocessing, feature services, and infra to avoid hidden data drift.

Monitoring signals to watch- Business/ML metrics: the critical metric (e.g., conversion rate, latency-weighted accuracy), secondary metrics (precision/recall, calibration), and cohort performance (new users, device types).- Data-quality signals: feature distributions, missing rates, schema drift, null/NaN counts, cardinality changes.- Infrastructure signals: inference latency, error rates, resource usage, queue/backpressure.- Observability: logs, traces, and real-time dashboards; collect per-user and per-cohort logs for root cause.

Automated rollback triggers- Predefined thresholds: absolute drop in critical metric > X% or statistically significant drop (p < 0.01) over a sliding window.- Data-quality alarms: feature distribution KL divergence > threshold, sudden spike in missing features.- Infrastructure alarms: 5xx error rate > Y%, latency > SLA.- Safety circuit: automated rollback action via orchestration (feature flag toggle or deployment rollback) if any trigger fires; require human-on-call notification with context.

Mitigation & postmortem- Fallback: route affected traffic immediately to previous stable model (hot standby).- Snapshot inputs/outputs for failed traffic for debugging.- Run targeted A/B by cohort to isolate cause; roll forward only after root cause fixed and re-validation.- Document decisions and adjust thresholds/monitoring to prevent recurrence.

This approach minimizes customer exposure, gives statistical confidence, and automates safe rollback while preserving data for diagnosis.

MediumTechnical

28 practiced

Write a PySpark transformation that computes per-user daily aggregates from an events dataset while minimizing shuffle. Assume events(partitioned by event_date) with columns (user_id, event_type, value). Show how you would use partitioning and map-side combiners to reduce shuffle volume.

Sample Answer

Approach: exploit the existing partitioning by event_date, do local (map-side) pre-aggregation within each partition, then perform a single keyed shuffle across user_id+event_date using combineByKey (or reduceByKey). Persist minimal columns and write output partitioned by event_date to avoid downstream shuffles.

python

from pyspark.sql import SparkSession
from pyspark.sql.functions import col
from pyspark import StorageLevel

spark = SparkSession.builder.getOrCreate()

# Read events already partitioned by event_date on storage (e.g., /data/events/event_date=YYYY-MM-DD/)
events = spark.read.parquet("/data/events/") \
    .select("event_date", "user_id", "event_type", "value")  # prune cols early

# Convert to RDD of ((event_date, user_id), (count, sum)) and do map-side combiner
rdd = events.rdd.map(lambda r: ((r.event_date, r.user_id), (1, float(r.value))))

# combineByKey does local aggregation (map-side) then a single shuffle per key
def create_combiner(v): return v
def merge_value(acc, v): return (acc[0] + v[0], acc[1] + v[1])
def merge_combiners(a, b): return (a[0] + b[0], a[1] + b[1])

agg_rdd = rdd.combineByKey(create_combiner, merge_value, merge_combiners)

# Convert back to DataFrame and write partitioned by event_date (no shuffle on write if using partitioning and bucketed sink)
agg_df = agg_rdd.map(lambda kv: (kv[0][0], kv[0][1], kv[1][0], kv[1][1])) \
    .toDF(["event_date", "user_id", "event_count", "event_value_sum"])

# Repartition by event_date only (cheap if data already on disk partitioned by date) and write
agg_df.repartition(col("event_date")).write.mode("overwrite").partitionBy("event_date").parquet("/data/agg/user_daily/")

Key points:- Early column pruning reduces data shuffled.- Using RDD.combineByKey (map-side combiner) aggregates within each mapper partition before the shuffle, reducing network traffic.- Keys include event_date so final shuffle groups per user per day; writing partitioned by event_date avoids future shuffles.- Persist/partition awareness: if reading many dates, process per-date slices (filter by event_date) to limit working set.

Complexity: single shuffle across unique (date,user) keys; network I/O proportional to number of distinct keys. Edge cases: skewed users (hot keys) — handle with salting; nulls in value; very large single partitions — use larger parallelism or split by date ranges. Alternatives: use DataFrame aggregations with shuffle partitions tuned (spark.sql.shuffle.partitions) or use mapInPandas for vectorized map-side aggregation.

Unlock Full Question Bank

Get access to hundreds of AI System Scalability interview questions and detailed answers.

Join thousands of developers preparing for their dream job.