diff --git a/README.md b/README.md index c66016b..7cec997 100644 --- a/README.md +++ b/README.md @@ -63,6 +63,12 @@ The model is trained from a random initialization until convergence, which is de 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. +### Dataset cache and sharded datagen + +`benchmark` creates or reuses datasets under `dataset_dir`. New datasets use the current physical-shard dataset format, which stores one volume and mask file per logical sample per DistConv 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 with `_shard000000` files. Changing the DistConv shard layout changes the dataset cache key, so a run either reuses a dataset with matching sharding metadata or generates a new one. Older full-volume caches are ignored. + Each `benchmark` invocation performs exactly one benchmark run, in a run folder created under `base_run_dir` set in the config file. Every run parameter must be single-valued; a list (e.g. `problem_scale: [6, 7]`) is rejected by name, since parameter sweeps are not supported. To compare parameter settings, launch one benchmark run per setting. For reproducibility, the run folder holds a copy of the benchmark config yml as `base_config.yaml` plus the fully merged `config.yaml` for that run. After the run completes, statistics from the run are stored in `train_stats.csv`. Additionally, users can inspect plots of the training and validation losses over time in ` 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}" + ) base_seed = int(config.seed) diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index 168f002..c8f66b1 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -48,6 +48,8 @@ # writer without stat-ing every volume file. _STALE_PROBE_MAX_DEPTH = 3 _STALE_PROBE_MAX_DIRS = 10000 +DEFAULT_DC_NUM_SHARDS = [1, 1, 1] +DEFAULT_DC_SHARD_DIMS = [2, 3, 4] # Bumped from 2 to 3 when instance point clouds moved from float64 to float32: # the storage layout is unchanged, but float32 voxel binning shifts a handful of # boundary voxels, so a float64-era dataset must not be reused as if it were @@ -55,11 +57,10 @@ # (grid_size - 1 - span)/2 to (grid_size - span)/2: every volume and mask # generated before that was misregistered by half a voxel (with the first # half-voxel of each axis clipped onto plane 0), so those datasets must be -# regenerated rather than reused. This version stamps new datasets, gates reuse -# below, and feeds the config_id hash, so an older dataset is neither matched nor -# scanned. The loader in data_loading.py keeps its own (lower) minimum-layout -# version and still reads v4 through the modern dense path. -DATASET_FORMAT_VERSION = 4 +# regenerated rather than reused. Bumped from 4 to 5 when generated volumes +# became physical shard files named by DistConv shard id, so dense v4 datasets +# are never mistaken for shard-suffixed datasets. +DATASET_FORMAT_VERSION = 5 INCLUDE_KEYS = [ "dataset_format_version", "n_categories", @@ -69,6 +70,8 @@ "variance_threshold", "n_fracts_per_vol", "val_split", + "dc_num_shards", + "dc_shard_dims", ] @@ -84,6 +87,16 @@ def canonicalize(input): return input +def _with_dataset_defaults(config: Dict[str, Any]) -> Dict[str, Any]: + """Return a config dict with dataset-version and default shard keys set.""" + + config = config.copy() + config["dataset_format_version"] = DATASET_FORMAT_VERSION + config.setdefault("dc_num_shards", list(DEFAULT_DC_NUM_SHARDS)) + config.setdefault("dc_shard_dims", list(DEFAULT_DC_SHARD_DIMS)) + return config + + def _get_required_keys_dict( config: Dict[str, Any], include_keys: list[str] ) -> Dict[str, Any]: @@ -91,6 +104,7 @@ def _get_required_keys_dict( Build a dict containing only the required keys. Raises KeyError if any required key is missing. """ + config = _with_dataset_defaults(config) missing = [key for key in include_keys if key not in config] if missing: raise KeyError( @@ -101,6 +115,41 @@ def _get_required_keys_dict( return canonicalize(required) +def _canonicalize_shard_layout(volume_config: Dict[str, Any]) -> Dict[str, Any]: + """Normalize shard layout ordering so equivalent 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 " + f"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 _volume_config( + config_dict: Dict[str, Any], canonicalize_shard_layout: bool = True +) -> Dict[str, Any]: + """Build the hashable dataset config subset.""" + + volume_config = _get_required_keys_dict( + config=config_dict, + include_keys=INCLUDE_KEYS, + ) + if canonicalize_shard_layout: + volume_config = _canonicalize_shard_layout(volume_config) + return volume_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] @@ -221,17 +270,6 @@ def _cleanup_stale_staging_dirs( anything outside ``base`` are never touched. Failures are logged and ignored: cleanup is opportunistic and must never break the decision it runs inside. - - "No sign of life" is the delicate part, because the staging directory of a - *running* job is visible here (unique staging names mean two live jobs never - share a directory, but they do share the base). Age alone is not enough: - generation is not bounded by a day -- at the larger scales it is measured in - days -- and after the first minutes it writes only at depth >= 2, so the top - of the tree stops changing while the job is perfectly healthy. Judging by - the top-level mtimes alone therefore let a concurrent same-config start - rmtree a live generation out from under its peers. ``_staging_dir_is_live`` - is the answer: an explicit heartbeat maintained by the writers, backed by a - bounded-depth mtime probe for directories that predate it. """ now = time.time() cutoff = now - max_age @@ -296,13 +334,6 @@ def _decide_reuse_or_generate( staging and final paths for a new generation. Making this decision in one place and broadcasting it prevents ranks from diverging when their views of the shared filesystem differ. - - The scan is deliberately forgiving: a candidate whose metadata is missing, - unreadable, or malformed is warned about and skipped rather than allowed to - raise. This function runs inside a window where every peer is already - waiting in the decision broadcast, so a crash here is a job-wide hang; a - poison directory (exactly what a killed job leaves behind) must never be - able to cause one. """ # Rank 0 is the only rank that touches this base, so this is also the one # safe place to reclaim staging dirs orphaned by earlier killed jobs. @@ -375,6 +406,15 @@ def _reraisable(exc: BaseException) -> BaseException | None: return exc if isinstance(exc, KeyboardInterrupt) else None +def _ensure_config_has_sharding(config: Namespace) -> None: + """Populate default sharding fields on minimal Namespace configs.""" + + if not hasattr(config, "dc_num_shards"): + config.dc_num_shards = list(DEFAULT_DC_NUM_SHARDS) + if not hasattr(config, "dc_shard_dims"): + config.dc_shard_dims = list(DEFAULT_DC_SHARD_DIMS) + + def get_dataset( config: Namespace, require_commit: bool = False, # default: ignore commit mismatches for reuse @@ -392,15 +432,20 @@ def get_dataset( comm = MPI.COMM_WORLD rank = comm.Get_rank() log = setup_mpi_logger(__file__, getattr(config, "verbose", 0)) + _ensure_config_has_sharding(config) root = Path(config.dataset_dir) - root.mkdir(exist_ok=True) + root.mkdir(parents=True, exist_ok=True) - # Get dict of required keys and compute config_id + # The physical dataset layout is defined by dc_num_shards/dc_shard_dims, + # matching the DistConv layout. The hash canonicalizes shard ordering so + # equivalent dimension/order spellings reuse the same cache, while metadata + # preserves the requested order used for shard ids and filenames. 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) @@ -413,10 +458,6 @@ def get_dataset( # same branch. Scanning the shared filesystem independently per rank lets # divergent views (stale metadata caches, a racing job's rename) strand some # ranks in the generation collectives while others return early. - # Everything rank 0 does here happens while the peers are already blocked in - # the broadcast below, so a rank-0 exception would strand the whole job. - # Any failure is therefore turned into an error sentinel that travels - # through the same broadcast and makes every rank raise the same error. interrupt = None if rank == 0: try: @@ -424,10 +465,6 @@ def get_dataset( base, config_id, commit, require_commit, log ) except BaseException as e: - # BaseException, not (Exception, SystemExit): a KeyboardInterrupt - # delivered to rank 0 alone (Ctrl-C on the launching terminal, a - # site watchdog SIGINT) would otherwise skip the broadcast and hang - # every peer -- the exact failure this guard exists to prevent. decision = ( "error", f"rank 0 failed to select a dataset under {base}: " @@ -439,8 +476,6 @@ def get_dataset( decision = comm.bcast(decision, root=0) if decision[0] == "error": - # Rank 0 keeps the abort signal it was actually given; the peers, which - # only ever saw the sentinel, report it as a generation failure. if interrupt is not None: raise interrupt raise RuntimeError(f"dataset selection failed: {decision[1]}") @@ -493,7 +528,7 @@ def get_dataset( meta = { "config_id": config_id, "dataset_format_version": DATASET_FORMAT_VERSION, - "config_subset": volume_config, + "config_subset": metadata_volume_config, "include_keys": INCLUDE_KEYS, "code_commit": commit, "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py index 997a73f..811c7be 100644 --- a/ScaFFold/datagen/volumegen.py +++ b/ScaFFold/datagen/volumegen.py @@ -17,7 +17,7 @@ import random import time from math import ceil -from typing import Dict +from typing import Callable, Dict import numpy as np from mpi4py import MPI @@ -25,8 +25,18 @@ from ScaFFold.datagen import layout from ScaFFold.utils.config_utils import Config from ScaFFold.utils.data_types import MASK_DTYPE, VOLUME_DTYPE +from ScaFFold.utils.spatial_sharding import ( + normalize_sharding, + shard_file_suffix, + shard_id_to_indices, + spatial_slices, + total_shards, +) from ScaFFold.utils.utils import setup_mpi_logger +DEFAULT_DC_NUM_SHARDS = (1, 1, 1) +DEFAULT_DC_SHARD_DIMS = (2, 3, 4) + # Liveness marker for the directory being generated into. Volume writing is the # long phase of a generation and it happens deep inside the tree # (``volumes//N.npy``), so the top of the staging directory can look @@ -185,6 +195,105 @@ def resolve_grid_size(config) -> int: return int(config.vol_size) +def _sharding_values(config): + num_shards = getattr(config, "dc_num_shards", DEFAULT_DC_NUM_SHARDS) + shard_dims = getattr(config, "dc_shard_dims", DEFAULT_DC_SHARD_DIMS) + return num_shards, shard_dims + + +def _physical_sharding(config): + """Return normalized physical sharding from the generation config.""" + + num_shards, shard_dims = _sharding_values(config) + return normalize_sharding(num_shards, shard_dims) + + +def _validate_generation_config(config): + """Validate sharded generation settings and return normalized layout data.""" + + num_shards, shard_dims = _physical_sharding(config) + n_total_shards = total_shards(num_shards) + grid_size = resolve_grid_size(config) + + return num_shards, shard_dims, n_total_shards, grid_size + + +def _voxelized_fractals_for_volume( + config, + curr_vol: np.ndarray, + fractal_colors: np.ndarray, + instances_dir: str, + grid_size: int, + point_cloud_loader: Callable[[str], np.ndarray] = load_np_ptcloud, +): + """Load and voxelize all fractals needed for one logical volume.""" + + n_fracts_per_vol = config.n_fracts_per_vol + voxelized_fractals = [] + + for curr_fract in range(n_fracts_per_vol): + curr_category = int(curr_vol[1 + 2 * curr_fract]) + curr_instance = int(curr_vol[1 + 2 * curr_fract + 1]) + fractal_color = fractal_colors[curr_category] + + point_cloud_path = os.path.join( + instances_dir, + f"{curr_category:06d}", + f"{curr_category:06d}_{curr_instance:04d}.npy", + ) + if point_cloud_loader is load_np_ptcloud and not os.path.exists( + point_cloud_path + ): + raise FileNotFoundError( + f"File {point_cloud_path} does not exist. " + "Ensure you have run 'scaffold generate_fractals ...'" + ) + + points = point_cloud_loader(point_cloud_path) + idx = points_to_voxel_indices(points, grid_size) + voxelized_fractals.append((curr_category, fractal_color, idx)) + + return voxelized_fractals + + +def _render_volume_shard(config, voxelized_fractals, shard_id: int): + """Render one physical shard from precomputed global voxel indices.""" + num_shards, shard_dims = _physical_sharding(config) + shard_indices = shard_id_to_indices(shard_id, num_shards) + slices = spatial_slices( + (config.vol_size, config.vol_size, config.vol_size), + shard_dims, + num_shards, + shard_indices, + ) + local_shape = tuple(s.stop - s.start for s in slices) + + volume = np.full((3, *local_shape), 0, dtype=VOLUME_DTYPE) + mask = np.full(local_shape, 0, dtype=MASK_DTYPE) + + for curr_category, fractal_color, idx in voxelized_fractals: + keep = np.ones(idx.shape[0], dtype=bool) + for axis, axis_slice in enumerate(slices): + keep &= idx[:, axis] >= axis_slice.start + keep &= idx[:, axis] < axis_slice.stop + + if not np.any(keep): + continue + + local_idx = idx[keep] + local_idx[:, 0] -= slices[0].start + local_idx[:, 1] -= slices[1].start + local_idx[:, 2] -= slices[2].start + d = local_idx[:, 0] + h = local_idx[:, 1] + w = local_idx[:, 2] + + volume[:, d, h, w] = fractal_color[:, None] + mask[d, h, w] = curr_category + 1 + + return volume, mask + + def main(config: Dict): # Initialize MPI comm = MPI.COMM_WORLD @@ -199,6 +308,7 @@ def main(config: Dict): volumes_contents_path = os.path.join(dataset_dir, "volumes_contents.csv") n_fracts_per_vol = config.n_fracts_per_vol + _, _, n_total_shards, grid_size = _validate_generation_config(config) random.seed(config.seed) # Python np.random.seed(config.seed) # NumPy @@ -260,6 +370,21 @@ def main(config: Dict): # Broadcast to all ranks volumes_contents = comm.bcast(volumes_contents, root=0) + # Each writer defensively creates the output dirs after root's broadcast. + for subdir in ["training", "validation"]: + os.makedirs(os.path.join(vol_path, subdir), exist_ok=True) + os.makedirs(os.path.join(mask_path, subdir), exist_ok=True) + + # Rank 0 creates shared metadata above; wait before local writer setup. + comm.Barrier() + + for subdir in ["training", "validation"]: + os.makedirs(os.path.join(vol_path, subdir), exist_ok=True) + os.makedirs(os.path.join(mask_path, subdir), exist_ok=True) + + # Wait until every rank has ensured the writer directories exist. + comm.Barrier() + # Determine train/val split globally so all ranks know where to save num_volumes = len(volumes_contents) random.seed(config.seed) # Reset seed to ensure all ranks get same split @@ -267,11 +392,11 @@ def main(config: Dict): random.sample(range(num_volumes), int(num_volumes * config.val_split / 100)) ) - # Work distribution - num_volumes = len(volumes_contents) - stride = ceil(num_volumes / size) + # Work distribution: each task renders one physical shard of one logical volume. + total_tasks = num_volumes * n_total_shards + stride = ceil(total_tasks / size) start_idx = rank * stride - end_idx = min(((rank + 1) * stride), num_volumes) + end_idx = min(((rank + 1) * stride), total_tasks) # Per-rank generation status. A rank that fails records the message here and # keeps executing the collective sequence below so every rank stays in step; @@ -282,21 +407,19 @@ def main(config: Dict): try: if start_idx >= end_idx: - log.debug("Rank %s given no volumes to generate", rank) + log.debug("Rank %s given no physical shard tasks to generate", rank) else: - volumes_contents_subset = volumes_contents[start_idx:end_idx] log.debug( - "Rank %s responsible for volumes %s through %s", + "Rank %s responsible for physical shard tasks %s through %s", rank, - volumes_contents_subset[0][0], - volumes_contents_subset[-1][0], + start_idx, + end_idx - 1, ) np.random.seed(config.seed) fractal_colors = np.random.rand(config.n_categories, 3) - grid_size = resolve_grid_size(config) # The instance library is keyed by seed (see ScaFFold.datagen # .layout), so a volume can only ever be built from point clouds # this run's seed produced. Resolved once, outside the loop. @@ -310,89 +433,76 @@ def main(config: Dict): heartbeat = StagingHeartbeat(dataset_dir) heartbeat.beat() start_time = time.time() - for i, curr_vol in enumerate(volumes_contents_subset): + n_generated_shards = 0 + cached_volume_idx = None + cached_global_vol_idx = None + cached_voxelized_fractals = None + + for i, task_idx in enumerate(range(start_idx, end_idx)): heartbeat.beat() if i % 10 == 0: - log.debug("Rank %s processing local volume %s", rank, i) - - volume = np.full( - (config.vol_size, config.vol_size, config.vol_size, 3), - 0, - dtype=VOLUME_DTYPE, - ) - mask = np.full( - (config.vol_size, config.vol_size, config.vol_size), - 0, - dtype=MASK_DTYPE, - ) - - global_vol_idx = curr_vol[0] - vol_seed = config.seed + int(global_vol_idx) - random.seed(vol_seed) - np.random.seed(vol_seed) - - for curr_fract in range(n_fracts_per_vol): - curr_category = curr_vol[1 + 2 * curr_fract] - curr_instance = curr_vol[1 + 2 * curr_fract + 1] - fractal_color = fractal_colors[curr_category] - - point_cloud_path = os.path.join( - instances_dir, - f"{curr_category:06d}", - f"{curr_category:06d}_{curr_instance:04d}.npy", + log.debug( + "Rank %s processing local physical shard task %s", rank, i ) - if not os.path.exists(point_cloud_path): - raise FileNotFoundError( - f"File {point_cloud_path} does not exist. " - "Ensure you have run 'scaffold generate_fractals ...'" - ) + volume_idx = task_idx // n_total_shards + shard_id = task_idx % n_total_shards - points = load_np_ptcloud(point_cloud_path) - voxel_idx = points_to_voxel_indices(points, grid_size) + if cached_volume_idx != volume_idx: + curr_vol = volumes_contents[volume_idx] + global_vol_idx = int(curr_vol[0]) + vol_seed = config.seed + global_vol_idx + random.seed(vol_seed) + np.random.seed(vol_seed) - assert voxel_idx.shape[1] == volume.ndim - 1, ( - f"voxel index width {voxel_idx.shape[1]} != volume spatial " - f"dims {volume.ndim - 1}" + cached_voxelized_fractals = _voxelized_fractals_for_volume( + config, + curr_vol, + fractal_colors, + instances_dir, + grid_size, ) + cached_volume_idx = volume_idx + cached_global_vol_idx = global_vol_idx - # Scatter only the occupied voxels: O(points) writes instead - # of two full-volume boolean-mask traversals per fractal. - rows, cols, depths = ( - voxel_idx[:, 0], - voxel_idx[:, 1], - voxel_idx[:, 2], - ) - volume[rows, cols, depths] = fractal_color - mask[rows, cols, depths] = curr_category + 1 + volume_to_save, mask_to_save = _render_volume_shard( + config, + cached_voxelized_fractals, + shard_id, + ) # Determine destination folder - subdir = "validation" if global_vol_idx in val_indices else "training" - # Tensors must logically be channels-first, later we will change striding/storage to channels-last on GPU (metadata will always stay channels-first). - volume_channels_first = volume.transpose((3, 0, 1, 2)) - volume_to_save = np.ascontiguousarray( - volume_channels_first, dtype=VOLUME_DTYPE + subdir = ( + "validation" if cached_global_vol_idx in val_indices else "training" ) - mask_to_save = np.ascontiguousarray(mask, dtype=MASK_DTYPE) + shard_suffix = shard_file_suffix(shard_id) - vol_file = os.path.join(vol_path, subdir, f"{global_vol_idx}.npy") + vol_file = os.path.join( + vol_path, subdir, f"{cached_global_vol_idx}{shard_suffix}.npy" + ) with open(vol_file, "wb") as f: np.save(f, volume_to_save) mask_file = os.path.join( - mask_path, subdir, f"{global_vol_idx}_mask.npy" + mask_path, + subdir, + f"{cached_global_vol_idx}{shard_suffix}_mask.npy", ) with open(mask_file, "wb") as f: np.save(f, mask_to_save) + n_generated_shards += 1 end_time = time.time() total_time = end_time - start_time if rank == 0: + shard_rate = n_generated_shards / total_time log.info( - "Rank 0 generated %s volumes in %.2f seconds | %.2f volumes per second", - len(volumes_contents_subset), + "Rank 0 generated %s volume shards from %s physical shard " + "tasks in %.2f seconds | %.2f shards per second", + n_generated_shards, + end_idx - start_idx, total_time, - len(volumes_contents_subset) / total_time, + shard_rate, ) except BaseException as e: # Capture the failure locally instead of letting it unwind past the diff --git a/ScaFFold/utils/data_loading.py b/ScaFFold/utils/data_loading.py index 8a75b1d..55fb6df 100644 --- a/ScaFFold/utils/data_loading.py +++ b/ScaFFold/utils/data_loading.py @@ -14,6 +14,7 @@ import hashlib import pickle +import re from dataclasses import dataclass from os import listdir from os.path import isfile, join, splitext @@ -27,9 +28,18 @@ from torch.utils.data import Dataset from ScaFFold.utils.data_types import MASK_DTYPE, VOLUME_DTYPE +from ScaFFold.utils.spatial_sharding import ( + chunk_slice, + normalize_sharding, + shard_file_suffix, + shard_indices_to_id, + total_shards, +) from ScaFFold.utils.utils import customlog DATASET_FORMAT_VERSION = 2 +PHYSICAL_SHARDED_DATASET_FORMAT_VERSION = 5 +MAX_SUPPORTED_DATASET_FORMAT_VERSION = PHYSICAL_SHARDED_DATASET_FORMAT_VERSION LEGACY_DATASET_FORMAT_VERSION = 1 META_FILENAME = "meta.yaml" @@ -67,19 +77,6 @@ def __post_init__(self): f"Invalid shard_index {shard_index} for shard_dim {shard_dim} with {num_shards} shards" ) - @staticmethod - def _chunk_slice(size: int, num_shards: int, shard_index: int) -> slice: - """Match torch.chunk-style uneven shard boundaries.""" - - chunk_size = (size + num_shards - 1) // num_shards - start = shard_index * chunk_size - if start >= size: - raise ValueError( - f"Empty local shard: dim size {size}, num_shards {num_shards}, shard_index {shard_index}" - ) - stop = min(size, start + chunk_size) - return slice(start, stop) - def slice_array( self, array: np.ndarray, axis_map: Dict[int, int], array_label: str ) -> np.ndarray: @@ -99,7 +96,7 @@ def slice_array( raise ValueError( f"Axis {axis} out of range for {array_label} with shape {array.shape}" ) - slices[axis] = self._chunk_slice(array.shape[axis], num_shards, shard_index) + slices[axis] = chunk_slice(array.shape[axis], num_shards, shard_index) return array[tuple(slices)] @@ -118,7 +115,30 @@ def __init__( self.mask_suffix = mask_suffix self.spatial_shard_spec = spatial_shard_spec self.dataset_root = self.images_dir.parents[1] - self.dataset_format_version = self._load_dataset_format_version() + self.dataset_meta = self._load_dataset_metadata() + self.dataset_format_version = int( + self.dataset_meta.get( + "dataset_format_version", LEGACY_DATASET_FORMAT_VERSION + ) + ) + if self.dataset_format_version > MAX_SUPPORTED_DATASET_FORMAT_VERSION: + raise RuntimeError( + f"Unsupported dataset format version {self.dataset_format_version}; " + f"expected <= {MAX_SUPPORTED_DATASET_FORMAT_VERSION}" + ) + self.physical_shards = ( + self.dataset_format_version >= PHYSICAL_SHARDED_DATASET_FORMAT_VERSION + ) + self.physical_num_shards, self.physical_shard_dims = ( + self._load_physical_sharding() + ) + self.physical_total_shards = ( + total_shards(self.physical_num_shards) if self.physical_shards else 1 + ) + self.shard_id = self._select_physical_shard_id() + self.shard_suffix = ( + shard_file_suffix(self.shard_id) if self.physical_shards else "" + ) # os.listdir order is filesystem/client dependent and explicitly # arbitrary. Sorting makes the index -> file mapping deterministic and @@ -131,7 +151,9 @@ def __init__( for file in listdir(images_dir) if isfile(join(images_dir, file)) and not file.startswith(".") ] - self.ids = sorted(splitext(file)[0] for file in image_files) + self.ids, self._image_paths = self._index_paths_by_id( + self.images_dir, image_files, "image" + ) if not self.ids: raise RuntimeError( f"No input file found in {images_dir}, make sure you put your images there" @@ -141,15 +163,12 @@ def __init__( # fetches then index these maps in O(1) instead of scanning (and # fnmatching) the whole directory on every call, which on a shared # filesystem turns each fetch into a burst of metadata traffic. - self._image_paths = self._index_paths_by_stem( - self.images_dir, image_files, "image" - ) mask_files = [ entry.name for entry in self.mask_dir.iterdir() if entry.is_file() and not entry.name.startswith(".") ] - self._mask_paths = self._index_paths_by_stem(self.mask_dir, mask_files, "mask") + _, self._mask_paths = self._index_paths_by_id(self.mask_dir, mask_files, "mask") # Belt-and-braces: when a process group is live, verify every rank built # the identical id list. Any residual divergence (e.g. inconsistent @@ -163,6 +182,12 @@ def __init__( self.mask_values = self._load_mask_values(data_dir) customlog(f"Unique mask values: {self.mask_values}") customlog(f"Dataset format version: {self.dataset_format_version}") + if self.physical_shards: + customlog( + f"Loading physical shard files with suffix {self.shard_suffix}; " + f"dc_num_shards={self.physical_num_shards}, " + f"dc_shard_dims={self.physical_shard_dims}" + ) # Masks are handed off in a signed 16-bit carrier (widened to long on # the compute device), so the largest class id the carrier will hold @@ -179,7 +204,7 @@ def _max_class_id(self): """Return the largest class id ``_to_mask_carrier`` will have to carry. The bound differs by format, and using the wrong one is unsafe in one - direction and needlessly strict in the other. v2+ masks ship *raw* + direction and needlessly strict in the other. Modern masks ship *raw* ``category + 1`` ids, and the per-split table lists only the categories present in that split -- so a sparse split can declare two classes while holding an id in the tens of thousands, which the class *count* check @@ -198,7 +223,7 @@ def _max_class_id(self): def _load_mask_values(self, data_dir): """Return the label-remap table for this split. - v2 datasets store dense class ids and never remap, so the per-split + Modern datasets store dense class ids and never remap, so the per-split pickle is loaded verbatim for bookkeeping. Legacy (v1) datasets remap raw voxel values by their position in this list; a per-split table would assign the same raw value different class ids across splits whenever a @@ -271,32 +296,54 @@ def _verify_ids_consistent_across_ranks(self, images_dir): def __len__(self): return len(self.ids) - @staticmethod - def _index_paths_by_stem(directory, filenames, label): - """Map each file's extension-less name to its full path. + def _id_from_filename(self, filename, label): + if not self.physical_shards: + stem = splitext(filename)[0] + if label == "mask": + if not stem.endswith(self.mask_suffix): + return None + return stem + return stem + + suffix = re.escape(self.shard_suffix) + if label == "mask": + pattern = re.compile( + rf"^(?P.+){suffix}{re.escape(self.mask_suffix)}\.npy$" + ) + else: + pattern = re.compile(rf"^(?P.+){suffix}\.npy$") + match = pattern.match(filename) + if match is None: + return None + return match.group("id") + + def _index_paths_by_id(self, directory, filenames, label): + """Map each logical sample id to its full path. - Raises if two files in ``directory`` share a stem, which would make the - stem an ambiguous key and silently pick one of them at fetch time. + Raises if two files in ``directory`` share an id, which would make the + id an ambiguous key and silently pick one of them at fetch time. """ paths = {} for filename in filenames: - stem = splitext(filename)[0] + sample_id = self._id_from_filename(filename, label) + if sample_id is None: + continue full_path = directory / filename - existing = paths.get(stem) + existing = paths.get(sample_id) if existing is not None: raise RuntimeError( - f"Ambiguous {label} id '{stem}' in {directory}: matches both " - f"{existing.name} and {filename}. Every id must map to exactly " - "one file." + f"Ambiguous {label} id '{sample_id}' in {directory}: matches " + f"both {existing.name} and {filename}. Every id must map to " + "exactly one file." ) - paths[stem] = full_path - return paths + paths[sample_id] = full_path + return sorted(paths), paths @staticmethod def _load_numpy_array(path, mmap_mode=None): return np.load(path, allow_pickle=False, mmap_mode=mmap_mode) - def _load_dataset_format_version(self): + def _load_dataset_metadata(self): """Determine which on-disk layout this dataset uses. Only a *missing* ``meta.yaml`` means legacy v1: those datasets predate @@ -309,7 +356,7 @@ def _load_dataset_format_version(self): """ meta_path = self.dataset_root / META_FILENAME if not meta_path.exists(): - return LEGACY_DATASET_FORMAT_VERSION + return {"dataset_format_version": LEGACY_DATASET_FORMAT_VERSION} try: with open(meta_path, "r") as meta_file: @@ -324,7 +371,7 @@ def _load_dataset_format_version(self): version = meta.get("dataset_format_version") if isinstance(meta, dict) else None try: - return int(version) + int(version) except (TypeError, ValueError): raise ValueError( f"Dataset metadata {meta_path} is missing a usable " @@ -333,6 +380,83 @@ def _load_dataset_format_version(self): "its layout. Repair the file or regenerate the dataset." ) from None + return meta + + def _load_physical_sharding(self): + """Load and normalize the physical shard layout from metadata.""" + + if not self.physical_shards: + return (), () + + config_subset = self.dataset_meta.get("config_subset") or {} + num_shards = config_subset.get("dc_num_shards") + shard_dims = config_subset.get("dc_shard_dims") + if num_shards is None or shard_dims is None: + raise RuntimeError( + "Physical dataset is missing shard metadata. Expected " + "config_subset.dc_num_shards/config_subset.dc_shard_dims in meta.yaml." + ) + + return normalize_sharding(num_shards, shard_dims) + + @staticmethod + def _layout_by_dim(num_shards, shard_dims): + """Map each sharded dimension to its shard count.""" + + return {int(dim): int(num) for num, dim in zip(num_shards, shard_dims)} + + def _physical_layout_matches_spatial_spec(self): + """Return whether dataset shards match the requested spatial layout.""" + + if self.spatial_shard_spec is None: + return False + return self._layout_by_dim( + self.physical_num_shards, self.physical_shard_dims + ) == self._layout_by_dim( + self.spatial_shard_spec.num_shards, + self.spatial_shard_spec.shard_dims, + ) + + def _physical_shard_id_for_spatial_spec(self): + """Return the physical shard id selected by the spatial shard spec.""" + + spec_indices_by_dim = { + int(dim): int(index) + for dim, index in zip( + self.spatial_shard_spec.shard_dims, + self.spatial_shard_spec.shard_indices, + ) + } + shard_indices = tuple( + spec_indices_by_dim[int(dim)] for dim in self.physical_shard_dims + ) + return shard_indices_to_id(shard_indices, self.physical_num_shards) + + def _select_physical_shard_id(self): + """Select the physical shard file this dataset instance should read.""" + + if not self.physical_shards: + return 0 + if self.spatial_shard_spec is None: + if self.physical_total_shards == 1: + return 0 + raise RuntimeError( + "Physical dataset has multiple shard files, but no SpatialShardSpec " + "was provided. Use a DistConv layout matching the dataset." + ) + if not self._physical_layout_matches_spatial_spec(): + raise RuntimeError( + "Physical dataset shard layout does not match the requested " + "DistConv layout. Physical dataset layout and DistConv layout " + "must match. " + f"dataset dc_num_shards={self.physical_num_shards}, " + f"dataset dc_shard_dims={self.physical_shard_dims}, " + f"dc_num_shards={self.spatial_shard_spec.num_shards}, " + f"dc_shard_dims={self.spatial_shard_spec.shard_dims}" + ) + + return self._physical_shard_id_for_spatial_spec() + @staticmethod def _prepare_legacy_image(img): return np.ascontiguousarray(img.transpose((3, 0, 1, 2)), dtype=VOLUME_DTYPE) @@ -368,7 +492,7 @@ def _prepare_optimized_mask(mask, materialize): return np.asarray(mask, dtype=MASK_DTYPE, order="C") def _slice_image_array(self, img): - if self.spatial_shard_spec is None: + if self.spatial_shard_spec is None or self.physical_shards: return img if self.dataset_format_version >= DATASET_FORMAT_VERSION: @@ -378,7 +502,7 @@ def _slice_image_array(self, img): return self.spatial_shard_spec.slice_array(img, axis_map, "image") def _slice_mask_array(self, mask): - if self.spatial_shard_spec is None: + if self.spatial_shard_spec is None or self.physical_shards: return mask axis_map = {2: 0, 3: 1, 4: 2} @@ -393,15 +517,18 @@ def _image_path(self, name): ) def _mask_path(self, name): - key = name + self.mask_suffix + key = name if self.physical_shards else name + self.mask_suffix try: return self._mask_paths[key] except KeyError: raise KeyError(f"No mask file found for the ID {key} in {self.mask_dir}") + def _materialize_shard_slice(self): + return self.spatial_shard_spec is not None and not self.physical_shards + def _load_prepared_mask(self, name): """Load, shard-slice, and label-prepare one mask as ``MASK_DTYPE``.""" - materialize = self.spatial_shard_spec is not None + materialize = self._materialize_shard_slice() # Memmap lets each rank slice out just its local shard without eagerly # reading the full sample into process memory first; the prepare step # then copies that slice out of the backing file. @@ -432,7 +559,7 @@ def load_mask_only(self, idx): def __getitem__(self, idx): name = self.ids[idx] - materialize = self.spatial_shard_spec is not None + materialize = self._materialize_shard_slice() mmap_mode = "r" if materialize else None img = self._load_numpy_array(self._image_path(name), mmap_mode=mmap_mode) img = self._slice_image_array(img) diff --git a/ScaFFold/utils/spatial_sharding.py b/ScaFFold/utils/spatial_sharding.py new file mode 100644 index 0000000..91664dc --- /dev/null +++ b/ScaFFold/utils/spatial_sharding.py @@ -0,0 +1,133 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +from math import prod +from typing import Iterable, Tuple + + +def normalize_sharding(num_shards: Iterable[int], shard_dims: Iterable[int]): + """Validate and normalize spatial sharding config.""" + + num_shards = tuple(int(x) for x in num_shards) + shard_dims = tuple(int(x) for x in shard_dims) + + if len(num_shards) != len(shard_dims): + raise ValueError( + f"num_shards {num_shards} must have same length as shard_dims {shard_dims}" + ) + if len(set(shard_dims)) != len(shard_dims): + raise ValueError(f"Shard dimensions must be unique: {shard_dims}") + + for num_shards_i, shard_dim_i in zip(num_shards, shard_dims): + if num_shards_i < 1: + raise ValueError(f"Invalid num_shards value {num_shards_i}") + if shard_dim_i not in (2, 3, 4): + raise ValueError( + f"Invalid shard_dim {shard_dim_i}: only 3D spatial dimensions 2, 3, and 4 are supported" + ) + + return num_shards, shard_dims + + +def total_shards(num_shards: Iterable[int]) -> int: + """Return the total number of shards in a multi-dimensional layout.""" + + return prod(tuple(int(x) for x in num_shards)) + + +def shard_id_to_indices(shard_id: int, num_shards: Iterable[int]) -> Tuple[int, ...]: + """Convert row-major linear shard id to multi-dimensional shard indices.""" + + num_shards = tuple(int(x) for x in num_shards) + total = total_shards(num_shards) + if shard_id < 0 or shard_id >= total: + raise ValueError( + f"shard_id {shard_id} out of range for num_shards={num_shards}" + ) + + indices = [] + linear_idx = int(shard_id) + stride = total + for num_shards_i in num_shards: + stride //= num_shards_i + indices.append(linear_idx // stride) + linear_idx %= stride + return tuple(indices) + + +def shard_indices_to_id(shard_indices: Iterable[int], num_shards: Iterable[int]) -> int: + """Convert multi-dimensional shard indices to row-major linear shard id.""" + + shard_indices = tuple(int(x) for x in shard_indices) + num_shards = tuple(int(x) for x in num_shards) + if len(shard_indices) != len(num_shards): + raise ValueError( + f"shard_indices {shard_indices} must match num_shards {num_shards}" + ) + + shard_id = 0 + stride = 1 + for shard_index_i, num_shards_i in zip( + reversed(shard_indices), reversed(num_shards) + ): + if shard_index_i < 0 or shard_index_i >= num_shards_i: + raise ValueError( + f"Invalid shard index {shard_index_i} for num_shards={num_shards}" + ) + shard_id += shard_index_i * stride + stride *= num_shards_i + return shard_id + + +def chunk_slice(size: int, num_shards: int, shard_index: int) -> slice: + """Match torch.chunk-style uneven shard boundaries.""" + + chunk_size = (size + num_shards - 1) // num_shards + start = shard_index * chunk_size + if start >= size: + raise ValueError( + f"Empty local shard: dim size {size}, num_shards {num_shards}, shard_index {shard_index}" + ) + stop = min(size, start + chunk_size) + return slice(start, stop) + + +def spatial_slices( + spatial_shape: Iterable[int], + shard_dims: Iterable[int], + num_shards: Iterable[int], + shard_indices: Iterable[int], +) -> Tuple[slice, slice, slice]: + """Return local D/H/W slices for DistConv spatial dims 2/3/4.""" + + spatial_shape = tuple(int(x) for x in spatial_shape) + if len(spatial_shape) != 3: + raise ValueError(f"Expected 3D spatial shape, got {spatial_shape}") + + slices = [slice(0, size) for size in spatial_shape] + for shard_dim, num_shards_i, shard_index_i in zip( + shard_dims, num_shards, shard_indices + ): + spatial_axis = int(shard_dim) - 2 + slices[spatial_axis] = chunk_slice( + spatial_shape[spatial_axis], int(num_shards_i), int(shard_index_i) + ) + + return tuple(slices) + + +def shard_file_suffix(shard_id: int) -> str: + """Return the filename suffix for a physical shard id.""" + + return f"_shard{int(shard_id):06d}" diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 40458ba..7192cdf 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -189,6 +189,13 @@ def main(kwargs_dict: dict = {}): ) log.info(f"rank={rank}, world_size={world_size}") + total_distconv_shards = math.prod(config.dc_num_shards) + if world_size % total_distconv_shards != 0: + raise ValueError( + f"world_size={world_size} must be divisible by total number of " + f"distconv shards={total_distconv_shards}" + ) + # Generate or retrieve dataset begin_code_region("get_dataset") dataset_dir = get_dataset( @@ -212,14 +219,6 @@ def main(kwargs_dict: dict = {}): group_norm_groups=config.group_norm_groups, ) # DDP + DistConv setup - # Ensure world_size is divisible by total distconv shards - total_distconv_shards = math.prod(config.dc_num_shards) - if world_size % total_distconv_shards != 0: - raise ValueError( - f"world_size={world_size} must be divisible by total number of " - f"distconv shards={total_distconv_shards}" - ) - ps = ParallelStrategy( num_shards=config.dc_num_shards, shard_dim=config.dc_shard_dims,