Add offline time-window indices for Yambda

Build stable indices and offsets during cache preparation when enabled. Reuse read-only slices for streaming windows and offsets for sample counts, with serialized construction across ranks.

Author: JacoCheung

recommendation/README.MD

@@ -383,6 +383,30 @@ env overrides:
 | `EVAL_EVERY_DATA_PCT` | `streaming_train_eval_loop.eval_every_data_pct` | 0.0 | eval cadence by fraction of train data (full ref / RCP: 0.001) |
 | `MIN_HISTORY` | `get_dataset.min_history` | 4086 | anchor-eligibility floor (0 = ~all users incl. cold-start) |
 
+#### Offline time-window indices
+
+To bucketize anchors by time window, set this in
+[`yambda_5b.gin`](generative_recommenders/dlrm_v4/train/gin/yambda_5b.gin):
+
+```gin
+get_dataset.use_offline_ts_indices = True
+```
+
+The default is `False`. When enabled, `ts_indices_*.npy` and `ts_offsets_*.npy`
+are **built automatically on the first run** during dataset initialization,
+alongside the positions cache. This also works with an existing cache folder
+that lacks these files. Later runs reuse them.
+
+For read-only data mounts, build them beforehand from an existing cache:
+
+```bash
+python3 -m generative_recommenders.dlrm_v4.datasets.offline_ts_indices \
+    --cache-dir "$DLRM_DATA_PATH/processed_5b/hstu_cache_L4086" \
+    --history-length 4086 --min-history 4086 --window-seconds 86400
+```
+
+Use the same history and window settings as training.
+
 ### Test data order
 
 The evaluation window is traversed once per evaluation, in the same in-window

recommendation/generative_recommenders/dlrm_v4/datasets/offline_ts_indices.py

@@ -0,0 +1,111 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+
+"""Bucketize Yambda anchors by time window and save TS indices and offsets."""
+
+import argparse
+import fcntl
+import os
+import tempfile
+from pathlib import Path
+
+import numpy as np
+
+
+def build_offline_ts_indices(
+    cache_dir, history_length, min_history=None, window_seconds=86400,
+    chunk_size=1_000_000,
+):
+    """Bucketize anchors by time window, preserving their order within each bucket."""
+    if history_length <= 0 or window_seconds <= 0 or chunk_size <= 0:
+        raise ValueError("history_length, window_seconds and chunk_size must be positive")
+    if min_history is not None and min_history < 0:
+        raise ValueError("min_history must be nonnegative")
+    cache_dir = Path(cache_dir)
+    tag = f"L{history_length}"
+    if min_history is not None and min_history != history_length:
+        tag += f"_m{min_history}"
+    # Keep derived files with their source cache so separate caches cannot collide.
+    indices_path = cache_dir / f"ts_indices_{tag}_W{window_seconds}.npy"
+    offsets_path = cache_dir / f"ts_offsets_{tag}_W{window_seconds}.npy"
+    # Complete caches bypass the build lock, as do training-time index reads.
+    if indices_path.exists() and offsets_path.exists():
+        return indices_path, offsets_path
+
+    # Each rank calls this during dataset initialization. Reuse the positions
+    # lock so only one rank builds missing output, avoiding duplicate work and
+    # concurrent writes to the same files without adding a separate TS lock.
+    with open(cache_dir / "_positions_lock", "a") as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        # Another rank may have finished building while this rank waited.
+        if indices_path.exists() and offsets_path.exists():
+            return indices_path, offsets_path
+        positions = np.load(cache_dir / f"positions_{tag}.npy", mmap_mode="r")
+        anchor_path = cache_dir / f"anchor_ts_{tag}.npy"
+        has_anchor_ts = anchor_path.exists()
+        timestamps = np.load(
+            anchor_path if has_anchor_ts else cache_dir / "flat_timestamps.npy", mmap_mode="r"
+        )
+        if any(array.ndim != 1 or array.dtype != np.int64 for array in (positions, timestamps)):
+            raise ValueError("positions and timestamps must be one-dimensional int64 arrays")
+        count = len(positions)
+        if has_anchor_ts and len(timestamps) != count:
+            raise ValueError("Anchor timestamps and positions have different lengths")
+
+        def timestamp_chunk(start):
+            stop = min(start + chunk_size, count)
+            return timestamps[start:stop] if has_anchor_ts else timestamps[positions[start:stop]]
+
+        t_min, t_max = 0, -1
+        for start in range(0, count, chunk_size):
+            chunk = timestamp_chunk(start)
+            t_min = int(chunk.min()) if start == 0 else min(t_min, int(chunk.min()))
+            t_max = int(chunk.max()) if start == 0 else max(t_max, int(chunk.max()))
+        if t_max - t_min > np.iinfo(np.int64).max:
+            raise ValueError("Timestamp span exceeds int64 capacity")
+        windows = (t_max - t_min) // window_seconds + 1 if count else 0
+        counts = np.zeros(windows, dtype=np.int64)
+        for start in range(0, count, chunk_size):
+            buckets = (timestamp_chunk(start) - t_min) // window_seconds
+            counts += np.bincount(buckets, minlength=windows)
+        offsets = np.concatenate(([0], np.cumsum(counts)))
+        cursors = offsets[:-1].copy()
+        with tempfile.TemporaryDirectory(prefix=".ts-indices-", dir=cache_dir) as tmp:
+            staging = Path(tmp)
+            indices = np.lib.format.open_memmap(
+                staging / indices_path.name, mode="w+", dtype=np.int64, shape=(count,)
+            )
+            try:
+                for start in range(0, count, chunk_size):
+                    buckets = (timestamp_chunk(start) - t_min) // window_seconds
+                    order = np.argsort(buckets, kind="stable")
+                    sizes = np.bincount(buckets, minlength=windows)
+                    begin = 0
+                    for bucket in np.flatnonzero(sizes):
+                        end = begin + int(sizes[bucket])
+                        cursor = int(cursors[bucket])
+                        indices[cursor:cursor + end - begin] = order[begin:end] + start
+                        cursors[bucket] += end - begin
+                        begin = end
+                indices.flush()
+            finally:
+                indices._mmap.close()
+            np.save(staging / offsets_path.name, offsets)
+            # Publish offsets last so an interrupted build cannot appear complete.
+            offsets_path.unlink(missing_ok=True)
+            os.replace(staging / indices_path.name, indices_path)
+            os.replace(staging / offsets_path.name, offsets_path)
+    return indices_path, offsets_path
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--cache-dir", required=True)
+    parser.add_argument("--history-length", type=int, required=True)
+    parser.add_argument("--min-history", type=int)
+    parser.add_argument("--window-seconds", type=int, default=86400)
+    parser.add_argument("--chunk-size", type=int, default=1_000_000)
+    print(*build_offline_ts_indices(**vars(parser.parse_args())), sep="\n")

recommendation/generative_recommenders/dlrm_v4/datasets/yambda.py

@@ -219,6 +219,7 @@ def __init__(
         streaming_shuffle_seed: int = 0,
         train_split_percentage: float = 1.0,
         split_salt: int = 0,
+        use_offline_ts_indices: bool = False,
         *args,
         **kwargs,
     ) -> None:
@@ -283,6 +284,9 @@ def __init__(
         self._active: Optional[np.ndarray] = None
         self.is_eval: bool = False
         self._anchor_ts: Optional[np.ndarray] = None
+        self._use_offline_ts_indices = use_offline_ts_indices
+        self._ts_indices: Optional[np.ndarray] = None
+        self._ts_offsets: Optional[np.ndarray] = None
         self._t_min: Optional[int] = None
         self._t_max: Optional[int] = None
         self._cache_dir: Optional[str] = cache_dir
@@ -315,6 +319,15 @@ def __init__(
         self._ensure_positions_built(
             cache_dir, self._positions_name, self._min_history
         )
+        if self._use_offline_ts_indices:
+            from generative_recommenders.dlrm_v4.datasets.offline_ts_indices import (
+                build_offline_ts_indices,
+            )
+
+            build_offline_ts_indices(
+                cache_dir, history_length, self._min_history,
+                window_seconds=self._streaming_window_seconds,
+            )
         self._positions: np.ndarray = _load_npy_readonly(
             os.path.join(cache_dir, self._positions_name)
         )
@@ -596,6 +609,25 @@ def _ensure_streaming_index(self) -> None:
         Multi-rank safe via an exclusive file lock + atomic rename; all ranks
         then mmap the result read-only (shared physical pages, ~0 anon).
         """
+        if self._use_offline_ts_indices:
+            if self._ts_indices is None:
+                directory = Path(self._cache_dir)
+                tag = self._positions_name[len("positions_"):-4]
+                tag += f"_W{self._streaming_window_seconds}"
+                indices = _load_npy_readonly(directory / f"ts_indices_{tag}.npy")
+                offsets = _load_npy_readonly(directory / f"ts_offsets_{tag}.npy")
+                if (
+                    indices.dtype != np.int64 or indices.ndim != 1
+                    or len(indices) != len(self._positions)
+                    or offsets.dtype != np.int64 or offsets.ndim != 1
+                    or len(offsets) == 0 or offsets[0] != 0
+                    or offsets[-1] != len(indices)
+                    or np.any(offsets[1:] < offsets[:-1])
+                ):
+                    raise ValueError("Invalid offline TS indices or offsets")
+                self._ts_offsets = offsets
+                self._ts_indices = indices
+            return
         if self._anchor_ts is not None:
             return
         import fcntl
@@ -626,9 +658,21 @@ def _ensure_streaming_index(self) -> None:
         self._t_min = int(self._anchor_ts.min())
         self._t_max = int(self._anchor_ts.max())
 
+    def _offline_window_bounds(self, start_ts: int, num_ts: int = 1) -> Tuple[int, int]:
+        windows = len(self._ts_offsets) - 1
+        start = min(windows, max(0, start_ts))
+        end = min(windows, max(0, start_ts + max(0, num_ts)))
+        return int(self._ts_offsets[start]), int(self._ts_offsets[end])
+
+    def _offline_window_indices(self, start_ts: int, num_ts: int = 1) -> np.ndarray:
+        lo, hi = self._offline_window_bounds(start_ts, num_ts)
+        return self._ts_indices[lo:hi]
+
     def num_windows(self) -> int:
         """Number of fixed-duration windows spanning [t_min, t_max]."""
         self._ensure_streaming_index()
+        if self._ts_offsets is not None:
+            return len(self._ts_offsets) - 1
         assert self._t_min is not None and self._t_max is not None
         span = self._t_max - self._t_min + 1
         w = self._streaming_window_seconds
@@ -653,6 +697,13 @@ def window_indices(
         O(N) mask here keeps only one ~26 GB array resident and is robust.
         """
         self._ensure_streaming_index()
+        if self._ts_indices is not None:
+            idx = self._offline_window_indices(ts)
+            do_sort = self._streaming_sort_within_window if sort_by_time is None else sort_by_time
+            if do_sort and idx.size:
+                timestamps = self.store.flat_timestamps[self._positions[idx]]
+                idx = idx[np.argsort(timestamps, kind="stable")]
+            return idx
         assert self._anchor_ts is not None and self._t_min is not None
         w = self._streaming_window_seconds
         lo = self._t_min + ts * w
@@ -778,6 +829,15 @@ def total_train_anchors(self, start_ts: int, num_ts: int) -> int:
         by a fraction of a window).
         """
         self._ensure_streaming_index()
+        if self._ts_indices is not None:
+            lo, hi = self._offline_window_bounds(start_ts, num_ts)
+            if self._train_split_percentage >= 1.0:
+                return hi - lo
+            indices = self._ts_indices[lo:hi]
+            return sum(
+                int(np.count_nonzero(~self._eval_anchor_mask(indices[start:start + 1_000_000])))
+                for start in range(0, len(indices), 1_000_000)
+            )
         assert self._anchor_ts is not None and self._t_min is not None
         if num_ts <= 0:
             return 0

recommendation/generative_recommenders/dlrm_v4/train/gin/yambda_5b.gin

@@ -528,6 +528,8 @@ nl/env_int.default = 3
 # to the dataset's available window count at runtime; override via $NUM_TRAIN_TS.
 get_dataset.streaming_window_seconds = 86400
 get_dataset.streaming_sort_within_window = False
+# Bucketize anchors by time window on first use and reuse the cached indices.
+get_dataset.use_offline_ts_indices = False
 # In-window shuffle to break user-major batching (consecutive sliding-window
 # anchors otherwise come from the same few users -> few unique embedding reads).
 # Diversity dial in [0,1] -- the AGREED, config-invariant benchmark knob. Maps to

recommendation/generative_recommenders/dlrm_v4/utils.py

@@ -1634,6 +1634,7 @@ def get_dataset(
     streaming_shuffle_seed: int = 0,
     train_split_percentage: float = 1.0,
     split_salt: int = 0,
+    use_offline_ts_indices: bool = False,
 ):
     """
     Get dataset class and configuration by name.
@@ -1777,6 +1778,7 @@ def get_dataset(
                 # streaming-train-eval; ignored by the default train-eval path).
                 "streaming_window_seconds": streaming_window_seconds,
                 "streaming_sort_within_window": streaming_sort_within_window,
+                "use_offline_ts_indices": use_offline_ts_indices,
                 # In-window shuffle diversity dial in [0,1]: K=round(frac*N) within-
                 # segment shuffle. 0=off/user-major, 1=full. Config-invariant and
                 # deterministic by (seed, ts).