NetKet 3.21.0__tar.gz → 3.22.0.dev1__tar.gz

{netket-3.21.0 → netket-3.22.0.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.21.0
+Version: 3.22.0.dev1
 Summary: Netket : Machine Learning toolbox for many-body quantum systems.
 Project-URL: homepage, https://www.netket.org
 Project-URL: documentation, https://netket.readthedocs.io/en/latest/#
@@ -20,7 +20,7 @@ Classifier: Topic :: Scientific/Engineering :: Physics
 Requires-Python: >=3.11
 Requires-Dist: einops>=0.8.1
 Requires-Dist: equinox>=0.11.3
-Requires-Dist: flax>=0.10.2
+Requires-Dist: flax>=0.10.6
 Requires-Dist: igraph>=0.10.0
 Requires-Dist: jax>=0.7.0
 Requires-Dist: numba>=0.59
@@ -50,14 +50,14 @@ Requires-Dist: ruff==0.14.1; extra == 'dev'
 Requires-Dist: wheel>=0.38.1; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: ipywidgets>=8.1.7; extra == 'docs'
-Requires-Dist: myst-nb<1.4,>=1.0; extra == 'docs'
-Requires-Dist: myst-parser<4.1.0,>=2.0.0; extra == 'docs'
+Requires-Dist: myst-nb<1.5,>=1.0; extra == 'docs'
+Requires-Dist: myst-parser<5.1.0,>=2.0.0; extra == 'docs'
 Requires-Dist: nbconvert>=7.16.6; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints<4.0,>=1.22; extra == 'docs'
 Requires-Dist: sphinx-book-theme~=1.1.0; extra == 'docs'
 Requires-Dist: sphinx-copybutton>=0.5.2; extra == 'docs'
 Requires-Dist: sphinx-design>=0.4.0; extra == 'docs'
-Requires-Dist: sphinx<9.1.0,>=8.2.3; extra == 'docs'
+Requires-Dist: sphinx<9.2.0,>=8.2.3; extra == 'docs'
 Requires-Dist: sphinxcontrib-fulltoc~=1.2.0; extra == 'docs'
 Requires-Dist: sphinxcontrib-jsmath~=1.0.1; extra == 'docs'
 Provides-Extra: extra

netket-3.22.0.dev1/netket/_src/callbacks/auto_chunk_size.py

@@ -0,0 +1,220 @@
+import jax
+from jax.errors import JaxRuntimeError
+
+from netket.utils import struct, timing, dispatch
+
+from netket._src.callbacks.base import AbstractCallback
+
+
+@dispatch.dispatch
+def get_forward_operator(driver):
+    """
+    Returns an operator whose expectation value computes the forward
+    pass for the driver.
+
+    This is a dispatch function: register new overloads with
+    ``@get_forward_operator.dispatch`` for custom driver types.
+
+    Example:
+       class VMC(AbstractVariationalDriver):
+            ...
+
+        @get_forward_operator.dispatch
+        def _get_forward_operator_vmc(driver: VMC):
+            return driver.hamiltonian
+    """
+    raise NotImplementedError(
+        f"get_forward_operator not implemented for {type(driver)}"
+    )
+
+
+class AutoChunkSize(AbstractCallback):
+    """
+    Automatically tunes the chunk size for the (i) sampler, (ii) forward pass and
+    (iii) backward pass of a driver and a variational state in order to avoid OOM errors.
+
+    This callback is useful when the optimal chunk size is not known in advance,
+    and the user wants to avoid manually tuning it.
+
+    The callback will try to run the sampler, forward pass and backward pass with
+    decreasing chunk sizes until it finds one that fits in memory. The chunk size
+    is halved on each failure, down to ``minimum_chunk_size``.
+
+    After the largest chunk size that works is found, the callback will print those
+    values and store them in the callback object. If you reuse the same callback object
+    in a subsequent run, it will reuse the same chunk sizes.
+
+    .. rubric:: Supporting custom drivers
+
+    To tune the forward pass, the callback needs to know which operator to use to
+    estimate the loss. This is determined by the :func:`get_forward_operator`
+    dispatch function, which is already registered for the built-in drivers
+    (:class:`~netket.driver.VMC`, :class:`~netket.driver.VMC_SR`,
+    :class:`~netket.driver.Infidelity_SR`).
+
+    If you implement a custom driver, you must register a new overload so that
+    ``AutoChunkSize`` knows how to probe the forward pass:
+
+    .. code-block:: python
+
+        from netket._src.callbacks.auto_chunk_size import get_forward_operator
+
+        class MyDriver(AbstractVariationalDriver):
+            def __init__(self, hamiltonian, ...):
+                self._ham = hamiltonian
+                ...
+
+        @get_forward_operator.dispatch
+        def _(driver: MyDriver):
+            return driver._ham
+
+    The overload should return an operator whose expectation value exercises the
+    same forward-pass code path as the actual optimisation step. The registration
+    is best placed at the bottom of the file that defines the driver class, so
+    it takes effect as soon as the driver is imported.
+    """
+
+    minimum_chunk_size: int = struct.field(pytree_node=False, serialize=False)
+    sampler_chunk_size: int | None = struct.field(pytree_node=False, serialize=False)
+    chunk_size: int | None = struct.field(pytree_node=False, serialize=False)
+    chunk_size_bwd: int | None = struct.field(pytree_node=False, serialize=False)
+
+    def __init__(
+        self,
+        sampler_chunk_size: int | None = None,
+        chunk_size: int | None = None,
+        chunk_size_bwd: int | None = None,
+        minimum_chunk_size: int = 1,
+    ):
+        """
+        Initialize the callback.
+
+        Args:
+            sampler_chunk_size: The initial chunk size to use for the sampler (default: None).
+                If None, the callback will try to find a chunk size that works starting from
+                the number of chains per rank.
+            chunk_size: The initial chunk size to use for the forward pass (default: None).
+                If None, the callback will try to find a chunk size that works starting from
+                the number of samples.
+            chunk_size_bwd: The initial chunk size to use for the backward pass (default: None).
+                If None, the callback will try to find a chunk size that works starting from
+                the number of samples.
+            minimum_chunk_size: The minimum chunk size to use (default: 1). If the chunk size
+                is decreased to this value and it still fails, the callback will raise an error.
+        """
+        self.sampler_chunk_size = sampler_chunk_size
+        self.chunk_size = chunk_size
+        self.chunk_size_bwd = chunk_size_bwd
+
+        self.minimum_chunk_size = minimum_chunk_size
+
+    def on_run_start(self, step, driver):
+        if jax.process_index() == 0:
+            print("Automatically tuning chunk size....")
+
+        # Try to check if sampling works
+        set_sampler_chunk_size = False
+        sampler_chunk_size = None
+
+        if hasattr(driver.state, "sampler"):
+            if hasattr(driver.state.sampler, "chunk_size"):
+                if self.sampler_chunk_size is not None:
+                    sampler_chunk_size = self.sampler_chunk_size
+                    driver.state.sampler = driver.state.sampler.replace(
+                        chunk_size=self.sampler_chunk_size
+                    )
+                while True:
+                    try:
+                        with timing.timed_scope(force=True) as timer:
+                            timer.block_until_ready(driver.state.samples)
+                    except JaxRuntimeError:
+                        set_sampler_chunk_size = True
+                        sampler_chunk_size = driver.state.sampler.chunk_size
+                        if sampler_chunk_size is None:
+                            sampler_chunk_size = driver.state.sampler.n_chains_per_rank
+
+                        print(
+                            f"- Sampling failed with {sampler_chunk_size = }, trying with {sampler_chunk_size // 2 }"
+                        )
+                        sampler_chunk_size = sampler_chunk_size // 2
+                        if sampler_chunk_size < self.minimum_chunk_size:
+                            raise ValueError("Sampling failed with minimum chunk size")
+                        driver.state.sampler = driver.state.sampler.replace(
+                            chunk_size=sampler_chunk_size
+                        )
+                    else:
+                        print(f"- Sampling succeeded with {sampler_chunk_size = }")
+                        break
+                self.sampler_chunk_size = sampler_chunk_size
+        # Try to check if forward pass works
+        set_vstate_chunk_size = False
+        vstate_chunk_size = None
+        if hasattr(driver.state, "chunk_size"):
+            operator = get_forward_operator(driver)
+            if self.chunk_size is not None:
+                vstate_chunk_size = self.chunk_size
+                driver.state.chunk_size = self.chunk_size
+            while True:
+                try:
+                    with timing.timed_scope(force=True) as timer:
+                        timer.block_until_ready(driver.estimate(operator))
+
+                except JaxRuntimeError:
+                    set_vstate_chunk_size = True
+                    vstate_chunk_size = driver.state.chunk_size
+                    if vstate_chunk_size is None:
+                        vstate_chunk_size = driver.state.n_samples
+
+                    print(
+                        f"- Forward pass failed with {vstate_chunk_size = }, trying with {vstate_chunk_size // 2 }"
+                    )
+                    vstate_chunk_size = vstate_chunk_size // 2
+                    if vstate_chunk_size < self.minimum_chunk_size:
+                        raise ValueError("Forward pass failed with minimum chunk size")
+
+                    driver.state.chunk_size = vstate_chunk_size
+                else:
+                    print(f"- Forward pass succeeded with {vstate_chunk_size = }")
+                    break
+            self.chunk_size = vstate_chunk_size
+        # backward
+        set_bwd_chunk_size = False
+        chunk_size_bwd = None
+        if hasattr(driver, "chunk_size_bwd"):
+            driver.reset_step()
+            if self.chunk_size_bwd is not None:
+                chunk_size_bwd = self.chunk_size_bwd
+                driver.chunk_size_bwd = self.chunk_size_bwd
+
+            while True:
+                try:
+                    with timing.timed_scope(force=True) as timer:
+                        timer.block_until_ready(driver.compute_loss_and_update())
+
+                except JaxRuntimeError:
+                    set_bwd_chunk_size = True
+                    chunk_size_bwd = driver.chunk_size_bwd
+                    if chunk_size_bwd is None:
+                        chunk_size_bwd = driver.state.n_samples
+
+                    print(
+                        f"- Optimisation step failed with {chunk_size_bwd = }, trying with {chunk_size_bwd // 2 }"
+                    )
+                    chunk_size_bwd = chunk_size_bwd // 2
+                    if chunk_size_bwd < self.minimum_chunk_size:
+                        raise ValueError(
+                            "Optimisation pass failed with minimum chunk size"
+                        )
+                    driver.chunk_size_bwd = chunk_size_bwd
+                else:
+                    print(f"- Optimisation step succeeded with {chunk_size_bwd = }")
+                    break
+            self.chunk_size_bwd = chunk_size_bwd
+
+        print("Chunk size tuning complete:")
+        if set_sampler_chunk_size:
+            print(f"  {sampler_chunk_size = }")
+        if set_vstate_chunk_size:
+            print(f"  {vstate_chunk_size = }")
+        if set_bwd_chunk_size:
+            print(f"  {chunk_size_bwd = }")

netket-3.22.0.dev1/netket/_src/callbacks/auto_slurm_requeue.py

@@ -0,0 +1,171 @@
+import subprocess
+import os
+from datetime import timedelta, datetime
+
+import jax
+
+from netket.utils import struct
+
+from netket._src.callbacks.base import AbstractCallback
+
+
+def get_time_left(job_id):
+    """Queries the remaining time for the current Slurm job using the SLURM_JOB_ID environment variable."""
+    # Retrieve the job ID from the environment variable
+    try:
+        # Run the squeue command to fetch the remaining time
+        result = subprocess.run(
+            ["squeue", "-h", "-j", str(job_id), "-O", "TimeLeft"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        # Extract the time left from the command's output
+        time_left_str = result.stdout.strip()
+        if not time_left_str or time_left_str == "UNLIMITED":
+            # print(f"Job {job_id} has unlimited or no time limit.")
+            return None
+
+        # Parse the time left dynamically
+        if "-" in time_left_str:
+            days, time_str = time_left_str.split("-")
+            days = int(days)
+        else:
+            days = 0
+            time_str = time_left_str
+
+        # Split the time string and pad missing fields with zeros
+        time_parts = list(map(int, time_str.split(":")))
+        while len(time_parts) < 3:  # Ensure [hours, minutes, seconds]
+            time_parts.insert(0, 0)
+
+        hours, minutes, seconds = time_parts
+
+        # Calculate total seconds and convert to timedelta
+        total_seconds = days * 86400 + hours * 3600 + minutes * 60 + seconds
+        return timedelta(seconds=total_seconds)
+
+    except subprocess.TimeoutExpired:
+        print("squeue command timed out.")
+    except subprocess.CalledProcessError as e:
+        print(f"Error querying squeue: {e}")
+    except ValueError as e:
+        print(f"Error parsing time left: {e}")
+
+    return None
+
+
+def is_requeueable(jobid):
+    """
+    Check if a Slurm job is requeueable.
+
+    Args:
+        jobid (str): The job ID to check.
+
+    Returns:
+        bool: True if the job is requeueable, False otherwise.
+    """
+    try:
+        # Run scontrol to get job details
+        result = subprocess.run(
+            ["scontrol", "show", "job", jobid],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+
+        # Parse the output for the Requeue field
+        for line in result.stdout.splitlines():
+            if "Requeue=" in line:
+                requeue_value = line.split("Requeue=")[1].split()[0]
+                return requeue_value == "1"
+
+        # If the Requeue field is not found
+        raise RuntimeError(f"Could not determine Requeue status for job {jobid}.")
+
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"Error while running scontrol: {e.stderr.strip()}")
+
+
+class AutoSlurmRequeue(AbstractCallback):
+    """
+    A callback that automatically requeues a Slurm job if it is about to run out
+    of time.
+
+    This callback should be used together with a form of checkpointing to ensure
+    that the job can be requeued without losing progress.
+    """
+
+    before: timedelta = struct.field(pytree_node=False, serialize=False)
+    max_requeue_count: int = struct.field(pytree_node=False, serialize=False)
+    _time_to_requeue: timedelta | None = struct.field(
+        pytree_node=False, serialize=False, default=None
+    )
+    _enabled: bool = struct.field(pytree_node=False, serialize=False, default=True)
+
+    def __init__(
+        self, before: timedelta = timedelta(minutes=5), max_requeue_count: int = 3
+    ):
+        """
+        Initialize the auto-requeue callback.
+
+        Args:
+            before: The time before the job ends to check for requeueing (default: 5 minute).
+                This should be a timedelta object or a number of seconds, and it should be at least
+                as long as the time it takes an iteration to run.
+            max_requeue_count: Maximum number of times the job should be requeued.
+        """
+        if isinstance(before, (int)):
+            before = timedelta(seconds=before)
+        if not isinstance(before, timedelta):
+            raise TypeError(
+                f"Expected before to be a timedelta or an int (seconds), but got {type(before)}"
+            )
+        if not isinstance(max_requeue_count, int) or max_requeue_count < 0:
+            raise ValueError("max_requeue_count must be a non-negative integer.")
+
+        self.before = before
+        self.max_requeue_count = max_requeue_count
+
+    def on_run_start(self, step, driver):
+        jobid = os.getenv("SLURM_JOB_ID")
+        requeue_count = int(os.getenv("SLURM_RESTART_COUNT", "0"))
+
+        if jobid is None:
+            print("SLURM_JOB_ID not found. Skipping auto-requeue check.")
+            self._enabled = False
+            return
+
+        if requeue_count >= self.max_requeue_count:
+            print(
+                f"Job has been requeued {requeue_count} times, exceeding the limit of {self.max_requeue_count}. Disabling auto-requeue."
+            )
+            self._enabled = False
+            return
+
+        if not is_requeueable(jobid):
+            raise RuntimeError(
+                "Job is not requeable. Use `--requeue` option when scheduling the job \n"
+                "for example use `sbatch --requeue file` or add #SLURM --requeue at the top"
+                "of the file."
+            )
+
+        time_left = get_time_left(jobid)
+
+        if time_left is None:
+            print("No time limit found. Skipping auto-requeue check.")
+            self._enabled = False
+            return
+
+        self._time_to_requeue = datetime.now() + time_left - self.before
+        self._enabled = True
+
+    def on_step_end(self, step, log_data, driver):
+        if not self._enabled:
+            return
+
+        if datetime.now() > self._time_to_requeue:
+            print("Reached requeue time. Requeueing job...")
+            if jax.process_index() == 0:
+                subprocess.run(["scontrol", "requeue", os.getenv("SLURM_JOB_ID")])
+            self._enabled = False

netket-3.22.0.dev1/netket/_src/callbacks/base.py

@@ -0,0 +1,126 @@
+from netket.utils import struct
+
+
+class StopRun(Exception):
+    """
+    Exception to be raised by callbacks to gracefully stop the optimisation loop.
+
+    Raise this exception (or a subclass of it) from any callback hook to stop
+    the :meth:`~netket.driver.AbstractVariationalDriver.run` loop early.
+    The driver will catch it, call
+    :meth:`~netket.callbacks.AbstractCallback.on_run_end` on all callbacks, and
+    then return normally — no traceback is printed and no exception propagates to
+    the caller.
+
+    The message passed to the exception will be printed to stdout.
+
+    Example::
+
+        class StopAfterConvergence(nk.callbacks.AbstractCallback):
+            def on_step_end(self, step, log_data, driver):
+                if log_data["Energy"].error_of_mean < 1e-4:
+                    raise nk.callbacks.StopRun("Energy converged.")
+    """
+
+    pass
+
+
+class AbstractCallback(struct.Pytree, mutable=True):
+    """
+    Abstract base class for callbacks in advanced variational drivers.
+
+    This class is a Pytree, so it can be used with JAX transformations and
+    automatically handles serialisation, but fields must be declared
+    with `struct.field(pytree_node=False)` as class attributes.
+
+    Subclass this class and override any of the hook methods to inject custom
+    logic at specific points of the optimisation loop.  All hook methods have
+    no-op default implementations, so you only need to override the ones you
+    need.
+
+    To stop the optimisation early from inside any hook, raise
+    :class:`~netket.callbacks.StopRun` (or a subclass of it).  The driver will
+    catch it, call :meth:`on_run_end` on all callbacks, and return normally.
+
+    For a full description of the run loop structure and every available hook,
+    including pseudocode showing exactly when each hook is called, see
+    :ref:`advanced_custom_callbacks`.
+    """
+
+    def __init_subclass__(cls, **kwargs):
+        from netket.errors import CallbackLegacyHookError
+
+        super().__init_subclass__(**kwargs)
+
+        old_methods = [
+            m for m in ("on_legacy_run", "on_parameter_update") if m in cls.__dict__
+        ]
+        if old_methods:
+            raise CallbackLegacyHookError(cls.__name__, old_methods)
+
+    def __init__(
+        self,
+    ):
+        pass
+
+    @property
+    def callback_order(self) -> int:
+        """
+        An integer representing the order in which this callback should be called.
+
+        Lower numbers are called first, and higher numbers are called later.
+
+        This can be redefined in subclasses to change the order in which callbacks are called.
+        (Default: 0, for all callbacks, 10 for loggers).
+        """
+        return 0
+
+    def on_run_start(self, step, driver):
+        pass
+
+    def on_step_start(self, step, log_data, driver):
+        pass
+
+    def on_compute_update_start(self, step, log_data, driver):
+        pass
+
+    def on_compute_update_end(self, step, log_data, driver) -> bool:
+        """
+        Callback called at the end of the compute update phase, after computing the loss and its gradient.
+
+        This is called before the parameters are updated, so it can be used to implement custom
+        logic for rejecting a step based on the computed loss or gradient.
+
+        Returns:
+            A boolean indicating whether to reject the step (i.e. repeat it with the same parameters).
+            If it returns None, it is treated as False.
+        """
+        return False
+
+    def before_parameter_update(self, step, log_data, driver):
+        """
+        Called after all update logic has been computed and the step has been accepted,
+        but before the driver applies the parameter update.
+
+        At this point:
+
+        - The loss and its gradient have been computed by :meth:`~netket.driver.AbstractVariationalDriver.compute_loss_and_update`.
+        - The step has been accepted (not rejected by :meth:`on_compute_update_end`).
+        - ``driver.step_count`` still refers to the *current* step — it has not yet been incremented.
+        - The variational state parameters have **not** yet changed.
+
+        This is the right place to estimate additional observables, add data to
+        ``log_data``, or take a snapshot of the state for logging.  Callbacks with a
+        lower :attr:`callback_order` run first, so observables callbacks (order 0) are
+        guaranteed to populate ``log_data`` before logger callbacks (order 10) read it.
+        """
+        pass
+
+    def on_step_end(self, step, log_data, driver):
+        pass
+
+    def on_run_end(self, step, driver):
+        pass
+
+    def on_run_error(self, step, error, driver):
+        pass

netket-3.22.0.dev1/netket/_src/callbacks/callback_list.py

@@ -0,0 +1,52 @@
+from typing import Any
+
+
+from netket._src.callbacks.base import AbstractCallback
+
+
+class CallbackList(AbstractCallback):
+    callbacks: tuple[Any]
+
+    def __init__(
+        self,
+        callbacks=None,
+    ):
+        if isinstance(callbacks, AbstractCallback):
+            callbacks = (callbacks,)
+        self.callbacks = tuple(sorted(callbacks, key=lambda x: x.callback_order))
+
+    def on_run_start(self, step, driver):
+        for callback in self.callbacks:
+            callback.on_run_start(step, driver)
+
+    def on_step_start(self, step, log_data, driver):
+        for callback in self.callbacks:
+            callback.on_step_start(step, log_data, driver)
+
+    def on_compute_update_start(self, step, log_data, driver):
+        for callback in self.callbacks:
+            callback.on_compute_update_start(step, log_data, driver)
+
+    def on_compute_update_end(self, step, log_data, driver) -> bool:
+        reject_step = False
+        for callback in self.callbacks:
+            _reject = callback.on_compute_update_end(step, log_data, driver)
+            if _reject is True:
+                reject_step = True
+        return reject_step
+
+    def before_parameter_update(self, step, log_data, driver):
+        for callback in self.callbacks:
+            callback.before_parameter_update(step, log_data, driver)
+
+    def on_step_end(self, step, log_data, driver):
+        for callback in self.callbacks:
+            callback.on_step_end(step, log_data, driver)
+
+    def on_run_end(self, step, driver):
+        for callback in self.callbacks:
+            callback.on_run_end(step, driver)
+
+    def on_run_error(self, step, error, driver):
+        for callback in self.callbacks:
+            callback.on_run_error(step, error, driver)

netket-3.22.0.dev1/netket/_src/callbacks/legacy_wrappers.py

@@ -0,0 +1,73 @@
+from typing import TYPE_CHECKING, Any, Callable
+from functools import wraps
+import copy
+
+from netket.vqs import VariationalState
+from netket.utils import struct
+
+from netket._src.callbacks.base import AbstractCallback, StopRun
+
+if TYPE_CHECKING:
+    from netket._src.driver.abstract_variational_driver import AbstractDriver
+
+LegacyCallbackT = Callable[[int, dict, "AbstractDriver"], bool]
+
+
+class LegacyCallbackWrapper(AbstractCallback):
+    """
+    Wraps a legacy callback function (i.e. a function that takes the step, log_data and driver as arguments)
+    into a callback that can be used with the new callback system.
+
+    The legacy callback should return a boolean indicating whether to continue the run (True) or stop it (False).
+    """
+
+    callback: LegacyCallbackT
+
+    def __init__(self, callback: LegacyCallbackT):
+        self.callback = callback
+
+    def before_parameter_update(self, step, log_data, driver):
+        do_continue = self.callback(step, log_data, driver)
+        if not do_continue:
+            raise StopRun(f"Legacy callback {self.callback} requested to stop the run.")
+
+
+class LegacyLoggerWrapper(AbstractCallback):
+    """
+    Wraps a legacy logger function (i.e. a function that takes the step, log_data and variational state as arguments)
+    into a callback that can be used with the new callback system.
+    """
+
+    logger: Any
+
+    _vstate: VariationalState | None = struct.field(serialize=False)
+    """
+    Used to store the variational state between when parameters are updated and
+    when we actually log it.
+    """
+
+    def __init__(self, logger: Any):
+        self.logger = logger
+        self._vstate = None
+
+    @property
+    @wraps(AbstractCallback.callback_order)
+    def callback_order(self) -> int:
+        return 10
+
+    def before_parameter_update(self, step, log_data, driver):
+        # Store the vstate internally so that if we modify the parameters
+        # we can still log the previous one.
+        self._vstate = copy.copy(driver.state)
+
+    def on_step_end(self, step, log_data, driver):
+        self.logger(step, log_data, self._vstate)
+        self._vstate = None
+
+    def on_run_end(self, step, driver):
+        self.logger.flush(driver.state)
+        self._vstate = None
+
+    def on_run_error(self, step, error, driver):
+        self.logger.flush(driver.state)
+        self._vstate = None