Resume sampling from existing ZarrTrace

Author: lucianopaz

pymc/backends/__init__.py

@@ -72,7 +72,7 @@
 from pymc.backends.arviz import predictions_to_inference_data, to_inference_data
 from pymc.backends.base import BaseTrace, IBaseTrace
 from pymc.backends.ndarray import NDArray
-from pymc.backends.zarr import ZarrTrace
+from pymc.backends.zarr import TraceAlreadyInitialized, ZarrTrace
 from pymc.blocking import PointType
 from pymc.model import Model
 from pymc.step_methods.compound import BlockedStep, CompoundStep
@@ -132,15 +132,41 @@ def init_traces(
 ) -> tuple[RunType | None, Sequence[IBaseTrace]]:
     """Initialize a trace recorder for each chain."""
     if isinstance(backend, ZarrTrace):
-        backend.init_trace(
-            chains=chains,
-            draws=expected_length - tune,
-            tune=tune,
-            step=step,
-            model=model,
-            vars=trace_vars,
-            test_point=initial_point,
-        )
+        try:
+            backend.init_trace(
+                chains=chains,
+                draws=expected_length - tune,
+                tune=tune,
+                step=step,
+                model=model,
+                vars=trace_vars,
+                test_point=initial_point,
+            )
+        except TraceAlreadyInitialized:
+            # Trace has already been initialized. We need to make sure that the
+            # tracked variable names and the number of chains match, and then resize
+            # the zarr groups to the desired number of draws and tune.
+            backend.assert_model_and_step_are_compatible(
+                step=step,
+                model=model,
+                vars=trace_vars,
+            )
+            assert backend.posterior.chain.size == chains, (
+                f"The requested number of chains {chains} does not match the number "
+                f"of chains stored in the trace ({backend.posterior.chain.size})."
+            )
+            vars, var_names = backend.parse_varnames(model=model, vars=trace_vars)
+            backend.link_model_and_step(
+                chains=chains,
+                draws=expected_length - tune,
+                tune=tune,
+                step=step,
+                model=model,
+                vars=vars,
+                var_names=var_names,
+                test_point=initial_point,
+            )
+            backend.resize(tune=tune, draws=expected_length - tune)
         return None, backend.straces
     if HAS_MCB and isinstance(backend, Backend):
         return init_chain_adapters(

pymc/sampling/mcmc.py

@@ -767,7 +767,7 @@ def joined_blas_limiter():
             UserWarning,
         )
     rngs = get_random_generator(random_seed).spawn(chains)
-    random_seed_list = [rng.integers(2**30) for rng in rngs]
+    random_seed_list: list[int] = [int(rng.integers(2**30)) for rng in rngs]
 
     if not discard_tuned_samples and not return_inferencedata and not isinstance(trace, ZarrTrace):
         warnings.warn(
@@ -993,11 +993,8 @@ def _sample_return(
     Final step of `pm.sampler`.
     """
     if isinstance(traces, ZarrTrace):
-        # Split warmup from posterior samples
-        traces.split_warmup_groups()
-
         # Set sampling time
-        traces.sampling_time = t_sampling
+        traces.sampling_time = traces.sampling_time + t_sampling
 
         # Compute number of actual draws per chain
         total_draws_per_chain = traces._sampling_state.draw_idx[:]
@@ -1226,7 +1223,7 @@ def _sample(
         callback=callback,
     )
     try:
-        for it, stats in enumerate(sampling_gen):
+        for it, stats in sampling_gen:
             progress_manager.update(
                 chain_idx=chain, is_last=False, draw=it, stats=stats, tuning=it > tune
             )
@@ -1251,7 +1248,7 @@ def _iter_sample(
     rng: np.random.Generator,
     model: Model | None = None,
     callback: SamplingIteratorCallback | None = None,
-) -> Iterator[list[dict[str, Any]]]:
+) -> Iterator[tuple[int, list[dict[str, Any]]]]:
     """Sample one chain with a generator (singleprocess).
 
     Parameters
@@ -1285,14 +1282,29 @@ def _iter_sample(
     step.set_rng(rng)
 
     point = start
-    if isinstance(trace, ZarrChain):
-        trace.link_stepper(step)
+    initial_draw_idx = 0
+    step.tune = bool(tune)
+    if hasattr(step, "reset_tuning"):
+        step.reset_tuning()
+    trace.link_stepper(step)
+    stored_draw_idx, stored_sampling_state = trace.get_stored_draw_and_state()
+    if stored_draw_idx > 0:
+        if stored_sampling_state is not None:
+            step.sampling_state = stored_sampling_state
+        else:
+            raise RuntimeError(
+                "Cannot use the supplied ZarrTrace to restart sampling because "
+                "it has no sampling_state information stored. You will have to "
+                "resample from scratch."
+            )
+        initial_draw_idx = stored_draw_idx
+        point = trace.get_mcmc_point()
+    else:
+        # Store initial point in trace
+        trace.set_mcmc_point(point)
 
     try:
-        step.tune = bool(tune)
-        if hasattr(step, "reset_tuning"):
-            step.reset_tuning()
-        for i in range(draws):
+        for i in range(initial_draw_idx, draws):
             if i == 0 and hasattr(step, "iter_count"):
                 step.iter_count = 0
             if i == tune:
@@ -1308,17 +1320,15 @@ def _iter_sample(
                     draw=Draw(chain, i == draws, i, i < tune, stats, point),
                 )
 
-            yield stats
+            yield i, stats
 
     except (KeyboardInterrupt, BaseException):
-        if isinstance(trace, ZarrChain):
-            trace.record_sampling_state(step=step)
+        trace.record_sampling_state(step=step)
         trace.close()
         raise
 
     else:
-        if isinstance(trace, ZarrChain):
-            trace.record_sampling_state(step=step)
+        trace.record_sampling_state(step=step)
         trace.close()
 
 
@@ -1377,7 +1387,6 @@ def _mp_sample(
 
     # We did draws += tune in pm.sample
     draws -= tune
-    zarr_chains: list[ZarrChain] | None = None
     zarr_recording = False
     if all(isinstance(trace, ZarrChain) for trace in traces):
         if isinstance(cast(ZarrChain, traces[0])._posterior.store, MemoryStore):
@@ -1388,7 +1397,6 @@ def _mp_sample(
                 "DirectoryStore or ZipStore"
             )
         else:
-            zarr_chains = cast(list[ZarrChain], traces)
             zarr_recording = True
 
     sampler = ps.ParallelSampler(
@@ -1403,7 +1411,9 @@ def _mp_sample(
         progressbar_theme=progressbar_theme,
         blas_cores=blas_cores,
         mp_ctx=mp_ctx,
-        zarr_chains=zarr_chains,
+        # We only need to pass the traces when zarr_recording is happening because
+        # it's the only backend that can resume sampling
+        traces=traces if zarr_recording else None,
     )
     try:
         try:

pymc/sampling/parallel.py

@@ -30,6 +30,7 @@
 from rich.theme import Theme
 from threadpoolctl import threadpool_limits
 
+from pymc.backends.base import IBaseTrace
 from pymc.backends.zarr import ZarrChain
 from pymc.blocking import DictToArrayBijection
 from pymc.exceptions import SamplingError
@@ -105,8 +106,9 @@ def __init__(
         rng_state: RandomGeneratorState,
         blas_cores,
         chain: int,
-        zarr_chains: list[ZarrChain] | bytes | None = None,
-        zarr_chains_is_pickled: bool = False,
+        traces: list[IBaseTrace] | bytes | None = None,
+        traces_is_pickled: bool = False,
+        zarr_recording: bool = False,
     ):
         # Because of https://github.com/numpy/numpy/issues/27727, we can't send
         # the rng instance to the child process because pickling (copying) looses
@@ -117,13 +119,12 @@ def __init__(
         self._step_method = step_method
         self._step_method_is_pickled = step_method_is_pickled
         self.chain = chain
-        self._zarr_recording = False
-        self._zarr_chain: ZarrChain | None = None
-        if zarr_chains_is_pickled:
-            self._zarr_chain = cloudpickle.loads(zarr_chains)[self.chain]
-        elif zarr_chains is not None:
-            self._zarr_chain = cast(list[ZarrChain], zarr_chains)[self.chain]
-        self._zarr_recording = self._zarr_chain is not None
+        self._zarr_recording = zarr_recording
+        self._trace: IBaseTrace | None = None
+        if traces_is_pickled:
+            self._trace = cloudpickle.loads(traces)[self.chain]
+        elif traces is not None:
+            self._trace = cast(list[IBaseTrace], traces)[self.chain]
 
         self._shared_point = shared_point
         self._rng = rng
@@ -165,7 +166,7 @@ def run(self):
 
     def _link_step_to_zarrchain(self):
         if self._zarr_recording:
-            self._zarr_chain.link_stepper(self._step_method)
+            self._trace.link_stepper(self._step_method)
 
     def _wait_for_abortion(self):
         while True:
@@ -194,6 +195,24 @@ def _start_loop(self):
 
         draw = 0
         tuning = True
+        if self._zarr_recording:
+            trace = self._trace
+            stored_draw_idx = trace._sampling_state.draw_idx[self.chain]
+            stored_sampling_state = trace._sampling_state.sampling_state[self.chain]
+            if stored_draw_idx > 0:
+                if stored_sampling_state is not None:
+                    self._step_method.sampling_state = stored_sampling_state
+                else:
+                    raise RuntimeError(
+                        "Cannot use the supplied ZarrTrace to restart sampling because "
+                        "it has no sampling_state information stored. You will have to "
+                        "resample from scratch."
+                    )
+                draw = stored_draw_idx
+                self._write_point(trace.get_mcmc_point())
+            else:
+                # Store starting point in trace's mcmc_point
+                trace.set_mcmc_point(self._point)
 
         msg = self._recv_msg()
         if msg[0] == "abort":
@@ -220,7 +239,7 @@ def _start_loop(self):
                 raise KeyboardInterrupt()
             elif msg[0] == "write_next":
                 if zarr_recording:
-                    self._zarr_chain.record(point, stats)
+                    self._trace.record(point, stats)
                 self._write_point(point)
                 is_last = draw + 1 == self._draws + self._tune
                 self._msg_pipe.send(("writing_done", is_last, draw, tuning, stats))
@@ -247,8 +266,9 @@ def __init__(
         start: dict[str, np.ndarray],
         blas_cores,
         mp_ctx,
-        zarr_chains: list[ZarrChain] | None = None,
-        zarr_chains_pickled: bytes | None = None,
+        traces: Sequence[IBaseTrace] | None = None,
+        traces_pickled: bytes | None = None,
+        zarr_recording: bool = False,
     ):
         self.chain = chain
         process_name = f"worker_chain_{chain}"
@@ -271,15 +291,15 @@ def __init__(
         self._readable = True
         self._num_samples = 0
 
-        zarr_chains_send: list[ZarrChain] | bytes | None = None
-        if zarr_chains_pickled is not None:
-            zarr_chains_send = zarr_chains_pickled
-        elif zarr_chains is not None:
+        traces_send: Sequence[IBaseTrace] | bytes | None = None
+        if traces_pickled is not None:
+            traces_send = traces_pickled
+        elif traces is not None:
             if mp_ctx.get_start_method() == "spawn":
                 raise ValueError(
-                    "please provide a pre-pickled zarr_chains when multiprocessing start method is 'spawn'"
+                    "please provide a pre-pickled traces when multiprocessing start method is 'spawn'"
                 )
-            zarr_chains_send = zarr_chains
+            traces_send = traces
 
         if step_method_pickled is not None:
             step_method_send = step_method_pickled
@@ -305,8 +325,9 @@ def __init__(
                 get_state_from_generator(rng),
                 blas_cores,
                 self.chain,
-                zarr_chains_send,
-                zarr_chains_pickled is not None,
+                traces_send,
+                traces_pickled is not None,
+                zarr_recording,
             ),
         )
         self._process.start()
@@ -429,7 +450,7 @@ def __init__(
         progressbar_theme: Theme | None = default_progress_theme,
         blas_cores: int | None = None,
         mp_ctx=None,
-        zarr_chains: list[ZarrChain] | None = None,
+        traces: Sequence[IBaseTrace] | None = None,
     ):
         if any(len(arg) != chains for arg in [rngs, start_points]):
             raise ValueError(f"Number of rngs and start_points must be {chains}.")
@@ -450,15 +471,14 @@ def __init__(
             mp_ctx = multiprocessing.get_context(mp_ctx)
 
         step_method_pickled = None
-        zarr_chains_pickled = None
-        self.zarr_recording = False
-        if zarr_chains is not None:
-            assert all(isinstance(zarr_chain, ZarrChain) for zarr_chain in zarr_chains)
-            self.zarr_recording = True
+        traces_pickled = None
+        self.zarr_recording = traces is not None and all(
+            isinstance(trace, ZarrChain) for trace in traces
+        )
         if mp_ctx.get_start_method() != "fork":
             step_method_pickled = cloudpickle.dumps(step_method, protocol=-1)
-            if zarr_chains is not None:
-                zarr_chains_pickled = cloudpickle.dumps(zarr_chains, protocol=-1)
+            if traces is not None:
+                traces_pickled = cloudpickle.dumps(traces, protocol=-1)
 
         self._samplers = [
             ProcessAdapter(
@@ -471,8 +491,9 @@ def __init__(
                 start,
                 blas_cores,
                 mp_ctx,
-                zarr_chains=zarr_chains,
-                zarr_chains_pickled=zarr_chains_pickled,
+                traces=traces,
+                traces_pickled=traces_pickled,
+                zarr_recording=self.zarr_recording,
             )
             for chain, rng, start in zip(range(chains), rngs, start_points)
         ]