pytorch-lightning

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.
14
import logging
15
import os
16
from unittest import mock
17

18
import pytest
19
from lightning.fabric.plugins.environments import KubeflowEnvironment
20

21

22
@mock.patch.dict(os.environ, {}, clear=True)
23
def test_default_attributes():
24
    """Test the default attributes when no environment variables are set."""
25
    env = KubeflowEnvironment()
26
    assert env.creates_processes_externally
27

28
    with pytest.raises(KeyError):
29
        # MASTER_ADDR is required
30
        env.main_address
31
    with pytest.raises(KeyError):
32
        # MASTER_PORT is required
33
        env.main_port
34
    with pytest.raises(KeyError):
35
        # WORLD_SIZE is required
36
        env.world_size()
37
    with pytest.raises(KeyError):
38
        # RANK is required
39
        env.global_rank()
40
    assert env.local_rank() == 0
41

42

43
@mock.patch.dict(
44
    os.environ,
45
    {
46
        "KUBERNETES_PORT": "tcp://127.0.0.1:443",
47
        "MASTER_ADDR": "1.2.3.4",
48
        "MASTER_PORT": "500",
49
        "WORLD_SIZE": "20",
50
        "RANK": "1",
51
    },
52
)
53
def test_attributes_from_environment_variables(caplog):
54
    """Test that the torchelastic cluster environment takes the attributes from the environment variables."""
55
    env = KubeflowEnvironment()
56
    assert env.main_address == "1.2.3.4"
57
    assert env.main_port == 500
58
    assert env.world_size() == 20
59
    assert env.global_rank() == 1
60
    assert env.local_rank() == 0
61
    assert env.node_rank() == 1
62
    # setter should be no-op
63
    with caplog.at_level(logging.DEBUG, logger="lightning.fabric.plugins.environments"):
64
        env.set_global_rank(100)
65
    assert env.global_rank() == 1
66
    assert "setting global rank is not allowed" in caplog.text
67

68
    caplog.clear()
69

70
    with caplog.at_level(logging.DEBUG, logger="lightning.fabric.plugins.environments"):
71
        env.set_world_size(100)
72
    assert env.world_size() == 20
73
    assert "setting world size is not allowed" in caplog.text
74

75

76
def test_detect_kubeflow():
77
    """Test that the KubeflowEnvironment does not support auto-detection."""
78
    with pytest.raises(NotImplementedError, match="can't be detected automatically"):
79
        KubeflowEnvironment.detect()
80