callbacks.py

import os
import sys
import time
from dataclasses import dataclass
from datetime import datetime
from enum import IntFlag
from logging import getLogger
from typing import *

import numpy as np

from .checkpoint import BaseCheckpoint, CheckpointManager
from .errors import NaNMetricError
from .formatting import MetricsFormatter, format_duration, format_as_asctime
from .logging_ import print_with_time
from .metrics import ScalarMetricsLogger, ScalarMetricCollector
from .mlstorage import ExperimentDoc
from .stateful import StatefulObjectGroup, StatefulObject
from .utils import NOT_SET

__all__ = [
    'CallbackData', 'Callback', 'CallbackList',
    'LoggerMode', 'LoggerCallback', 'StopOnNaN',
    'BaseTrainCallback', 'BaseCheckpointCallback',
    'AutoCheckpoint', 'EarlyStopping',
]


@dataclass
class CallbackData(object):
    """
    Data carried by a cycle begin/end event from :class:`Callback`.
    """

    __slots__ = ('stage', 'index', 'size', 'start_timestamp',
                 'end_timestamp', 'exc_time', 'metrics')

    stage: 'Stage'
    """The stage that calls the callback."""

    index: Optional[int]
    """Index of the epoch or batch, start from 1."""

    size: Optional[int]
    """The size of the batch."""

    start_timestamp: float
    """Start timestamp of the stage/epoch/batch."""

    end_timestamp: Optional[float]
    """End timestamp of the stage/epoch/batch, available at the cycle end."""

    exc_time: Optional[float]
    """Execution time of the stage/epoch/batch, available at the cycle end."""

    metrics: Optional[Dict[str, Any]]
    """Metrics dict, available at the cycle end."""


class Callback(object):
    """Base class of a callback for a machine learning stage."""

    priority: int = 0
    """
    The priority of the callback.  Smaller priority indicates the callback
    should be called earlier than other callbacks with larger priorities.
    """

    ###########
    # metrics #
    ###########
    def on_metrics(self, data: CallbackData):
        pass  # pragma: no cover

    ##################
    # general events #
    ##################
    def on_stage_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_stage_end(self, data: CallbackData):
        pass  # pragma: no cover

    def on_epoch_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_epoch_end(self, data: CallbackData):
        pass  # pragma: no cover

    def on_batch_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_batch_end(self, data: CallbackData):
        pass  # pragma: no cover

    ################
    # train events #
    ################
    def on_train_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_train_end(self, data: CallbackData):
        pass  # pragma: no cover

    def on_train_epoch_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_train_epoch_end(self, data: CallbackData):
        pass  # pragma: no cover

    def on_train_batch_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_train_batch_end(self, data: CallbackData):
        pass  # pragma: no cover

    #####################
    # validation events #
    #####################
    def on_validation_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_validation_end(self, data: CallbackData):
        pass  # pragma: no cover

    def on_validation_batch_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_validation_batch_end(self, data: CallbackData):
        pass  # pragma: no cover

    ###############
    # test events #
    ###############
    def on_test_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_test_end(self, data: CallbackData):
        pass  # pragma: no cover

    def on_test_batch_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_test_batch_end(self, data: CallbackData):
        pass  # pragma: no cover

    ##################
    # predict events #
    ##################
    def on_predict_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_predict_end(self, data: CallbackData):
        pass  # pragma: no cover

    def on_predict_batch_begin(self, data: CallbackData):
        pass  # pragma: no cover

    def on_predict_batch_end(self, data: CallbackData):
        pass  # pragma: no cover


class CallbackList(Sequence[Callback]):
    """
    A callback list, which maintains the orders of callbacks according to
    their priority.
    """

    _SORTED = object()

    def __init__(self,
                 callbacks: Optional[Iterator[Callback]] = None,
                 *,
                 _sorted=None):
        if callbacks is not None:
            if _sorted is not self._SORTED:
                callbacks = sorted(callbacks, key=lambda cb: cb.priority)
            else:
                callbacks = list(callbacks)
        self._callbacks = callbacks

    def __len__(self) -> int:
        return len(self._callbacks)

    def __iter__(self):
        return iter(self._callbacks)

    def __eq__(self, other):
        return isinstance(other, CallbackList) and \
            self._callbacks == other._callbacks

    def __getitem__(self, item):
        return self._callbacks[item]

    def __delitem__(self, item):
        del self._callbacks[item]

    def __copy__(self):
        return self.clone()

    def clone(self) -> 'CallbackList':
        return CallbackList(self._callbacks, _sorted=self._SORTED)

    def add(self, callback: Callback):
        """
        Add a callback to this list, respecting the `priority`.

        Args:
            callback: The callback object.
        """
        i = len(self._callbacks) - 1
        while i > -1:
            if self._callbacks[i].priority <= callback.priority:
                break
            i -= 1
        self._callbacks.insert(i + 1, callback)

    def remove(self, callback: Callback):
        """
        Remove a callback from this list.

        Args:
            callback: The callback to be removed.

        Raises:
            ValueError: If `callback` is not present.
        """
        self._callbacks.remove(callback)


@dataclass
class _LoggerContext(object):
    """The context of an open stage in :class:`LoggerCallback`."""

    __slots__ = ('stage', 'progress', 'metrics_collector', 'batch_metrics',
                 'last_console_log_time', 'last_remote_push_time')

    stage: 'Stage'
    progress: Dict[str, Any]
    metrics_collector: ScalarMetricsLogger
    """
    Metrics logger to accumulate the mean and std of metrics.  This logger
    will be cleared at the beginning when `on_epoch_begin` is called.

    For validation, test and predict, this should effectively accumulate
    the metrics throughout the whole stage, since the `on_epoch_begin`
    callback will never be called.
    """
    batch_metrics: Dict[str, Any]
    """
    The current batch metrics.  Will be cleared after each batch.
    """
    last_console_log_time: float
    """Last time that the logs have been written to console."""
    last_remote_push_time: float
    """Last time that the logs have been pushed to remote."""

    @staticmethod
    def new_context(stage) -> '_LoggerContext':
        now_time = time.time()
        return _LoggerContext(
            stage=stage,
            progress={},
            metrics_collector=ScalarMetricsLogger(),
            batch_metrics={},
            # set these two log times as the current time, such that these
            # logs will not be written immediately after the stage begins.
            last_console_log_time=now_time,
            last_remote_push_time=now_time,
        )

    def update_metrics(self,
                       metrics: Mapping[str, Any],
                       replace: bool = False,
                       batch_size: Optional[float] = None) -> None:
        """
        Update the epoch metrics logger and batch metrics dict (if a batch
        is currently active) according to `metrics`.

        Args:
            metrics: The batch, epoch or stage metrics from stage callback.
            replace: Whether to replace the epoch/stage metrics
                instead of updating them.
            batch_size: The batch size information from stage callback.
        """
        # We expect the metrics to be scalars.  If not, we shall take average.
        raw_metrics = {}
        averaged_metrics = {}
        if metrics:
            for key, val in metrics.items():
                key = self.stage.type.add_metric_prefix(key)
                if np.shape(val) == ():
                    averaged_metrics[key] = val
                else:
                    raw_metrics[key] = val

            updater = self.metrics_collector.replace \
                if replace else self.metrics_collector.update
            updater(raw_metrics)
            updater(averaged_metrics, weight=batch_size or 1.)

        # if inside a batch, update the batch metrics dict
        if self.stage.batch.is_active:
            self.batch_metrics.update(averaged_metrics)
            self.batch_metrics.update(
                {k: np.mean(v) for k, v in raw_metrics.items()})

    def copy_metrics_from_nested_context(self, ctx: '_LoggerContext'):
        """
        Copy the final metrics from nested stage.

        Args:
            ctx: The nested stage context.
        """
        # obtain the final metrics from the nested context
        nested_metrics = ctx.metrics_collector.to_json(mean_only=True)

        # if currently a batch is active, update the batch metrics
        if self.stage.batch.is_active:
            self.batch_metrics.update(nested_metrics)

        # update the final metrics
        self.metrics_collector.update(nested_metrics)

    def next_epoch(self):
        """Reset the internal states and enter the next epoch."""
        self.metrics_collector.clear()
        self.batch_metrics.clear()

    def next_batch(self):
        """Reset the internal states and enter the next batch."""
        self.batch_metrics.clear()


def _console_writer(s: str) -> None:
    sys.stdout.write(s)
    sys.stdout.flush()


def _print_log(console_writer: Optional[Callable[[str], None]],
               text: str,
               nl: bool = True,
               show_time: bool = True):
    if console_writer is not None:
        if show_time:
            time_str = format_as_asctime(datetime.now())
            text = f'[{time_str}] {text}'
        if nl:
            text += '\n'
        console_writer(text)


class LoggerMode(IntFlag):
    """Integer flags of logger mode."""

    NONE = 0x0

    LOG_START_END = 0x1
    """Log at the stage start/end."""

    LOG_EVERY_EPOCH = 0x2
    """Log at every epoch."""

    LOG_MAJOR = LOG_START_END | LOG_EVERY_EPOCH
    """Log at the stage start/end and at every epoch."""

    LOG_EVERY_FEW_BATCHES = 0x4
    """Log after every few batches."""

    LOG_EVERY_FEW_SECONDS = 0x8
    """Log every few seconds."""

    DEFAULT = LOG_MAJOR | LOG_EVERY_FEW_SECONDS
    """Default log mode."""

    def check_integrity(self):
        if LoggerMode.LOG_EVERY_FEW_SECONDS in self and \
                LoggerMode.LOG_EVERY_FEW_BATCHES in self:
            raise ValueError(
                '`LOG_EVERY_FEW_SECONDS` and `LOG_EVERY_FEW_BATCHES` '
                'cannot be both enabled.'
            )


class LoggerCallback(Callback):
    """
    Callback that logs training/testing/predicting progress and metrics
    to console and to MLStorage server.

    For performance considerations, batch metrics and progress information
    will be written to console every ``console_log_interval`` seconds,
    and sent to server every ``remote_log_interval`` seconds.

    The progress info will be stored as `progress.<stage.type>` field, and
    the batch metrics will be stored in `progress.<stage.type>.batch_metrics`
    field.  Stages with different types thus will not override the progress
    information and batch metrics of each other.

    Batch metrics will be accumulated by :class:`MetricsLogger`, and reported
    at the end of the epoch.  If the epoch callback provides metrics with the
    same names as the batch metrics, the epoch metrics will override the batch
    metrics.  These metrics are the epoch metrics.

    Moreover, metrics provided by the stage end callback are the stage metrics.
    Epoch metrics and stage metrics will be stored in the `result` field.
    For epoch and stage metrics of the train stage, the metrics will be saved
    as-is; but for other stages, the metrics names will be enforced to have
    the following prefix:

    *  validation stage: "val_" or "valid_".
    *  test stage: "test_"
    *  predict stage: "pred_" or "predict_"

    For nested stages (e.g., validation stage inside a train stage), the
    progress and metrics of the inner stages will not be written to the
    console, but will indeed be sent to the server.
    """

    # a sufficiently large value, should run after almost all callbacks
    priority = 999999

    _ctx_stack: List[_LoggerContext]
    """The stack of :class:`_LoggerContext`, one for every :class:`Stage`."""

    console_mode: LoggerMode
    """The console logger mode."""

    console_writer: Optional[Callable[[str], None]]
    """The console writer."""

    console_log_batch_freq: int
    """Write batch progress and metrics every this number of batches."""

    console_log_interval: float
    """Write batch progress and metrics every this number of seconds."""

    remote_doc: Optional[ExperimentDoc]
    """The :class:`ExperimentDoc`, where to push progress and metrics."""

    remote_push_interval: float
    """Push updates to the remote every this number of seconds."""

    enabled: bool
    """Whether or not this :class:`LoggerCallback` is enabled?"""

    def __init__(self,
                 console_mode: LoggerMode = LoggerMode.DEFAULT,
                 console_writer: Optional[
                     Callable[[str], None]] = _console_writer,
                 console_log_batch_freq: int = 100,
                 console_log_interval: float = 10.,
                 remote_doc: Optional[ExperimentDoc] = NOT_SET,
                 remote_push_interval: float = 60.,
                 metrics_formatter: MetricsFormatter = MetricsFormatter()):
        """
        Construct a new :class:`LoggerCallback`.

        Args:
            console_mode: Mode of console log.
            console_writer: The console writer.
            console_log_batch_freq: Log to console every this number of batches,
                if `LOG_EVERY_FEW_BATCHES` is enabled in `console_mode`.
            console_log_interval: Log to console every this number of seconds,
                if `LOG_EVERY_FEW_SECONDS` is enabled in `console_mode`.
            remote_doc: The remote doc object, where to push updates.
            remote_push_interval: Push to remote every this number of seconds.
            metrics_formatter: The metrics formatter.
        """
        # check the argument
        console_mode.check_integrity()

        # get the remote document according to the context and the environment
        if remote_doc is NOT_SET:
            remote_doc = ExperimentDoc.default_doc()

        self._ctx_stack = []
        self.console_mode = console_mode
        self.console_writer = console_writer
        self.console_log_batch_freq = console_log_batch_freq
        self.console_log_interval = console_log_interval
        self.remote_doc = remote_doc
        self.remote_push_interval = remote_push_interval
        self.metrics_formatter = metrics_formatter
        self._enabled = self.remote_doc is not None or bool(self.console_mode)

    @property
    def enabled(self) -> bool:
        """Whether or not this logger callback is enabled?"""
        return self._enabled

    @property
    def in_nested_stage(self) -> bool:
        """Whether or not this logger callback is in nested stage?"""
        return len(self._ctx_stack) > 1

    @property
    def ctx(self) -> _LoggerContext:
        """Get the current active context (at the top of context stack)."""
        return self._ctx_stack[-1]

    @property
    def stage(self) -> 'Stage':
        """Get the current active stage (at the top of context stack)."""
        return self.ctx.stage

    def _should_write_start_end_console_log(self) -> bool:
        return (not self.in_nested_stage and
                LoggerMode.LOG_START_END in self.console_mode)

    def _should_write_epoch_console_log(self) -> bool:
        return (not self.in_nested_stage and
                LoggerMode.LOG_EVERY_EPOCH in self.console_mode)

    def _should_write_batch_console_log(self,
                                        batch_id: int,
                                        end_timestamp: float) -> bool:
        # if we are now in a nested stage, we shall never write batch log
        if self.in_nested_stage:
            return False

        # if the epoch log is enabled, and this is the final batch, we
        # shall write epoch log instead of the batch log.
        if (LoggerMode.LOG_EVERY_EPOCH in self.console_mode and
                self.stage.epoch is not None and
                batch_id == self.stage.batch.total):
            return False

        # if the best validation mark is True, write batch log
        if self.stage.best_validation_mark:
            return True

        # ordinary checks for the batch
        if (LoggerMode.LOG_EVERY_FEW_BATCHES in self.console_mode and
                batch_id % self.console_log_batch_freq == 0):
            return True

        if (LoggerMode.LOG_EVERY_FEW_SECONDS in self.console_mode and
                end_timestamp - self.ctx.last_console_log_time >=
                self.console_log_interval):
            return True

        # no log is required to be written now
        return False

    def _should_push_batch_remote_log(self,
                                      batch_id: int,
                                      end_timestamp: float) -> bool:
        return (end_timestamp - self.ctx.last_remote_push_time >=
                self.remote_push_interval)

    def _push_to_remote(self, result: Optional[Dict[str, Any]] = None):
        payload = {
            f'progress.{self.stage.name}': self.ctx.progress
        }
        if result:
            payload['result'] = result
        self.remote_doc.update(payload)
        self.ctx.last_remote_push_time = time.time()

    def _write_stage_or_epoch_end_console_log(
            self,
            result_dict: Optional[Dict[str, Any]],
            prefix: str = '',
            suffix: str = '',
            show_time: bool = False,
            is_stage_end: bool = False) -> None:
        # first, compose the log line
        buf = []
        # - <prefix>
        if prefix:
            buf.append(prefix)
        # - <metrics>
        if result_dict:
            result_str = self.metrics_formatter.format(
                result_dict,
                sep=(': ', ' - '),
                known_names=self.stage.known_metrics
            )
            buf.append(result_str)
        # - <suffix>
        if suffix:
            buf.append(suffix)
        log_line = ' - '.join(buf)
        # - " (*)" mark
        if not is_stage_end and self.stage.best_validation_mark:
            log_line += ' (*)'

        # then, print the log
        _print_log(self.console_writer, log_line, show_time=show_time)
        self.ctx.last_console_log_time = time.time()

    def _batch_console_head(self, batch=None) -> str:
        # the batch counter
        max_batch = str(self.ctx.progress.get('max_batch', ''))
        if batch is None:
            batch = str(self.ctx.progress.get('batch', ''))
        if max_batch:
            return f'{batch:>{len(max_batch)}s}/{max_batch}'
        return batch

    def _update_progress_time_info(self, end_time: Optional[float]):
        # update elapsed
        if end_time is not None:
            self.ctx.progress['elapsed'] = end_time - self.stage.start_timestamp

        # update eta
        eta = self.stage.get_eta()
        if eta is not None and eta > 1e-7:
            self.ctx.progress['eta'] = eta
        else:
            self.ctx.progress.pop('eta', None)

    def on_metrics(self, data: CallbackData):
        if not data.stage.is_active and self.remote_doc is not None:
            # some immediate metrics outside a loop context
            m_logger = ScalarMetricsLogger()
            m_logger.update(data.metrics)
            payload = {'result': m_logger.to_json()}
            self.remote_doc.update(payload)

    def on_stage_begin(self, data: CallbackData):
        self._ctx_stack.append(_LoggerContext.new_context(data.stage))

        # write console log
        if self._should_write_start_end_console_log():
            _print_log(
                self.console_writer,
                f'{self.stage.name.capitalize()} started',
                show_time=True
            )

        # start the remote doc worker if this is the first stage
        if len(self._ctx_stack) == 1 and self.remote_doc is not None:
            self.remote_doc.start_worker()

    def on_stage_end(self, data: CallbackData):
        try:
            # set the progress info
            self._update_progress_time_info(data.end_timestamp)

            # replace the epoch metrics with stage metrics, if provided
            if data.metrics:
                self.ctx.update_metrics(data.metrics, replace=True)

            # obtain the stage result dict
            stage_result = self.ctx.metrics_collector.to_json()

            # write the console log
            if self._should_write_start_end_console_log():
                log_prefix = f'{self.stage.name.capitalize()} finished'
                if data.exc_time is not None:
                    elapsed_str = format_duration(
                        data.exc_time, precision=1, count_down=True)
                    log_prefix += f' in {elapsed_str}'

                self._write_stage_or_epoch_end_console_log(
                    result_dict=stage_result,
                    prefix=log_prefix,
                    suffix='',
                    show_time=True,
                    is_stage_end=True,
                )

            # push to remote
            if self.remote_doc is not None:
                self._push_to_remote(stage_result)

        finally:
            # pop this stage
            if len(self._ctx_stack) > 1:
                self._ctx_stack[-2].copy_metrics_from_nested_context(self.ctx)

            self._ctx_stack.pop()

            # stop the remote doc worker if there is no context left
            if not self._ctx_stack and self.remote_doc is not None:
                self.remote_doc.stop_worker()

    def on_epoch_begin(self, data: CallbackData):
        # set the progress info
        self.ctx.progress['epoch'] = data.index
        if data.stage.epoch.total is not None:
            self.ctx.progress['max_epoch'] = data.stage.epoch.total

        # set the context to enter next epoch
        self.ctx.next_epoch()

        # write epoch beginning log
        if self._should_write_epoch_console_log():
            _print_log(self.console_writer, f'Epoch {data.stage.epoch}',
                       show_time=False)

    def on_epoch_end(self, data: CallbackData):
        # set the progress info
        self._update_progress_time_info(data.end_timestamp)
        if data.exc_time is not None:
            self.ctx.progress['epoch_time'] = data.exc_time

        # We use the metric values provided in `data.metrics` as the final
        # metric values for the epoch, to replace any batch metrics.
        self.ctx.update_metrics(data.metrics, replace=True)
        epoch_result = self.ctx.metrics_collector.to_json()

        # write the console log
        if self._should_write_epoch_console_log():
            # log_prefix: total number of executed batches + exc_time
            batch = self.ctx.progress.get('batch', None)
            log_prefix = f'{batch} iters'
            if data.exc_time:
                elapsed_str = format_duration(
                    data.exc_time, precision=1, count_down=True)
                log_prefix += f' in {elapsed_str}'

            # eta
            eta = self.stage.get_eta()
            if eta is not None:
                # just to be consistent with the format of batch logs
                log_prefix += f' - eta {format_duration(eta, count_down=True)}'

            self._write_stage_or_epoch_end_console_log(
                result_dict=epoch_result,
                prefix=log_prefix,
                suffix='',
            )

        # push to remote log
        if self.remote_doc is not None:
            self._push_to_remote(epoch_result)

    def on_batch_begin(self, data: CallbackData):
        self.ctx.progress['batch'] = data.index
        if data.stage.batch.total is not None:
            self.ctx.progress['max_batch'] = data.stage.batch.total
        self.ctx.progress.pop('batch_metrics', None)

        # set the context to enter next batch
        self.ctx.next_batch()

    def on_batch_end(self, data: CallbackData):
        # update the progress info
        self._update_progress_time_info(data.end_timestamp)
        if data.exc_time is not None:
            self.ctx.progress['batch_time'] = data.exc_time

        # update the metrics
        self.ctx.update_metrics(data.metrics, batch_size=data.size)

        # obtain the results of the batch
        batch_result = self.ctx.batch_metrics

        # Copy the batch metrics to the progress dict.
        # This assignment will be cleared at the beginning of the next batch.
        self.ctx.progress['batch_metrics'] = batch_result

        # write logs to console
        if self._should_write_batch_console_log(data.index,
                                                data.end_timestamp):
            buf = [self._batch_console_head()]
            if 'eta' in self.ctx.progress:
                eta_str = format_duration(self.ctx.progress["eta"],
                                          count_down=True)
                buf.append(f'eta {eta_str}')
            if batch_result:
                result_str = self.metrics_formatter.format(
                    batch_result,
                    sep=(': ', ' - '),
                    known_names=self.stage.known_metrics,
                )
                buf.append(result_str)
            log_line = ' - '.join(buf)
            if self.stage.best_validation_mark:
                log_line += ' (*)'
            _print_log(self.console_writer, log_line, show_time=False)
            self.ctx.last_console_log_time = time.time()

        # push the logs to remote
        if self.remote_doc is not None and \
                self._should_push_batch_remote_log(data.index,
                                                   data.end_timestamp):
            self._push_to_remote(batch_result)


class StopOnNaN(Callback):
    """
    Callback that raises :class:`NaNMetricError` whenever an NaN metric
    has been encountered.
    """

    # its priority should be even larger than the LoggerCallback, such that
    # the NaN metrics would be printed before exiting on NaNs
    priority = LoggerCallback.priority + 1

    def _check_metrics(self, metrics: Optional[Mapping[str, Any]]):
        if metrics:
            for key, val in metrics.items():
                if np.isnan(val):
                    raise NaNMetricError(key)

    def on_batch_end(self, data: CallbackData):
        self._check_metrics(data.metrics)

    def on_epoch_end(self, data: CallbackData):
        self._check_metrics(data.metrics)

    def on_stage_end(self, data: CallbackData):
        self._check_metrics(data.metrics)


class BaseTrainCallback(Callback):
    """
    Base callback class for train stages.

    Binds to the first train stage.  If the first stage that this callback
    encounters is not a train stage, then an error will be raised.

    If a subclass need to override :meth:`on_stage_begin()` or
    :meth:`on_stage_end()`, they should call the parent's method.
    For other event callbacks, they need to verify whether or not ``data.stage``
    equals to ``self.stage``.
    """

    stage: Optional['Stage'] = None
    """The current active train stage."""

    def on_stage_begin(self, data: CallbackData):
        if self.stage is None:
            if data.stage.type != StageType.TRAIN:
                raise RuntimeError(
                    f'The outer stage of `{self.__class__.__qualname__}` must '
                    f'be a train stage: got {data.stage.name} stage '
                    f'{data.stage!r}')
            self.stage = data.stage  # bind to this train stage

    def on_stage_end(self, data: CallbackData):
        if data.stage == self.stage:
            self.stage = None  # unbind from the current stage


class BaseCheckpointCallback(BaseTrainCallback):
    """
    Base class for checkpoint callbacks.

    Checkpoint callbacks are train callbacks, which will only work for a
    train stage.  Sub-classes should check ``if self.stage == data.stage``
    in any overrided method.
    """

    STAGE_STATE_KEY: str = '__stage'
    """State key that stores the stage states."""

    checkpoint: BaseCheckpoint
    """The checkpoint object."""

    root_dir: str
    """The root directory, where to save checkpoints."""

    state_objects: Dict[str, StatefulObject]
    """The state objects to be saved along with checkpoints."""

    max_checkpoints_to_keep: Optional[int]
    """
    Maximum number of checkpoints to keep.
    :obj:`None` means that all checkpoints will be kept.
    """

    save_stage_state: bool
    """Whether or not to save the stage state?"""

    checkpoint_manager: Optional[CheckpointManager] = None
    """The checkpoint manager instance."""

    def __init__(self,
                 checkpoint: BaseCheckpoint,
                 root_dir: str,
                 state_objects: Optional[Mapping[str, StatefulObject]] = None,
                 max_checkpoints_to_keep: Optional[int] = None,
                 save_stage_state: bool = True):
        """
        Construct a new :class:`BaseCheckpointCallback`.

        Args:
            checkpoint: The checkpoint object.
            root_dir: The root directory, where to save checkpoints.
            state_objects: The state objects to be saved along with checkpoints.
            max_checkpoints_to_keep: Maximum number of checkpoints to keep.
                Defaults to :obj:`None`, where all checkpoints will be kept.
            save_stage_state: Whether or not to save stage state?
        """
        # check the argument `state_objects`
        state_objects = {k: state_objects[k] for k in (state_objects or ())}
        for k in state_objects:
            if k == self.STAGE_STATE_KEY:
                raise ValueError(f'State object key {k!r} is reserved.')
            v = state_objects[k]
            if not isinstance(v, StatefulObject):
                raise ValueError(f'The item {k!r} in `state_objects` is not '
                                 f'a StatefulObject: got {v!r}')

        # memorize the argument
        self.checkpoint = checkpoint
        self.root_dir = os.path.abspath(root_dir)
        self.state_objects = state_objects
        self.max_checkpoints_to_keep = max_checkpoints_to_keep
        self.save_stage_state = save_stage_state

    def on_stage_begin(self, data: CallbackData):
        super().on_stage_begin(data)
        if data.stage == self.stage:
            if self.save_stage_state:
                self.state_objects[self.STAGE_STATE_KEY] = \
                    data.stage.state_proxy()
            self.checkpoint_manager = CheckpointManager(
                checkpoint=self.checkpoint,
                state_objects=StatefulObjectGroup(self.state_objects),
                root_dir=self.root_dir,
                max_to_keep=self.max_checkpoints_to_keep,
            )

    def on_stage_end(self, data: CallbackData):
        super().on_stage_end(data)
        if self.stage is None and self.checkpoint_manager is not None:
            self.checkpoint_manager = None
            self.state_objects.pop(self.STAGE_STATE_KEY, None)

    def make_checkpoint(self):
        epoch = self.stage.epoch.index
        batch = self.stage.batch.index
        ckpt_name = f'epoch-{epoch}-batch-{batch}'
        ckpt_path = self.checkpoint_manager.save(ckpt_name)
        getLogger(__name__).debug('Saved to checkpoint: %s', ckpt_path)


class AutoCheckpoint(BaseCheckpointCallback):
    """
    Callback to save train checkpoints automatically.
    """

    # priority is one larger than that of `EarlyStopping`, should run after it
    priority = 1000

    interval: Optional[float]
    """If not :obj:`None`, will save checkpoint every this number of seconds."""

    epoch_freq: Optional[int]
    """If not :obj:`None`, will save checkpoint every this number of epochs."""

    batch_freq: Optional[int]
    """If not :obj:`None`, will save checkpoint every this number of batches."""

    restore_checkpoint: Union[str, bool]
    """
    If :obj:`True`, restore the latest saved checkpoint from `root_dir` when
    train begins.  If a str, treat it as the path of a checkpoint, and restore
    it when train begins.
    """

    last_checkpoint_time: float
    """The timestamp when the last checkpoint was saved."""

    def __init__(self,
                 checkpoint: BaseCheckpoint,
                 root_dir: str,
                 *,
                 interval: Optional[float] = None,
                 epoch_freq: Optional[int] = None,
                 batch_freq: Optional[int] = None,
                 state_objects: Optional[Mapping[str, StatefulObject]] = None,
                 max_checkpoints_to_keep: Optional[int] = None,
                 restore_checkpoint: Union[str, bool] = True):
        """
        Construct a new :class:`AutoCheckpoint`.

        Args:
            checkpoint: The checkpoint object.
            root_dir: The root directory, where to save checkpoints.
            interval: If not :obj:`None`, will save checkpoint every this
                number of seconds.  One and only one of `interval`, `epoch_freq`
                and `batch_freq` can be not :obj:`None`.
            epoch_freq: If not :obj:`None`, will save checkpoint every this
                number of epochs.
            batch_freq: If not :obj:`None`, will save checkpoint every this
                number of batches.
            state_objects: The state objects to be saved along with checkpoints.
            max_checkpoints_to_keep: Maximum number of checkpoints to keep.
                Defaults to :obj:`None`, where all checkpoints will be kept.
            restore_checkpoint: If :obj:`True`, restore the latest saved
                checkpoint from `root_dir` when train begins.
                If a str, treat it as the path of a checkpoint, and restore
                it when train begins.
        """
        not_none_count = (
            int(interval is not None) + int(epoch_freq is not None) +
            int(batch_freq is not None)
        )
        if not_none_count != 1:
            raise ValueError('One and only one of `interval`, `epoch_freq` '
                             'and `batch_freq` should be specified.')
        if not isinstance(restore_checkpoint, str) and \
                restore_checkpoint not in (True, False):
            raise TypeError(f'`restore_checkpoint` must be a str or a bool: '