code: def multi_node_snapshot(comm, snapshot, replica_sets):
      """Create trainer extension for multi-node snapshots
  
      Provides generis multi-node snapshot saving and auto-load feature
      at multi-node environment, leveraging power of single-node
      snapshot.
  
      In many cases snapshot target may differ, e.g. only trainer of
      rank 0 process often has extensions such as ``LogReport`` and so
      on, to not confuse terminal output. Just loading at one process
      and broadcasting it to other processes does not work in that case.
  
      This wrapper addresses that issue by defining sets of replicas
      where within the set the target object is replicated and supposed
      to be same among processes. For example, a trainer example, only
      the trainer at rank ``0`` has special extensions and others
      doesn't::
  
          trainer = Trainer(updater)
          if comm.rank == 0:
              trainer.extend(extensions.DumpGraph('main/loss'))
              trainer.extend(extensions.LogReport())
              trainer.extend(extensions.PrintReport(
                  ['epoch', 'main/loss', 'validation/main/loss',
                   'main/accuracy', 'validation/main/accuracy', 'elapsed_time']))
              trainer.extend(extensions.ProgressBar())
  
      This case can be described with two replica sets, where each set
      can be represented as single integer that indicates rank number,
      or iterable set/list/generator of integers like this::
  
          replica_sets = [[0], range(1, comm.size)]
  
      Here the first replica set is described as ``[0]``, or simply in
      short just ``0``, and the second replica set is ``range(1,
      comm.size)``, representing rest of processes other than ``0``. The
      remaining list can be omitted. Thus in that case, it can be
      simplified more::
  
          replica_sets = [0,]
  
      In this case, the snapshot will be saved at rank ``0`` process and
      at rank ``1`` process. The latter represents the replica set of
      ``range(1, comm.size)`` . In this case autoloading at
      initialization of snapshot extension works after the restart
      cleanly, even though the size of the communicator differs.
  
      Once the replica sets are defined, it can be easily extended::
  
          replica_sets = [0,]
          snapshot = multi_node_snapshot(comm, extensions.snapshot(),
                                         replica_sets)
          trainer.extend(snapshot, trigger=(1, 'epoch'))
  
  
      More example tuples of replica set representation follows:
  
      ===================== ===== ==============================================
      code                  nproc actual sets
      ===================== ===== ==============================================
      ``[0]``               ``4`` ``[{0}, {1, 2, 3}]``
      ``[0, 1]``            ``4`` ``[{0}, {1}, {2, 3}]``
      ``[0, 1], [2, 3]]``   ``4`` ``[{0, 1}, {2, 3}]``
      ``[]``                ``4`` ``[{0, 1, 2, 3}]``
      ``[range(0, 8, 2)]``  ``8`` ``[set(range(0, 8, 2)), set(range(1, 8, 2))]``
      ===================== ===== ==============================================
  
      Args:
          comm (ChainerMN communicator): communicater object
          snapshot: Snapshot extension object obtained via
                :meth:`~chainer.training.extensions.snapshot` .
          replica_sets: list of replica set definition, where
                a replica set can be defined by single integer
                as rank number, or iterable integers.
  
      Returns:
          Trainer extension that wraps ``snapshot`` and properly
          controles number of snapshots.
  
      """
      return _MultiNodeSnapshot(comm, snapshot, replica_sets)
  
file_path: chainer/chainermn/extensions/_multi_node_snapshot.py
imports: import io
  from chainer.serializers import load_npz
  from chainer.serializers import save_npz
  from chainer.training.extension import Extension
  from chainer.training.extensions._snapshot import _find_latest_snapshot
label_desc: Use image node_id 1 for calling a global function with example usage: multi_node_snapshot(comm, snapshot, replica_sets) and returns: _MultiNodeSnapshot