def _test_idist_methods_in_hvd_context(backend, device):
# We explicitly set _model as _SerialModel
# then call idist.* methods and check that they give correct values
import horovod.torch as hvd
from ignite.distributed.utils import _SerialModel, _set_model
hvd.init()
_set_model(_SerialModel())
ws = hvd.size()
rank = hvd.rank()
local_rank = hvd.local_rank()
if torch.cuda.is_available():
torch.cuda.set_device(local_rank)
_test_distrib_config(local_rank,
backend=backend,
ws=ws,
true_device=device,
rank=rank)
hvd.shutdown()
def _destroy_mnodes_dist_context():
dist.barrier()
dist.destroy_process_group()
from ignite.distributed.utils import _SerialModel, _set_model
# We need to set synced model to initial state
_set_model(_SerialModel())
def test_idist_methods_in_xla_context():
# We explicitly set _model as _SerialModel
# then call idist.* methods and check that they give correct values
from ignite.distributed.utils import _set_model, _SerialModel
_set_model(_SerialModel())
_test_distrib_config(local_rank=0, backend="xla-tpu", ws=1, true_device="xla", rank=0)
def _test_sync(cls):
from ignite.distributed.utils import _SerialModel, _set_model
_set_model(_SerialModel())
sync()
from ignite.distributed.utils import _model
assert isinstance(_model, cls), f"{type(_model)} vs {cls}"
def _test_idist_methods_in_native_context(backend, device, local_rank):
# We explicitly set _model as _SerialModel
# then call idist.* methods and check that they give correct values
from ignite.distributed.utils import _SerialModel, _set_model
_set_model(_SerialModel())
ws = dist.get_world_size()
rank = dist.get_rank()
_test_distrib_config(local_rank, backend=backend, ws=ws, true_device=device, rank=rank)
def _test_sync(cls):
from ignite.distributed.utils import _set_model, _SerialModel
_set_model(_SerialModel())
sync()
from ignite.distributed.utils import _model
assert isinstance(_model, cls), "{} vs {}".format(type(_model), cls)
def _test_idist_methods_in_xla_context_in_child_proc(index):
# We explicitly set _model as _SerialModel
# then call idist.* methods and check that they give correct values
from ignite.distributed.utils import _set_model, _SerialModel
_set_model(_SerialModel())
import torch_xla.core.xla_model as xm
_test_distrib_config(
local_rank=index, backend="xla-tpu", ws=xm.xrt_world_size(), true_device="xla", rank=xm.get_ordinal()
)
def _test_idist_methods_in_native_context_set_local_rank(backend, device, local_rank):
# We explicitly set _model as _SerialModel
# then call idist.* methods and check that they give correct values
from ignite.distributed.utils import _SerialModel, _set_model
_set_model(_SerialModel())
lrank = int(os.environ["LOCAL_RANK"])
del os.environ["LOCAL_RANK"]
ws = dist.get_world_size()
rank = dist.get_rank()
idist.set_local_rank(local_rank)
_test_distrib_config(local_rank=local_rank, backend=backend, ws=ws, true_device=device, rank=rank)
os.environ["LOCAL_RANK"] = str(lrank)
