feat(data): preprocessing mode enum + prefetch_size config + lazy_cached memmap (#4173)

Author: w4nderlust

ludwig/backend/ray.py

@@ -714,7 +714,11 @@ def to_tensors(df: pd.DataFrame) -> pd.DataFrame:
         from ludwig.data.dataframe.dask import tensor_extension_casting
 
         with tensor_extension_casting(False):
-            predictions = dataset.ds.map_batches(to_tensors, batch_format="pandas").map_batches(
+            # Apply lazy-decode transforms before casting to tensors.  Without this,
+            # lazy audio/image columns still contain file-path strings at predict time,
+            # causing cast_as_tensor_dtype to raise a TypeError.
+            ds = dataset._with_lazy_decode(dataset.ds)
+            predictions = ds.map_batches(to_tensors, batch_format="pandas").map_batches(
                 batch_predictor,
                 batch_size=self.batch_size,
                 compute=ray.data.ActorPoolStrategy(),

ludwig/data/batcher/random_access.py

@@ -41,8 +41,15 @@ class RandomAccessBatcher(Batcher):
     is a cheap numpy slice.
 
     ``PandasDataset.initialize_batcher`` automatically sets ``prefetch_size``
-    to a non-zero value when any lazy column (audio/image path arrays) is
-    present, so callers don't need to set this manually.
+    from each feature's preprocessing config (``None`` → 4 for ``lazy`` /
+    ``lazy_cached`` features, 0 for ``eager``).  Users can override this via
+    the ``prefetch_size`` field in the feature's preprocessing config, or by
+    passing ``prefetch_size`` directly to ``initialize_batcher``.
+
+    After the first epoch completes in ``lazy_cached`` mode, ``set_epoch``
+    checks ``dataset.is_fully_cached()`` and resets ``prefetch_size`` to 0 —
+    memmap reads (~0.1 ms/batch) are fast enough that background pipelining
+    adds no measurable benefit and avoids unnecessary thread overhead.
     """
 
     def __init__(
@@ -234,6 +241,13 @@ def last_batch(self) -> bool:
         return False
 
     def set_epoch(self, epoch: int, batch_size: int) -> None:
+        """Reset state for a new epoch; disables prefetch when the dataset is fully cached.
+
+        If the dataset exposes ``is_fully_cached()`` and it returns ``True``
+        (i.e. all ``lazy_cached`` columns have finished their first-pass decode),
+        ``prefetch_size`` is set to 0 so subsequent epochs read directly from
+        the memmap without spinning up a producer thread.
+        """
         if self._prefetch_size > 0:
             self._stop_async()
 
@@ -244,6 +258,12 @@ def set_epoch(self, epoch: int, batch_size: int) -> None:
         self.sampler.set_epoch(epoch)
         self.sample_it = iter(self.sampler)
 
+        # After epoch 1, if all lazy columns are decoded and cached in memmaps,
+        # disable prefetch — memmap reads are fast enough that pipelining adds
+        # no measurable benefit and avoids unnecessary thread overhead.
+        if self._prefetch_size > 0 and hasattr(self.dataset, "is_fully_cached") and self.dataset.is_fully_cached():
+            self._prefetch_size = 0
+
         if self._prefetch_size > 0:
             self._start_async_epoch()
 

ludwig/data/dataset/pandas.py

@@ -160,14 +160,31 @@ def __init__(self, dataset, features, data_cache_fp, training_set_metadata=None)
         # During preprocessing, lazy features stored file paths as strings.  Here
         # we wrap those path arrays with the appropriate per-sample decode function
         # so that batches are decoded on demand rather than all at once.
+        self._auto_prefetch_size = 0
         if training_set_metadata is not None:
             self._init_lazy_columns(features, training_set_metadata)
 
         self.size = len(list(self.dataset.values())[0])
 
+    def _decoded_cache_path(self, proc_col: str, n: int, sample_shape: tuple) -> str:
+        """Return the path for the decoded numpy memmap file for a given column.
+
+        Placed next to the Parquet cache when ``data_cache_fp`` is set, otherwise
+        falls back to ``~/.cache/ludwig/lazy_media/``.
+        """
+        flat_shape = "_".join(str(d) for d in sample_shape)
+        fname = f"{proc_col}_decoded_n{n}_{flat_shape}_f32.npy"
+        if self.data_cache_fp:
+            return os.path.join(os.path.dirname(self.data_cache_fp), fname)
+        from ludwig.data.lazy_utils import get_default_lazy_cache_dir
+
+        return str(get_default_lazy_cache_dir() / fname)
+
     def _init_lazy_columns(self, features: dict, training_set_metadata: dict) -> None:
-        """Replace string-path arrays for lazy features with LazyColumn objects."""
-        from ludwig.data.lazy_utils import LazyColumn
+        """Replace string-path arrays for lazy features with LazyColumn or CachedLazyColumn objects."""
+        from ludwig.data.lazy_utils import CachedLazyColumn, LazyColumn
+
+        max_prefetch = 0
 
         for proc_col, feat_cfg in features.items():
             if proc_col not in self.dataset:
@@ -179,6 +196,11 @@ def _init_lazy_columns(self, features: dict, training_set_metadata: dict) -> Non
 
             paths = self.dataset[proc_col]
             feat_type = feat_cfg.get("type", "")
+            mode = feat_meta.get("mode", "lazy")
+            # Per-feature prefetch override (None means auto → 4 for lazy modes)
+            feat_prefetch = feat_meta.get("prefetch_size")
+            effective_prefetch = feat_prefetch if feat_prefetch is not None else 4
+            max_prefetch = max(max_prefetch, effective_prefetch)
 
             if feat_type == "audio" and "lazy_audio_params" in feat_meta:
                 from ludwig.features.audio_feature import AudioFeatureMixin
@@ -198,7 +220,16 @@ def _init_lazy_columns(self, features: dict, training_set_metadata: dict) -> Non
                 import torch as _torch
 
                 _audio_workers = max(1, (os.cpu_count() or 4) // max(1, _torch.get_num_threads()))
-                self.dataset[proc_col] = LazyColumn(paths, decode_fn, max_workers=_audio_workers)
+
+                if mode == "lazy_cached":
+                    # decode_fn returns (max_length, feature_dim) — match that order.
+                    sample_shape = (p["max_length"], p["feature_dim"])
+                    cache_path = self._decoded_cache_path(proc_col, len(paths), sample_shape)
+                    self.dataset[proc_col] = CachedLazyColumn(
+                        paths, decode_fn, cache_path, sample_shape, max_workers=_audio_workers
+                    )
+                else:
+                    self.dataset[proc_col] = LazyColumn(paths, decode_fn, max_workers=_audio_workers)
 
             elif feat_type == "image" and "lazy_image_params" in feat_meta:
                 import torch
@@ -220,7 +251,15 @@ def _init_lazy_columns(self, features: dict, training_set_metadata: dict) -> Non
                     channel_class_map=channel_class_map,
                     default_image=default_image,
                 )
-                self.dataset[proc_col] = LazyColumn(paths, decode_fn)
+
+                if mode == "lazy_cached":
+                    sample_shape = tuple(p["default_image_shape"])
+                    cache_path = self._decoded_cache_path(proc_col, len(paths), sample_shape)
+                    self.dataset[proc_col] = CachedLazyColumn(paths, decode_fn, cache_path, sample_shape)
+                else:
+                    self.dataset[proc_col] = LazyColumn(paths, decode_fn)
+
+        self._auto_prefetch_size = max_prefetch
 
     def to_df(self, features: Iterable[BaseFeature] | None = None) -> DataFrame:
         """Convert the dataset to a Pandas DataFrame.
@@ -271,6 +310,19 @@ def _has_lazy_columns(self) -> bool:
 
         return any(is_lazy_column(v) for v in self.dataset.values())
 
+    def is_fully_cached(self) -> bool:
+        """Return ``True`` when every ``CachedLazyColumn`` in this dataset has finished its first-pass decode.
+
+        Returns ``False`` if there are no ``CachedLazyColumn`` instances (i.e. the dataset
+        uses plain ``LazyColumn`` or eager mode).
+        """
+        from ludwig.data.lazy_utils import is_cached_lazy_column
+
+        cached_cols = [v for v in self.dataset.values() if is_cached_lazy_column(v)]
+        if not cached_cols:
+            return False
+        return all(col.is_fully_cached() for col in cached_cols)
+
     @contextlib.contextmanager
     def initialize_batcher(
         self,
@@ -282,12 +334,29 @@ def initialize_batcher(
         augmentation_pipeline=None,
         prefetch_size: int | None = None,
     ) -> Batcher:
+        """Yield a :class:`RandomAccessBatcher` configured for this dataset.
+
+        Parameters
+        ----------
+        prefetch_size:
+            Number of batches to pipeline in a background thread while the GPU
+            processes the current batch.  ``None`` (default) uses
+            ``self._auto_prefetch_size``, which is derived from the
+            ``prefetch_size`` field in each feature's preprocessing config
+            (``None`` → 4 for ``lazy``/``lazy_cached`` features, 0 for
+            ``eager``).  Pass ``0`` to disable prefetch entirely.
+
+            For ``lazy_cached`` mode, ``RandomAccessBatcher.set_epoch`` will
+            automatically reset ``prefetch_size`` to 0 after epoch 1 once
+            ``dataset.is_fully_cached()`` returns ``True`` — memmap reads are
+            fast enough that background pipelining adds no measurable benefit.
+        """
         # When the dataset contains lazy columns (audio / image file paths that
         # are decoded per-batch), automatically enable prefetch so the GPU is
         # not idle during decode.  Callers can override by passing prefetch_size
         # explicitly (0 to disable, N to set a specific depth).
         if prefetch_size is None:
-            prefetch_size = 4 if self._has_lazy_columns() else 0
+            prefetch_size = self._auto_prefetch_size if self._has_lazy_columns() else 0
 
         sampler = DistributedSampler(
             len(self), shuffle=should_shuffle, random_seed=random_seed, distributed=distributed

ludwig/data/dataset/ray.py

@@ -133,7 +133,9 @@ def _with_lazy_decode(self, ds: RayNativeDataset) -> RayNativeDataset:
         for proc_col, feat_cfg in self.features.items():
             feature_name = feat_cfg.get("name") or feat_cfg.get("column") or proc_col
             feat_meta = self.training_set_metadata.get(feature_name, {})
-            if not isinstance(feat_meta, dict) or not feat_meta.get("lazy"):
+            if not isinstance(feat_meta, dict) or (
+                not feat_meta.get("lazy") and feat_meta.get("mode", "eager") == "eager"
+            ):
                 continue
 
             feat_type = feat_cfg.get("type", "")

ludwig/data/lazy_utils.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import os
+import threading
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Callable
@@ -31,6 +32,17 @@
 class LazyColumn:
     """Array-like wrapper that decodes file paths on demand per batch.
 
+    Decoding runs in a ``ThreadPoolExecutor`` at batch-slice time, keeping peak
+    memory bounded to ``batch_size`` samples at any one time.  The thread pool
+    overlaps with the GPU forward pass when ``RandomAccessBatcher`` prefetch is
+    enabled.
+
+    .. warning::
+        For audio FBANK features, each ``decode_fn`` call already spawns PyTorch's
+        internal thread pool.  Creating many ``LazyColumn`` workers in parallel will
+        over-subscribe CPUs (up to 5× slowdown).  ``PandasDataset._init_lazy_columns``
+        caps ``max_workers`` at ``cpu_count // torch_num_threads`` for this case.
+
     Parameters
     ----------
     paths:
@@ -97,9 +109,151 @@ def __repr__(self) -> str:
         return f"LazyColumn(n={len(self._paths)}, decode_fn={self._decode_fn.__name__!r})"
 
 
+def _select(paths: np.ndarray, indices) -> tuple[list, list, bool]:
+    """Normalise any index type to ``(paths_list, int_indices_list, is_scalar)``."""
+    scalar = isinstance(indices, (int, np.integer))
+    if not scalar and isinstance(indices, np.ndarray) and indices.ndim == 0:
+        scalar = True
+
+    if scalar:
+        idx = int(indices)
+        return [paths[idx]], [idx], True
+
+    if isinstance(indices, slice):
+        int_indices = list(range(*indices.indices(len(paths))))
+    elif isinstance(indices, np.ndarray) and indices.dtype == bool:
+        int_indices = list(np.where(indices)[0])
+    elif isinstance(indices, np.ndarray):
+        int_indices = indices.tolist()
+    else:
+        int_indices = list(indices)
+
+    return [paths[i] for i in int_indices], int_indices, False
+
+
+class CachedLazyColumn:
+    """Like :class:`LazyColumn`, but writes decoded arrays to a numpy memmap for reuse.
+
+    **First pass (epoch 1):** behaves identically to :class:`LazyColumn` — decodes via
+    ``ThreadPoolExecutor`` — but also writes each decoded sample to a ``np.memmap`` file.
+    A per-sample boolean array tracks which indices have been written.  When every sample
+    has been written, the memmap is flushed and an empty ``.done`` sentinel file is created
+    next to the memmap.
+
+    **Subsequent passes (epoch 2+):** ``is_fully_cached()`` returns ``True`` and
+    ``__getitem__`` reads directly from the memmap (~0.1 ms/batch), bypassing the thread
+    pool entirely.  ``RandomAccessBatcher.set_epoch`` detects this via
+    ``dataset.is_fully_cached()`` and automatically disables prefetch.
+
+    If the ``.done`` file already exists when the object is constructed (e.g. a resumed
+    run), the memmap is opened in read-only mode immediately — no decode occurs at all.
+
+    Parameters
+    ----------
+    paths:
+        1-D array or list of file paths.
+    decode_fn:
+        Callable ``path -> np.ndarray`` (same contract as :class:`LazyColumn`).
+    cache_path:
+        Full path for the memmap file (e.g. ``/data/audio_proc_decoded_n1000_8_23_f32.npy``).
+    sample_shape:
+        Shape of a single decoded sample, e.g. ``(max_length, feature_dim)``.
+    dtype:
+        Element dtype for the memmap. Default ``np.float32``.
+    max_workers:
+        Thread-pool size for parallel decode on cache-miss. Default ``_DEFAULT_MAX_WORKERS``.
+    """
+
+    def __init__(
+        self,
+        paths: np.ndarray | list,
+        decode_fn: Callable[[str], np.ndarray],
+        cache_path: str,
+        sample_shape: tuple,
+        dtype=np.float32,
+        max_workers: int = _DEFAULT_MAX_WORKERS,
+    ) -> None:
+        self._paths = np.asarray(paths, dtype=object)
+        self._decode_fn = decode_fn
+        self._cache_path = cache_path
+        self._done_path = cache_path + ".done"
+        self._sample_shape = sample_shape
+        self._dtype = dtype
+        self._max_workers = max_workers
+        self._n = len(self._paths)
+        self._written = np.zeros(self._n, dtype=bool)
+        self._lock = threading.Lock()
+        self._memmap = None
+        self._fully_cached = os.path.exists(self._done_path)
+        if self._fully_cached:
+            self._written[:] = True
+            self._memmap = np.memmap(cache_path, dtype=dtype, mode="r", shape=(self._n, *sample_shape))
+
+    # ------------------------------------------------------------------
+    # numpy-compatible interface
+    # ------------------------------------------------------------------
+
+    def __len__(self) -> int:
+        return self._n
+
+    def __getitem__(self, indices) -> np.ndarray:
+        _, int_indices, scalar = _select(self._paths, indices)
+
+        if self._fully_cached:
+            result = np.array(self._memmap[int_indices])
+            return result[0] if scalar else result
+
+        # Decode any samples not yet in the cache.
+        need_decode = [i for i in int_indices if not self._written[i]]
+
+        if need_decode:
+            paths_list = [self._paths[i] for i in need_decode]
+            with ThreadPoolExecutor(max_workers=min(self._max_workers, len(paths_list))) as ex:
+                decoded = list(ex.map(self._decode_fn, paths_list))
+
+            with self._lock:
+                if self._memmap is None:
+                    os.makedirs(os.path.dirname(self._cache_path) or ".", exist_ok=True)
+                    self._memmap = np.memmap(
+                        self._cache_path, dtype=self._dtype, mode="w+", shape=(self._n, *self._sample_shape)
+                    )
+                for i, arr in zip(need_decode, decoded):
+                    if not self._written[i]:
+                        self._memmap[i] = arr
+                        self._written[i] = True
+
+                if self._written.all() and not self._fully_cached:
+                    self._fully_cached = True
+                    self._memmap.flush()
+                    Path(self._done_path).touch()
+
+        result = np.array(self._memmap[int_indices])
+        return result[0] if scalar else result
+
+    @property
+    def dtype(self):
+        return object
+
+    @property
+    def shape(self):
+        return (self._n,)
+
+    def is_fully_cached(self) -> bool:
+        """Return True once every sample has been decoded and written to the memmap."""
+        return self._fully_cached
+
+    def __repr__(self) -> str:
+        return f"CachedLazyColumn(n={self._n}, cached={self._fully_cached})"
+
+
 def is_lazy_column(col) -> bool:
-    """Return True if *col* is a ``LazyColumn`` instance."""
-    return isinstance(col, LazyColumn)
+    """Return True if *col* is a ``LazyColumn`` or ``CachedLazyColumn`` instance."""
+    return isinstance(col, (LazyColumn, CachedLazyColumn))
+
+
+def is_cached_lazy_column(col) -> bool:
+    """Return True if *col* is a ``CachedLazyColumn`` instance."""
+    return isinstance(col, CachedLazyColumn)
 
 
 def get_default_lazy_cache_dir() -> Path: