pytorch-lightning

1
import os
2
import sys
3
from unittest import mock
4

5
import pytest
6
from lightning.fabric.utilities.rank_zero import _get_rank
7

8

9
@pytest.mark.parametrize(
10
    ("env_vars", "expected"),
11
    [
12
        ({"RANK": "0"}, 1),
13
        ({"SLURM_PROCID": "0"}, 1),
14
        ({"LOCAL_RANK": "0"}, 1),
15
        ({"JSM_NAMESPACE_RANK": "0"}, 1),
16
        ({}, 1),
17
        ({"RANK": "1"}, None),
18
        ({"SLURM_PROCID": "2"}, None),
19
        ({"LOCAL_RANK": "3"}, None),
20
        ({"JSM_NAMESPACE_RANK": "4"}, None),
21
    ],
22
)
23
def test_rank_zero_known_environment_variables(env_vars, expected, monkeypatch):
24
    """Test that rank environment variables are properly checked for rank_zero_only."""
25
    with mock.patch.dict(os.environ, env_vars):
26
        # force module reload to re-trigger the rank_zero_only.rank global computation
27
        monkeypatch.delitem(sys.modules, "lightning_utilities.core.rank_zero", raising=False)
28
        monkeypatch.delitem(sys.modules, "lightning.fabric.utilities.rank_zero", raising=False)
29
        from lightning.fabric.utilities.rank_zero import rank_zero_only
30

31
        @rank_zero_only
32
        def foo():
33
            return 1
34

35
        assert foo() == expected
36

37

38
@pytest.mark.parametrize(
39
    ("environ", "expected_rank"),
40
    [
41
        ({"JSM_NAMESPACE_RANK": "3"}, 3),
42
        ({"JSM_NAMESPACE_RANK": "3", "SLURM_PROCID": "2"}, 2),
43
        ({"JSM_NAMESPACE_RANK": "3", "SLURM_PROCID": "2", "LOCAL_RANK": "1"}, 1),
44
        ({"JSM_NAMESPACE_RANK": "3", "SLURM_PROCID": "2", "LOCAL_RANK": "1", "RANK": "0"}, 0),
45
    ],
46
)
47
def test_rank_zero_priority(environ, expected_rank):
48
    """Test the priority in which the rank gets determined when multiple environment variables are available."""
49
    with mock.patch.dict(os.environ, environ):
50
        assert _get_rank() == expected_rank
51