fix(datasets): address review findings on the WebDataset shard path

Publication is now atomic where it claimed to be: shard names carry a
generation token derived from the packed contents, so a re-pack writes
alongside the shards the published index still references, the index is
swapped with os.replace, and the previous generation is removed only after
that commit. Deriving the token from the contents keeps re-packing unchanged
data reproducible rather than churning the directory.

Evaluation now applies the same rank-level shard guard as training: a rank
left without a shard never reaches validation_step while populated ranks do,
which deadlocks the process group rather than failing.

A train/val category_ids mismatch is rejected instead of silently evaluating
remapped labels against raw-trained predictions, pack_coco_to_shards
validates category_ids for library callers, and a split name carrying a glob
metacharacter is refused before it can match another split shards.

Each change has a regression test.

Author: JESUSROYETH

docs/learn/train/dataset-formats.md

@@ -414,14 +414,16 @@ Repeat once per split, into the same `--output-dir`:
 
 ```
 coco-shards/
-├── train-000000.tar        # ~100 MB each, override with --max-shard-mb
-├── train-000001.tar
+├── train-a1b2c3d4-000000.tar   # ~100 MB each, override with --max-shard-mb
+├── train-a1b2c3d4-000001.tar
 ├── ...
 ├── train-index.json        # shard list, sample count, categories
-├── val-000000.tar
+├── val-5e6f7a8b-000000.tar
 └── val-index.json
 ```
 
+The hex segment in a shard name is a generation token derived from the packed contents. It lets a re-pack write its shards alongside the ones the published index still points at, so the index swap is the only moment the pack changes; because it is derived from the contents rather than random, re-packing unchanged data reproduces the same names instead of churning the directory.
+
 Image bytes are copied verbatim — no re-encode — so a packed split decodes to exactly the same pixels as the directory it came from. A `.json` member next to each image carries that image's `image_id`, `file_name` and annotation list, segmentation polygons included.
 
 `--category-ids` chooses the label space, and the index records the choice so the loader never has to guess:

src/rfdetr/datasets/webdataset_io.py

@@ -39,6 +39,7 @@
 from __future__ import annotations
 
 import argparse
+import hashlib
 import io
 import json
 import os
@@ -112,7 +113,9 @@ def _validate_split_name(split: str) -> str:
         *split*, unchanged, once validated.
 
     Raises:
-        ValueError: If *split* is empty, or contains a path separator or a ``.``/``..`` segment.
+        ValueError: If *split* is empty, contains a path separator or a ``.``/``..`` segment, or carries a
+            glob metacharacter. The last one matters because the name is also matched against existing shard
+            files: a split called ``*`` would match — and then delete — the shards of every other split.
 
     Examples:
         >>> _validate_split_name("train")
@@ -124,15 +127,23 @@ def _validate_split_name(split: str) -> str:
     """
     if not split or "/" in split or "\\" in split or split in (".", ".."):
         raise ValueError(f"split {split!r} must not contain a path separator or '..'.")
+    forbidden = sorted(set("*?[]") & set(split))
+    if forbidden:
+        raise ValueError(
+            f"split {split!r} must not contain {''.join(forbidden)}: the name is matched against existing shard "
+            "files, so a glob metacharacter would reach other splits' shards."
+        )
     return split
 
 
-def _shard_name(split: str, index: int) -> str:
+def _shard_name(split: str, index: int, generation: str | None = None) -> str:
     """Return the file name of shard *index* of *split*.
 
     Args:
         split: Split name the shard belongs to.
         index: Zero-based shard number.
+        generation: Token identifying one pack run. Included in the name so that re-packing a split writes new
+            files instead of overwriting the ones the published index still points at.
 
     Returns:
         Shard file name.
@@ -143,9 +154,13 @@ def _shard_name(split: str, index: int) -> str:
     Examples:
         >>> _shard_name("train", 7)
         'train-000007.tar'
+        >>> _shard_name("train", 7, "a1b2c3d4")
+        'train-a1b2c3d4-000007.tar'
     """
     _validate_split_name(split)
-    return f"{split}-{index:06d}.tar"
+    if generation is None:
+        return f"{split}-{index:06d}.tar"
+    return f"{split}-{generation}-{index:06d}.tar"
 
 
 def index_name(split: str) -> str:
@@ -357,6 +372,56 @@ def _annotations_by_image(coco_data: dict[str, Any]) -> dict[Any, list[dict[str,
     return grouped
 
 
+def _pack_generation(payload: dict[str, Any], shard_sizes: list[int]) -> str:
+    """Derive a short generation token from what was packed.
+
+    Deterministic on purpose: the same split packed twice produces the same token, so re-packing unchanged data
+    reproduces the same file names rather than churning the directory. Any change to the sample count, the
+    category set, the per-shard sample counts or the shard byte sizes changes the token, which is what keeps a
+    re-pack from writing over the shards a currently published index still references.
+
+    Args:
+        payload: The index mapping for this pack, excluding the shard names it is about to name.
+        shard_sizes: Byte size of each staged shard, in order.
+
+    Returns:
+        Eight hex characters.
+
+    Examples:
+        >>> a = _pack_generation({"split": "train", "num_samples": 2}, [10, 20])
+        >>> a == _pack_generation({"split": "train", "num_samples": 2}, [10, 20])
+        True
+        >>> a == _pack_generation({"split": "train", "num_samples": 3}, [10, 20])
+        False
+    """
+    material = json.dumps({"index": payload, "sizes": shard_sizes}, sort_keys=True).encode("utf-8")
+    return hashlib.sha256(material).hexdigest()[:8]
+
+
+def _published_shard_names(destination: Path, split: str) -> set[str]:
+    """Return the shard names the currently published index of *split* references.
+
+    Reading the index rather than globbing keeps cleanup scoped to what this split actually published: the
+    directory legitimately holds other splits, and a half-written generation from a crashed run must not be
+    mistaken for a live one.
+
+    Args:
+        destination: Directory holding the pack.
+        split: Split whose published index to read.
+
+    Returns:
+        Shard file names from the published index, or an empty set when the split has no readable index yet.
+
+    Examples:
+        >>> _published_shard_names(Path("/nonexistent"), "train")
+        set()
+    """
+    try:
+        return set(read_shard_index(destination, split).shards)
+    except (WebDatasetSplitUnavailableError, ValueError, KeyError, json.JSONDecodeError):
+        return set()
+
+
 def pack_coco_to_shards(
     image_dir: str | Path,
     annotations_file: str | Path,
@@ -399,6 +464,10 @@ def pack_coco_to_shards(
     """
     if max_shard_bytes <= 0:
         raise ValueError(f"max_shard_bytes must be > 0, got {max_shard_bytes}.")
+    if category_ids not in ("remap", "raw"):
+        # The CLI constrains this with `choices`, but a library caller can pass anything. Without this the pack
+        # succeeds and only fails later, at read time, when ShardIndex.from_json rejects the policy it wrote.
+        raise ValueError(f"category_ids must be 'remap' or 'raw', got {category_ids!r}.")
     _validate_split_name(split)
 
     image_root = Path(image_dir)
@@ -494,6 +563,20 @@ def pack_coco_to_shards(
             tar = None
             shard_sample_counts.append(shard_samples)
 
+        provisional = list(shard_names)
+        shard_sizes = [(work_dir / name).stat().st_size for name in provisional]
+        generation = _pack_generation(
+            {
+                "split": split,
+                "num_samples": written,
+                "categories": list(coco_data.get("categories", ())),
+                "annotated_category_ids": sorted(annotated_ids),
+                "category_ids": category_ids,
+                "samples_per_shard": shard_sample_counts,
+            },
+            shard_sizes,
+        )
+        shard_names = [_shard_name(split, position, generation) for position in range(len(provisional))]
         index = ShardIndex(
             split=split,
             shards=tuple(shard_names),
@@ -509,12 +592,20 @@ def pack_coco_to_shards(
         # this run does not reproduce (e.g. it produced fewer of them) are removed too, so no stale, unindexed
         # shard is left behind. The index is written last, so a reader can never observe an index whose shard
         # list is only partially on disk.
-        stale_shards = {path.name for path in destination.glob(f"{split}-*.tar")} - set(shard_names)
-        for name in shard_names:
-            (work_dir / name).replace(destination / name)
-        for name in stale_shards:
-            (destination / name).unlink()
-        (destination / index_name(split)).write_bytes(index_bytes)
+        # Shard names carry this run's generation token, so moving them in cannot overwrite a shard the
+        # published index still points at. Publication is therefore: move this generation's shards in, swap the
+        # index with os.replace (atomic on POSIX), then drop the shards only the previous index referenced. A
+        # reader that opened the old index keeps a complete pack until that last step; one that opens the new
+        # index sees this generation in full. The window where a reader loses files it is mid-way through is
+        # narrowed to the cleanup, not the whole re-pack.
+        previous = _published_shard_names(destination, split)
+        for staged, published in zip(provisional, shard_names):
+            (work_dir / staged).replace(destination / published)
+        staged_index = destination / f".{index_name(split)}.{generation}"
+        staged_index.write_bytes(index_bytes)
+        staged_index.replace(destination / index_name(split))
+        for name in previous - set(shard_names):
+            (destination / name).unlink(missing_ok=True)
     finally:
         if tar is not None:
             tar.close()
@@ -894,12 +985,24 @@ def build_webdataset_loader(
         A ``DataLoader`` over *dataset*.
 
     Raises:
-        ValueError: If a planned epoch would leave a worker with no shard to read.
+        ValueError: If the split has fewer shards than distributed ranks, or if a planned epoch would leave a
+            worker with no shard to read.
     """
     _require_webdataset()
     ranks = _distributed_world_size() if world_size is None else world_size
+    shard_count = len(dataset.index.shards)
+    if ranks > shard_count:
+        # An empty DataLoader *worker* is harmless; an entirely empty *rank* is not. split_by_node would leave
+        # that rank with no batches, so it never enters validation_step/test_step while the populated ranks do —
+        # and their DDP forward and sync_dist=True logging then wait on a rank that never arrives. This applies
+        # to evaluation as much as to training, so it is checked before the fixed_epoch branch below.
+        raise ValueError(
+            f"Split {dataset.index.split!r} has {shard_count} shard(s) for {ranks} rank(s): a rank with no shard "
+            "never reaches the step function while the others do, which deadlocks the process group. Re-pack "
+            "with a smaller --max-shard-mb so the split has at least one shard per rank."
+        )
     if fixed_epoch:
-        shards = len(dataset.index.shards)
+        shards = shard_count
         slots = ranks * max(1, num_workers)
         if slots > shards:
             raise ValueError(
@@ -1019,7 +1122,23 @@ def build_webdataset(image_set: str, args: Any, resolution: int) -> WebDatasetDe
         keypoint_flip_pairs=None,
     )
 
-    cat2label = None if is_train else read_shard_index(root, "train").cat2label()
+    if is_train:
+        cat2label = None
+    else:
+        # Both indexes are read so a policy mismatch is refused rather than silently honoured. Passing None for a
+        # "raw" train split would let this split derive its own mapping from its own index, so a train split
+        # packed "raw" beside a val split packed "remap" would evaluate remapped labels against raw-trained
+        # predictions — wrong numbers, no error.
+        train_index = read_shard_index(root, "train")
+        split_index = read_shard_index(root, split)
+        if split_index.category_ids != train_index.category_ids:
+            raise ValueError(
+                f"Split {split!r} was packed with category_ids={split_index.category_ids!r} but the train split "
+                f"was packed with {train_index.category_ids!r}. The two label spaces do not match, so evaluation "
+                "would score predictions against different class indices than training used. Re-pack both splits "
+                "with the same --category-ids."
+            )
+        cat2label = train_index.cat2label()
     logger.info("Building WebDataset %s dataset at resolution %d from %s", image_set, resolution, root)
     return WebDatasetDetection(
         root,

tests/datasets/test_webdataset_io.py

@@ -36,7 +36,9 @@
     ShardIndex,
     WebDatasetDetection,
     WebDatasetSplitUnavailableError,
+    _pack_generation,
     _shard_url,
+    _validate_split_name,
     build_webdataset,
     build_webdataset_loader,
     index_name,
@@ -223,6 +225,68 @@ def test_partially_orphaned_annotations_are_rejected(self, tmp_path: Path) -> No
         with pytest.raises(ValueError, match="matching no image"):
             pack_coco_to_shards(image_dir, annotations, tmp_path / "shards")
 
+    @pytest.mark.parametrize(
+        "category_ids",
+        [pytest.param("contiguous", id="unknown-word"), pytest.param("", id="empty")],
+    )
+    def test_unknown_category_policy_is_rejected_before_anything_is_written(
+        self, tmp_path: Path, category_ids: str
+    ) -> None:
+        """The CLI's choices= does not protect a library caller; without this the pack only fails at read time."""
+        image_dir, annotations = _build_coco_split(tmp_path, count=2)
+        shard_dir = tmp_path / "shards"
+        with pytest.raises(ValueError, match="category_ids"):
+            pack_coco_to_shards(image_dir, annotations, shard_dir, split="train", category_ids=category_ids)
+        assert not shard_dir.exists() or not list(shard_dir.iterdir())
+
+    @pytest.mark.parametrize(
+        "split",
+        [pytest.param("*", id="star"), pytest.param("tr?in", id="question"), pytest.param("tr[ai]n", id="bracket")],
+    )
+    def test_glob_metacharacters_in_split_are_rejected(self, split: str) -> None:
+        """A split used as a glob would match, and then delete, other splits' shards."""
+        with pytest.raises(ValueError, match="glob metacharacter|must not contain"):
+            _validate_split_name(split)
+
+    def test_repacking_one_split_leaves_another_splits_shards_alone(self, tmp_path: Path) -> None:
+        shard_dir = tmp_path / "shards"
+        val_images, val_annotations = _build_coco_split(tmp_path, count=4, subdir="other")
+        val_index = pack_coco_to_shards(val_images, val_annotations, shard_dir, split="val")
+        val_bytes = {name: (shard_dir / name).read_bytes() for name in val_index.shards}
+
+        train_images, train_annotations = _build_coco_split(tmp_path, count=6, subdir="tr")
+        pack_coco_to_shards(train_images, train_annotations, shard_dir, split="train")
+        pack_coco_to_shards(train_images, train_annotations, shard_dir, split="train", max_shard_bytes=4096)
+
+        assert read_shard_index(shard_dir, "val").shards == val_index.shards
+        for name, payload in val_bytes.items():
+            assert (shard_dir / name).read_bytes() == payload
+
+    def test_a_repack_never_leaves_the_index_pointing_at_a_missing_shard(self, tmp_path: Path) -> None:
+        """Publication order must keep the index and the shards it names consistent at every step."""
+        shard_dir = tmp_path / "shards"
+        first_images, first_annotations = _build_coco_split(tmp_path, count=8, subdir="first")
+        pack_coco_to_shards(first_images, first_annotations, shard_dir, split="train", max_shard_bytes=4096)
+
+        second_images, second_annotations = _build_coco_split(tmp_path, count=3, subdir="second")
+        pack_coco_to_shards(second_images, second_annotations, shard_dir, split="train", max_shard_bytes=4096)
+
+        index = read_shard_index(shard_dir, "train")
+        assert index.num_samples == 3
+        for name in index.shards:
+            assert (shard_dir / name).exists()
+        # Nothing from the first generation is left behind unindexed.
+        on_disk = {path.name for path in shard_dir.glob("train-*.tar")}
+        assert on_disk == set(index.shards)
+
+    def test_generation_is_content_derived_so_repacking_is_reproducible(self, tmp_path: Path) -> None:
+        image_dir, annotations = _build_coco_split(tmp_path, count=4)
+        first = pack_coco_to_shards(image_dir, annotations, tmp_path / "a", split="train")
+        second = pack_coco_to_shards(image_dir, annotations, tmp_path / "b", split="train")
+        assert first.shards == second.shards
+        assert _pack_generation({"x": 1}, [2]) == _pack_generation({"x": 1}, [2])
+        assert _pack_generation({"x": 1}, [2]) != _pack_generation({"x": 2}, [2])
+
     @pytest.mark.parametrize("max_shard_bytes", [pytest.param(0, id="zero"), pytest.param(-1, id="negative")])
     def test_non_positive_shard_size_is_rejected(self, tmp_path: Path, max_shard_bytes: int) -> None:
         image_dir, annotations = _build_coco_split(tmp_path, count=1)
@@ -664,6 +728,28 @@ def test_evaluation_loader_reports_no_length(self, tmp_path: Path) -> None:
         with pytest.raises(TypeError):
             len(loader)
 
+    @pytest.mark.parametrize(
+        "fixed_epoch",
+        [pytest.param(True, id="training"), pytest.param(False, id="evaluation")],
+    )
+    def test_fewer_shards_than_ranks_is_rejected_on_both_paths(self, tmp_path: Path, fixed_epoch: bool) -> None:
+        """A rank with no shard never reaches the step function while the others do, which deadlocks DDP.
+
+        Evaluation used to skip this check, so only training was protected.
+        """
+        shard_dir = _pack(tmp_path, count=8)
+        dataset = WebDatasetDetection(shard_dir, "train", transforms=None)
+        shards = len(dataset.index.shards)
+        with pytest.raises(ValueError, match="rank"):
+            build_webdataset_loader(
+                dataset,
+                batch_size=2,
+                collate_fn=_count_collate,
+                num_workers=1,
+                fixed_epoch=fixed_epoch,
+                world_size=shards + 1,
+            )
+
     def test_more_workers_than_shards_is_rejected_for_training(self, tmp_path: Path) -> None:
         image_dir, annotations = _build_coco_split(tmp_path, count=16)
         shard_dir = tmp_path / "shards"
@@ -832,6 +918,25 @@ def test_both_resize_pipelines_produce_model_ready_tensors(
         assert image.ndim == 3
         assert target["boxes"].shape[-1] == 4
 
+    def test_category_policy_mismatch_between_splits_is_rejected(self, tmp_path: Path) -> None:
+        """Train packed raw beside val packed remap would score two different label spaces against each other."""
+        shard_dir = tmp_path / "shards"
+        train_images, train_annotations = _build_coco_split(tmp_path, count=6, subdir="tr")
+        pack_coco_to_shards(train_images, train_annotations, shard_dir, split="train", category_ids="raw")
+        val_images, val_annotations = _build_coco_split(tmp_path, count=4, subdir="va")
+        pack_coco_to_shards(val_images, val_annotations, shard_dir, split="val", category_ids="remap")
+        with pytest.raises(ValueError, match="category_ids"):
+            build_webdataset("val", self._namespace(shard_dir), 224)
+
+    def test_matching_raw_policy_across_splits_is_accepted(self, tmp_path: Path) -> None:
+        shard_dir = tmp_path / "shards"
+        train_images, train_annotations = _build_coco_split(tmp_path, count=6, subdir="tr2")
+        pack_coco_to_shards(train_images, train_annotations, shard_dir, split="train", category_ids="raw")
+        val_images, val_annotations = _build_coco_split(tmp_path, count=4, subdir="va2")
+        pack_coco_to_shards(val_images, val_annotations, shard_dir, split="val", category_ids="raw")
+        dataset = build_webdataset("val", self._namespace(shard_dir), 224)
+        assert dataset.cat2label is None
+
     def test_build_dataset_routes_the_webdataset_format(self, tmp_path: Path) -> None:
         shard_dir = _pack(tmp_path, count=4)
         dataset = build_dataset("train", self._namespace(shard_dir), 224)