pytorch-lightning
93 строки · 3.2 Кб
1# Copyright The Lightning AI team.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14import logging
15import os
16import re
17from unittest import mock
18
19import pytest
20from lightning.fabric.plugins.environments import TorchElasticEnvironment
21
22
23@mock.patch.dict(os.environ, {}, clear=True)
24def test_default_attributes():
25"""Test the default attributes when no environment variables are set."""
26env = TorchElasticEnvironment()
27assert env.creates_processes_externally
28assert env.main_address == "127.0.0.1"
29assert env.main_port == 12910
30with pytest.raises(KeyError):
31# world size is required to be passed as env variable
32env.world_size()
33with pytest.raises(KeyError):
34# local rank is required to be passed as env variable
35env.local_rank()
36assert env.node_rank() == 0
37
38
39@mock.patch.dict(
40os.environ,
41{
42"MASTER_ADDR": "1.2.3.4",
43"MASTER_PORT": "500",
44"WORLD_SIZE": "20",
45"RANK": "1",
46"LOCAL_RANK": "2",
47"GROUP_RANK": "3",
48},
49)
50def test_attributes_from_environment_variables(caplog):
51"""Test that the torchelastic cluster environment takes the attributes from the environment variables."""
52env = TorchElasticEnvironment()
53assert env.main_address == "1.2.3.4"
54assert env.main_port == 500
55assert env.world_size() == 20
56assert env.global_rank() == 1
57assert env.local_rank() == 2
58assert env.node_rank() == 3
59# setter should be no-op
60with caplog.at_level(logging.DEBUG, logger="lightning.fabric.plugins.environments"):
61env.set_global_rank(100)
62assert env.global_rank() == 1
63assert "setting global rank is not allowed" in caplog.text
64
65caplog.clear()
66
67with caplog.at_level(logging.DEBUG, logger="lightning.fabric.plugins.environments"):
68env.set_world_size(100)
69assert env.world_size() == 20
70assert "setting world size is not allowed" in caplog.text
71
72
73def test_detect():
74"""Test the detection of a torchelastic environment configuration."""
75with mock.patch.dict(os.environ, {}, clear=True):
76assert not TorchElasticEnvironment.detect()
77
78with mock.patch.dict(
79os.environ,
80{
81"TORCHELASTIC_RUN_ID": "",
82},
83):
84assert TorchElasticEnvironment.detect()
85
86
87@mock.patch.dict(os.environ, {"WORLD_SIZE": "8"})
88def test_validate_user_settings():
89"""Test that the environment can validate the number of devices and nodes set in Fabric/Trainer."""
90env = TorchElasticEnvironment()
91env.validate_settings(num_devices=4, num_nodes=2)
92with pytest.raises(ValueError, match=re.escape("the product (2 * 2) does not match the world size (8)")):
93env.validate_settings(num_devices=2, num_nodes=2)
94