BUG: Ensure seed sequences are restored through pickling (numpy#26260)

bashtage · web-flow · commit d567754fae4f · 2024-04-16T12:10:15.000+02:00
Explicity store and restore seed sequence closes numpy#26234 --- * BUG: Ensure seed sequences are restored through pickling Explicity store and restore seed sequence closes numpy#26234 * CLN: Simplify refactor Make more use of set and getstate to avoid changes in the pickling functions * BUG: Correct behavior for legacy pickles Add test for legacy pickles Include pickles for tests * MAINT: Correct types for pickle related functions * REF: Switch from string to type * REF: Swtich to returning bit generators Explicitly return bit generator rather than ctor
diff --git a/numpy/random/_generator.pyi b/numpy/random/_generator.pyi
@@ -68,9 +68,12 @@ class Generator:
     def __init__(self, bit_generator: BitGenerator) -> None: ...
     def __repr__(self) -> str: ...
     def __str__(self) -> str: ...
-    def __getstate__(self) -> dict[str, Any]: ...
-    def __setstate__(self, state: dict[str, Any]) -> None: ...
-    def __reduce__(self) -> tuple[Callable[[str], Generator], tuple[str], dict[str, Any]]: ...
+    def __getstate__(self) -> None: ...
+    def __setstate__(self, state: dict[str, Any] | None) -> None: ...
+    def __reduce__(self) -> tuple[
+        Callable[[BitGenerator], Generator],
+        tuple[BitGenerator],
+        None]: ...
     @property
     def bit_generator(self) -> BitGenerator: ...
     def spawn(self, n_children: int) -> list[Generator]: ...
diff --git a/numpy/random/_generator.pyx b/numpy/random/_generator.pyx
@@ -214,17 +214,19 @@ cdef class Generator:
 
     # Pickling support:
     def __getstate__(self):
-        return self.bit_generator.state
+        return None
 
-    def __setstate__(self, state):
-        self.bit_generator.state = state
+    def __setstate__(self, bit_gen):
+        if isinstance(bit_gen, dict):
+            # Legacy path
+            # Prior to 2.0.x only the state of the underlying bit generator
+            # was preserved and any seed sequence information was lost
+            self.bit_generator.state = bit_gen
 
     def __reduce__(self):
-        ctor, name_tpl, state = self._bit_generator.__reduce__()
-
         from ._pickle import __generator_ctor
-        # Requirements of __generator_ctor are (name, ctor)
-        return __generator_ctor, (name_tpl[0], ctor), state
+        # Requirements of __generator_ctor are (bit_generator, )
+        return __generator_ctor, (self._bit_generator, ), None
 
     @property
     def bit_generator(self):
diff --git a/numpy/random/_pickle.py b/numpy/random/_pickle.py
@@ -1,3 +1,4 @@
+from .bit_generator import BitGenerator
 from .mtrand import RandomState
 from ._philox import Philox
 from ._pcg64 import PCG64, PCG64DXSM
@@ -14,27 +15,30 @@
                  }
 
 
-def __bit_generator_ctor(bit_generator_name='MT19937'):
+def __bit_generator_ctor(bit_generator: str | type[BitGenerator] = 'MT19937'):
     """
     Pickling helper function that returns a bit generator object
 
     Parameters
     ----------
-    bit_generator_name : str
-        String containing the name of the BitGenerator
+    bit_generator : type[BitGenerator] or str
+        BitGenerator class or string containing the name of the BitGenerator
 
     Returns
     -------
-    bit_generator : BitGenerator
+    BitGenerator
         BitGenerator instance
     """
-    if bit_generator_name in BitGenerators:
-        bit_generator = BitGenerators[bit_generator_name]
+    if isinstance(bit_generator, type):
+        bit_gen_class = bit_generator
+    elif bit_generator in BitGenerators:
+        bit_gen_class = BitGenerators[bit_generator]
     else:
-        raise ValueError(str(bit_generator_name) + ' is not a known '
-                                                   'BitGenerator module.')
+        raise ValueError(
+            str(bit_generator) + ' is not a known BitGenerator module.'
+        )
 
-    return bit_generator()
+    return bit_gen_class()
 
 
 def __generator_ctor(bit_generator_name="MT19937",
@@ -44,8 +48,9 @@ def __generator_ctor(bit_generator_name="MT19937",
 
     Parameters
     ----------
-    bit_generator_name : str
-        String containing the core BitGenerator's name
+    bit_generator_name : str or BitGenerator
+        String containing the core BitGenerator's name or a
+        BitGenerator instance
     bit_generator_ctor : callable, optional
         Callable function that takes bit_generator_name as its only argument
         and returns an instantized bit generator.
@@ -55,6 +60,9 @@ def __generator_ctor(bit_generator_name="MT19937",
     rg : Generator
         Generator using the named core BitGenerator
     """
+    if isinstance(bit_generator_name, BitGenerator):
+        return Generator(bit_generator_name)
+    # Legacy path that uses a bit generator name and ctor
     return Generator(bit_generator_ctor(bit_generator_name))
 
 
@@ -76,5 +84,6 @@ def __randomstate_ctor(bit_generator_name="MT19937",
     rs : RandomState
         Legacy RandomState using the named core BitGenerator
     """
-
+    if isinstance(bit_generator_name, BitGenerator):
+        return RandomState(bit_generator_name)
     return RandomState(bit_generator_ctor(bit_generator_name))
diff --git a/numpy/random/bit_generator.pyi b/numpy/random/bit_generator.pyi
@@ -92,11 +92,17 @@ class SeedSequence(ISpawnableSeedSequence):
 class BitGenerator(abc.ABC):
     lock: Lock
     def __init__(self, seed: None | _ArrayLikeInt_co | SeedSequence = ...) -> None: ...
-    def __getstate__(self) -> dict[str, Any]: ...
-    def __setstate__(self, state: dict[str, Any]) -> None: ...
+    def __getstate__(self) -> tuple[dict[str, Any], ISeedSequence]: ...
+    def __setstate__(
+            self, state_seed_seq: dict[str, Any] | tuple[dict[str, Any], ISeedSequence]
+    ) -> None: ...
     def __reduce__(
         self,
-    ) -> tuple[Callable[[str], BitGenerator], tuple[str], tuple[dict[str, Any]]]: ...
+    ) -> tuple[
+        Callable[[str], BitGenerator],
+        tuple[str],
+        tuple[dict[str, Any], ISeedSequence]
+    ]: ...
     @abc.abstractmethod
     @property
     def state(self) -> Mapping[str, Any]: ...
diff --git a/numpy/random/bit_generator.pyx b/numpy/random/bit_generator.pyx
@@ -537,14 +537,27 @@ cdef class BitGenerator():
 
     # Pickling support:
     def __getstate__(self):
-        return self.state
+        return self.state, self._seed_seq
 
-    def __setstate__(self, state):
-        self.state = state
+    def __setstate__(self, state_seed_seq):
+
+        if isinstance(state_seed_seq, dict):
+            # Legacy path
+            # Prior to 2.0.x only the state of the underlying bit generator
+            # was preserved and any seed sequence information was lost
+            self.state = state_seed_seq
+        else:
+            self._seed_seq = state_seed_seq[1]
+            self.state = state_seed_seq[0]
 
     def __reduce__(self):
         from ._pickle import __bit_generator_ctor
-        return __bit_generator_ctor, (self.state['bit_generator'],), self.state
+
+        return (
+            __bit_generator_ctor,
+            (type(self), ),
+            (self.state, self._seed_seq)
+        )
 
     @property
     def state(self):
diff --git a/numpy/random/meson.build b/numpy/random/meson.build
@@ -139,6 +139,9 @@ py.install_sources(
     'tests/data/philox-testset-2.csv',
     'tests/data/sfc64-testset-1.csv',
     'tests/data/sfc64-testset-2.csv',
+    'tests/data/sfc64_np126.pkl.gz',
+    'tests/data/generator_pcg64_np126.pkl.gz',
+    'tests/data/generator_pcg64_np121.pkl.gz',
   ],
   subdir: 'numpy/random/tests/data'
 )
diff --git a/numpy/random/mtrand.pyi b/numpy/random/mtrand.pyi
@@ -73,7 +73,7 @@ class RandomState:
     def __str__(self) -> str: ...
     def __getstate__(self) -> dict[str, Any]: ...
     def __setstate__(self, state: dict[str, Any]) -> None: ...
-    def __reduce__(self) -> tuple[Callable[[str], RandomState], tuple[str], dict[str, Any]]: ...
+    def __reduce__(self) -> tuple[Callable[[BitGenerator], RandomState], tuple[BitGenerator], dict[str, Any]]: ...
     def seed(self, seed: None | _ArrayLikeFloat_co = ...) -> None: ...
     @overload
     def get_state(self, legacy: Literal[False] = ...) -> dict[str, Any]: ...
diff --git a/numpy/random/mtrand.pyx b/numpy/random/mtrand.pyx
@@ -205,10 +205,13 @@ cdef class RandomState:
         self.set_state(state)
 
     def __reduce__(self):
-        ctor, name_tpl, _ = self._bit_generator.__reduce__()
-
         from ._pickle import __randomstate_ctor
-        return __randomstate_ctor, (name_tpl[0], ctor), self.get_state(legacy=False)
+        # The third argument containing the state is required here since
+        # RandomState contains state information in addition to the state
+        # contained in the bit generator that described the gaussian
+        # generator. This argument is passed to __setstate__ after the
+        # Generator is created.
+        return __randomstate_ctor, (self._bit_generator, ), self.get_state(legacy=False)
 
     cdef _initialize_bit_generator(self, bit_generator):
         self._bit_generator = bit_generator
diff --git a/numpy/random/tests/data/generator_pcg64_np121.pkl.gz b/numpy/random/tests/data/generator_pcg64_np121.pkl.gz
diff --git a/numpy/random/tests/data/generator_pcg64_np126.pkl.gz b/numpy/random/tests/data/generator_pcg64_np126.pkl.gz
diff --git a/numpy/random/tests/data/sfc64_np126.pkl.gz b/numpy/random/tests/data/sfc64_np126.pkl.gz
diff --git a/numpy/random/tests/test_direct.py b/numpy/random/tests/test_direct.py
@@ -298,6 +298,24 @@ def test_pickle(self):
         aa = pickle.loads(pickle.dumps(ss))
         assert_equal(ss.state, aa.state)
 
+    def test_pickle_preserves_seed_sequence(self):
+        # GH 26234
+        # Add explicit test that bit generators preserve seed sequences
+        import pickle
+
+        bit_generator = self.bit_generator(*self.data1['seed'])
+        ss = bit_generator.seed_seq
+        bg_plk = pickle.loads(pickle.dumps(bit_generator))
+        ss_plk = bg_plk.seed_seq
+        assert_equal(ss.state, ss_plk.state)
+        assert_equal(ss.pool, ss_plk.pool)
+
+        bit_generator.seed_seq.spawn(10)
+        bg_plk = pickle.loads(pickle.dumps(bit_generator))
+        ss_plk = bg_plk.seed_seq
+        assert_equal(ss.state, ss_plk.state)
+        assert_equal(ss.n_children_spawned, ss_plk.n_children_spawned)
+
     def test_invalid_state_type(self):
         bit_generator = self.bit_generator(*self.data1['seed'])
         with pytest.raises(TypeError):
@@ -349,8 +367,9 @@ def test_getstate(self):
         bit_generator = self.bit_generator(*self.data1['seed'])
         state = bit_generator.state
         alt_state = bit_generator.__getstate__()
-        assert_state_equal(state, alt_state)
-
+        assert isinstance(alt_state, tuple)
+        assert_state_equal(state, alt_state[0])
+        assert isinstance(alt_state[1], SeedSequence)
 
 class TestPhilox(Base):
     @classmethod
@@ -413,6 +432,7 @@ def test_advange_large(self):
         assert state["state"] == advanced_state
 
 
+
 class TestPCG64DXSM(Base):
     @classmethod
     def setup_class(cls):
@@ -502,6 +522,29 @@ def setup_class(cls):
         cls.invalid_init_types = [(3.2,), ([None],), (1, None)]
         cls.invalid_init_values = [(-1,)]
 
+    def test_legacy_pickle(self):
+        # Pickling format was changed in 2.0.x
+        import gzip
+        import pickle
+
+        expected_state = np.array(
+            [
+                9957867060933711493,
+                532597980065565856,
+                14769588338631205282,
+                13
+            ],
+            dtype=np.uint64
+        )
+
+        base_path = os.path.split(os.path.abspath(__file__))[0]
+        pkl_file = os.path.join(base_path, "data", f"sfc64_np126.pkl.gz")
+        with gzip.open(pkl_file) as gz:
+            sfc = pickle.load(gz)
+
+        assert isinstance(sfc, SFC64)
+        assert_equal(sfc.state["state"]["state"], expected_state)
+
 
 class TestDefaultRNG:
     def test_seed(self):
diff --git a/numpy/random/tests/test_generator_mt19937.py b/numpy/random/tests/test_generator_mt19937.py
@@ -1,3 +1,4 @@
+import os.path
 import sys
 import hashlib
 
@@ -2738,10 +2739,50 @@ def test_generator_ctor_old_style_pickle():
     rg = np.random.Generator(np.random.PCG64DXSM(0))
     rg.standard_normal(1)
     # Directly call reduce which is used in pickling
-    ctor, args, state_a = rg.__reduce__()
+    ctor, (bit_gen, ), _ = rg.__reduce__()
     # Simulate unpickling an old pickle that only has the name
-    assert args[:1] == ("PCG64DXSM",)
-    b = ctor(*args[:1])
-    b.bit_generator.state = state_a
+    assert bit_gen.__class__.__name__ == "PCG64DXSM"
+    print(ctor)
+    b = ctor(*("PCG64DXSM",))
+    print(b)
+    b.bit_generator.state = bit_gen.state
     state_b = b.bit_generator.state
-    assert state_a == state_b
+    assert bit_gen.state == state_b
+
+
+def test_pickle_preserves_seed_sequence():
+    # GH 26234
+    # Add explicit test that bit generators preserve seed sequences
+    import pickle
+
+    rg = np.random.Generator(np.random.PCG64DXSM(20240411))
+    ss = rg.bit_generator.seed_seq
+    rg_plk = pickle.loads(pickle.dumps(rg))
+    ss_plk = rg_plk.bit_generator.seed_seq
+    assert_equal(ss.state, ss_plk.state)
+    assert_equal(ss.pool, ss_plk.pool)
+
+    rg.bit_generator.seed_seq.spawn(10)
+    rg_plk = pickle.loads(pickle.dumps(rg))
+    ss_plk = rg_plk.bit_generator.seed_seq
+    assert_equal(ss.state, ss_plk.state)
+
+
+@pytest.mark.parametrize("version", [121, 126])
+def test_legacy_pickle(version):
+    # Pickling format was changes in 1.22.x and in 2.0.x
+    import pickle
+    import gzip
+
+    base_path = os.path.split(os.path.abspath(__file__))[0]
+    pkl_file = os.path.join(
+        base_path, "data", f"generator_pcg64_np{version}.pkl.gz"
+    )
+    with gzip.open(pkl_file) as gz:
+        rg = pickle.load(gz)
+    state = rg.bit_generator.state['state']
+
+    assert isinstance(rg, Generator)
+    assert isinstance(rg.bit_generator, np.random.PCG64)
+    assert state['state'] == 35399562948360463058890781895381311971
+    assert state['inc'] == 87136372517582989555478159403783844777
diff --git a/numpy/random/tests/test_randomstate.py b/numpy/random/tests/test_randomstate.py
@@ -2052,8 +2052,8 @@ def test_randomstate_ctor_old_style_pickle():
     # Directly call reduce which is used in pickling
     ctor, args, state_a = rs.__reduce__()
     # Simulate unpickling an old pickle that only has the name
-    assert args[:1] == ("MT19937",)
-    b = ctor(*args[:1])
+    assert args[0].__class__.__name__ == "MT19937"
+    b = ctor(*("MT19937",))
     b.set_state(state_a)
     state_b = b.get_state(legacy=False)
 
diff --git a/numpy/typing/tests/data/pass/random.py b/numpy/typing/tests/data/pass/random.py
@@ -911,9 +911,7 @@
 
 def_gen.__str__()
 def_gen.__repr__()
-def_gen_state: dict[str, Any]
-def_gen_state = def_gen.__getstate__()
-def_gen.__setstate__(def_gen_state)
+def_gen.__setstate__(dict(def_gen.bit_generator.state))
 
 # RandomState
 random_st: np.random.RandomState = np.random.RandomState()
diff --git a/numpy/typing/tests/data/reveal/random.pyi b/numpy/typing/tests/data/reveal/random.pyi
@@ -953,9 +953,7 @@ assert_type(def_gen.shuffle(D_2D, axis=1), None)
 assert_type(np.random.Generator(pcg64), np.random.Generator)
 assert_type(def_gen.__str__(), str)
 assert_type(def_gen.__repr__(), str)
-def_gen_state = def_gen.__getstate__()
-assert_type(def_gen_state, dict[str, Any])
-assert_type(def_gen.__setstate__(def_gen_state), None)
+assert_type(def_gen.__setstate__(dict(def_gen.bit_generator.state)), None)
 
 # RandomState
 random_st: np.random.RandomState = np.random.RandomState()
