ML Engineer
Distributed Systems & Big Data
Large-scale ML requires distributed data processing and distributed training. You need to understand how Spark computes lazily across a cluster, how Kafka streams data in real time, and how PyTorch DDP synchronizes gradients across GPUs โ plus when each tool is appropriate.
Core Theory
- 1MapReduce paradigm: Map phase applies a function to each record independently (embarrassingly parallel). Shuffle phase groups outputs by key across nodes. Reduce phase aggregates each group. Example: word count โ Map emits (word, 1) per word, Reduce sums counts per word. Foundation for Spark, Hive, and all distributed SQL engines.
- 2Apache Spark: distributed computation engine built on MapReduce ideas. Key abstraction: RDD (Resilient Distributed Dataset) โ immutable, partitioned collection with lazy evaluation. DataFrames provide SQL-like API with Catalyst optimizer. DAG execution: Spark builds a directed acyclic graph of transformations and executes lazily when an action (collect, write, count) is called.
- 3Spark execution model: transformations (map, filter, join) are lazy โ build the DAG. Actions (count, collect, save) trigger execution. Narrow transformations (map, filter) process one partition independently. Wide transformations (groupBy, join) require shuffle โ data moves across nodes, expensive. Minimize shuffles: filter early, broadcast small tables, use bucketing for repeated joins.
- 4Apache Kafka: distributed message queue / streaming platform. Topics are partitioned logs โ producers write to partitions, consumers read from partitions. Consumer groups allow parallel consumption (one consumer per partition). Kafka retains messages for a configurable period, allowing replay. Used for: feature pipelines, real-time event streaming, ML training data collection.
- 5PyTorch Distributed Data Parallel (DDP): each GPU has a full copy of the model. Each GPU processes a different mini-batch subset. After backward pass, gradients are all-reduced (averaged) across all GPUs using the NCCL backend. Effective batch size = single-GPU batch ร number of GPUs. Linear scaling rule: multiply LR by (num_gpus / base_num_gpus) when scaling batch size.
- 6FSDP (Fully Sharded Data Parallel): shards model parameters, gradients, and optimizer states across GPUs. Each GPU holds only 1/N of everything. Before each forward/backward layer, parameters are gathered (all-gather) from all GPUs, then immediately freed after use. ZeRO Stage 3 equivalent in PyTorch. Enables training models 10-20ร larger than DDP allows.
- 7Columnar storage formats: Parquet and ORC store data column-by-column instead of row-by-row. For analytical queries that access a few columns from many rows (SELECT AVG(feature_1) FROM training_data), columnar storage avoids reading irrelevant columns โ often 10-50ร less I/O. Built-in compression per column (delta encoding for timestamps, run-length encoding for low-cardinality). Standard format for ML training data at scale.
- 8Data warehouse vs data lake: data warehouse (BigQuery, Redshift, Snowflake) โ structured, SQL-queryable, optimized for analytics, schema-on-write. Data lake (S3, GCS, HDFS) โ any format, cheap storage, schema-on-read, supports unstructured data. Delta Lake / Iceberg add ACID transactions to data lakes. ML typically reads from data lake (raw features) and writes results to data warehouse (metrics, aggregated labels).
Key Formulas
| 1 | # โโโ PyTorch DDP (Distributed Data Parallel) โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 2 | import torch |
| 3 | import torch.nn as nn |
| 4 | import torch.distributed as dist |
| 5 | from torch.nn.parallel import DistributedDataParallel as DDP |
| 6 | from torch.utils.data.distributed import DistributedSampler |
| 7 | |
| 8 | def train_ddp(rank: int, world_size: int, n_epochs: int = 3): |
| 9 | """ |
| 10 | Called on each GPU process via torch.multiprocessing.spawn. |
| 11 | rank = this GPU's index (0, 1, ..., world_size-1) |
| 12 | world_size = total number of GPUs |
| 13 | """ |
| 14 | # 1. Initialize process group |
| 15 | dist.init_process_group( |
| 16 | backend="nccl", # GPU-to-GPU communication |
| 17 | init_method="env://", # reads MASTER_ADDR, MASTER_PORT, RANK, WORLD_SIZE |
| 18 | world_size=world_size, |
| 19 | rank=rank |
| 20 | ) |
| 21 | torch.cuda.set_device(rank) |
| 22 | |
| 23 | # 2. Create model and wrap in DDP |
| 24 | model = nn.Sequential(nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 1)) |
| 25 | model = model.to(rank) |
| 26 | model = DDP(model, device_ids=[rank]) # gradients all-reduced automatically |
| 27 | |
| 28 | # 3. DistributedSampler ensures each GPU sees different data |
| 29 | dataset = torch.utils.data.TensorDataset( |
| 30 | torch.randn(1000, 128), torch.randn(1000, 1) |
| 31 | ) |
| 32 | sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank) |
| 33 | loader = torch.utils.data.DataLoader(dataset, batch_size=32, sampler=sampler) |
| 34 | |
| 35 | optimizer = torch.optim.Adam(model.parameters(), lr=1e-3 * world_size) # linear LR scaling |
| 36 | loss_fn = nn.MSELoss() |
| 37 | |
| 38 | for epoch in range(n_epochs): |
| 39 | sampler.set_epoch(epoch) # ensures different shuffles per epoch |
| 40 | total_loss = 0.0 |
| 41 | for X, y in loader: |
| 42 | X, y = X.to(rank), y.to(rank) |
| 43 | optimizer.zero_grad() |
| 44 | loss = loss_fn(model(X), y) |
| 45 | loss.backward() # DDP all-reduces gradients here automatically |
| 46 | optimizer.step() |
| 47 | total_loss += loss.item() |
| 48 | |
| 49 | if rank == 0: # only print on rank 0 to avoid duplicate logs |
| 50 | print(f"Epoch {epoch+1}: loss = {total_loss/len(loader):.4f}") |
| 51 | |
| 52 | dist.destroy_process_group() |
| 53 | |
| 54 | # Launch: torchrun --nproc_per_node=4 train.py |
| 55 | # Or: torch.multiprocessing.spawn(train_ddp, args=(4, 3), nprocs=4) |
| 56 | |
| 57 | # โโโ Spark DataFrame feature engineering sketch โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 58 | SPARK_CODE = """ |
| 59 | from pyspark.sql import SparkSession |
| 60 | from pyspark.sql import functions as F |
| 61 | from pyspark.sql.window import Window |
| 62 | |
| 63 | spark = SparkSession.builder.appName("FeatureEngineering").getOrCreate() |
| 64 | |
| 65 | # Read from S3 Parquet (lazy โ no computation yet) |
| 66 | df = spark.read.parquet("s3://bucket/raw-events/") |
| 67 | |
| 68 | # Compute rolling 30-day user feature (wide transformation = shuffle) |
| 69 | w = Window.partitionBy("user_id").orderBy("event_ts").rangeBetween(-30*86400, 0) |
| 70 | df_features = df.withColumn( |
| 71 | "user_30d_purchases", |
| 72 | F.count(F.when(df.event_type == "purchase", 1)).over(w) |
| 73 | ) |
| 74 | |
| 75 | # Filter early to reduce shuffle size (push predicates down) |
| 76 | df_features = df_features.filter(F.col("event_date") >= "2024-01-01") |
| 77 | |
| 78 | # Join with item metadata (broadcast small table to avoid shuffle) |
| 79 | item_meta = spark.read.parquet("s3://bucket/item-metadata/") # small table |
| 80 | df_joined = df_features.join(F.broadcast(item_meta), "item_id", "left") |
| 81 | |
| 82 | # Write as partitioned Parquet for efficient downstream reads |
| 83 | df_joined.write.partitionBy("event_date").parquet("s3://bucket/features/") |
| 84 | # Partitioned by date โ downstream queries on a date range skip irrelevant files |
| 85 | """ |
| 86 | print(SPARK_CODE) |
DDP training loop: every GPU runs the same code, processes different data via DistributedSampler, and DDP automatically all-reduces gradients after backward(). Only rank 0 prints logs. The Spark code shows the key patterns: lazy evaluation (no computation until write), window functions for rolling features, early filtering to reduce shuffle size, and broadcast joins for small tables.
Worked Interview Problems
3 problemsProblem
A Spark job that should take 5 minutes is taking 50 minutes. Walk through your debugging process.
Solution
Check the Spark UI (port 4040): look at the stage timeline. Is one stage taking 80% of the time? If so, that's your bottleneck. Common culprit: a shuffle (GroupBy, Join) creating too many or too few partitions.
Check data skew: in the slow stage, look at the task duration distribution. If 1 task out of 200 takes 40 minutes while others take 10 seconds, you have key skew โ one value (e.g., user_id='unknown') has millions of rows. Fix: salt the key (add random suffix for hot keys) or filter them separately.
Check partitioning: how many partitions does your RDD/DataFrame have? Too few (e.g., 2) โ underparallelism. Too many (e.g., 10,000 with small files) โ task overhead dominates. Ideal: 2-4 partitions per CPU core, each ~128MB. Use repartition(n) or coalesce(n).
Check for unnecessary shuffles: are you doing groupBy() on a column you could have filtered first? Push filter() before groupBy(). Are you joining two large tables? Ensure the smaller one is broadcast-joined if it fits in memory (<10GB).
Check serialization and UDFs: Python UDFs are slow (Python process per row). Replace with Spark SQL functions (F.col, F.when) or use pandas UDFs (vectorized). A single Python UDF can be 10-100ร slower than equivalent Spark SQL.
Answer
Debug order: Spark UI stage timeline โ data skew (task duration distribution) โ partition count โ unnecessary shuffles โ Python UDF usage. Most Spark performance problems are caused by: (1) data skew, (2) too few partitions, (3) avoidable shuffle (missing filter before groupBy), (4) Python UDFs.
Common Mistakes
Calling model.collect() or model.toPandas() on a large Spark DataFrame โ this pulls all data to the driver node, causing OOM. Only collect small result DataFrames (aggregated stats, not raw feature data).
Not setting sampler.set_epoch(epoch) in DDP training โ this causes all GPUs to see the same data order in every epoch, eliminating the randomness that prevents overfitting.
Forgetting dist.destroy_process_group() at the end of DDP training โ this leaks NCCL resources and can cause subsequent training runs to fail or hang.
Using .cache() or .persist() on every Spark DataFrame by default โ caching is only beneficial if the DataFrame is used multiple times. Unnecessary caching wastes cluster memory and can actually slow jobs by triggering earlier materialization.
Not using broadcast joins for small tables: if one side of a join fits in memory (<10GB), spark.sql.autoBroadcastJoinThreshold controls automatic broadcasting. Explicitly using F.broadcast() is safer and more readable.