scatter_dataset_test.py

import chainermn
from chainer.datasets import get_cifar10
import pickle

comm = chainermn.create_communicator('pure_nccl')

if comm.rank == 0:
    train, test = get_cifar10()
else:
    train, test = None, None

try:
    train = chainermn.scatter_dataset(train, comm, shuffle=True)
    test = chainermn.scatter_dataset(test, comm, shuffle=True)
except (pickle.UnpicklingError, EOFError) as e:
    print("**************************", flush=True)
    print("Rank {} Error occured!!! ".format(comm.rank), flush=True)
    print("{}: {}".format(type(e), str(e)), flush=True)
    print("**************************", flush=True)
    comm.mpi_comm.Abort(1)
else:
    if comm.rank == 0:
        print("OK")