Sharded dataset generation

README.md

@@ -59,6 +59,12 @@ The model is trained from a random initialization until convergence, which is de
 1. Once fractal generation completes, run the benchmark:  
     `torchrun-hpc -N 1 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark -c ScaFFold/configs/benchmark_default.yml`
 
+### Dataset cache and sharded datagen
+
+`benchmark` creates or reuses datasets under `dataset_dir`. New datasets are written in the v3 format, which stores one volume and mask file per logical sample per physical shard. The physical layout is controlled by `dc_num_shards` and `dc_shard_dims`; for example, `dc_num_shards: [1, 1, 2]` writes two physical shards per logical volume, with filenames such as `120_shard000000.npy` and `120_shard000001.npy`. Datasets are generated with the same sharding configuration used for model training.
+
+Unsharded runs use `dc_num_shards: [1, 1, 1]` and are still stored as v3 datasets with `_shard000000` files. ScaFFold only reuses matching v3 caches; older full-volume caches are ignored and a new v3 dataset is generated instead.
+
 ScaFFold benchmark training always uses PyTorch distributed execution with DistConv spatial parallelism. For a singleton run, launch one distributed rank rather than disabling distributed execution.
 
 `benchmark` creates a folder for the benchmark run(s) at `base_run_dir` set in the config file. For reproducibility, we store a copy of the benchmark run config yml. Within each run subfolder, `benchmark` creates a yml config for that specific run.
@@ -168,6 +174,8 @@ For n  in n_volumes:
     3. Save volume and mask  to files
 ```
 
+In the current v3 dataset format, this save step writes each logical sample as one or more physical shard files, matching the requested `dc_num_shards` layout. The dataloader then reads only the shard file needed by the current DistConv rank instead of loading a full volume and slicing it locally.
+
 ### Performance Profiling
 
 #### 1. Profiling with the PyTorch Profiler

ScaFFold/cli.py

@@ -115,6 +115,11 @@ def main():
     benchmark_parser.add_argument(
         "--base-run-dir", type=str, help="Subfolder of $(pwd) in which to run jobs."
     )
+    benchmark_parser.add_argument(
+        "--dataset-dir",
+        type=str,
+        help="Directory in which to store and query benchmark datasets.",
+    )
     benchmark_parser.add_argument(
         "--fract-base-dir",
         type=str,
@@ -161,6 +166,23 @@ def main():
         type=int,
         help="Number of DataLoader worker processes per rank.",
     )
+    benchmark_parser.add_argument(
+        "--dataloader-pin-memory",
+        type=int,
+        choices=[0, 1],
+        help="Whether DataLoader should pin host memory before device transfer.",
+    )
+    benchmark_parser.add_argument(
+        "--dataloader-persistent-workers",
+        type=int,
+        choices=[0, 1],
+        help="Whether DataLoader workers should persist across epochs.",
+    )
+    benchmark_parser.add_argument(
+        "--dataloader-prefetch-factor",
+        type=int,
+        help="Number of batches prefetched by each DataLoader worker.",
+    )
     benchmark_parser.add_argument(
         "--optimizer",
         type=str,

ScaFFold/datagen/category_search.py

@@ -198,8 +198,10 @@ def main(config: Config) -> None:
     log = setup_mpi_logger(__file__, getattr(config, "verbose", 0))
 
     datagen_batch_size = int(getattr(config, "datagen_batch_size", 10000))
-    if datagen_batch_size <= 0:
-        raise ValueError("datagen_batch_size must be positive")
+    if datagen_batch_size < 1:
+        raise ValueError(
+            f"datagen_batch_size must be positive, got {datagen_batch_size}"
+        )
 
     # FIXME anything else to ensure determinism?
     np.random.seed(config.seed + rank)

ScaFFold/datagen/get_dataset.py

@@ -21,7 +21,7 @@
 import time
 from argparse import Namespace
 from pathlib import Path
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 import yaml
 from mpi4py import MPI
@@ -30,7 +30,7 @@
 from ScaFFold.utils.utils import setup_mpi_logger
 
 META_FILENAME = "meta.yaml"
-DATASET_FORMAT_VERSION = 2
+DATASET_FORMAT_VERSION = 3
 INCLUDE_KEYS = [
     "dataset_format_version",
     "n_categories",
@@ -40,6 +40,8 @@
     "variance_threshold",
     "n_fracts_per_vol",
     "val_split",
+    "dc_num_shards",
+    "dc_shard_dims",
 ]
 
 
@@ -72,11 +74,45 @@ def _get_required_keys_dict(
     return canonicalize(required)
 
 
+def _canonicalize_v3_shard_layout(volume_config: Dict[str, Any]) -> Dict[str, Any]:
+    """Normalize shard layout ordering so equivalent v3 layouts share cache IDs."""
+
+    canonical_config = volume_config.copy()
+    num_shards = canonical_config["dc_num_shards"]
+    shard_dims = canonical_config["dc_shard_dims"]
+    if len(num_shards) != len(shard_dims):
+        raise ValueError(
+            f"dc_num_shards {num_shards} must have same length as dc_shard_dims {shard_dims}"
+        )
+
+    shard_layout = sorted(
+        (int(shard_dim), int(num_shard))
+        for num_shard, shard_dim in zip(num_shards, shard_dims)
+    )
+    canonical_config["dc_shard_dims"] = [shard_dim for shard_dim, _ in shard_layout]
+    canonical_config["dc_num_shards"] = [num_shard for _, num_shard in shard_layout]
+    return canonical_config
+
+
 def _hash_volume_config(volume_config: Dict[str, Any]) -> str:
     s = json.dumps(volume_config, separators=(",", ":"), sort_keys=True).encode()
     return hashlib.sha256(s).hexdigest()[:12]
 
 
+def _volume_config(config_dict, canonicalize_shard_layout=True):
+    """Build the hashable v3 dataset config subset."""
+
+    versioned_config = config_dict.copy()
+    versioned_config["dataset_format_version"] = DATASET_FORMAT_VERSION
+    volume_config = _get_required_keys_dict(
+        config=versioned_config,
+        include_keys=INCLUDE_KEYS,
+    )
+    if canonicalize_shard_layout:
+        volume_config = _canonicalize_v3_shard_layout(volume_config)
+    return volume_config
+
+
 def _git_commit_short(log) -> str:
     try:
         return (
@@ -101,6 +137,37 @@ def _git_commit_short(log) -> str:
         return "no-commit-id"
 
 
+def _find_reusable_dataset(
+    root: Path,
+    config_id: str,
+    commit: str,
+    require_commit: bool,
+) -> Optional[Path]:
+    """Find the newest reusable dataset matching format, config, and commit."""
+
+    base = root / config_id
+    if not base.exists():
+        return None
+
+    candidates = sorted(
+        (p for p in base.iterdir() if p.is_dir()), key=lambda p: p.name, reverse=True
+    )
+    for dataset_path in candidates:
+        meta_path = dataset_path / META_FILENAME
+        if not meta_path.exists():
+            continue
+        meta = yaml.safe_load(meta_path.read_text()) or {}
+        if meta.get("config_id") != config_id:
+            continue
+        if meta.get("dataset_format_version") != DATASET_FORMAT_VERSION:
+            continue
+        if require_commit and meta.get("code_commit") != commit:
+            continue
+        return dataset_path
+
+    return None
+
+
 def get_dataset(
     config: Namespace,
     require_commit: bool = False,  # default: ignore commit mismatches for reuse
@@ -122,40 +189,33 @@ def get_dataset(
     root = Path(config.dataset_dir)
     root.mkdir(exist_ok=True)
 
-    # Get dict of required keys and compute config_id
+    # V3 is the current physical-shard format. The physical dataset layout is
+    # defined by dc_num_shards/dc_shard_dims, matching the DistConv layout.
     config_dict = vars(config).copy()
-    config_dict["dataset_format_version"] = DATASET_FORMAT_VERSION
-    volume_config = _get_required_keys_dict(
-        config=config_dict, include_keys=INCLUDE_KEYS
+    volume_config = _volume_config(config_dict)
+    metadata_volume_config = _volume_config(
+        config_dict,
+        canonicalize_shard_layout=False,
     )
     config_id = _hash_volume_config(volume_config)
     commit = _git_commit_short(log)
 
-    base = root / config_id
-    base.mkdir(parents=True, exist_ok=True)
-
-    # Try to reuse latest candidate dataset
-    candidates = sorted(
-        (p for p in base.iterdir() if p.is_dir()), key=lambda p: p.name, reverse=True
+    # Reuse only matching v3 physical-shard datasets.
+    dataset_path = _find_reusable_dataset(
+        root,
+        config_id,
+        commit,
+        require_commit,
     )
-    for dataset_path in candidates:
-        meta_path = dataset_path / META_FILENAME
-        if not meta_path.exists():
-            continue
-        meta = yaml.safe_load(meta_path.read_text())
-        if meta.get("config_id") != config_id:
-            continue
-        if meta.get("dataset_format_version", 1) != DATASET_FORMAT_VERSION:
-            continue
-        if require_commit and meta.get("code_commit") != commit:
-            continue
-        # If we pass the above checks, this dataset can be reused
-        log.info("Reusing existing dataset at %s", dataset_path)
+    if dataset_path is not None:
+        log.info("Reusing existing v3 dataset at %s", dataset_path)
         return dataset_path
 
     # Otherwise, generate a new dataset
+    base = root / config_id
     log.info("No valid existing dataset found at %s. Generating new dataset.", base)
     if rank == 0:
+        base.mkdir(parents=True, exist_ok=True)
         ts = time.strftime("%Y%m%d-%H%M%S")
         dest = base / f"{ts}__{commit}"
         tmp = base / f".tmp_{ts}"
@@ -178,34 +238,52 @@ def get_dataset(
 
     # Check that all ranks succeeded in volumegen, then sync
     all_ok = comm.allreduce(1 if ok else 0, op=MPI.MIN) == 1
-    comm.Barrier()
+    errs = comm.gather(err, root=0)
 
-    errs = comm.gather(err, root=0) if not all_ok else None
-    if not all_ok:
-        if rank == 0:
+    failure_msg = None
+    if rank == 0:
+        if not all_ok:
             try:
                 shutil.rmtree(tmp, ignore_errors=True)
             except Exception:
                 pass
             msgs = "; ".join(e for e in errs if e)
-            raise RuntimeError(f"dataset generation failed: {msgs or 'unknown error'}")
-        raise RuntimeError("dataset generation failed on another rank")
+            failure_msg = f"dataset generation failed: {msgs or 'unknown error'}"
+
+    failure_msg = comm.bcast(failure_msg, root=0)
+    if failure_msg:
+        raise RuntimeError(failure_msg)
 
     # rank 0 has file write + move
+    finalize_err = ""
     if rank == 0:
-        # Write to tmp, then move, so readers never see half-written dataset
-        meta = {
-            "config_id": config_id,
-            "dataset_format_version": DATASET_FORMAT_VERSION,
-            "config_subset": volume_config,
-            "include_keys": INCLUDE_KEYS,
-            "code_commit": commit,
-            "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
-        }
-        (tmp / META_FILENAME).write_text(
-            yaml.safe_dump(meta, sort_keys=True, default_flow_style=False)
-        )
-        tmp.rename(dest)
+        try:
+            # Write to tmp, then move, so readers never see half-written dataset
+            meta = {
+                "config_id": config_id,
+                "dataset_format_version": DATASET_FORMAT_VERSION,
+                "config_subset": metadata_volume_config,
+                "include_keys": INCLUDE_KEYS,
+                "code_commit": commit,
+                "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+            }
+            (tmp / META_FILENAME).write_text(
+                yaml.safe_dump(meta, sort_keys=True, default_flow_style=False)
+            )
+            tmp.rename(dest)
+        except Exception as e:
+            finalize_err = (
+                f"dataset finalization failed: rank 0: {type(e).__name__}: {e}"
+            )
+
+    finalize_err = comm.bcast(finalize_err, root=0)
+    if finalize_err:
+        if rank == 0:
+            try:
+                shutil.rmtree(tmp, ignore_errors=True)
+            except Exception:
+                pass
+        raise RuntimeError(finalize_err)
 
     # ensure the rename is visible everywhere before returning
     comm.Barrier()