def select_strategy(self) -> Strategy:
if isinstance(self.distributed_backend, Accelerator) and self.distributed_backend.strategy is not None:
plugin = self.distributed_backend.strategy
elif self.use_ddp2:
plugin = DDP2Strategy(parallel_devices=self.parallel_devices, cluster_environment=self.cluster_environment)
elif self.use_ddp and self.use_deepspeed:
plugin = DeepSpeedStrategy(
cluster_environment=self.select_cluster_environment(), parallel_devices=self.parallel_devices
)
elif self.use_ddp:
use_slurm_ddp = self.use_ddp and self._is_slurm_managing_tasks()
use_torchelastic_ddp = self.use_ddp and TorchElasticEnvironment.detect()
use_kubeflow_ddp = self.use_ddp and KubeflowEnvironment.detect()
use_ddp_spawn = self._strategy_type == _StrategyType.DDP_SPAWN
use_ddp_cpu_spawn = use_ddp_spawn and self.use_cpu
use_tpu_spawn = self.use_tpu and self._strategy_type == _StrategyType.TPU_SPAWN
use_ddp_cpu_torch_elastic = use_ddp_cpu_spawn and TorchElasticEnvironment.detect()
use_ddp_cpu_kubeflow = use_ddp_cpu_spawn and KubeflowEnvironment.detect()
use_ddp_cpu_slurm = use_ddp_cpu_spawn and self._is_slurm_managing_tasks()
use_ddp_sharded = self._strategy_type == _StrategyType.DDP_SHARDED
use_ddp_sharded_spawn = self._strategy_type == _StrategyType.DDP_SHARDED_SPAWN
use_ddp_fully_sharded = self._strategy_type == _StrategyType.DDP_FULLY_SHARDED
if use_tpu_spawn:
ddp_strategy_cls = TPUSpawnStrategy
elif use_ddp_sharded:
ddp_strategy_cls = DDPShardedStrategy
elif use_ddp_sharded_spawn:
ddp_strategy_cls = DDPSpawnShardedStrategy
elif (
use_ddp_cpu_slurm
or use_slurm_ddp
or use_ddp_cpu_torch_elastic
or use_torchelastic_ddp
or use_kubeflow_ddp
or use_ddp_cpu_kubeflow
):
ddp_strategy_cls = DDPStrategy
elif use_ddp_spawn or use_ddp_cpu_spawn:
ddp_strategy_cls = DDPSpawnStrategy
elif use_ddp_fully_sharded:
ddp_strategy_cls = DDPFullyShardedStrategy
else:
ddp_strategy_cls = DDPStrategy
plugin = ddp_strategy_cls(
parallel_devices=self.parallel_devices, cluster_environment=self.cluster_environment
)
elif self.use_dp:
plugin = DataParallelStrategy(parallel_devices=self.parallel_devices)
elif self.use_horovod:
plugin = HorovodStrategy(parallel_devices=self.parallel_devices)
elif self.use_tpu and isinstance(self.tpu_cores, list):
plugin = SingleTPUStrategy(self.tpu_id)
elif self.use_ipu:
plugin = IPUStrategy(parallel_devices=self.parallel_devices)
else:
single_gpu_ordinal = device_parser.determine_root_gpu_device(self.parallel_device_ids)
plugin = SingleDeviceStrategy(device=single_gpu_ordinal if self.use_gpu else "cpu")
return plugin
def select_training_type_plugin(self) -> TrainingTypePlugin:
if self.use_ddp2:
plugin = DDP2Plugin(
parallel_devices=self.parallel_devices,
num_nodes=self.num_nodes,
cluster_environment=self.cluster_environment,
sync_batchnorm=self.sync_batchnorm,
)
elif self.use_ddp and self.use_deepspeed:
plugin = DeepSpeedPlugin(
num_nodes=self.num_nodes,
cluster_environment=self.select_cluster_environment(),
parallel_devices=self.parallel_devices
)
elif self.use_ddp:
use_slurm_ddp = self.use_ddp and self.is_slurm_managing_tasks
use_torchelastic_ddp = self.use_ddp and TorchElasticEnvironment.is_using_torchelastic()
use_ddp_spawn = self._distrib_type == DistributedType.DDP_SPAWN
use_ddp_cpu_spawn = self.use_ddp and self.on_cpu
use_tpu_spawn = self.on_tpu and self._distrib_type == DistributedType.TPU_SPAWN
use_ddp_cpu_torch_elastic = use_ddp_cpu_spawn and TorchElasticEnvironment.is_using_torchelastic()
use_ddp_cpu_slurm = use_ddp_cpu_spawn and self.is_slurm_managing_tasks
use_ddp_sharded = self._distrib_type == DistributedType.DDP_SHARDED
use_ddp_sharded_spawn = self._distrib_type == DistributedType.DDP_SHARDED_SPAWN
# TODO: decouple from TE
# ddp script mode uses the same flags as TE
if os.environ.get("PL_IN_DDP_SUBPROCESS", False):
use_torchelastic_ddp = False
if use_tpu_spawn:
ddp_plugin_cls = TPUSpawnPlugin
elif use_ddp_sharded:
ddp_plugin_cls = DDPShardedPlugin
elif use_ddp_sharded_spawn:
ddp_plugin_cls = DDPSpawnShardedPlugin
elif use_ddp_cpu_slurm or use_slurm_ddp or use_ddp_cpu_torch_elastic or use_torchelastic_ddp:
ddp_plugin_cls = DDPPlugin
elif use_ddp_spawn or use_ddp_cpu_spawn:
ddp_plugin_cls = DDPSpawnPlugin
else:
ddp_plugin_cls = DDPPlugin
plugin = ddp_plugin_cls(
parallel_devices=self.parallel_devices,
num_nodes=self.num_nodes,
cluster_environment=self.cluster_environment,
sync_batchnorm=self.sync_batchnorm,
)
elif self.use_dp:
plugin = DataParallelPlugin(parallel_devices=self.parallel_devices)
elif self.use_horovod:
plugin = HorovodPlugin(parallel_devices=self.parallel_devices)
elif self.on_tpu and isinstance(self.tpu_cores, list):
plugin = SingleTPUPlugin(self.tpu_id)
else:
single_gpu_ordinal = device_parser.determine_root_gpu_device(self.parallel_device_ids)
plugin = SingleDevicePlugin(device=torch.device(f"cuda:{single_gpu_ordinal}" if self.on_gpu else "cpu"))
return plugin
def select_cluster_environment(self) -> ClusterEnvironment:
if self._cluster_environment is not None:
return self._cluster_environment
if self.is_slurm_managing_tasks:
env = SLURMEnvironment()
elif TorchElasticEnvironment.is_using_torchelastic():
env = TorchElasticEnvironment()
else:
env = LightningEnvironment()
return env
def test_detect_after_1_9_1():
"""Test the detection of a torchelastic environment configuration after 1.9.1."""
with mock.patch.dict(os.environ, {}):
assert not TorchElasticEnvironment.detect()
with mock.patch.dict(
os.environ,
{
"TORCHELASTIC_RUN_ID": "",
},
):
assert TorchElasticEnvironment.detect()
def select_cluster_environment(self):
if self.cluster_environment is not None:
return self.cluster_environment
if self.is_slurm_managing_tasks:
env = SLURMEnvironment()
elif self.is_using_torchelastic:
env = TorchElasticEnvironment()
# TODO: decouple DDP from TE
# maybe introduce a DefaultEnvironment?
os.environ["PL_IN_DDP_SUBPROCESS"] = "1"
else:
# TODO: maybe introduce a DefaultEnvironment?
env = TorchElasticEnvironment()
return env
def test_default_attributes():
""" Test the default attributes when no environment variables are set. """
env = TorchElasticEnvironment()
assert env.creates_children()
assert env.master_address() == "127.0.0.1"
assert env.master_port() == 12910
assert env.world_size() is None
with pytest.raises(KeyError):
# local rank is required to be passed as env variable
env.local_rank()
assert env.node_rank() == 0
def test_detect_before_1_9_1():
"""Test the detection of a torchelastic environment configuration before 1.9.1."""
with mock.patch.dict(os.environ, {}):
assert not TorchElasticEnvironment.detect()
with mock.patch.dict(
os.environ,
{
"RANK": "",
"GROUP_RANK": "",
"LOCAL_RANK": "",
"LOCAL_WORLD_SIZE": "",
},
):
assert TorchElasticEnvironment.detect()
def select_cluster_environment(self) -> ClusterEnvironment:
if self._cluster_environment is not None:
return self._cluster_environment
if self._is_slurm_managing_tasks():
env = SLURMEnvironment()
rank_zero_info("Multiprocessing is handled by SLURM.")
elif TorchElasticEnvironment.is_using_torchelastic():
env = TorchElasticEnvironment()
elif KubeflowEnvironment.is_using_kubeflow():
env = KubeflowEnvironment()
elif LSFEnvironment.is_using_lsf():
env = LSFEnvironment()
else:
env = LightningEnvironment()
return env
def environment_combinations():
expected = dict(global_rank=3, local_rank=1, node_rank=1, world_size=4)
# Lightning
variables = {
"CUDA_VISIBLE_DEVICES": "0,1,2,4",
"LOCAL_RANK": "1",
"NODE_RANK": "1",
"WORLD_SIZE": "8"
}
environment = LightningEnvironment()
yield environment, variables, expected
# SLURM
variables = {
"CUDA_VISIBLE_DEVICES": "0,1,2,4",
"SLURM_JOB_NAME": "SOME_NAME",
"SLURM_LOCALID": "1",
"SLURM_NODEID": "1",
"SLURM_PROCID": "3",
"SLURM_NTASKS": "4",
}
environment = SLURMEnvironment()
yield environment, variables, expected
# TorchElastic
variables = {
"CUDA_VISIBLE_DEVICES": "0,1,2,4",
"LOCAL_RANK": "1",
"GROUP_RANK": "1",
"RANK": "3",
"WORLD_SIZE": "4",
"LOCAL_WORLD_SIZE": "2",
}
environment = TorchElasticEnvironment()
yield environment, variables, expected
def _check_strategy_and_fallback(self) -> None:
"""Checks edge cases when the strategy selection was a string input, and we need to fall back to a
different choice depending on other parameters or the environment."""
# current fallback and check logic only apply to user pass in str config and object config
# TODO this logic should apply to both str and object config
strategy_flag = "" if isinstance(self._strategy_flag, Strategy) else self._strategy_flag
if strategy_flag in ("ddp_spawn", "ddp_spawn_find_unused_parameters_false") and (
TorchElasticEnvironment.detect() or KubeflowEnvironment.detect() or self._is_slurm_managing_tasks()
):
strategy_flag = "ddp"
if strategy_flag in ("dp", "ddp2") and self._accelerator_flag == "cpu":
rank_zero_warn(f"{strategy_flag!r} is not supported on CPUs, hence setting `strategy='ddp'`.")
strategy_flag = "ddp"
if (
strategy_flag in DDPFullyShardedNativeStrategy.get_registered_strategies()
or isinstance(self._strategy_flag, DDPFullyShardedNativeStrategy)
) and self._accelerator_flag != "gpu":
raise MisconfigurationException(
f"You selected strategy to be `{DDPFullyShardedNativeStrategy.strategy_name}`, "
"but GPU accelerator is not used."
)
if strategy_flag:
self._strategy_flag = strategy_flag
def _check_strategy_and_fallback(self) -> None:
"""Checks edge cases when the strategy selection was a string input, and we need to fall back to a
different choice depending on other parameters or the environment."""
# current fallback and check logic only apply to user pass in str config and object config
# TODO this logic should apply to both str and object config
strategy_flag = "" if isinstance(self._strategy_flag,
Strategy) else self._strategy_flag
if strategy_flag in ("ddp_spawn",
"ddp_spawn_find_unused_parameters_false") and (
TorchElasticEnvironment.detect()
or KubeflowEnvironment.detect()
or self._is_slurm_managing_tasks()):
strategy_flag = "ddp"
if strategy_flag == "dp" and self._accelerator_flag == "cpu":
rank_zero_warn(
f"{strategy_flag!r} is not supported on CPUs, hence setting `strategy='ddp'`."
)
strategy_flag = "ddp"
if (strategy_flag
in DDPFullyShardedNativeStrategy.get_registered_strategies() or
isinstance(self._strategy_flag, DDPFullyShardedNativeStrategy)
) and self._accelerator_flag not in ("cuda", "gpu"):
raise MisconfigurationException(
f"You selected strategy to be `{DDPFullyShardedNativeStrategy.strategy_name}`, "
"but GPU accelerator is not used.")
if strategy_flag in _DDP_FORK_ALIASES and "fork" not in torch.multiprocessing.get_all_start_methods(
):
raise ValueError(
f"You selected `Trainer(strategy='{strategy_flag}')` but process forking is not supported on this"
f" platform. We recommed `Trainer(strategy='ddp_spawn')` instead."
)
if strategy_flag:
self._strategy_flag = strategy_flag
def select_cluster_environment(self) -> ClusterEnvironment:
if self._cluster_environment is not None:
return self._cluster_environment
if self.is_slurm_managing_tasks:
env = SLURMEnvironment()
# TODO: decouple DDP from SLURM
# refactor and let generic cluster env hold the information about who spawns the processes
os.environ["PL_IN_DDP_SUBPROCESS"] = "1"
elif self.is_using_torchelastic:
env = TorchElasticEnvironment()
# TODO: decouple DDP from TE
# refactor and let generic cluster env hold the information about who spawns the processes
os.environ["PL_IN_DDP_SUBPROCESS"] = "1"
else:
# TODO: maybe introduce a DefaultEnvironment?
env = TorchElasticEnvironment()
return env
def test_sync_batchnorm_ddp(tmpdir):
seed_everything(234)
set_random_master_port()
# define datamodule and dataloader
dm = MNISTDataModule()
dm.prepare_data()
dm.setup(stage=None)
train_dataloader = dm.train_dataloader()
model = SyncBNModule()
bn_outputs = []
# shuffle is false by default
for batch_idx, batch in enumerate(train_dataloader):
x, _ = batch
_, out_bn = model.forward(x, batch_idx)
bn_outputs.append(out_bn)
# get 3 steps
if batch_idx == 2:
break
bn_outputs = [x.cuda() for x in bn_outputs]
# reset datamodule
# batch-size = 16 because 2 GPUs in DDP
dm = MNISTDataModule(batch_size=16, dist_sampler=True)
dm.prepare_data()
dm.setup(stage=None)
model = SyncBNModule(gpu_count=2, bn_targets=bn_outputs)
trainer = Trainer(gpus=2,
num_nodes=1,
accelerator='ddp_spawn',
max_epochs=1,
max_steps=3,
sync_batchnorm=True,
num_sanity_val_steps=0,
replace_sampler_ddp=False,
plugins=[
DDPSpawnPlugin(
parallel_devices=[
torch.device("cuda", 0),
torch.device("cuda", 1)
],
num_nodes=1,
sync_batchnorm=True,
cluster_environment=TorchElasticEnvironment(),
find_unused_parameters=True)
])
trainer.fit(model, dm)
assert trainer.state == TrainerState.FINISHED, "Sync batchnorm failing with DDP"
def test_attributes_from_environment_variables():
""" Test that the torchelastic cluster environment takes the attributes from the environment variables. """
env = TorchElasticEnvironment()
assert env.master_address() == "1.2.3.4"
assert env.master_port() == 500
assert env.world_size() == 20
assert env.local_rank() == 2
assert env.node_rank() == 3
def is_using_torchelastic(self) -> bool:
"""
.. deprecated:: v1.3
Will be removed in v1.5.0.
Returns:
``True`` if the current process was launched using the torchelastic command.
"""
rank_zero_deprecation(
"The property `AcceleratorConnector.is_using_torchelastic` was deprecated in v1.3"
" and will be removed in 1.5. Use `TorchElasticEnvironment.is_using_torchelastic()` instead.",
)
return TorchElasticEnvironment.is_using_torchelastic()
def test_default_attributes():
"""Test the default attributes when no environment variables are set."""
env = TorchElasticEnvironment()
assert env.creates_processes_externally
assert env.main_address == "127.0.0.1"
assert env.main_port == 12910
with pytest.raises(KeyError):
# world size is required to be passed as env variable
env.world_size()
with pytest.raises(KeyError):
# local rank is required to be passed as env variable
env.local_rank()
assert env.node_rank() == 0
def test_v1_4_0_deprecated_lightning_distributed_data_parallel(tmpdir):
model = BoringModel()
trainer = Trainer(
default_root_dir=tmpdir,
fast_dev_run=True,
gpus=2,
accelerator="ddp_spawn",
plugins=[
CustomDDPPlugin(
parallel_devices=[torch.device("cuda", 0), torch.device("cuda", 1)],
cluster_environment=TorchElasticEnvironment(),
)
]
)
trainer.fit(model)
def parse_gpu_ids(
gpus: Optional[Union[int, str, List[int]]]) -> Optional[List[int]]:
"""
Parses the GPU ids given in the format as accepted by the
:class:`~pytorch_lightning.trainer.Trainer`.
Args:
gpus: An int -1 or string '-1' indicate that all available GPUs should be used.
A list of ints or a string containing list of comma separated integers
indicates specific GPUs to use.
An int 0 means that no GPUs should be used.
Any int N > 0 indicates that GPUs [0..N) should be used.
Returns:
a list of gpus to be used or ``None`` if no GPUs were requested
If no GPUs are available but the value of gpus variable indicates request for GPUs
then a MisconfigurationException is raised.
"""
# Check that gpus param is None, Int, String or List
_check_data_type(gpus)
# Handle the case when no gpus are requested
if gpus is None or isinstance(gpus, int) and gpus == 0:
return None
if _compare_version("pytorch_lightning", operator.ge,
"1.5") and isinstance(gpus,
str) and gpus.strip() == "0":
# TODO: in v1.5 combine this with the above if statement
return None
# We know user requested GPUs therefore if some of the
# requested GPUs are not available an exception is thrown.
gpus = _normalize_parse_gpu_string_input(gpus)
gpus = _normalize_parse_gpu_input_to_list(gpus)
if not gpus:
raise MisconfigurationException(
"GPUs requested but none are available.")
if TorchElasticEnvironment.is_using_torchelastic(
) and len(gpus) != 1 and len(_get_all_available_gpus()) == 1:
# omit sanity check on torchelastic as by default shows one visible GPU per process
return gpus
return _sanitize_gpu_ids(gpus)
def parse_gpu_ids(gpus: Optional[Union[int, str, List[int]]]) -> Optional[List[int]]:
"""
Parses the GPU ids given in the format as accepted by the
:class:`~pytorch_lightning.trainer.Trainer`.
Args:
gpus: An int -1 or string '-1' indicate that all available GPUs should be used.
A list of unique ints or a string containing list of comma separated unique integers
indicates specific GPUs to use.
An int 0 means that no GPUs should be used.
Any int N > 0 indicates that GPUs [0..N) should be used.
Returns:
a list of gpus to be used or ``None`` if no GPUs were requested
Raises:
MisconfigurationException:
If no GPUs are available but the value of gpus variable indicates request for GPUs
"""
# Check that gpus param is None, Int, String or List
_check_data_type(gpus)
# Handle the case when no gpus are requested
if gpus is None or (isinstance(gpus, int) and gpus == 0) or str(gpus).strip() in ("0", "[]"):
return None
# We know user requested GPUs therefore if some of the
# requested GPUs are not available an exception is thrown.
gpus = _normalize_parse_gpu_string_input(gpus)
gpus = _normalize_parse_gpu_input_to_list(gpus)
if not gpus:
raise MisconfigurationException("GPUs requested but none are available.")
if TorchElasticEnvironment.detect() and len(gpus) != 1 and len(_get_all_available_gpus()) == 1:
# omit sanity check on torchelastic as by default shows one visible GPU per process
return gpus
# Check that gpus are unique. Duplicate gpus are not supported by the backend.
_check_unique(gpus)
return _sanitize_gpu_ids(gpus)
def _check_strategy_and_fallback(self) -> None:
"""Checks edge cases when the strategy selection was a string input, and we need to fall back to a
different choice depending on other parameters or the environment."""
# current fallback and check logic only apply to user pass in str config and object config
# TODO this logic should apply to both str and object config
strategy_flag = "" if isinstance(self._strategy_flag,
Strategy) else self._strategy_flag
if strategy_flag == "ddp_cpu":
if _TPU_AVAILABLE:
raise MisconfigurationException(
"`accelerator='ddp_cpu'` is not supported on TPU machines. "
"Learn more: https://github.com/PyTorchLightning/pytorch-lightning/issues/7810"
)
if self._devices_flag == 1 and self._num_nodes_flag > 1:
strategy_flag = DDPStrategy.strategy_name
else:
strategy_flag = "ddp_spawn"
if self._accelerator_flag == "gpu":
rank_zero_warn(
"You requested one or more GPUs, but set `accelerator='ddp_cpu'`. Training will not use GPUs."
)
self._accelerator_flag = "cpu"
self.accelerator = CPUAccelerator()
if strategy_flag in ("ddp_spawn",
"ddp_spawn_find_unused_parameters_false") and (
TorchElasticEnvironment.detect()
or KubeflowEnvironment.detect()
or self._is_slurm_managing_tasks()):
strategy_flag = "ddp"
if strategy_flag in ("dp", "ddp2") and self._accelerator_flag == "cpu":
rank_zero_warn(
f"{strategy_flag!r} is not supported on CPUs, hence setting `strategy='ddp'`."
)
strategy_flag = "ddp"
if strategy_flag:
self._strategy_flag = strategy_flag
def select_training_type_plugin(self) -> TrainingTypePlugin:
if (isinstance(self.distributed_backend, Accelerator)
and self.distributed_backend.training_type_plugin is not None):
plugin = self.distributed_backend.training_type_plugin
elif self.use_ddp2:
plugin = DDP2Plugin(parallel_devices=self.parallel_devices,
cluster_environment=self.cluster_environment)
elif self.use_ddp and self.use_deepspeed:
plugin = DeepSpeedPlugin(
cluster_environment=self.select_cluster_environment(),
parallel_devices=self.parallel_devices)
elif self.use_ddp:
use_slurm_ddp = self.use_ddp and self.is_slurm_managing_tasks
use_torchelastic_ddp = self.use_ddp and TorchElasticEnvironment.is_using_torchelastic(
)
use_kubeflow_ddp = self.use_ddp and KubeflowEnvironment.is_using_kubeflow(
)
use_ddp_spawn = self._distrib_type == DistributedType.DDP_SPAWN
use_ddp_cpu_spawn = use_ddp_spawn and self.use_cpu
use_tpu_spawn = self.use_tpu and self._distrib_type == DistributedType.TPU_SPAWN
use_ddp_cpu_torch_elastic = use_ddp_cpu_spawn and TorchElasticEnvironment.is_using_torchelastic(
)
use_ddp_cpu_kubeflow = use_ddp_cpu_spawn and KubeflowEnvironment.is_using_kubeflow(
)
use_ddp_cpu_slurm = use_ddp_cpu_spawn and self.is_slurm_managing_tasks
use_ddp_sharded = self._distrib_type == DistributedType.DDP_SHARDED
use_ddp_sharded_spawn = self._distrib_type == DistributedType.DDP_SHARDED_SPAWN
use_ddp_fully_sharded = self._distrib_type == DistributedType.DDP_FULLY_SHARDED
if use_tpu_spawn:
ddp_plugin_cls = TPUSpawnPlugin
elif use_ddp_sharded:
ddp_plugin_cls = DDPShardedPlugin
elif use_ddp_sharded_spawn:
ddp_plugin_cls = DDPSpawnShardedPlugin
elif (use_ddp_cpu_slurm or use_slurm_ddp
or use_ddp_cpu_torch_elastic or use_torchelastic_ddp
or use_kubeflow_ddp or use_ddp_cpu_kubeflow):
ddp_plugin_cls = DDPPlugin
elif use_ddp_spawn or use_ddp_cpu_spawn:
ddp_plugin_cls = DDPSpawnPlugin
elif use_ddp_fully_sharded:
ddp_plugin_cls = DDPFullyShardedPlugin
else:
ddp_plugin_cls = DDPPlugin
plugin = ddp_plugin_cls(
parallel_devices=self.parallel_devices,
cluster_environment=self.cluster_environment)
elif self.use_dp:
plugin = DataParallelPlugin(parallel_devices=self.parallel_devices)
elif self.use_horovod:
plugin = HorovodPlugin(parallel_devices=self.parallel_devices)
elif self.use_tpu and isinstance(self.tpu_cores, list):
plugin = SingleTPUPlugin(self.tpu_id)
elif self.use_ipu:
plugin = IPUPlugin(parallel_devices=self.parallel_devices)
else:
single_gpu_ordinal = device_parser.determine_root_gpu_device(
self.parallel_device_ids)
plugin = SingleDevicePlugin(device=torch.device(
f"cuda:{single_gpu_ordinal}" if self.use_gpu else "cpu"))
return plugin
def test_attributes_from_environment_variables(caplog):
"""Test that the torchelastic cluster environment takes the attributes from the environment variables."""
env = TorchElasticEnvironment()
assert env.main_address == "1.2.3.4"
assert env.main_port == 500
assert env.world_size() == 20
assert env.global_rank() == 1
assert env.local_rank() == 2
assert env.node_rank() == 3
# setter should be no-op
with caplog.at_level(logging.DEBUG,
logger="pytorch_lightning.plugins.environments"):
env.set_global_rank(100)
assert env.global_rank() == 1
assert "setting global rank is not allowed" in caplog.text
caplog.clear()
with caplog.at_level(logging.DEBUG,
logger="pytorch_lightning.plugins.environments"):
env.set_world_size(100)
assert env.world_size() == 20
assert "setting world size is not allowed" in caplog.text