Optimize exposure index training IO

dataset.py

@@ -45,7 +45,7 @@ def _monthly_exposure_columns() -> list[str]:
 class ExposureCache:
     """Random-access view over files produced by prepare_exposure_cache.py."""
 
-    def __init__(self, cache_dir: str | Path, row_group_cache_size: int = 16):
+    def __init__(self, cache_dir: str | Path, row_group_cache_size: int = 4):
         cache_dir = Path(cache_dir)
         self.cache_dir = cache_dir
         manifest_path = cache_dir / "exposure_manifest.json"
@@ -129,6 +129,21 @@ class ExposureCache:
 
         self._key_to_index: dict[tuple[int, int, int], int] | None = None
 
+    def locality_key(self, indices: np.ndarray) -> tuple[int, int]:
+        """Return a stable parquet locality key for sampler-side batching."""
+        indices = np.asarray(indices, dtype=np.int64)
+        valid = indices[indices >= 0]
+        if len(valid) == 0:
+            return (2**31 - 1, 2**31 - 1)
+        if self.storage != "parquet_index":
+            return (0, int(valid[0] // 1024))
+        file_ids = np.asarray(self.daily_file_ids[valid], dtype=np.int64)
+        row_groups = np.asarray(self.daily_row_groups[valid], dtype=np.int64)
+        groups = np.stack([file_ids, row_groups], axis=1)
+        unique_groups, counts = np.unique(groups, axis=0, return_counts=True)
+        best = unique_groups[int(np.argmax(counts))]
+        return (int(best[0]), int(best[1]))
+
     def build_age_index(self, birth_date_by_eid: dict[int, np.datetime64]) -> None:
         keys: dict[tuple[int, int, int], int] = {}
         eids = np.asarray(self.eids, dtype=np.int64)
@@ -159,18 +174,67 @@ class ExposureCache:
         return out
 
     def daily_window(self, index: int) -> np.ndarray:
-        if index < 0:
-            return np.full(DAILY_EXPOSURE_SHAPE, np.nan, dtype=np.float32)
-        if self.storage == "dense_npy":
-            return np.asarray(self.daily[index], dtype=np.float32)
-        return self._parquet_window("daily", index)
+        return self.daily_windows(np.asarray([index], dtype=np.int64))[0]
 
     def monthly_window(self, index: int) -> np.ndarray:
-        if index < 0:
-            return np.full(MONTHLY_EXPOSURE_SHAPE, np.nan, dtype=np.float32)
+        return self.monthly_windows(np.asarray([index], dtype=np.int64))[0]
+
+    def daily_windows(self, indices: np.ndarray) -> np.ndarray:
+        return self._windows("daily", indices)
+
+    def monthly_windows(self, indices: np.ndarray) -> np.ndarray:
+        return self._windows("monthly", indices)
+
+    def _windows(
+        self,
+        kind: Literal["daily", "monthly"],
+        indices: np.ndarray,
+    ) -> np.ndarray:
+        indices = np.asarray(indices, dtype=np.int64)
+        shape = DAILY_EXPOSURE_SHAPE if kind == "daily" else MONTHLY_EXPOSURE_SHAPE
+        out = np.full((len(indices), shape[0], shape[1]), np.nan, dtype=np.float32)
+        valid_pos = np.nonzero(indices >= 0)[0]
+        if len(valid_pos) == 0:
+            return out
+
+        valid_indices = indices[valid_pos]
         if self.storage == "dense_npy":
-            return np.asarray(self.monthly[index], dtype=np.float32)
-        return self._parquet_window("monthly", index)
+            source = self.daily if kind == "daily" else self.monthly
+            out[valid_pos] = np.asarray(source[valid_indices], dtype=np.float32)
+            return out
+
+        if kind == "daily":
+            file_ids = np.asarray(self.daily_file_ids[valid_indices], dtype=np.int64)
+            row_groups = np.asarray(self.daily_row_groups[valid_indices], dtype=np.int64)
+            row_in_groups = np.asarray(self.daily_row_in_groups[valid_indices], dtype=np.int64)
+            columns = _daily_exposure_columns()
+        else:
+            file_ids = np.asarray(self.monthly_file_ids[valid_indices], dtype=np.int64)
+            row_groups = np.asarray(self.monthly_row_groups[valid_indices], dtype=np.int64)
+            row_in_groups = np.asarray(
+                self.monthly_row_in_groups[valid_indices],
+                dtype=np.int64,
+            )
+            columns = _monthly_exposure_columns()
+
+        group_keys = np.stack([file_ids, row_groups], axis=1)
+        unique_groups, inverse = np.unique(group_keys, axis=0, return_inverse=True)
+        for group_idx, (file_id, row_group) in enumerate(unique_groups):
+            group_pos = np.nonzero(inverse == group_idx)[0]
+            frame = self._read_parquet_row_group(
+                kind,
+                int(file_id),
+                int(row_group),
+                columns,
+            )
+            row_values = frame.iloc[row_in_groups[group_pos]].reindex(columns=columns)
+            values = (
+                row_values.to_numpy(dtype=np.float32, copy=True)
+                .reshape(len(group_pos), shape[1], shape[0])
+                .transpose(0, 2, 1)
+            )
+            out[valid_pos[group_pos]] = values
+        return out
 
     def _parquet_window(self, kind: Literal["daily", "monthly"], index: int) -> np.ndarray:
         if kind == "daily":
@@ -308,13 +372,17 @@ class _ExpoBaseDataset(Dataset):
         include_no_event_in_uts_target: bool = False,
         exposure_cache_dir: str | Path | None = None,
         mask_onset_exposure: bool = False,
+        exposure_row_group_cache_size: int = 4,
     ) -> None:
         self.data_prefix = data_prefix
         self.labels_file = labels_file
         self.no_event_interval_years = float(no_event_interval_years)
         self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target)
         self.exposure_cache = (
-            ExposureCache(exposure_cache_dir)
+            ExposureCache(
+                exposure_cache_dir,
+                row_group_cache_size=exposure_row_group_cache_size,
+            )
             if exposure_cache_dir is not None
             else None
         )
@@ -450,14 +518,14 @@ class _ExpoBaseDataset(Dataset):
         if self.exposure_cache is None:
             raise RuntimeError("Exposure cache is not enabled")
 
-        daily = np.stack(
-            [self.exposure_cache.daily_window(int(idx)) for idx in exposure_index],
-            axis=0,
-        ).astype(np.float32, copy=True)
-        monthly = np.stack(
-            [self.exposure_cache.monthly_window(int(idx)) for idx in exposure_index],
-            axis=0,
-        ).astype(np.float32, copy=True)
+        daily = self.exposure_cache.daily_windows(exposure_index).astype(
+            np.float32,
+            copy=False,
+        )
+        monthly = self.exposure_cache.monthly_windows(exposure_index).astype(
+            np.float32,
+            copy=False,
+        )
 
         if self.mask_onset_exposure:
             daily[:, 0, :] = np.nan
@@ -485,6 +553,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
         include_no_event_in_uts_target: bool = False,
         exposure_cache_dir: str | Path | None = None,
         mask_onset_exposure: bool = False,
+        exposure_row_group_cache_size: int = 4,
     ) -> None:
         super().__init__(
             data_prefix=data_prefix,
@@ -493,6 +562,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
             include_no_event_in_uts_target=include_no_event_in_uts_target,
             exposure_cache_dir=exposure_cache_dir,
             mask_onset_exposure=mask_onset_exposure,
+            exposure_row_group_cache_size=exposure_row_group_cache_size,
         )
 
         self.samples: List[Dict] = []
@@ -581,6 +651,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
         validation_query_seed: int = 42,
         exposure_cache_dir: str | Path | None = None,
         mask_onset_exposure: bool = False,
+        exposure_row_group_cache_size: int = 4,
     ) -> None:
         if split not in {"train", "valid", "test"}:
             raise ValueError(f"split must be train/valid/test, got {split!r}")
@@ -592,6 +663,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
             include_no_event_in_uts_target=include_no_event_in_uts_target,
             exposure_cache_dir=exposure_cache_dir,
             mask_onset_exposure=mask_onset_exposure,
+            exposure_row_group_cache_size=exposure_row_group_cache_size,
         )
 
         self.split = split

train_next_step.py

@@ -12,13 +12,13 @@ import json
 import logging
 import math
 import time
-from typing import Any, Dict
+from typing import Any, Dict, Iterator, List
 
 import numpy as np
 import torch
 from torch.nn.utils import clip_grad_norm_
 from torch.optim import AdamW
-from torch.utils.data import DataLoader, RandomSampler
+from torch.utils.data import DataLoader, RandomSampler, Sampler
 from tqdm.auto import tqdm
 
 from dataset import HealthDataset, collate_fn
@@ -53,6 +53,59 @@ EXPOSURE_INPUT_KEYS = (
 )
 
 
+class ExposureLocalityBatchSampler(Sampler[List[int]]):
+    """Randomized batches with within-buffer sorting by exposure parquet locality."""
+
+    def __init__(
+        self,
+        data_source,
+        batch_size: int,
+        buffer_size: int,
+        seed: int,
+        drop_last: bool = False,
+    ) -> None:
+        self.data_source = data_source
+        self.batch_size = int(batch_size)
+        self.buffer_size = max(int(buffer_size), self.batch_size)
+        self.seed = int(seed)
+        self.drop_last = bool(drop_last)
+        self.epoch = 0
+
+    def __iter__(self) -> Iterator[List[int]]:
+        n = len(self.data_source)
+        generator = torch.Generator().manual_seed(self.seed + self.epoch)
+        self.epoch += 1
+        shuffled = torch.randperm(n, generator=generator).tolist()
+        for start in range(0, n, self.buffer_size):
+            buffer = shuffled[start:start + self.buffer_size]
+            buffer.sort(key=self._locality_key)
+            for batch_start in range(0, len(buffer), self.batch_size):
+                batch = buffer[batch_start:batch_start + self.batch_size]
+                if len(batch) < self.batch_size and self.drop_last:
+                    continue
+                yield batch
+
+    def __len__(self) -> int:
+        n = len(self.data_source)
+        if self.drop_last:
+            return n // self.batch_size
+        return (n + self.batch_size - 1) // self.batch_size
+
+    def _locality_key(self, local_idx: int) -> tuple[int, int, int]:
+        dataset = self.data_source
+        raw_idx = int(local_idx)
+        if hasattr(dataset, "dataset") and hasattr(dataset, "indices"):
+            raw_idx = int(dataset.indices[local_idx])
+            dataset = dataset.dataset
+        sample = getattr(dataset, "samples", [])[raw_idx]
+        exposure_index = sample.get("exposure_index")
+        exposure_cache = getattr(dataset, "exposure_cache", None)
+        if exposure_index is None or exposure_cache is None:
+            return (2**31 - 1, 2**31 - 1, raw_idx)
+        file_id, row_group = exposure_cache.locality_key(exposure_index)
+        return (file_id, row_group, raw_idx)
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Train DeepHealth with next-token/point supervision")
@@ -83,6 +136,15 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--exposure_conv_kernel_size", type=int, default=7)
     parser.add_argument("--exposure_mlp_ratio", type=float, default=4.0)
     parser.add_argument("--no_exposure_gate", action="store_true")
+    parser.add_argument(
+        "--exposure_row_group_cache_size",
+        type=int,
+        default=4,
+        help=(
+            "Number of parquet exposure row groups cached per DataLoader worker "
+            "when using indexed exposure storage."
+        ),
+    )
 
     parser.add_argument("--target_mode", type=str, default="uts",
                         choices=["delphi2m", "uts"])
@@ -106,6 +168,21 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--patience", type=int, default=15)
     parser.add_argument("--min_lr_ratio", type=float, default=0.1)
     parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument(
+        "--prefetch_factor",
+        type=int,
+        default=4,
+        help="DataLoader batches prefetched per worker when num_workers > 0.",
+    )
+    parser.add_argument(
+        "--exposure_locality_buffer_size",
+        type=int,
+        default=4096,
+        help=(
+            "Training-only shuffle buffer sorted by exposure parquet locality. "
+            "Set 0 to use the standard RandomSampler."
+        ),
+    )
     parser.add_argument("--device", type=str, default="cuda")
     parser.add_argument("--progress_interval", type=int, default=20)
 
@@ -116,6 +193,12 @@ def parse_args() -> argparse.Namespace:
     )
     if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
         raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
+    if args.num_workers > 0 and args.prefetch_factor <= 0:
+        raise ValueError("prefetch_factor must be positive when num_workers > 0")
+    if args.exposure_row_group_cache_size < 0:
+        raise ValueError("exposure_row_group_cache_size must be non-negative")
+    if args.exposure_locality_buffer_size < 0:
+        raise ValueError("exposure_locality_buffer_size must be non-negative")
     if args.target_mode == "uts":
         args.readout_name = args.readout_name or "same_time_group_end"
         args.include_no_event_in_uts_target = True
@@ -364,6 +447,10 @@ def build_metadata(
         "exposure_conv_kernel_size": int(args.exposure_conv_kernel_size),
         "exposure_mlp_ratio": float(args.exposure_mlp_ratio),
         "exposure_use_gate": not bool(args.no_exposure_gate),
+        "exposure_row_group_cache_size": int(args.exposure_row_group_cache_size),
+        "num_workers": int(args.num_workers),
+        "prefetch_factor": int(args.prefetch_factor),
+        "exposure_locality_buffer_size": int(args.exposure_locality_buffer_size),
         "split_sizes": {
             "train": int(len(train_subset)),
             "val": int(len(val_subset)),
@@ -394,6 +481,13 @@ def main() -> None:
     logger.info(f"Device: {device}")
     logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
     logger.info(f"exposure_cache_dir={args.exposure_cache_dir}")
+    logger.info(
+        "DataLoader IO: "
+        f"num_workers={args.num_workers}, "
+        f"prefetch_factor={args.prefetch_factor if args.num_workers > 0 else None}, "
+        f"exposure_row_group_cache_size={args.exposure_row_group_cache_size}, "
+        f"exposure_locality_buffer_size={args.exposure_locality_buffer_size}"
+    )
 
     dataset = HealthDataset(
         data_prefix=args.data_prefix,
@@ -402,6 +496,7 @@ def main() -> None:
         include_no_event_in_uts_target=args.include_no_event_in_uts_target,
         exposure_cache_dir=args.exposure_cache_dir,
         mask_onset_exposure=args.mask_onset_exposure,
+        exposure_row_group_cache_size=args.exposure_row_group_cache_size,
     )
     if args.train_eid_file and args.val_eid_file and args.test_eid_file:
         train_subset, val_subset, test_subset = split_dataset_by_eid_files(
@@ -430,35 +525,50 @@ def main() -> None:
         f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
     )
 
-    train_loader = DataLoader(
-        train_subset,
-        batch_size=args.batch_size,
-        sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)),
-        collate_fn=collate_fn,
-        num_workers=args.num_workers,
-        pin_memory=device.type == "cuda",
-        persistent_workers=args.num_workers > 0,
-        prefetch_factor=2 if args.num_workers > 0 else None,
+    dataloader_kwargs = {
+        "collate_fn": collate_fn,
+        "num_workers": args.num_workers,
+        "pin_memory": device.type == "cuda",
+        "persistent_workers": args.num_workers > 0,
+        "prefetch_factor": args.prefetch_factor if args.num_workers > 0 else None,
+    }
+    use_locality_sampler = (
+        args.exposure_cache_dir is not None
+        and args.exposure_locality_buffer_size > 0
     )
+    if use_locality_sampler:
+        logger.info("Using exposure-locality batch sampler for training")
+        train_loader = DataLoader(
+            train_subset,
+            batch_sampler=ExposureLocalityBatchSampler(
+                train_subset,
+                batch_size=args.batch_size,
+                buffer_size=args.exposure_locality_buffer_size,
+                seed=args.seed,
+            ),
+            **dataloader_kwargs,
+        )
+    else:
+        train_loader = DataLoader(
+            train_subset,
+            batch_size=args.batch_size,
+            sampler=RandomSampler(
+                train_subset,
+                generator=torch.Generator().manual_seed(args.seed),
+            ),
+            **dataloader_kwargs,
+        )
     val_loader = DataLoader(
         val_subset,
         batch_size=args.batch_size,
         shuffle=False,
-        collate_fn=collate_fn,
-        num_workers=args.num_workers,
-        pin_memory=device.type == "cuda",
-        persistent_workers=args.num_workers > 0,
-        prefetch_factor=2 if args.num_workers > 0 else None,
+        **dataloader_kwargs,
     )
     test_loader = DataLoader(
         test_subset,
         batch_size=args.batch_size,
         shuffle=False,
-        collate_fn=collate_fn,
-        num_workers=args.num_workers,
-        pin_memory=device.type == "cuda",
-        persistent_workers=args.num_workers > 0,
-        prefetch_factor=2 if args.num_workers > 0 else None,
+        **dataloader_kwargs,
     )
 
     model = build_model(args, dataset).to(device)