Use eid-aligned exposure cache

train_next_step.py

@@ -102,8 +102,8 @@ class ExposureLocalityBatchSampler(Sampler[List[int]]):
         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)
+        block_id, block_offset = exposure_cache.locality_key(exposure_index)
+        return (block_id, block_offset, raw_idx)
 
 
 def parse_args() -> argparse.Namespace:
@@ -136,16 +136,6 @@ 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"])
     parser.add_argument("--readout_name", type=str, default=None,
@@ -195,8 +185,6 @@ def parse_args() -> argparse.Namespace:
         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":
@@ -447,7 +435,6 @@ 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),
@@ -485,7 +472,6 @@ def main() -> None:
         "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}"
     )
 
@@ -496,7 +482,6 @@ 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(