# Ludwig Lazy / Streaming Preprocessing Plan

## Problem Statement

Ludwig's preprocessing pipeline is **not lazy**: before the training loop starts, it decodes and
transforms every sample in the dataset into in-memory NumPy arrays. This causes:

1. **OOM on media datasets** — 1 000 images at original resolution × their tensor size saturates
   RAM before a single batch reaches the GPU.
1. **Fixed preprocessing bottleneck** — transforms run single-threaded on the CPU ahead of training,
   blocking the GPU.
1. **No online statistics** — normalization stats (mean/std) require a full pass that materialises
   the entire dataset.
1. **Impossible streaming** — datasets that exceed RAM (or that arrive over a network) cannot be
   used at all today.

______________________________________________________________________

## How Other Frameworks Solve This

### PyTorch / torchvision

`Dataset.__getitem__` is the canonical lazy primitive. Each DataLoader worker calls it
independently, so only `batch_size × num_workers × prefetch_factor` samples are alive at once.
Transforms (resize, normalize, augment) are composed in the Dataset constructor and run inside
worker processes, overlapping with the GPU forward pass.

Statistics (ImageNet mean/std) are pre-computed offline and shipped as constants. For custom
datasets a single Welford accumulation scan (O(1) memory) is standard.

### TensorFlow `tf.data`

Declarative lazy pipeline: `.map(fn, num_parallel_calls=AUTOTUNE)` → `.batch()` →
`.prefetch(AUTOTUNE)`. Nothing is executed until the training iterator pulls a batch.
`.cache()` optionally persists the decoded dataset to disk after the first epoch.

### HuggingFace `datasets`

Arrow IPC memory-map: only the pages actually accessed are loaded from disk. Audio/image columns
store raw bytes or file paths — decoded waveforms are produced per-sample at access time via
`cast_column("audio", Audio())`. Passing `decode=False` gives raw bytes with zero decoding cost.

### MONAI (3D medical imaging)

Lazy resampling accumulates sequential spatial transforms (rotate → crop → resize) as pending
homogeneous-matrix operations and fuses them into a single interpolation, eliminating intermediate
materialisation and reducing quality loss.

### torchaudio

Transforms (`MelSpectrogram`, `Resample`, `MFCC`) are `nn.Module` subclasses — composable,
JIT-scriptable, and GPU-movable. Applying them post-batch on GPU is 10-20× faster than per-sample
on CPU.

______________________________________________________________________

## Design Principles for Ludwig's New Pipeline

1. **Store paths, not arrays.** The Arrow dataset (or in-memory DataFrame) holds file paths /
   metadata, never decoded tensors. Decode happens inside the DataLoader worker.
1. **Transforms as `nn.Module`.** All feature-specific transforms become PyTorch modules:
   composable, testable, batchable, GPU-ready.
1. **Online statistics with Welford's algorithm.** Replace full-dataset decode-then-accumulate
   with a single streaming pass that needs O(1) memory per feature.
1. **Prefetching and parallelism.** DataLoader `num_workers > 0` + `prefetch_factor` provides
   CPU/GPU overlap for free once decode is inside `__getitem__`.
1. **Backward compatibility.** Existing configs that work today must continue to work. New
   behaviour is opt-in at the feature level via `lazy: true` or automatic for audio/image features.

______________________________________________________________________

## Implementation Phases

______________________________________________________________________

### Phase 0 — Foundations (no user-facing change)

**Goal:** establish shared infrastructure that later phases build on.

#### 0-A. Welford online stats accumulator

- Add `ludwig/data/statistics.py` with `WelfordAccumulator` (supports merge across shards).
- API: `acc.update(x: np.ndarray)`, `acc.result() → (mean, std, min, max, count)`,
  `acc.merge(other)`.
- Unit-test with known distributions.

#### 0-B. Feature transform protocol

- Define `FeatureTransform(nn.Module)` base class in `ludwig/features/transforms.py`.
- Existing preprocessing logic stays unchanged for now; this just establishes the interface.
- `__call__(self, x: Tensor) → Tensor`; serialisable via `torch.save`.

#### 0-C. DataLoader worker-safe file handle pool

- Research whether Ludwig's existing multiprocessing setup needs `worker_init_fn` to reset
  open file handles or RNG state (same issue as HuggingFace's `IterableDataset` shard split).
- Document the policy; no code change yet.

**Exit criteria:** `WelfordAccumulator` tests pass; `FeatureTransform` ABC exists.

______________________________________________________________________

### Phase 1 — Lazy Audio (highest OOM risk)

**Goal:** audio columns decode inside the DataLoader worker, not during preprocessing.

#### 1-A. `AudioFeature` path-only mode

- After the preprocessing step, the Arrow/DataFrame column holds **file paths** (strings), not
  decoded tensors.
- Remove `_process_in_memory` eager decode; replace with a lightweight path-validity check.
- The existing `read_audio_files_to_dict` path (used today) becomes a fallback for non-lazy mode.

#### 1-B. `AudioDataset.__getitem__` decode

- Create `ludwig/data/datasets/audio_dataset.py` (or extend `LudwigDataset`).
- In `__getitem__`, call `torchaudio.load(path)` → resample → pad/trim to target length.
- Apply `FeatureTransform` chain (mel spectrogram, normalization).

#### 1-C. Online stats for audio normalization

- Replace current two-pass (decode all → compute stats) with a streaming Welford scan:
  `torchaudio.load` one file at a time, update accumulator, discard tensor.
- Triggered once at `prepare_data` time; stats stored in the feature config cache.

#### 1-D. Integration test

- Run `pytest tests/integration_tests/test_audio_feature.py` — no OOM on the medium-size
  audio fixtures.
- Add a test that processes 10 000 short clips and checks RSS stays under 2 GB.

**Exit criteria:** audio smoke tests pass on datasets with >10 h of audio; RSS bounded by
`batch_size × clip_length × sr × 4 bytes`.

______________________________________________________________________

### Phase 2 — Lazy Image

**Goal:** image columns store paths; decode + augment happens per-sample in DataLoader workers.

#### 2-A. `ImageFeature` path-only mode

- Preprocessing stores file paths (already partially the case for CSV/HDF5 workflows).
- Remove upfront `_process_in_memory` full-batch decode; keep only metadata inference
  (height, width, channels) via a single sampled read.

#### 2-B. `ImageDataset.__getitem__` decode

- `PIL.Image.open(path).convert("RGB")` → `torchvision.transforms` pipeline.
- Apply `FeatureTransform` (resize, center crop, normalize).
- Multi-channel and TIFF support via existing `read_image_from_path` helper.

#### 2-C. Online stats for image normalization

- Welford accumulator over pixel values per channel; one forward pass, batch-level updates
  (read image → reshape to (H×W, C) → `acc.update`).
- Offer `use_imagenet_stats: true` shortcut (ships precomputed constants) for the common case.

#### 2-D. Augmentation support

- `ImageFeature` config gains an optional `augmentation:` block (equivalent to
  `torchvision.transforms.v2`).
- Augmentation runs inside `__getitem__` → free via DataLoader parallelism.

**Exit criteria:** image smoke tests pass; openfake and synthia complete without OOM at
original resolution.

______________________________________________________________________

### Phase 3 — Lazy Text / Tabular (lower priority)

**Goal:** tokenisation and numerical transforms move into the DataLoader worker.

#### 3-A. Text tokenisation in `__getitem__`

- Currently: tokenise all text upfront into integer id arrays → store in HDF5.
- New: store raw strings; tokenise on-the-fly in `__getitem__`.
- Tradeoff: tokenisation is cheap but repeated; cache tokenised output to Arrow after first epoch
  if `cache_encoder_embeddings: true`.

#### 3-B. Numerical normalisation as `nn.Module`

- `NormalizationTransform(mean, std)` replaces the current in-place column mutation.
- Stats computed once via Welford in `prepare_data`; transform applied in `__getitem__`.

#### 3-C. Category embedding stays eager

- Category integer encoding is O(1) per sample and zero-memory; no change needed.

**Exit criteria:** existing tabular test suite passes unchanged.

______________________________________________________________________

### Phase 4 — Streaming Dataset Source

**Goal:** support HuggingFace streaming datasets (`streaming=True`) as a first-class source,
enabling datasets that exceed local disk space.

#### 4-A. `HFStreamingDataset(IterableDataset)`

- Wraps a `datasets.IterableDataset` in a PyTorch `IterableDataset`.
- Shard splitting: `worker_init_fn` slices the stream across DataLoader workers (HF provides
  `dataset.shard(num_shards=N, index=i)`).
- Shuffle buffer: configurable per-dataset (small for media, large for text).

#### 4-B. Online stats for streaming sources

- Welford accumulator updated during the stats-scan pre-training pass (a full iteration over
  the stream without gradient accumulation).
- Stats cached to disk keyed by dataset id + revision + split.

#### 4-C. No-shuffle streaming mode

- For datasets too large even for a shuffle buffer, expose `shuffle: false` (warn user).

**Exit criteria:** `AudioSet`, `VoxPopuli`, `LibriSpeech` stream through without touching disk;
stats computed in one pass.

______________________________________________________________________

### Phase 5 — MONAI-style Transform Fusion (stretch goal)

**Goal:** fuse adjacent spatial transforms (resize → crop → normalize) into a single tensor
operation to eliminate intermediate materialisation.

#### 5-A. Transform graph analysis

- At `setup` time, walk the `FeatureTransform` chain; identify composable sequences.
- A sequence is composable when all transforms implement `as_matrix() → Tensor[4×4]`.

#### 5-B. Single-pass spatial transform

- Fuse the homogeneous matrices; apply via `F.grid_sample` once.
- Benchmark vs. sequential PIL transforms on ImageNet validation set.

#### 5-C. Augmentation randomness

- Random transforms (flip, rotation jitter) break deterministic fusion; keep them as
  post-fusion `nn.Module` steps.

**Exit criteria:** ≥ 15% wall-clock speedup on image-only training runs vs. Phase 2 baseline.

______________________________________________________________________

## Data-Flow Diagram (target state)

```
Disk / HF Hub
    │
    ▼
┌──────────────────────────────────────────────────────┐
│  prepare_data  (single process, once per dataset)    │
│                                                      │
│  • Build Arrow metadata table (paths, labels, ids)   │
│  • Welford scan → normalization stats                │
│  • Write stats to feature config cache              │
└───────────────────────────┬──────────────────────────┘
                            │ Arrow file / paths
                            ▼
┌──────────────────────────────────────────────────────┐
│  LudwigDataset.__getitem__  (per DataLoader worker)  │
│                                                      │
│  • Read path from Arrow row                          │
│  • torchaudio.load / PIL.Image.open                  │
│  • FeatureTransform chain (resize, mel spec, norm)   │
│  • Return sample dict of tensors                     │
└───────────────────────────┬──────────────────────────┘
                            │ batch of tensors (pinned RAM)
                            ▼
┌──────────────────────────────────────────────────────┐
│  GPU training loop                                   │
│  • H2D transfer (pin_memory → non-blocking)          │
│  • Forward pass                                      │
│  • Gradient accumulation                             │
└──────────────────────────────────────────────────────┘
```

______________________________________________________________________

## Key Files to Create / Modify

| File                                           | Action                                              |
| ---------------------------------------------- | --------------------------------------------------- |
| `ludwig/data/statistics.py`                    | NEW — `WelfordAccumulator`                          |
| `ludwig/features/transforms.py`                | NEW — `FeatureTransform` base + concrete transforms |
| `ludwig/features/audio_feature.py`             | MODIFY — remove eager decode, add lazy path         |
| `ludwig/features/image_feature.py`             | MODIFY — remove eager decode, add lazy path         |
| `ludwig/data/datasets/base_dataset.py`         | MODIFY — add `__getitem__` decode hooks             |
| `ludwig/data/datasets/audio_dataset.py`        | NEW (or extend base)                                |
| `ludwig/data/datasets/image_dataset.py`        | NEW (or extend base)                                |
| `ludwig/data/datasets/hf_streaming_dataset.py` | NEW — Phase 4                                       |
| `tests/unit/data/test_statistics.py`           | NEW — Welford unit tests                            |
| `tests/integration_tests/test_lazy_audio.py`   | NEW                                                 |
| `tests/integration_tests/test_lazy_image.py`   | NEW                                                 |

______________________________________________________________________

## Risks and Mitigations

| Risk                                                               | Mitigation                                                                                        |
| ------------------------------------------------------------------ | ------------------------------------------------------------------------------------------------- |
| DataLoader worker forking copies open file handles / CUDA contexts | `worker_init_fn` resets handles; delay CUDA init until after fork                                 |
| Welford online stats diverge from current batch stats              | Unit-test Welford against numpy batch stats on synthetic data                                     |
| Per-sample disk I/O is slower than batched HDF5 read               | Benchmark; add optional post-first-epoch disk cache for decoded tensors                           |
| Existing HDF5 preprocessing cache invalidated                      | Introduce cache version key; fall back to eager mode on cache miss                                |
| `IterableDataset` with shuffle buffer OOMs on large-media streams  | Default shuffle buffer scales with declared feature size; audio → 500, image → 200, text → 10 000 |

______________________________________________________________________

## Open Questions

1. **HDF5 cache vs. Arrow:** Should the decoded-tensor cache (post-Phase 2) write Arrow or HDF5?
   Arrow supports mmap natively; HDF5 has chunked access. Arrow preferred for new code.
1. **Normalization stat freshness:** If the dataset changes between runs, cached Welford stats
   are stale. Hash the dataset id + revision + split into the cache key.
1. **Multi-output features:** When one column feeds multiple output features (e.g. image →
   classification + detection), the decode must happen once. Ensure `LudwigDataset` returns a
   shared tensor, not independent copies.
1. **Tokenisation caching:** Tokenising every sample every epoch wastes CPU. Cache to Arrow after
   the first epoch under `${cache_dir}/<dataset_hash>/tokenized/`.