Class MultiNodeCheckpointManager

caffe2/python/checkpoint.py:432–634

Coordinates checkpointing and checkpointing across multiple nodes. Each of `init`, `load` and `save` will build TaskGroups which will trigger checkpointing on each of the nodes involved in a distributed job. Args: db_prefix: The prefix used to construct full db name. Since

class MultiNodeCheckpointManager:
    """
    Coordinates checkpointing and checkpointing across multiple nodes.
    Each of `init`, `load` and `save` will build TaskGroups which will
    trigger checkpointing on each of the nodes involved in a distributed job.

    Args:
        db_prefix: The prefix used to construct full db name. Since `absolute_path`
            is set to True, this will be used as db_name in SaveOp.
        db_type: Type of database to use for storing checkpoint.
        metadata_handler: An optional object capable of reading/writing
            checkpoint info in storage of choice.
    """
    def __init__(self, db_prefix, db_type, metadata_handler=None):
        self._node_managers = None
        self._db_prefix = db_prefix
        self._db_type = db_type
        self._metadata_handler = metadata_handler
        self._path_prefix = None
        self._path_type = None

    def _task_group(self, func, *args, **kw):
        assert self._node_managers is not None, 'init must be called first.'
        with TaskGroup(WorkspaceType.GLOBAL) as task_group:
            for node, manager in self._node_managers:
                with Node(node):
                    func(manager, *args, **kw)
            return task_group

    """
    Args:
        nodes: An array of nodes where this checkpoint manager is running.
        retrieve_from_epoch: Set to a number to load blobs from this epoch.
        path_prefix: Used to construct db name or path where checkpoint files are
            stored.
        path_type: Indicate the type of path where checkpoint files are stored.
    """
    def init(
        self, nodes, retrieve_from_epoch=None, path_prefix=None, path_type=None
    ):
        if self._node_managers is not None:
            assert [node for node, _ in self._node_managers] == nodes
            return TaskGroup(WorkspaceType.GLOBAL)
        self._node_managers = []
        for node in nodes:
            with Node(node):
                manager = CheckpointManager(
                    db_prefix=self._db_prefix,
                    node_name=str(node),
                    db_type=self._db_type)
                self._node_managers.append((node, manager))
        return self._task_group(
            CheckpointManager.init,
            nodes=[node],
            retrieve_from_epoch=retrieve_from_epoch,
            path_prefix=path_prefix,
            path_type=path_type)