From d2405fd3b5691b18b86f1e469e676e2505b78469 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Fri, 10 Jul 2026 00:33:00 +0200 Subject: [PATCH 1/4] init without precommit --- tests/test_sampler.py | 76 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 75 insertions(+), 1 deletion(-) diff --git a/tests/test_sampler.py b/tests/test_sampler.py index d37baf21..b0c59dfa 100644 --- a/tests/test_sampler.py +++ b/tests/test_sampler.py @@ -3,16 +3,18 @@ from __future__ import annotations 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 +881,75 @@ 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 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 "sequential": + return SequentialSampler(chunk_size=10, preload_nchunks=2, batch_size=5) + 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. + return DistributedSampler(inner, dist_info=lambda: (0, 1)) + case _: + raise ValueError(f"unknown sampler kind {kind!r}") + + +@pytest.mark.parametrize("kind", ["random", "sequential", "class", "distributed"]) +def test_sampler_is_serializable_and_reproducible(kind: str): + """Every sampler pickles and its random *state* (not just the seed) round-trips.""" + n_obs = 100 + + def advance_pickle_indices(seed: int) -> list[int]: + sampler = _make_sampler(kind, seed, n_obs) + collect_indices(sampler, n_obs) # advance the rng one pass + restored = pickle.loads(pickle.dumps(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 pickle + collect_indices(sampler, n_obs) + restored = pickle.loads(pickle.dumps(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 + + if restored.shuffle: + # tests for the controlled stochastic behavior + # without these tests, passing this test becomes + # trivial if there is a seed collapse etc. + + # check if the state is not reset: an advanced-then-saved + # stochastic sampler must not match a fresh one + assert restored_indices != fresh_indices + # check if saved state tracks the seed: identical for a + # same-seed run, different otherwise + assert restored_indices == advance_pickle_indices(seed=0) + assert restored_indices != advance_pickle_indices(seed=1) From 64846b208125b6c29557cb644a24306a2e4ecad1 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Tue, 14 Jul 2026 14:47:15 +0200 Subject: [PATCH 2/4] add deepcopy test and remove sequentialsampler test --- tests/test_sampler.py | 48 +++++++++++++++++++++---------------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/tests/test_sampler.py b/tests/test_sampler.py index b0c59dfa..634a30f8 100644 --- a/tests/test_sampler.py +++ b/tests/test_sampler.py @@ -2,6 +2,7 @@ from __future__ import annotations +import copy import math import pickle import sys @@ -884,7 +885,7 @@ def test_wraps_sequential_sampler(self, make_distributed_sampler: Callable[..., # ============================================================================= -# Serialization / reproducibility (all implemented samplers) +# Serialization / reproducibility (all implemented stochastic samplers) # ============================================================================= @@ -893,8 +894,6 @@ def _make_sampler(kind: str, seed: int, n_obs: int) -> Sampler: match kind: case "random": return RandomSampler(chunk_size=10, preload_nchunks=2, batch_size=5, rng=np.random.default_rng(seed)) - case "sequential": - return SequentialSampler(chunk_size=10, preload_nchunks=2, batch_size=5) 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)) @@ -909,30 +908,41 @@ def _make_sampler(kind: str, seed: int, n_obs: int) -> Sampler: 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. + # so the resulting sampler stays picklable/copyable. return DistributedSampler(inner, dist_info=lambda: (0, 1)) case _: raise ValueError(f"unknown sampler kind {kind!r}") -@pytest.mark.parametrize("kind", ["random", "sequential", "class", "distributed"]) -def test_sampler_is_serializable_and_reproducible(kind: str): - """Every sampler pickles and its random *state* (not just the seed) round-trips.""" +# 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_pickle_indices(seed: int) -> list[int]: + 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 = pickle.loads(pickle.dumps(sampler)) + 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 pickle + # advance the rng by one full pass, then snapshot via the round-trip collect_indices(sampler, n_obs) - restored = pickle.loads(pickle.dumps(sampler)) + restored = round_trip(sampler) assert isinstance(restored, type(sampler)) restored_indices, _, _ = collect_indices(restored, n_obs) @@ -940,16 +950,6 @@ def advance_pickle_indices(seed: int) -> list[int]: assert restored_indices, "sampler produced no indices" assert restored_indices == original_indices - - if restored.shuffle: - # tests for the controlled stochastic behavior - # without these tests, passing this test becomes - # trivial if there is a seed collapse etc. - - # check if the state is not reset: an advanced-then-saved - # stochastic sampler must not match a fresh one - assert restored_indices != fresh_indices - # check if saved state tracks the seed: identical for a - # same-seed run, different otherwise - assert restored_indices == advance_pickle_indices(seed=0) - assert restored_indices != advance_pickle_indices(seed=1) + assert restored_indices != fresh_indices + assert restored_indices == advance_round_trip_indices(seed=0) + assert restored_indices != advance_round_trip_indices(seed=1) From 4d8b3b6c4056a2d596a17e6eafcd8224199f8330 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Wed, 15 Jul 2026 15:21:50 +0200 Subject: [PATCH 3/4] remove unnecessary extra advance hence the helper func --- tests/test_sampler.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/tests/test_sampler.py b/tests/test_sampler.py index 634a30f8..cc9833f6 100644 --- a/tests/test_sampler.py +++ b/tests/test_sampler.py @@ -930,26 +930,23 @@ def test_sampler_is_serializable_and_reproducible(kind: str, round_trip: Callabl """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 + # advance the rng one full pass before snapshotting collect_indices(sampler, n_obs) restored = round_trip(sampler) assert isinstance(restored, type(sampler)) restored_indices, _, _ = collect_indices(restored, n_obs) + assert len(restored_indices) > 0, "sampler produced no indices" original_indices, _, _ = collect_indices(sampler, n_obs) - assert restored_indices, "sampler produced no indices" + # round-trip preserved the live rng *state*: restored keeps producing what the original does... assert restored_indices == original_indices + # ...and that state had genuinely advanced past a fresh seed-0 sampler -- otherwise we'd only be + # proving the seed round-tripped, not the state. assert restored_indices != fresh_indices - assert restored_indices == advance_round_trip_indices(seed=0) - assert restored_indices != advance_round_trip_indices(seed=1) + # a different seed must produce a different sequence (guards against an ignored/hard-coded rng) + fresh_indices_seed_1, _, _ = collect_indices(_make_sampler(kind, seed=1, n_obs=n_obs), n_obs) + assert fresh_indices != fresh_indices_seed_1 From d569be8619ec5e216bb4d2b7717729ab456d91a5 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Sat, 18 Jul 2026 20:11:48 +0200 Subject: [PATCH 4/4] implement eq, deepcopy, and not copy. Also add tests --- src/annbatch/abc/sampler.py | 64 ++++++++++++++++++++++++++++++++++++- tests/test_sampler.py | 29 ++++++++++++++++- 2 files changed, 91 insertions(+), 2 deletions(-) diff --git a/src/annbatch/abc/sampler.py b/src/annbatch/abc/sampler.py index 028f70bf..12b031f5 100644 --- a/src/annbatch/abc/sampler.py +++ b/src/annbatch/abc/sampler.py @@ -2,8 +2,9 @@ from __future__ import annotations +import copy from abc import ABC, abstractmethod -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, NoReturn, Self import numpy as np @@ -15,6 +16,32 @@ from annbatch.types import LoadRequest +def _attr_equal(a: object, b: object) -> bool: + """Structural equality for a single value stored in a sampler's ``__dict__``. + + Handles the container types samplers keep as state -- most importantly the + :class:`numpy.random.Generator`, whose equality must compare the live *bit + generator state* (not object identity) so a round-tripped sampler counts as + equal to its source. numpy arrays and pandas objects (which return + element-wise ``==``) are compared structurally, and nested samplers dispatch + back to :meth:`Sampler.__eq__`. + """ + if isinstance(a, np.random.Generator) or isinstance(b, np.random.Generator): + return ( + isinstance(a, np.random.Generator) + and isinstance(b, np.random.Generator) + and a.bit_generator.state == b.bit_generator.state + ) + if isinstance(a, Sampler) or isinstance(b, Sampler): + return a == b + if isinstance(a, np.ndarray) or isinstance(b, np.ndarray): + return isinstance(a, np.ndarray) and isinstance(b, np.ndarray) and bool(np.array_equal(a, b)) + # pandas DataFrame/Series/Index/Categorical all expose a structural `.equals` + if hasattr(a, "equals") and hasattr(b, "equals") and type(a) is type(b): + return bool(a.equals(b)) + return bool(a == b) + + class Sampler(ABC): """Base sampler class. @@ -24,6 +51,41 @@ class Sampler(ABC): _mask: slice = slice(0, None) _rng: np.random.Generator | None = None + def __eq__(self, other: object) -> bool: + """Two samplers are equal iff they have the same type and the same state. + + State includes the random number generator's *bit generator state*, so a + sampler equals a pickle/deepcopy round-trip of itself but not a fresh (or + differently advanced) sampler built from the same seed. + """ + if type(self) is not type(other): + return NotImplemented + if self.__dict__.keys() != other.__dict__.keys(): + return False + return all(_attr_equal(self.__dict__[key], other.__dict__[key]) for key in self.__dict__) + + def __copy__(self) -> NoReturn: + """Refuse shallow copies -- they would share the rng with the original. + + A shallow copy keeps the same :class:`numpy.random.Generator` object, so + advancing one sampler would advance the "copy" too, silently breaking the + independence a copy is meant to provide. Use :func:`copy.deepcopy` instead. + """ + raise TypeError( + f"{type(self).__name__} does not support shallow copying: a shallow copy would share the " + "random number generator with the original, so the two would not sample independently. " + "Use copy.deepcopy() instead." + ) + + def __deepcopy__(self, memo: dict[int, Any]) -> Self: + """Deep-copy every attribute into an independent sampler (rng included).""" + cls = type(self) + new = cls.__new__(cls) + memo[id(self)] = new + for key, value in self.__dict__.items(): + setattr(new, key, copy.deepcopy(value, memo)) + return new + @property def mask(self) -> slice: """The observation range this sampler operates on.""" diff --git a/tests/test_sampler.py b/tests/test_sampler.py index cc9833f6..7e502efe 100644 --- a/tests/test_sampler.py +++ b/tests/test_sampler.py @@ -937,6 +937,7 @@ def test_sampler_is_serializable_and_reproducible(kind: str, round_trip: Callabl collect_indices(sampler, n_obs) restored = round_trip(sampler) assert isinstance(restored, type(sampler)) + assert restored == sampler restored_indices, _, _ = collect_indices(restored, n_obs) assert len(restored_indices) > 0, "sampler produced no indices" @@ -948,5 +949,31 @@ def test_sampler_is_serializable_and_reproducible(kind: str, round_trip: Callabl # proving the seed round-tripped, not the state. assert restored_indices != fresh_indices # a different seed must produce a different sequence (guards against an ignored/hard-coded rng) - fresh_indices_seed_1, _, _ = collect_indices(_make_sampler(kind, seed=1, n_obs=n_obs), n_obs) + sampler_seed_1 = _make_sampler(kind, seed=1, n_obs=n_obs) + fresh_indices_seed_1, _, _ = collect_indices(sampler_seed_1, n_obs) assert fresh_indices != fresh_indices_seed_1 + assert sampler != sampler_seed_1 + + +@pytest.mark.parametrize("kind", ["random", "class", "distributed"]) +@pytest.mark.parametrize("advance", [False, True], ids=["unchanged", "rng-advanced"]) +def test_sampler_eq(kind: str, advance: bool): + """``__eq__`` holds across a state-preserving copy and breaks once a sampler's rng advances.""" + n_obs = 100 + sampler = _make_sampler(kind, seed=0, n_obs=n_obs) + other = copy.deepcopy(sampler) + if advance: + collect_indices(other, n_obs) + + assert (sampler == other) is not advance + # a different seed and a non-sampler object are never equal + assert sampler != _make_sampler(kind, seed=1, n_obs=n_obs) + assert sampler != object() + + +@pytest.mark.parametrize("kind", ["random", "class", "distributed"]) +def test_sampler_shallow_copy_is_rejected(kind: str): + """A shallow copy would share the rng, so ``copy.copy`` must raise rather than silently alias state.""" + sampler = _make_sampler(kind, seed=0, n_obs=100) + with pytest.raises(TypeError, match="deepcopy"): + copy.copy(sampler)