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

15
from lightning.fabric.strategies import STRATEGY_REGISTRY
16

17

18
def test_strategy_registry_with_new_strategy():
19
    class TestStrategy:
20
        strategy_name = "test_strategy"
21

22
        def __init__(self, param1, param2):
23
            self.param1 = param1
24
            self.param2 = param2
25

26
    strategy_name = "test_strategy"
27
    strategy_description = "Test Strategy"
28

29
    # TODO(fabric): Registering classes that do not inherit from Strategy should not be allowed
30
    STRATEGY_REGISTRY.register(strategy_name, TestStrategy, description=strategy_description, param1="abc", param2=123)
31

32
    assert strategy_name in STRATEGY_REGISTRY
33
    assert STRATEGY_REGISTRY[strategy_name]["description"] == strategy_description
34
    assert STRATEGY_REGISTRY[strategy_name]["init_params"] == {"param1": "abc", "param2": 123}
35
    assert STRATEGY_REGISTRY[strategy_name]["strategy_name"] == "test_strategy"
36
    assert isinstance(STRATEGY_REGISTRY.get(strategy_name), TestStrategy)
37

38
    STRATEGY_REGISTRY.remove(strategy_name)
39
    assert strategy_name not in STRATEGY_REGISTRY
40

41

42
def test_available_strategies_in_registry():
43
    expected = {
44
        "ddp",
45
        "deepspeed",
46
        "deepspeed_stage_1",
47
        "deepspeed_stage_1_offload",
48
        "deepspeed_stage_2",
49
        "deepspeed_stage_2_offload",
50
        "deepspeed_stage_3",
51
        "deepspeed_stage_3_offload",
52
        "deepspeed_stage_3_offload_nvme",
53
        "ddp_spawn",
54
        "ddp_fork",
55
        "ddp_notebook",
56
        "single_tpu",  # legacy
57
        "single_xla",
58
        "xla",
59
        "xla_fsdp",
60
        "dp",
61
        "fsdp",
62
        "fsdp_cpu_offload",
63
    }
64
    assert set(STRATEGY_REGISTRY.available_strategies()) == expected
65