fix: bugs related to rng, mask and deepcopy by selmanozleyen · Pull Request #285 · scverse/annbatch

selmanozleyen · 2026-09-23T15:22:29Z

DistributedSampler kept the sampler it wrapped by reference and re-sharded its mask on every n_batches/validate/_sample call, so wrapping left the caller's own sampler rank-sharded and re-seeded. Shard a copy. split_given_size is np.split, which returns views, and every caller reshuffles one row-id buffer between windows, so an already-yielded load request's splits changed underneath the caller. Copy the pieces out in the helper, where all the samplers route through. ClassSampler also hoisted the split list out of its window loop, so that moves back in. A mask assigned mid-pass had no effect on that pass, whose slices are all drawn when it starts. Raise, but only when the range actually moves, since DistributedSampler re-assigns the same shard on every call. class_weights as a pandas Series converted positionally and dropped the index, silently moving weights onto other classes. Reject it. `rng or np.random.default_rng()` swapped a falsy generator for a fresh unseeded one, so rng=0 lost reproducibility without a word while rng=42 failed loudly at the first draw. Both constructors now route through resolve_rng, which defaults only None and rejects a seed with a message saying what to pass instead. Also correct the ClassSampler summary, which said a class is drawn per batch when it is drawn per lcm(chunk_size, batch_size) rows. The Implementation section already had the rule; the summary contradicted it.

codecov · 2026-09-23T15:27:53Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.62%. Comparing base (b8385d7) to head (cc5e553).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #285      +/-   ##
==========================================
+ Coverage   91.30%   91.62%   +0.32%     
==========================================
  Files          18       18              
  Lines        1495     1517      +22     
==========================================
+ Hits         1365     1390      +25     
+ Misses        130      127       -3

Files with missing lines	Coverage Δ
src/annbatch/samplers/_chunk_sampler.py	`100.00% <100.00%> (ø)`
...nnbatch/samplers/_class_samplers/_class_sampler.py	`100.00% <100.00%> (+1.14%)`	⬆️
.../annbatch/samplers/_class_samplers/_rle_manager.py	`97.59% <100.00%> (+1.29%)`	⬆️
...amplers/_class_samplers/_weighted_class_sampler.py	`100.00% <100.00%> (ø)`
src/annbatch/samplers/_distributed_sampler.py	`96.22% <100.00%> (+0.07%)`	⬆️
src/annbatch/samplers/_random_sampler.py	`100.00% <ø> (ø)`
src/annbatch/samplers/_utils.py	`97.82% <100.00%> (+0.39%)`	⬆️
src/annbatch/utils.py	`74.40% <100.00%> (+0.20%)`	⬆️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

ilan-gold · 2026-09-28T14:16:08Z

-        self._sampler = sampler
-        if sampler.rng is not None:
-            sampler.rng = _spawn_worker_rng(sampler.rng, self._rank)
+        # a copy: n_batches/validate/_sample re-shard this sampler's mask on every call


What does "reshard" mean?

it was lingo for paritioning the data from the existing code. on distributed sampler

Ok I see now - these update the mask. My other comment touched on this point - I don't think this is necessarily needed. At least not n_batches

ilan-gold · 2026-09-28T14:17:18Z

-        if sampler.rng is not None:
-            sampler.rng = _spawn_worker_rng(sampler.rng, self._rank)
+        # a copy: n_batches/validate/_sample re-shard this sampler's mask on every call
+        self._sampler = copy.deepcopy(sampler)


Ok, right, so this is basically what #250 is supposed to fix, right? Our own copy method should handle this sort of thing? Or?

No because there is no copy call in the previous code that hooks into anything. Making our classes deepcopyable explicitly doesn't also guarantee we deepcopied everything before. Or that every consumer of that class deepcopied it before using it. We can't have that in python. This is not only special to rng's, it's also for any non primitive we pass. We could also observe such a bug in classes for example. A caller who passed it might modify it later.

Perhaps we can make a repo wide rule for samplers, that they should have full ownership of the attributes they have by deepcopy. But I think we should make it limited to our own samplers. Loader still works with any Sampler. DistributedSampler wouldn't work for example with a sampler that has a file open (extreme example I know but still ilustrates the ontology difference).

On #250, I think this is why we should just raise on copy.copy()/__copy__ and allow only copy.deepcopy()/__deepcopy__.

Making our classes deepcopyable explicitly doesn't also guarantee we deepcopied everything before.

Yeah, I just meant we should do this, not that it was being done. We should make our classes have explicit deep copy behavior instead of relying on whatever python thinks we want.

On #250, I think this is why we should just raise on copy.copy()/copy and allow only copy.deepcopy()/deepcopy.

Right, this is the sort of thing I am getting at. The rest of your comment I agree with.

ilan-gold · 2026-09-28T14:18:21Z

 def split_given_size(a: np.ndarray, size: int) -> list[np.ndarray]:
-    """Wrapper around `np.split` to split up an array into `size` chunks"""
-    return np.split(a, np.arange(size, len(a), size))
+    """Wrapper around `np.split` to split up an array into `size` chunks.
+
+    `np.split` returns views, so the pieces are copied: callers reshuffle one row-id buffer
+    between windows, which would rewrite the splits they already yielded.
+    """
+    return [piece.copy() for piece in np.split(a, np.arange(size, len(a), size))]


Are we sure every single place needs a copy? Maybe we should add a copy argument and keep things localized to the places that actually need a copy? WDYT?

Yeah I think a copy argument makes sense. The only legal consumer of Sampler is Loader, we don't even need to copy atm.

The only reasonable case would be when someone wants to get the full schedule ahead. Having an explicit arg would help document this. But I am not sure about copy. I'd perhaps call it reuse_splits:bool = True or something to make it clearer.

I think copy is clearer and widely known.

But it could be confused with other things for this reason no? Like could it also mean it copies the inner sampler? and not deep copy? Also there isn't another case where something is copied here

Also it would be a user facing arg right?

I think you misunderstood what I was saying. I was saying split_given_size gets a copy argument that is used when corrupting the outputs is a real concern. But when we internally control the output of split_given_size, we don't need copy. I wasn't saying the Sampler' functions need a copy` argument that copies splits - the splits the samplers yield should not be corruptible views, your bug description was correct. Sorry for the misunderstanding.

ilan-gold · 2026-09-28T14:23:27Z

+    # `len(loader)` mid-epoch goes through DistributedSampler.n_batches, which re-assigns the
+    # same shard; re-assigning the range a pass is already reading is a no-op, not a move


Looking at

annbatch/src/annbatch/samplers/_distributed_sampler.py

Lines 130 to 141 in b8385d7

def _shard_mask(self, n_obs: int) -> slice:

"""Return the contiguous observation slice for this rank."""

per_rank = n_obs // self._world_size

if self._enforce_equal_batches:

per_rank = per_rank // self._sampler.batch_size * self._sampler.batch_size

rank_start = self._rank * per_rank

rank_stop = rank_start + per_rank

return slice(rank_start, rank_stop)

def n_batches(self, n_obs: int) -> int:

self._sampler.mask = self._shard_mask(n_obs)

return self._sampler.n_batches(n_obs)

, I think n_batches should just not set the mask and instead answer the query of how many batches are yielded for a given number of observations. No?

I agree but can it be a separate PR?

This, and the presence of the same behavior in validate and _sample seem pretty important - I can't remember what the need of that side-effect is. I'm wondering if there is a way we can mitigate the bugs here but simply not having that side effect.

ilan-gold · 2026-10-02T16:30:03Z

+- `class_weights` given as a {class}`pandas.Series` now raises {class}`TypeError`. It used to be read positionally, discarding the index and attaching the weights to the wrong classes; pass `class_weights.reindex(classes.categories).to_numpy()` instead.
+- `rng` must now be a {class}`numpy.random.Generator` or `None`; anything else raises {class}`TypeError`. A falsy seed such as `rng=0` used to be silently swapped for a fresh unseeded generator, losing reproducibility; pass `np.random.default_rng(0)` instead.


We just released a minor version, so let's not have to do breaking changes.. The first one is a bug since it corrupted data and the second can be smoothed out with a warning.

ilan-gold · 2026-10-02T16:34:15Z

 def split_given_size(a: np.ndarray, size: int) -> list[np.ndarray]:
-    """Wrapper around `np.split` to split up an array into `size` chunks"""
-    return np.split(a, np.arange(size, len(a), size))
+    """Wrapper around `np.split` to split up an array into `size` chunks.
+
+    `np.split` returns views, so the pieces are copied: callers reshuffle one row-id buffer
+    between windows, which would rewrite the splits they already yielded.
+    """
+    return [piece.copy() for piece in np.split(a, np.arange(size, len(a), size))]


I think you misunderstood what I was saying. I was saying split_given_size gets a copy argument that is used when corrupting the outputs is a real concern. But when we internally control the output of split_given_size, we don't need copy. I wasn't saying the Sampler' functions need a copy` argument that copies splits - the splits the samplers yield should not be corruptible views, your bug description was correct. Sorry for the misunderstanding.

ilan-gold · 2026-10-02T16:36:47Z

+    # `len(loader)` mid-epoch goes through DistributedSampler.n_batches, which re-assigns the
+    # same shard; re-assigning the range a pass is already reading is a no-op, not a move


This, and the presence of the same behavior in validate and _sample seem pretty important - I can't remember what the need of that side-effect is. I'm wondering if there is a way we can mitigate the bugs here but simply not having that side effect.

selmanozleyen requested a review from ilan-gold September 23, 2026 15:22

selmanozleyen added the skip-gpu-ci Whether gpu ci should be skipped label Sep 23, 2026

Merge branch 'main' into fix/review-findings

62ba9a2

ilan-gold reviewed Sep 28, 2026

View reviewed changes

selmanozleyen added 2 commits October 2, 2026 14:24

Clarify open-pass tests, changelog

7858450

Add user-facing copy option for splits

cc5e553

selmanozleyen force-pushed the fix/review-findings branch from 1b54d72 to cc5e553 Compare October 2, 2026 13:16

ilan-gold reviewed Oct 2, 2026

View reviewed changes

		# `len(loader)` mid-epoch goes through DistributedSampler.n_batches, which re-assigns the
		# same shard; re-assigning the range a pass is already reading is a no-op, not a move

	def _shard_mask(self, n_obs: int) -> slice:
	"""Return the contiguous observation slice for this rank."""
	per_rank = n_obs // self._world_size
	if self._enforce_equal_batches:
	per_rank = per_rank // self._sampler.batch_size * self._sampler.batch_size
	rank_start = self._rank * per_rank
	rank_stop = rank_start + per_rank
	return slice(rank_start, rank_stop)

	def n_batches(self, n_obs: int) -> int:
	self._sampler.mask = self._shard_mask(n_obs)
	return self._sampler.n_batches(n_obs)

		- `class_weights` given as a {class}`pandas.Series` now raises {class}`TypeError`. It used to be read positionally, discarding the index and attaching the weights to the wrong classes; pass `class_weights.reindex(classes.categories).to_numpy()` instead.
		- `rng` must now be a {class}`numpy.random.Generator` or `None`; anything else raises {class}`TypeError`. A falsy seed such as `rng=0` used to be silently swapped for a fresh unseeded generator, losing reproducibility; pass `np.random.default_rng(0)` instead.

Conversation

selmanozleyen commented Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

selmanozleyen commented Sep 23, 2026 •

edited

Loading

codecov Bot commented Sep 23, 2026 •

edited

Loading