Repository: Lightning-AI/pytorch-lightning
File: tests/tests_fabric/utilities/test_distributed.py (Python)

Code before:
    with (
        mock.patch("lightning.fabric.utilities.distributed._distributed_is_initialized", return_value=True),
        mock.patch("lightning.fabric.utilities.distributed.torch.distributed") as dist_mock,
    ):
        barrier.__enter__()
        dist_mock.new_group.assert_called_once()
        assert barrier.barrier == barrier.group.monitored_barrier
        assert barrier.barrier.call_count == 0
        barrier()
        assert barrier.barrier.call_count == 1
        barrier.__exit__(None, None, None)
        assert barrier.barrier.call_count == 2
        dist_mock.destroy_process_group.assert_called_once()


@mock.patch("lightning.fabric.utilities.distributed.atexit")
@mock.patch("lightning.fabric.utilities.distributed.torch.distributed.init_process_group")
def test_init_dist_connection_registers_destruction_handler(_, atexit_mock):
    _init_dist_connection(LightningEnvironment(), "nccl")
    atexit_mock.register.assert_called_once_with(_destroy_dist_connection)
    atexit_mock.reset_mock()
    _init_dist_connection(LightningEnvironment(), "gloo")
    atexit_mock.register.assert_not_called()


def test_get_default_process_group_backend_for_device():
    # register a custom backend for test
    torch.utils.rename_privateuse1_backend("pcu")

    def mock_backend(store, group_rank, group_size, timeout):
        pass

    torch.distributed.Backend.register_backend(
        "pccl",
        lambda store, group_rank, group_size, timeout: mock_backend(store, group_rank, group_size, timeout),
        devices=["pcu"],
    )

    # test that the default backend is correctly set for each device
    devices = [torch.device("cpu"), torch.device("cuda:0"), torch.device("pcu:0")]
    backends = ["gloo", "nccl", "pccl"]
    for device, backend in zip(devices, backends):
        assert _get_default_process_group_backend_for_device(device) == backend


@RunIf(min_torch="2.4")
def test_is_dtensor(monkeypatch):
    from torch.distributed._tensor import DTensor

    assert _is_dtensor(Mock(spec=DTensor))
    assert not _is_dtensor(torch.zeros(2, 2))

Review comment (suggestion):
```suggestion
def test_get_default_process_group_backend_for_device():
    """Test that each device type maps to its correct default process group backend."""
```

Code after:
    with (
        mock.patch("lightning.fabric.utilities.distributed._distributed_is_initialized", return_value=True),
        mock.patch("lightning.fabric.utilities.distributed.torch.distributed") as dist_mock,
    ):
        barrier.__enter__()
        dist_mock.new_group.assert_called_once()
        assert barrier.barrier == barrier.group.monitored_barrier
        assert barrier.barrier.call_count == 0
        barrier()
        assert barrier.barrier.call_count == 1
        barrier.__exit__(None, None, None)
        assert barrier.barrier.call_count == 2
        dist_mock.destroy_process_group.assert_called_once()


@mock.patch("lightning.fabric.utilities.distributed.atexit")
@mock.patch("lightning.fabric.utilities.distributed.torch.distributed.init_process_group")
def test_init_dist_connection_registers_destruction_handler(_, atexit_mock):
    _init_dist_connection(LightningEnvironment(), "nccl")
    atexit_mock.register.assert_called_once_with(_destroy_dist_connection)
    atexit_mock.reset_mock()
    _init_dist_connection(LightningEnvironment(), "gloo")
    atexit_mock.register.assert_not_called()


def test_get_default_process_group_backend_for_device():
    """Test that each device type maps to its correct default process group backend."""
    # register a custom backend for test
    torch.utils.rename_privateuse1_backend("pcu")

    def mock_backend(store, group_rank, group_size, timeout):
        pass

    torch.distributed.Backend.register_backend(
        "pccl",
        lambda store, group_rank, group_size, timeout: mock_backend(store, group_rank, group_size, timeout),
        devices=["pcu"],
    )

    # test that the default backend is correctly set for each device
    devices = [torch.device("cpu"), torch.device("cuda:0"), torch.device("pcu:0")]
    backends = ["gloo", "nccl", "pccl"]
    for device, backend in zip(devices, backends):
        assert _get_default_process_group_backend_for_device(device) == backend


@RunIf(min_torch="2.4")
def test_is_dtensor(monkeypatch):
    from torch.distributed._tensor import DTensor

    assert _is_dtensor(Mock(spec=DTensor))