ML Engineer: Comprehensive Technical Prep Guide

from dataclasses import dataclass, field, asdict
from typing import Optional, List
from pathlib import Path

@dataclass
class TrainingConfig:
    model_name: str = "resnet50"
    learning_rate: float = 1e-4
    batch_size: int = 32
    num_epochs: int = 100
    checkpoint_dir: Path = Path("./checkpoints")
    gpus: List[int] = field(default_factory=lambda: [0, 1, 2, 3])
    mixed_precision: bool = True
    
    # Easy to add new fields later — nothing breaks
    wandb_project: Optional[str] = None

# Create with defaults or overrides
config = TrainingConfig(learning_rate=3e-4)
config.checkpoint_dir.mkdir(parents=True, exist_ok=True)

# Convert to dict (for serialization, logging, etc.)
as_dict = asdict(config)

# BAD: loads entire file into memory
def read_all_lines(path):
    with open(path) as f:
        return f.readlines()  # all lines in memory at once

# GOOD: yields one line at a time
def read_lines(path):
    with open(path) as f:
        for line in f:
            yield line.strip()

# Process without loading entire file
for line in read_lines("huge_dataset.txt"):
    process(line)

# Generic batching utility — works on any iterable
from itertools import islice

def batched(iterable, batch_size):
    """Yield successive batches from an iterable."""
    it = iter(iterable)
    while True:
        batch = list(islice(it, batch_size))
        if not batch:
            break
        yield batch

# Now you can batch ANYTHING — files, database rows, API results
for batch in batched(read_lines("data.txt"), batch_size=1000):
    process_batch(batch)  # list of 1000 lines

# List comprehension — builds the entire list in memory
squares = [x**2 for x in range(10_000_000)]  # ~80MB in memory

# Generator expression — computes on demand
squares = (x**2 for x in range(10_000_000))  # almost no memory
total = sum(squares)  # processes one at a time

from contextlib import contextmanager
import time

@contextmanager
def timer(label=""):
    """Time a code block. Works even if the code throws an error."""
    start = time.perf_counter()
    try:
        yield  # code inside the 'with' block runs here
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label}: {elapsed:.3f}s")

# Usage
with timer("data loading"):
    data = load_dataset()
# prints "data loading: 2.341s"

from functools import wraps

def retry(max_attempts=3, delay=1.0):
    """Retry a function if it fails. Useful for flaky I/O."""
    def decorator(fn):
        @wraps(fn)  # preserves original function's name/docstring
        def wrapper(*args, **kwargs):
            for attempt in range(max_attempts):
                try:
                    return fn(*args, **kwargs)
                except Exception as e:
                    if attempt == max_attempts - 1:
                        raise
                    time.sleep(delay * (2 ** attempt))  # exponential backoff
        return wrapper
    return decorator

@retry(max_attempts=3)
def save_checkpoint(model, path):
    torch.save(model.state_dict(), path)

from pathlib import Path

root = Path("/data/experiments")
model_dir = root / "models" / "resnet"     # joins paths with /
model_dir.mkdir(parents=True, exist_ok=True)

# Find files
checkpoints = sorted(model_dir.glob("*.pt"))       # non-recursive
all_zarr    = list(root.rglob("*.zarr"))            # recursive

# File info
for p in checkpoints:
    print(f"{p.stem}: {p.stat().st_size / 1e6:.1f} MB")
    # p.stem = filename without extension
    # p.suffix = ".pt"
    # p.name = "model_v3.pt"
    # p.parent = the directory containing it

from typing import Dict, List, Optional, Tuple, Union, Callable

def load_data(
    path: str,
    batch_size: int = 32,
    transform: Optional[Callable] = None,
) -> List[Dict[str, torch.Tensor]]:
    ...

import numpy as np

# From Python data
a = np.array([1, 2, 3])                        # 1D
b = np.array([[1, 2], [3, 4]])                  # 2D

# Filled arrays
np.zeros((3, 4))                                # all zeros
np.ones((2, 3, 4), dtype=np.float32)            # all ones, specific type
np.full((3, 3), fill_value=7.0)                 # all sevens
np.empty((5, 5))                                # uninitialized (fast, garbage values)
np.eye(4)                                       # 4x4 identity matrix

# Ranges
np.arange(0, 10, 0.5)                           # [0, 0.5, 1, ..., 9.5]
np.linspace(0, 1, num=100)                      # 100 evenly spaced points

# Random (modern API)
rng = np.random.default_rng(seed=42)
rng.standard_normal((3, 4))                     # normal distribution
rng.uniform(0, 1, size=(100,))                  # uniform [0, 1)
rng.integers(0, 10, size=(5, 5))                # random ints

arr = np.zeros((2, 3, 4), dtype=np.float32)

arr.shape      # (2, 3, 4) — dimensions
arr.ndim       # 3 — number of dimensions
arr.dtype      # float32 — data type of each element
arr.size       # 24 — total number of elements
arr.nbytes     # 96 — total bytes (24 elements × 4 bytes each)
arr.strides    # (48, 16, 4) — bytes to jump per dimension

a = np.arange(10)

# VIEWS (share memory, modifications propagate)
b = a[2:5]           # slicing creates a view
b[0] = 99            # a[2] is now 99 too!
c = a.reshape(2, 5)  # reshape usually creates a view
d = a.T              # transpose of 2D array is a view

# COPIES (independent memory)
e = a[2:5].copy()    # explicit copy
f = a[[0, 3, 7]]     # fancy indexing always copies
g = a[a > 5]         # boolean indexing always copies

arr = np.arange(24)

arr.reshape(2, 3, 4)     # new shape — must preserve total elements
arr.reshape(6, -1)        # -1 means "infer this dimension" → (6, 4)
arr.reshape(-1)           # flatten to 1D

# Transpose
m = np.zeros((3, 4))
m.T                       # (4, 3) — transpose (view, no copy)

# For 3D+ arrays, specify axis order
t = np.zeros((2, 3, 4))
t.transpose(0, 2, 1)      # (2, 4, 3) — swap last two dims

# Add/remove dimensions
a = np.array([1, 2, 3])   # shape (3,)
a[np.newaxis, :]           # shape (1, 3) — add batch dimension
a[:, np.newaxis]           # shape (3, 1) — add feature dimension
np.expand_dims(a, axis=0)  # same as newaxis

b = np.zeros((1, 3, 1, 4))
np.squeeze(b)              # removes all size-1 dims → (3, 4)
np.squeeze(b, axis=0)      # remove specific dim → (3, 1, 4)

# Concatenate vs Stack
x = np.ones((3, 4))
y = np.zeros((3, 4))
np.concatenate([x, y], axis=0)  # (6, 4) — join along existing axis
np.stack([x, y], axis=0)        # (2, 3, 4) — creates NEW axis

# Example: Normalize data (subtract mean per feature)
data = np.random.randn(1000, 64)   # 1000 samples, 64 features
mean = data.mean(axis=0)            # shape (64,)

# Broadcasting: (1000, 64) - (64,) → (64,) becomes (1, 64) → (1000, 64)
centered = data - mean  # works! each row gets the mean subtracted

# Example: Outer product
col = np.array([[1], [2], [3]])     # shape (3, 1)
row = np.array([10, 20, 30])        # shape (3,)
# Broadcasting: (3, 1) * (3,) → (3, 1) * (1, 3) → (3, 3)
outer = col * row

# COMMON GOTCHA — shapes that DON'T broadcast:
# (4, 3) + (4,)  → ERROR! trailing dims 3 vs 4 don't match
# Fix: reshape (4,) to (4, 1)

arr = np.arange(20).reshape(4, 5)
# [[ 0,  1,  2,  3,  4],
#  [ 5,  6,  7,  8,  9],
#  [10, 11, 12, 13, 14],
#  [15, 16, 17, 18, 19]]

# Basic slicing (returns VIEWS)
arr[1:3]             # rows 1-2
arr[:, 2:]           # all rows, columns 2 onward
arr[::2]             # every other row (step=2)
arr[::-1]            # reversed rows

# Boolean indexing (returns COPIES)
mask = arr > 10
arr[mask]             # flat array of elements where mask is True
arr[arr % 3 == 0]     # elements divisible by 3

# Fancy indexing (returns COPIES)
arr[[0, 2, 3]]        # select rows 0, 2, 3
arr[[0, 2], [1, 3]]   # elements at (0,1) and (2,3) → array([1, 13])

# np.where — conditional selection
np.where(arr > 10, arr, 0)  # keep values > 10, replace rest with 0

# SLOW — Python loop
def normalize_slow(data):
    result = np.empty_like(data)
    for i in range(len(data)):
        result[i] = (data[i] - data.mean()) / data.std()
    return result

# FAST — vectorized (100x faster)
def normalize_fast(data):
    return (data - data.mean()) / data.std()

# Einstein summation — powerful shorthand for complex operations
A = np.random.randn(3, 4)
B = np.random.randn(4, 5)

np.einsum('ij,jk->ik', A, B)     # matrix multiply (same as A @ B)
np.einsum('ii->', A[:3, :3])      # trace (sum of diagonal)
np.einsum('ij->j', A)             # sum over rows (same as A.sum(axis=0))

# Batch operations
batch_A = np.random.randn(8, 3, 4)
batch_B = np.random.randn(8, 4, 5)
np.einsum('bij,bjk->bik', batch_A, batch_B)  # batch matmul

Tensor: [[1, 2, 3],
         [4, 5, 6]]

Shape: (2, 3)
In memory: [1, 2, 3, 4, 5, 6]  ← just a flat array

Strides: (3, 1)
  - To move to the next row: skip 3 elements
  - To move to the next column: skip 1 element

To access element [1, 2]:
  offset = 1×3 + 2×1 = 5  →  memory[5] = 6  ✓

Original:      Transposed (no data copy!):
Shape (2, 3)   Shape (3, 2)
Stride (3, 1)  Stride (1, 3)
Memory: [1, 2, 3, 4, 5, 6]  ← SAME memory for both!

t = torch.tensor([[1, 2, 3], [4, 5, 6]])
t.is_contiguous()       # True — standard layout
t.stride()              # (3, 1)

t2 = t.T
t2.is_contiguous()      # False — transposed, strides are swapped
t2.stride()             # (1, 3)

# view() requires contiguous memory — it will fail on t2
# t2.view(-1)  → RuntimeError!

# Solutions:
t2.contiguous().view(-1)  # make a contiguous copy first
t2.reshape(-1)            # reshape handles non-contiguous automatically

torch.float32  (torch.float)   — default. 4 bytes per element.
torch.float64  (torch.double)  — double precision. 8 bytes. Rarely needed in ML.
torch.float16  (torch.half)    — half precision. 2 bytes. Faster but overflows easily.
torch.bfloat16                 — "brain" float. 2 bytes. Same RANGE as float32 
                                  but less precision. PREFERRED for training large models.
torch.int32    (torch.int)     — 32-bit integers
torch.int64    (torch.long)    — 64-bit integers. Used for indices/labels.
torch.bool                     — True/False

import torch

# From data
t = torch.tensor([1.0, 2.0, 3.0])
t = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)

# Filled tensors
torch.zeros(3, 4)
torch.ones(2, 3)
torch.full((3, 3), fill_value=3.14)
torch.empty(5, 5)          # uninitialized (fast, garbage values)
torch.eye(4)               # identity matrix

# Ranges
torch.arange(0, 10, 2)     # [0, 2, 4, 6, 8]
torch.linspace(0, 1, 50)   # 50 points from 0 to 1

# Random
torch.rand(3, 4)            # uniform [0, 1)
torch.randn(3, 4)           # normal distribution (mean=0, std=1)
torch.randint(0, 10, (3, 4)) # random integers

# Like another tensor (same shape, dtype, device)
x = torch.randn(3, 4, device='cuda')
y = torch.zeros_like(x)    # same shape/dtype/device as x

import numpy as np

# NumPy → PyTorch (SHARES memory)
np_arr = np.array([1.0, 2.0, 3.0])
t = torch.from_numpy(np_arr)
t[0] = 99
print(np_arr[0])   # 99.0 — both changed!

# PyTorch → NumPy
t = torch.randn(3, 4)
n = t.numpy()           # shares memory (CPU tensor only)
n = t.detach().cpu().numpy()  # safe way for GPU tensors
# .detach() = remove from computation graph
# .cpu() = move to CPU
# .numpy() = convert to numpy

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Move existing tensor
t = torch.randn(3, 4)
t_gpu = t.to(device)          # generic (works for any device)
t_gpu = t.cuda()               # shorthand for GPU
t_cpu = t_gpu.cpu()            # back to CPU

# Create directly on GPU
t = torch.randn(3, 4, device=device)

# Pin memory (important for DataLoader)
# "Pinned" memory = page-locked RAM that GPU can access directly
# Makes CPU→GPU transfer faster
t = torch.randn(3, 4).pin_memory()
t_gpu = t.to(device, non_blocking=True)  # async transfer

t = torch.arange(24)

# Reshape
t.view(2, 3, 4)       # requires contiguous, never copies
t.reshape(2, 3, 4)    # always works, may copy
t.view(-1, 4)          # infer first dim → (6, 4)

# Transpose / Permute
m = torch.randn(3, 4)
m.T                    # transpose 2D (same as m.t())
t3d = torch.randn(2, 3, 4)
t3d.permute(0, 2, 1)   # (2, 4, 3) — reorder dimensions

# Squeeze / Unsqueeze
t = torch.randn(1, 3, 1, 4)
t.squeeze()             # (3, 4) — remove ALL size-1 dims
t.squeeze(0)            # (3, 1, 4) — remove specific dim

t2 = torch.randn(3, 4)
t2.unsqueeze(0)         # (1, 3, 4) — add batch dim
t2.unsqueeze(-1)        # (3, 4, 1) — add trailing dim

# Concatenate / Stack
a = torch.randn(3, 4)
b = torch.randn(3, 4)
torch.cat([a, b], dim=0)    # (6, 4) — join along existing dim
torch.stack([a, b], dim=0)  # (2, 3, 4) — creates new dim

# Expand (no copy) vs Repeat (copies)
t = torch.randn(1, 3)
t.expand(4, 3)          # (4, 3) — reads same memory 4 times
t.repeat(4, 1)          # (4, 3) — actually copies data 4 times

# Basic gradient computation
x = torch.tensor([2.0, 3.0], requires_grad=True)
y = (x ** 2).sum()    # y = x[0]^2 + x[1]^2 = 4 + 9 = 13
y.backward()           # compute dy/dx
print(x.grad)          # tensor([4., 6.])  — dy/dx = 2x

# In a training loop
optimizer.zero_grad()          # reset gradients from previous step
output = model(input_data)     # forward pass (builds computation graph)
loss = criterion(output, target)
loss.backward()                # backward pass (computes gradients)
optimizer.step()               # update parameters using gradients

# Disable gradient tracking (for inference)
with torch.no_grad():
    predictions = model(test_data)
    # no graph built, no memory for intermediate activations
    # MUCH faster and uses less memory

# Detach from graph
z = x.detach()    # new tensor, no grad tracking, shares data

from torch.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()
    
    # Forward pass in mixed precision
    with autocast(device_type='cuda'):
        output = model(data)       # computed in float16/bfloat16
        loss = criterion(output, target)
    
    # Backward pass with gradient scaling
    # (prevents float16 gradients from underflowing to zero)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    """
    You must implement:
    - __len__: how many samples total
    - __getitem__: load and return one sample by index
    """
    def __init__(self, data_dir, transform=None):
        self.data_dir = Path(data_dir)
        # Build an index of all files (lightweight — don't load data yet)
        self.file_list = sorted(self.data_dir.glob("*.npy"))
        self.transform = transform
    
    def __len__(self):
        return len(self.file_list)
    
    def __getitem__(self, idx):
        # Load SINGLE sample (called by DataLoader in parallel)
        data = np.load(self.file_list[idx])
        tensor = torch.from_numpy(data).float()
        
        if self.transform:
            tensor = self.transform(tensor)
        
        return tensor, idx  # often return (data, label) tuple

dataset = CustomDataset("/data/samples")

loader = DataLoader(
    dataset,
    batch_size=32,          # samples per batch
    shuffle=True,           # randomize order each epoch
    num_workers=4,          # parallel loading processes
    pin_memory=True,        # faster CPU→GPU transfer
    drop_last=True,         # drop incomplete final batch
    prefetch_factor=2,      # batches loaded ahead per worker
    persistent_workers=True, # keep workers alive between epochs
)

# Training loop
for epoch in range(num_epochs):
    for batch_data, batch_labels in loader:
        batch_data = batch_data.to(device, non_blocking=True)
        # ... train

from torch.nn.utils.rnn import pad_sequence

def collate_variable_length(batch):
    """
    Custom collate for variable-length sequences.
    Pads shorter sequences with zeros to match the longest.
    """
    sequences = [item[0] for item in batch]  # list of tensors, different lengths
    labels = torch.tensor([item[1] for item in batch])
    lengths = torch.tensor([len(s) for s in sequences])
    
    # Pad all sequences to the length of the longest
    padded = pad_sequence(sequences, batch_first=True, padding_value=0)
    
    # Create attention mask: 1 where real data, 0 where padding
    mask = torch.arange(padded.size(1)).unsqueeze(0) < lengths.unsqueeze(1)
    
    return {"sequences": padded, "labels": labels, "mask": mask, "lengths": lengths}

loader = DataLoader(dataset, batch_size=32, collate_fn=collate_variable_length)

class StreamingDataset(torch.utils.data.IterableDataset):
    def __init__(self, shard_paths):
        self.shard_paths = shard_paths  # list of file paths
    
    def __iter__(self):
        # Split shards across workers (if using num_workers > 0)
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            shards = self.shard_paths
        else:
            # Each worker gets a subset of shards
            per_worker = len(self.shard_paths) // worker_info.num_workers
            start = worker_info.id * per_worker
            shards = self.shard_paths[start:start + per_worker]
        
        for path in shards:
            data = np.load(path, mmap_mode='r')  # memory-mapped
            for i in range(len(data)):
                yield torch.from_numpy(data[i].copy()).float()

num_workers:
  - Start with num_CPU_cores // num_GPUs
  - Too few → GPU starves waiting for data
  - Too many → CPU contention, worker overhead
  - Profile: if GPU utilization < 90%, increase workers

pin_memory:
  - Always True when training on GPU
  - Pairs with non_blocking=True on .to(device)

prefetch_factor:
  - Default 2 is usually fine
  - Increase for high-latency storage (NFS, cloud)

persistent_workers:
  - True if Dataset __init__ is expensive
  - Avoids respawning workers every epoch

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DistributedSampler

def main():
    # torchrun sets these environment variables automatically
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    
    # Initialize process group
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    
    # Model — wrap with DDP
    model = MyModel().to(device)
    model = DDP(model, device_ids=[local_rank])
    # Note: to access the original model, use model.module
    
    # Data — DistributedSampler splits data across GPUs
    dataset = MyDataset(data_path)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = DataLoader(
        dataset,
        batch_size=32,
        sampler=sampler,    # replaces shuffle=True
        num_workers=4,
        pin_memory=True,
    )
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    
    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # CRITICAL — ensures different shuffling each epoch
        
        for data, target in loader:
            data = data.to(device, non_blocking=True)
            target = target.to(device, non_blocking=True)
            
            optimizer.zero_grad()
            loss = compute_loss(model(data), target)
            loss.backward()      # gradients are all-reduced automatically by DDP
            optimizer.step()
        
        # Save checkpoint — ONLY on rank 0
        if rank == 0:
            torch.save({
                "epoch": epoch,
                "model_state_dict": model.module.state_dict(),  # .module!
                "optimizer_state_dict": optimizer.state_dict(),
            }, f"checkpoint_{epoch}.pt")
    
    dist.destroy_process_group()

# Single machine, 4 GPUs
torchrun --standalone --nproc_per_node=4 train.py

# 2 machines, 4 GPUs each (run on each machine)
torchrun --nnodes=2 --nproc_per_node=4 \
         --node_rank=0 --master_addr=10.0.0.1 --master_port=29500 \
         train.py

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

# Wrap model with FSDP
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # shard everything
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
    ),
    device_id=local_rank,
)

accumulation_steps = 4
# effective_batch_size = batch_size × accumulation_steps × world_size

for step, (data, target) in enumerate(loader):
    loss = model(data, target) / accumulation_steps  # scale loss
    
    if (step + 1) % accumulation_steps != 0:
        with model.no_sync():  # skip gradient sync (DDP optimization)
            loss.backward()
    else:
        loss.backward()        # gradients synced here
        optimizer.step()
        optimizer.zero_grad()

Without chunking:                With chunking (1000×1000):
┌──────────────────────┐         ┌────┬────┬────┬─ ...
│                      │         │ C0 │ C1 │ C2 │
│  One giant block     │         ├────┼────┼────┤
│  Must read ALL       │         │ C3 │ C4 │ C5 │
│  to access anything  │         ├────┼────┼────┤
│                      │         │ C6 │ C7 │ C8 │
└──────────────────────┘         └────┴────┴────┘
                                 100 separate chunks
                                 Each compressed independently
                                 Read only what you need

import h5py
import numpy as np

# ── Writing ──
with h5py.File("experiment.h5", "w") as f:
    # Create groups (like directories)
    microscopy = f.create_group("microscopy")
    
    # Create a chunked, compressed dataset
    images = microscopy.create_dataset(
        "images",
        shape=(10000, 512, 512),       # 10000 images, 512×512 each
        dtype=np.float32,
        chunks=(1, 512, 512),          # one image per chunk
        compression="gzip",
        compression_opts=4,            # compression level (1-9)
    )
    
    # Write data
    images[0] = np.random.randn(512, 512)      # write single image
    images[10:20] = np.random.randn(10, 512, 512)  # write a batch
    
    # Metadata (key-value attributes on any group or dataset)
    f.attrs["experiment_name"] = "imaging_v3"
    images.attrs["pixel_size_um"] = 0.65

# ── Reading ──
with h5py.File("experiment.h5", "r") as f:
    # Navigate hierarchy
    print(list(f.keys()))                 # ['microscopy']
    print(list(f["microscopy"].keys()))   # ['images']
    
    # Read a single image (only that chunk is loaded from disk)
    img = f["microscopy/images"][42]
    
    # Read a slice
    batch = f["microscopy/images"][100:132]
    
    # Inspect without loading
    ds = f["microscopy/images"]
    print(f"Shape: {ds.shape}, Chunks: {ds.chunks}, Compression: {ds.compression}")

experiment.zarr/
├── .zgroup                          # JSON: group metadata
├── microscopy/
│   ├── .zgroup
│   └── images/
│       ├── .zarray                  # JSON: shape, chunks, dtype, compressor
│       ├── 0.0                      # chunk at position (0, 0)
│       ├── 0.1                      # chunk at position (0, 1)
│       ├── 1.0                      # chunk at position (1, 0)
│       └── ...                      # each chunk is a separate file

import zarr
import numpy as np

# ── Writing ──
root = zarr.open("experiment.zarr", mode="w")

# Create group
microscopy = root.create_group("microscopy")

# Create chunked, compressed array
images = microscopy.create_array(
    "images",
    shape=(10000, 512, 512),
    chunks=(10, 512, 512),    # 10 images per chunk
    dtype=np.float32,
)

# Write data
images[0:10] = np.random.randn(10, 512, 512).astype(np.float32)

# Attributes (metadata)
images.attrs["pixel_size_um"] = 0.65

# ── Reading ──
root = zarr.open("experiment.zarr", mode="r")
img = root["microscopy/images"][42]           # reads only the chunk containing image 42
batch = root["microscopy/images"][100:132]

# Info
arr = root["microscopy/images"]
print(f"Shape: {arr.shape}, Chunks: {arr.chunks}, Dtype: {arr.dtype}")

HDF5:
  ✓ Multi-language support (C, Java, Fortran, MATLAB)
  ✓ Single file (easy to move around)
  ✓ Mature, battle-tested (20+ years)
  ✗ Not cloud-native (single file = can't parallelize I/O well)
  ✗ Limited parallel write support

Zarr:
  ✓ Cloud-native (each chunk is a separate object in S3/GCS)
  ✓ Parallel reads AND writes (different chunks = no conflict)
  ✓ Better compression options (any numcodecs codec)
  ✓ No GIL issues for compression/decompression
  ✗ Creates many small files (can be slow on local filesystems)
  ✗ Less multi-language support (mostly Python)

Rule of thumb:
  - Cloud storage → Zarr
  - Multi-language interop → HDF5
  - Parallel writes from many processes → Zarr
  - Small experiments, local files → either works

import tensorstore as ts

spec = {
    'driver': 'zarr',
    'kvstore': {
        'driver': 'file',
        'path': '/data/experiment.zarr/images',
    },
}
dataset = ts.open(spec, open=True).result()
data = dataset[0:10].read().result()  # returns numpy array

class ZarrDataset(torch.utils.data.Dataset):
    def __init__(self, zarr_path, key="images"):
        self.data = zarr.open(zarr_path, mode="r")[key]
    
    def __len__(self):
        return self.data.shape[0]
    
    def __getitem__(self, idx):
        # Only reads the chunk containing this sample
        return torch.from_numpy(self.data[idx].astype(np.float32))

loader = DataLoader(ZarrDataset("data.zarr"), batch_size=32, num_workers=8)

import mlflow

mlflow.set_tracking_uri("http://mlflow-server:5000")
mlflow.set_experiment("model-training")

with mlflow.start_run(run_name="finetune-v3"):
    # Log hyperparameters (logged once)
    mlflow.log_params({
        "model": "resnet50",
        "lr": 1e-4,
        "batch_size": 32,
        "num_gpus": 4,
    })
    
    # Log metrics over time (logged per step/epoch)
    for epoch in range(100):
        train_loss = train_one_epoch(model, loader)
        val_loss = evaluate(model, val_loader)
        
        mlflow.log_metrics({
            "train_loss": train_loss,
            "val_loss": val_loss,
        }, step=epoch)
    
    # Log files as artifacts
    mlflow.log_artifact("config.yaml")                   # any file
    mlflow.log_artifact("checkpoint.pt", "checkpoints")  # into a subdirectory
    
    # Log the final model
    mlflow.pytorch.log_model(model, "model")

def save_checkpoint(model, optimizer, scheduler, epoch, loss, path):
    """Save complete training state."""
    torch.save({
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scheduler_state_dict": scheduler.state_dict() if scheduler else None,
        "loss": loss,
        "rng_state": torch.random.get_rng_state(),
        "cuda_rng_state": torch.cuda.get_rng_state_all(),
    }, path)

def load_checkpoint(model, optimizer, scheduler, path, device):
    """Resume training from checkpoint."""
    ckpt = torch.load(path, map_location=device, weights_only=False)
    
    model.load_state_dict(ckpt["model_state_dict"])
    optimizer.load_state_dict(ckpt["optimizer_state_dict"])
    if scheduler and ckpt.get("scheduler_state_dict"):
        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
    
    # Restore RNG state for reproducibility
    torch.random.set_rng_state(ckpt["rng_state"])
    if ckpt.get("cuda_rng_state"):
        torch.cuda.set_rng_state_all(ckpt["cuda_rng_state"])
    
    return ckpt["epoch"] + 1  # next epoch to train

class CheckpointManager:
    """Keep only the N most recent checkpoints to save disk space."""
    
    def __init__(self, directory, max_kept=5):
        self.directory = Path(directory)
        self.directory.mkdir(parents=True, exist_ok=True)
        self.max_kept = max_kept
    
    def save(self, state_dict, epoch):
        path = self.directory / f"ckpt_{epoch:04d}.pt"
        tmp = path.with_suffix(".tmp")
        torch.save(state_dict, tmp)  # write to temp file
        tmp.rename(path)             # atomic rename (no partial files)
        self._cleanup()
        return path
    
    def latest(self):
        ckpts = sorted(self.directory.glob("ckpt_*.pt"))
        return ckpts[-1] if ckpts else None
    
    def _cleanup(self):
        ckpts = sorted(self.directory.glob("ckpt_*.pt"))
        for old in ckpts[:-self.max_kept]:
            old.unlink()

import torch.distributed.checkpoint as dcp

# Save — ALL ranks participate (each saves its shard)
state = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
dcp.save(state, checkpoint_id="/checkpoints/step_10000")
# Creates:
# /checkpoints/step_10000/
#   ├── .metadata      (shard mapping)
#   ├── __0_0.distcp    (rank 0's shard)
#   ├── __1_0.distcp    (rank 1's shard)
#   └── ...

# Load — reshards automatically!
# Save on 8 GPUs → load on 4 GPUs: DCP handles it
state = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
dcp.load(state, checkpoint_id="/checkpoints/step_10000")
model.load_state_dict(state["model"])

# Start simple: "Process some data with normalization"

def normalize(data, mean=0, std=1):
    return (data - mean) / (std + 1e-8)

def process(data, config):
    return normalize(data, config.get("mean", 0), config.get("std", 1))

# New requirement: "Now also support log transform and clipping"
# With a registry, you just ADD functions. Nothing existing changes.

TRANSFORMS = {
    "normalize": lambda data, **kw: (data - kw.get("mean", 0)) / (kw.get("std", 1) + 1e-8),
    "log": lambda data, **kw: np.log1p(np.clip(data, 0, None)),
    "clip": lambda data, **kw: np.clip(data, kw.get("low", 0), kw.get("high", 1)),
}

def process(data, steps):
    """steps = [{"name": "clip", "low": 0, "high": 100}, {"name": "normalize"}]"""
    for step in steps:
        name = step["name"]
        params = {k: v for k, v in step.items() if k != "name"}
        data = TRANSFORMS[name](data, **params)
    return data

# Another requirement: "Now support custom user-defined transforms"
# Registry is already a dict — users just add to it:
TRANSFORMS["custom_scale"] = lambda data, **kw: data * kw.get("factor", 1.0)

from abc import ABC, abstractmethod

class DataReader(ABC):
    @abstractmethod
    def read(self, key, slices=None): ...
    
    @abstractmethod
    def keys(self): ...

class HDF5Reader(DataReader):
    def __init__(self, path):
        self.f = h5py.File(path, "r")
    def read(self, key, slices=None):
        return self.f[key][slices] if slices else self.f[key][:]
    def keys(self):
        return list(self.f.keys())

class ZarrReader(DataReader):
    def __init__(self, path):
        self.store = zarr.open(path, "r")
    def read(self, key, slices=None):
        return self.store[key][slices] if slices else self.store[key][:]
    def keys(self):
        return list(self.store.keys())

# Factory
def open_data(path):
    if path.endswith(".h5"):
        return HDF5Reader(path)
    return ZarrReader(path)

# Caller doesn't care about the format
reader = open_data(some_path)
data = reader.read("images", slice(0, 32))