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

@@ -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,33 @@ def _iter_sample(
     step.set_rng(rng)
 
     point = start
+    initial_draw_idx = 0
+    step.tune = bool(tune)
+    if hasattr(step, "reset_tuning"):
+        step.reset_tuning()
     if isinstance(trace, ZarrChain):
         trace.link_stepper(step)
+        stored_draw_idx = trace._sampling_state.draw_idx[chain]
+        stored_sampling_state = trace._sampling_state.sampling_state[chain]
+        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):
+            diverging = False
+
             if i == 0 and hasattr(step, "iter_count"):
                 step.iter_count = 0
             if i == tune:
@@ -1308,7 +1324,7 @@ 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):

pymc/sampling/parallel.py

@@ -194,6 +194,24 @@ def _start_loop(self):
 
         draw = 0
         tuning = True
+        if self._zarr_recording:
+            trace = self._zarr_chain
+            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":

pymc/sampling/population.py

@@ -19,7 +19,7 @@
 
 from collections.abc import Iterator, Sequence
 from copy import copy
-from typing import TypeAlias
+from typing import TypeAlias, cast
 
 import cloudpickle
 import numpy as np
@@ -37,7 +37,7 @@
     PopulationArrayStepShared,
     StatsType,
 )
-from pymc.step_methods.compound import StepMethodState
+from pymc.step_methods.compound import CompoundStepState, StepMethodState
 from pymc.step_methods.metropolis import DEMetropolis
 from pymc.util import CustomProgress
 
@@ -54,7 +54,7 @@ def _sample_population(
     *,
     initial_points: Sequence[PointType],
     draws: int,
-    start: Sequence[PointType],
+    start: list[PointType],
     rngs: Sequence[np.random.Generator],
     step: BlockedStep | CompoundStep,
     tune: int,
@@ -151,7 +151,9 @@ def warn_population_size(
 class PopulationStepper:
     """Wraps population of step methods to step them in parallel with single or multiprocessing."""
 
-    def __init__(self, steppers, parallelize: bool, progressbar: bool = True):
+    def __init__(
+        self, steppers, parallelize: bool, progressbar: bool = True, first_draw_idx: int = 0
+    ):
         """Use multiprocessing to parallelize chains.
 
         Falls back to sequential evaluation if multiprocessing fails.
@@ -330,7 +332,7 @@ def _prepare_iter_population(
     *,
     draws: int,
     step,
-    start: Sequence[PointType],
+    start: list[PointType],
     parallelize: bool,
     traces: Sequence[BaseTrace],
     tune: int,
@@ -376,10 +378,35 @@ def _prepare_iter_population(
         raise ValueError("Argument `draws` should be above 0.")
 
     # The initialization of traces, samplers and points must happen in the right order:
+    # 0. previous sampling state is loaded if possible
     # 1. population of points is created
     # 2. steppers are initialized and linked to the points object
     # 3. a PopulationStepper is configured for parallelized stepping
 
+    # 0. load sampling state and start point from traces if possible
+    first_draw_idx = 0
+    stored_sampling_states: Sequence[StepMethodState | CompoundStepState] | None = None
+    can_resume_sampling = False
+    if isinstance(traces[0], ZarrChain):
+        # All traces share the same store. This lets us load the past sampling states and draw
+        # indices for all chain
+        stored_draw_idxs = traces[0]._sampling_state.draw_idx[:]
+        stored_sampling_states = cast(
+            Sequence[StepMethodState | CompoundStepState],
+            traces[0]._sampling_state.sampling_state[:],
+        )
+        can_resume_sampling = (
+            all(stored_draw_idxs > 0)
+            and all(stored_draw_idxs == stored_draw_idxs[0])
+            and all(sampling_state is not None for sampling_state in stored_sampling_states)
+        )
+        for chain, trace in enumerate(traces):
+            trace = cast(ZarrChain, trace)
+            if can_resume_sampling:
+                start[chain] = trace.get_mcmc_point()
+            else:
+                trace.set_mcmc_point(start[chain])
+
     # 1. create a population (points) that tracks each chain
     # it is updated as the chains are advanced
     population = [start[c] for c in range(nchains)]
@@ -401,15 +428,25 @@ def _prepare_iter_population(
         for sm in chainstep.methods if isinstance(step, CompoundStep) else [chainstep]:
             if isinstance(sm, PopulationArrayStepShared):
                 sm.link_population(population, c)
+        if can_resume_sampling:
+            chainstep.sampling_state = cast(Sequence[CompoundStepState], stored_sampling_states)[c]
         steppers.append(chainstep)
 
     # 3. configure the PopulationStepper (expensive call)
-    popstep = PopulationStepper(steppers, parallelize, progressbar=progressbar)
+    popstep = PopulationStepper(
+        steppers, parallelize, progressbar=progressbar, first_draw_idx=first_draw_idx
+    )
 
     # Because the preparations above are expensive, the actual iterator is
     # in another method. This way the progbar will not be disturbed.
     return _iter_population(
-        draws=draws, tune=tune, popstep=popstep, steppers=steppers, traces=traces, points=population
+        draws=draws,
+        tune=tune,
+        popstep=popstep,
+        steppers=steppers,
+        traces=traces,
+        points=population,
+        first_draw_idx=first_draw_idx,
     )
 
 
@@ -421,6 +458,7 @@ def _iter_population(
     steppers,
     traces: Sequence[BaseTrace],
     points,
+    first_draw_idx=0,
 ) -> Iterator[int]:
     """Iterate a ``PopulationStepper``.
 
@@ -450,7 +488,7 @@ def _iter_population(
     try:
         with popstep:
             # iterate draws of all chains
-            for i in range(draws):
+            for i in range(first_draw_idx, draws):
                 # this call steps all chains and returns a list of (point, stats)
                 # the `popstep` may interact with subprocesses internally
                 updates = popstep.step(i == tune, points)