def test__native_dist_model_warning_index_less_localrank( local_rank, world_size): assert _NativeDistModel.create_from_context() is None dist.init_process_group("nccl", "tcp://0.0.0.0:2222", world_size=world_size, rank=local_rank) dist.barrier() # We deliberately incorrectly set cuda device to 0 torch.cuda.set_device(0) model = _NativeDistModel.create_from_context() assert isinstance(model, _NativeDistModel), f"{type(model)} vs _NativeDistModel" if local_rank == 1: with pytest.warns( UserWarning, match=r"Current device index is less than current local rank." ): model.device() dist.destroy_process_group()
def _test__native_dist_model_create_from_context_no_local_rank(): if "LOCAL_RANK" in os.environ: del os.environ["LOCAL_RANK"] from ignite.distributed.comp_models.base import ComputationModel if ComputationModel._ext_local_rank is not None: ComputationModel._ext_local_rank = None with pytest.warns(UserWarning, match=r"Local rank information for native distributed setting will be initialized"): _NativeDistModel.create_from_context()
def _test__native_dist_model_create_from_context_dist(local_rank, rank, world_size, true_backend, true_device): assert _NativeDistModel.create_from_context() is None dist.init_process_group(true_backend, "tcp://0.0.0.0:2222", world_size=world_size, rank=rank) dist.barrier() true_conf = { "device": true_device, "local_rank": local_rank, "rank": rank, "world_size": world_size, "node_index": 0, "nnodes": 1, "nproc_per_node": world_size, } _test__native_dist_model_create_from_context_env_local_rank(true_conf) _test__native_dist_model_create_from_context_set_local_rank(true_conf) dist.destroy_process_group()
def _test__native_dist_model_create_from_context_env_local_rank(true_conf): import os remove_lrank = False if "LOCAL_RANK" not in os.environ: os.environ["LOCAL_RANK"] = str(true_conf["local_rank"]) remove_lrank = True model = _NativeDistModel.create_from_context() _assert_model(model, true_conf) if remove_lrank: del os.environ["LOCAL_RANK"]
def _test__native_dist_model_create_from_context_set_local_rank(true_conf): from ignite.distributed.comp_models.base import ComputationModel lrank = None if "LOCAL_RANK" in os.environ: lrank = os.environ["LOCAL_RANK"] del os.environ["LOCAL_RANK"] ComputationModel._ext_local_rank = true_conf["local_rank"] model = _NativeDistModel.create_from_context() _assert_model(model, true_conf) ComputationModel._ext_local_rank = None if lrank is not None: os.environ["LOCAL_RANK"] = lrank