Repository: dask/dask
Pull request: Speed up dask.bag.Bag.random_sample
Closes #10351
```python
%timeit random_state_data_python(1, 0)
%timeit random_state_data_python(10, 0)
%timeit random_state_data_python(100, 0)
%timeit random_state_data_python(1000, 0)
```
Before:
```
265 µs ± 1.72 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
2.64 ms ± 16.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
26.5 ms ± 146 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
272 ms ± 1.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```
After:
```
131 µs ± 1.47 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
234 µs ± 1.21 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
1.26 ms ± 9.68 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
20.4 ms ± 137 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

Problem statement:
Parallel random_state_data_python
I've noticed my workers spend a lot of time on dask/bag/core.py random_sample computing the list comprehension in random_state_data_python. Currently this has to happen serially, as python.random.Random increments the state serially. However there are other sources of randomness (like SystemRandom) that can be accessed in parallel, but don't generate reproducible sequences.

Is reproducibility required here? Could we switch to SystemRandom (or similar) and parallelize this function?

Hints from issue discussion:
I transferred this to the dask/dask repo since it doesn't require a distributed cluster.

My take on this is that for most practical purposes this is not required to be reproducible. At the very least the default is not strictly reproducible if the user does not provide a random seed. Accepting a random seed and ensuring reproducibility in this case is a useful feature, e.g. for testing and I would prefer keeping this.

However, I'm not entirely sure what level of parallelism you are looking for. I would be slightly concerned if we parallelized within `random_state_data_python`. What did you have in mind?
```python
def random_state_data_python(n, random_state=None):
    """Return a list of tuples that can be passed to
    ``random.Random.setstate``.

    Parameters
    ----------
    n : int
        Number of tuples to return.
    random_state : int or ``random.Random``, optional
        If an int, is used to seed a new ``random.Random``.
    """
    if not isinstance(random_state, Random):
        random_state = Random(random_state)

    maxuint32 = 1 << 32
    return [
        (
            3,
            tuple(random_state.randint(0, maxuint32) for i in range(624)) + (624,),
            None,
        )
        for i in range(n)
    ]
```
I'm finding I'm spending a bunch of time in that list comprehension, so yes I was imagining we parallelize that loop. What's the concern?
> What's the concern?

Primarily I'm concerned about complexity around managing a threadpool for this. I wouldn't want to initialize a new pool for this for various reasons and passing an existing one to this function also feels cumbersome.

Instead, I would prefer us thinking about whether we should not modify our algorithm here. We're generating `NPartitions * 624` random numbers to create a deterministic `RandomState`, see https://github.com/dask/dask/pull/1687 to fix the issue here https://github.com/dask/dask/issues/1684
there was also an alternative proposal to fix this issue here (https://github.com/dask/dask/pull/1689) which was rejected because it had inferior statistical properties about reproducibility. We may want to revisit this decision before considering to parallelize this. If we went back to simply generating seeds this would no longer be an issue.

cc @jcrist @shoyer @crusaderky 

I see no reason not to accelerate this function as a wrapper around, or variant of, `dask.utils.random_state_data`, and only fall back to pure Python if numpy isn't installed. This would mean that the random sequence would differ when numpy is installed from when it isn't, but I doubt this will be a problem for anybody.

numpy is 12~18x faster:
```python
%timeit random_state_data_python(1, 0)
265 µs ± 1.72 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
%timeit [a.tolist() for a in random_state_data(1, 0)]
134 µs ± 4.17 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

%timeit random_state_data_python(10, 0)
2.64 ms ± 16.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit [a.tolist() for a in random_state_data(10, 0)]
225 µs ± 780 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%timeit random_state_data_python(100, 0)
26.5 ms ± 146 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit [a.tolist() for a in random_state_data(100, 0)]
1.37 ms ± 26.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%timeit random_state_data_python(1000, 0)
272 ms ± 1.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit [a.tolist() for a in random_state_data(1000, 0)]
16.8 ms ± 144 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```

Commit message: doctest

Patch to review:
diff --git a/dask/bag/core.py b/dask/bag/core.py
index 502f19e53a5..170fb5a0b9c 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -675,9 +675,9 @@ def random_sample(self, prob, random_state=None):
         >>> import dask.bag as db
         >>> b = db.from_sequence(range(5))
         >>> list(b.random_sample(0.5, 43))
-        [0, 3, 4]
+        [0, 1, 4]
         >>> list(b.random_sample(0.5, 43))
-        [0, 3, 4]
+        [0, 1, 4]
         """
         if not 0 <= prob <= 1:
             raise ValueError("prob must be a number in the interval [0, 1]")
@@ -2530,7 +2530,9 @@ def random_sample(x, state_data, prob):
             yield i
 
 
-def random_state_data_python(n, random_state=None):
+def random_state_data_python(
+    n: int, random_state: int | Random | None = None
+) -> list[tuple[int, tuple[int, ...], None]]:
     """Return a list of tuples that can be passed to
     ``random.Random.setstate``.
 
@@ -2540,19 +2542,37 @@ def random_state_data_python(n, random_state=None):
         Number of tuples to return.
     random_state : int or ``random.Random``, optional
         If an int, is used to seed a new ``random.Random``.
-    """
-    if not isinstance(random_state, Random):
-        random_state = Random(random_state)
 
+    See Also
+    --------
+    dask.utils.random_state_data
+    """
     maxuint32 = 1 << 32
-    return [
-        (
-            3,
-            tuple(random_state.randint(0, maxuint32) for i in range(624)) + (624,),
-            None,
-        )
-        for i in range(n)
-    ]
+
+    try:
+        import numpy as np
+
+        if isinstance(random_state, Random):
+            random_state = random_state.randint(0, maxuint32)
+        np_rng = np.random.RandomState(random_state)
+
+        random_data = np_rng.bytes(624 * n * 4)  # `n * 624` 32-bit integers
+        arr = np.frombuffer(random_data, dtype=np.uint32).reshape((n, -1))
+        return [(3, tuple(row) + (624,), None) for row in arr.tolist()]
+
+    except ImportError:
+        # Pure python (much slower)
+        if not isinstance(random_state, Random):
+            random_state = Random(random_state)
+
+        return [
+            (
+                3,
+                tuple(random_state.randint(0, maxuint32) for _ in range(624)) + (624,),
+                None,
+            )
+            for _ in range(n)
+        ]
 
 
 def split(seq, n):


Code review comments:
diff_hunk: @@ -2540,19 +2542,37 @@ def random_state_data_python(n, random_state=None):
         Number of tuples to return.
     random_state : int or ``random.Random``, optional
         If an int, is used to seed a new ``random.Random``.
-    """
-    if not isinstance(random_state, Random):
-        random_state = Random(random_state)
 
+    See Also
+    --------
+    dask.utils.random_state_data
+    """
     maxuint32 = 1 << 32
-    return [
-        (
-            3,
-            tuple(random_state.randint(0, maxuint32) for i in range(624)) + (624,),
-            None,
-        )
-        for i in range(n)
-    ]
+
+    try:
+        import numpy as np
+
+        if isinstance(random_state, Random):
+            random_state = random_state.randint(0, maxuint32)
+        np_rng = np.random.RandomState(random_state)
path: dask/bag/core.py
text: @user1:
numpy (and dask) `RandomState` is frozen and wont be getting any updates. Cant we use the new `Generator` interface here?

And if all we want is a bunch of reproducible parallel streams, perhaps we can move `_spawn_bitgens` from `da.random` to `dask.utils` and use that. Bonus would be getting rid of all those `624`s (Mersenne Twister detail) in the code.

Do let me know if I am not making any sense.

@user1:
On second look, perhaps just using `default_rng` instead of `RandomState` in line 2557 is enough. `624` is too deep in the stack to handle.

Feel free to ignore. Sorry

@author:
Sorry for the wait. I replaced RandomState with default_rng. Do you think this is sufficient to merge this PR?

@user1:
yes, good for merge afaic. thank you