chore: sampler serialization tests

tests/test_sampler.py

@@ -2,17 +2,20 @@
 
 from __future__ import annotations
 
+import copy
 import math
+import pickle
 import sys
 from functools import partial
 from typing import TYPE_CHECKING
 from unittest.mock import MagicMock, patch
 
 import numpy as np
+import pandas as pd
 import pytest
 
 from annbatch.abc import Sampler
-from annbatch.samplers import DistributedSampler, RandomSampler, SequentialSampler
+from annbatch.samplers import ClassSampler, DistributedSampler, RandomSampler, SequentialSampler
 from annbatch.samplers._utils import WorkerInfo
 
 if TYPE_CHECKING:
@@ -879,3 +882,74 @@ def test_wraps_sequential_sampler(self, make_distributed_sampler: Callable[...,
             for j in range(i + 1, world_size):
                 assert set(all_indices[i]).isdisjoint(set(all_indices[j]))
         assert set().union(*all_indices) == set(range(n_obs))
+
+
+# =============================================================================
+# Serialization / reproducibility (all implemented stochastic samplers)
+# =============================================================================
+
+
+def _make_sampler(kind: str, seed: int, n_obs: int) -> Sampler:
+    """Build the sampler identified by ``kind``, seeded with ``seed``, sized for ``n_obs`` obs."""
+    match kind:
+        case "random":
+            return RandomSampler(chunk_size=10, preload_nchunks=2, batch_size=5, rng=np.random.default_rng(seed))
+        case "class":
+            # n_obs obs, 5 classes, each a contiguous run of n_obs // 5 (>= chunk_size)
+            classes = pd.Categorical(np.repeat(np.arange(5), n_obs // 5))
+            return ClassSampler(
+                chunk_size=10,
+                preload_nchunks=2,
+                batch_size=5,
+                classes=classes,
+                num_samples=50,
+                rng=np.random.default_rng(seed),
+            )
+        case "distributed":
+            inner = RandomSampler(chunk_size=10, preload_nchunks=2, batch_size=5, rng=np.random.default_rng(seed))
+            # dist_info is only called at construction (rank 0 of 1) and is not stored,
+            # so the resulting sampler stays picklable/copyable.
+            return DistributedSampler(inner, dist_info=lambda: (0, 1))
+        case _:
+            raise ValueError(f"unknown sampler kind {kind!r}")
+
+
+# A round-trip duplicates a sampler into an *independent* object with the same state.
+# ``copy.deepcopy`` -- not shallow ``copy.copy`` -- is the copy analog of pickle: a shallow
+# copy would share the ``rng`` object with the original, breaking the independence these
+# assertions rely on.
+_ROUND_TRIPS = {
+    "pickle": lambda sampler: pickle.loads(pickle.dumps(sampler)),
+    "deepcopy": copy.deepcopy,
+}
+
+
+@pytest.mark.parametrize("round_trip", _ROUND_TRIPS.values(), ids=_ROUND_TRIPS.keys())
+@pytest.mark.parametrize("kind", ["random", "class", "distributed"])
+def test_sampler_is_serializable_and_reproducible(kind: str, round_trip: Callable[[Sampler], Sampler]):
+    """Every sampler survives a pickle/deepcopy round-trip with its random *state* (not just the seed) preserved."""
+    n_obs = 100
+
+    def advance_round_trip_indices(seed: int) -> list[int]:
+        sampler = _make_sampler(kind, seed, n_obs)
+        collect_indices(sampler, n_obs)  # advance the rng one pass
+        restored = round_trip(sampler)
+        indices, _, _ = collect_indices(restored, n_obs)
+        return indices
+
+    sampler = _make_sampler(kind, seed=0, n_obs=n_obs)
+    fresh_indices, _, _ = collect_indices(_make_sampler(kind, seed=0, n_obs=n_obs), n_obs)
+
+    # advance the rng by one full pass, then snapshot via the round-trip
+    collect_indices(sampler, n_obs)
+    restored = round_trip(sampler)
+    assert isinstance(restored, type(sampler))
+
+    restored_indices, _, _ = collect_indices(restored, n_obs)
+    original_indices, _, _ = collect_indices(sampler, n_obs)
+    assert restored_indices, "sampler produced no indices"
+
+    assert restored_indices == original_indices
+    assert restored_indices != fresh_indices
+    assert restored_indices == advance_round_trip_indices(seed=0)
+    assert restored_indices != advance_round_trip_indices(seed=1)