Add tests for multi-node data gen partitioning logic

Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>

Author: fynnsu

scripts/data_generation_offline.py

@@ -26,6 +26,10 @@
 from safetensors import safe_open
 from tqdm import tqdm
 
+from speculators.data_generation.offline import (
+    get_existing_hidden_state_indices,
+    get_indices_to_process,
+)
 from speculators.data_generation.vllm_client import (
     DEFAULT_MAX_RETRIES,
     DEFAULT_REQUEST_TIMEOUT,
@@ -184,77 +188,6 @@ def parse_args():
     return parser.parse_args()
 
 
-def get_existing_hidden_state_indices(output_path: Path) -> list[int]:
-    """Find existing `hs_i.safetensors` files (where i is the file index)"""
-
-    existing_file_indices = []
-
-    if not output_path.exists():
-        return existing_file_indices
-
-    for file_path in output_path.iterdir():
-        if file_path.name.startswith("hs_") and file_path.name.endswith(".safetensors"):
-            index_str = file_path.stem[3:]  # Remove "hs_" prefix
-            try:
-                file_index = int(index_str)
-                existing_file_indices.append(file_index)
-            except ValueError:
-                continue
-
-    return sorted(existing_file_indices)
-
-
-def get_indices_to_process(
-    num_samples: int,
-    max_samples: int | None,
-    existing: list[int],
-    world_size: int,
-    rank: int,
-) -> list[int]:
-    """Determines which indices should be processed. If max_samples is None
-    returns all dataset indices not in existing. Otherwise gets the first
-    `max_samples - len(existing)` samples not already in existing.
-
-    Args:
-        num_samples: Total size of preprocessed dataset
-        max_samples: (Optional) limit for number of samples to process
-        existing: list of ids that have already been processed
-        world_size: Number of nodes to generate on
-        rank: The rank of the local node
-
-    Returns:
-        list of dataset indices to process
-    """
-
-    if len(existing) >= num_samples:
-        logger.info("All samples already processed!")
-        return []
-
-    target = min(max_samples, num_samples) if max_samples else num_samples
-
-    chunk_size = target // world_size
-    remainder = target % world_size
-    # Distribute remainder across the first `remainder` ranks so chunks differ
-    # by at most 1.
-    start = rank * chunk_size + min(rank, remainder)
-    end = start + chunk_size + (1 if rank < remainder else 0)
-
-    existing_s = set(existing)
-    to_process = [i for i in range(start, end) if i not in existing_s]
-
-    if not to_process:
-        logger.info("All samples for this rank already processed!")
-        return []
-
-    if len(existing_s & set(range(start, end))) > 0:
-        logger.info(
-            f"Found {len(existing_s & set(range(start, end)))} existing samples"
-            f" for rank {rank}."
-        )
-
-    return to_process
-
-
 def check_safetensors_file(path: Path, tokens: list[int]):
     with safe_open(path, "pt") as f:
         t_ids = f.get_tensor("token_ids").tolist()

src/speculators/data_generation/offline.py

@@ -0,0 +1,75 @@
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+def get_existing_hidden_state_indices(output_path: Path) -> list[int]:
+    """Find existing `hs_i.safetensors` files (where i is the file index)"""
+
+    existing_file_indices = []
+
+    if not output_path.exists():
+        return existing_file_indices
+
+    for file_path in output_path.iterdir():
+        if file_path.name.startswith("hs_") and file_path.name.endswith(".safetensors"):
+            index_str = file_path.stem[3:]  # Remove "hs_" prefix
+            try:
+                file_index = int(index_str)
+                existing_file_indices.append(file_index)
+            except ValueError:
+                continue
+
+    return sorted(existing_file_indices)
+
+
+def get_indices_to_process(
+    num_samples: int,
+    max_samples: int | None,
+    existing: list[int],
+    world_size: int,
+    rank: int,
+) -> list[int]:
+    """Determines which indices should be processed. If max_samples is None
+    returns all dataset indices not in existing. Otherwise gets the first
+    `max_samples - len(existing)` samples not already in existing.
+
+    Args:
+        num_samples: Total size of preprocessed dataset
+        max_samples: (Optional) limit for number of samples to process
+        existing: list of ids that have already been processed
+        world_size: Number of nodes to generate on
+        rank: The rank of the local node
+
+    Returns:
+        list of dataset indices to process
+    """
+
+    if len(existing) >= num_samples:
+        logger.info("All samples already processed!")
+        return []
+
+    target = min(max_samples, num_samples) if max_samples else num_samples
+
+    chunk_size = target // world_size
+    remainder = target % world_size
+    # Distribute remainder across the first `remainder` ranks so chunks differ
+    # by at most 1.
+    start = rank * chunk_size + min(rank, remainder)
+    end = start + chunk_size + (1 if rank < remainder else 0)
+
+    existing_s = set(existing)
+    to_process = [i for i in range(start, end) if i not in existing_s]
+
+    if not to_process:
+        logger.info("All samples for this rank already processed!")
+        return []
+
+    if len(existing_s & set(range(start, end))) > 0:
+        logger.info(
+            f"Found {len(existing_s & set(range(start, end)))} existing samples"
+            f" for rank {rank}."
+        )
+
+    return to_process

tests/unit/data_generation/__init__.py

@@ -0,0 +1,117 @@
+import pytest
+
+from speculators.data_generation.offline import (
+    get_existing_hidden_state_indices,
+    get_indices_to_process,
+)
+
+
+# ===== get_indices_to_process Tests =====
+
+
+class TestGetIndicesToProcess:
+    def test_single_node_no_max_samples(self):
+        result = get_indices_to_process(10, None, [], world_size=1, rank=0)
+        assert result == list(range(10))
+
+    def test_single_node_with_max_samples(self):
+        result = get_indices_to_process(10, 5, [], world_size=1, rank=0)
+        assert result == [0, 1, 2, 3, 4]
+
+    def test_single_node_max_samples_exceeds_num_samples(self):
+        result = get_indices_to_process(5, 10, [], world_size=1, rank=0)
+        assert result == list(range(5))
+
+    def test_single_node_with_existing(self):
+        result = get_indices_to_process(10, None, [2, 5, 7], world_size=1, rank=0)
+        assert result == [0, 1, 3, 4, 6, 8, 9]
+
+    def test_all_samples_already_processed(self):
+        result = get_indices_to_process(5, None, list(range(5)), world_size=1, rank=0)
+        assert result == []
+
+    def test_multi_node_even_split(self):
+        r0 = get_indices_to_process(10, None, [], world_size=2, rank=0)
+        r1 = get_indices_to_process(10, None, [], world_size=2, rank=1)
+        assert r0 == [0, 1, 2, 3, 4]
+        assert r1 == [5, 6, 7, 8, 9]
+
+    def test_multi_node_uneven_split(self):
+        r0 = get_indices_to_process(10, None, [], world_size=3, rank=0)
+        r1 = get_indices_to_process(10, None, [], world_size=3, rank=1)
+        r2 = get_indices_to_process(10, None, [], world_size=3, rank=2)
+        assert r0 == [0, 1, 2, 3]
+        assert r1 == [4, 5, 6]
+        assert r2 == [7, 8, 9]
+
+    def test_multi_node_no_overlap_and_full_coverage(self):
+        num_samples = 17
+        world_size = 4
+        all_indices = []
+        for rank in range(world_size):
+            chunk = get_indices_to_process(
+                num_samples, None, [], world_size=world_size, rank=rank
+            )
+            all_indices.extend(chunk)
+        assert sorted(all_indices) == list(range(num_samples))
+        assert len(all_indices) == len(set(all_indices))
+
+    def test_multi_node_with_max_samples(self):
+        r0 = get_indices_to_process(100, 10, [], world_size=2, rank=0)
+        r1 = get_indices_to_process(100, 10, [], world_size=2, rank=1)
+        assert r0 == [0, 1, 2, 3, 4]
+        assert r1 == [5, 6, 7, 8, 9]
+
+    def test_multi_node_with_existing(self):
+        result = get_indices_to_process(10, None, [1, 3], world_size=2, rank=0)
+        assert result == [0, 2, 4]
+
+    def test_multi_node_rank_fully_processed(self):
+        result = get_indices_to_process(
+            10, None, [0, 1, 2, 3, 4], world_size=2, rank=0
+        )
+        assert result == []
+
+    def test_existing_exceeds_num_samples(self):
+        result = get_indices_to_process(5, None, list(range(10)), world_size=1, rank=0)
+        assert result == []
+
+
+# ===== get_existing_hidden_state_indices Tests =====
+
+
+class TestGetExistingHiddenStateIndices:
+    def test_nonexistent_directory(self, tmp_path):
+        result = get_existing_hidden_state_indices(tmp_path / "nonexistent")
+        assert result == []
+
+    def test_empty_directory(self, tmp_path):
+        result = get_existing_hidden_state_indices(tmp_path)
+        assert result == []
+
+    def test_finds_safetensor_files(self, tmp_path):
+        (tmp_path / "hs_0.safetensors").touch()
+        (tmp_path / "hs_3.safetensors").touch()
+        (tmp_path / "hs_7.safetensors").touch()
+        result = get_existing_hidden_state_indices(tmp_path)
+        assert result == [0, 3, 7]
+
+    def test_ignores_non_numeric_suffixes(self, tmp_path):
+        (tmp_path / "hs_0.safetensors").touch()
+        (tmp_path / "hs_abc.safetensors").touch()
+        (tmp_path / "hs_.safetensors").touch()
+        result = get_existing_hidden_state_indices(tmp_path)
+        assert result == [0]
+
+    def test_ignores_unrelated_files(self, tmp_path):
+        (tmp_path / "hs_0.safetensors").touch()
+        (tmp_path / "other_file.txt").touch()
+        (tmp_path / "hs_1.pt").touch()
+        result = get_existing_hidden_state_indices(tmp_path)
+        assert result == [0]
+
+    def test_results_are_sorted(self, tmp_path):
+        for i in [9, 2, 5, 0]:
+            (tmp_path / f"hs_{i}.safetensors").touch()
+        result = get_existing_hidden_state_indices(tmp_path)
+        assert result == [0, 2, 5, 9]