lightning/tests/tests_lite/test_connector.py

# Copyright The PyTorch Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License

import os
from re import escape
from typing import Any, Dict
from unittest import mock

import pytest
import torch
import torch.distributed
from tests_lite.helpers.runif import RunIf

import lightning_lite
from lightning_lite.accelerators import TPUAccelerator
from lightning_lite.accelerators.accelerator import Accelerator
from lightning_lite.accelerators.cpu import CPUAccelerator
from lightning_lite.accelerators.cuda import CUDAAccelerator
from lightning_lite.accelerators.mps import MPSAccelerator
from lightning_lite.connector import _Connector
from lightning_lite.plugins import DoublePrecision, NativeMixedPrecision, Precision, TPUPrecision
from lightning_lite.plugins.environments import (
    KubeflowEnvironment,
    LightningEnvironment,
    SLURMEnvironment,
    TorchElasticEnvironment,
)
from lightning_lite.plugins.io import TorchCheckpointIO
from lightning_lite.strategies import (
    DataParallelStrategy,
    DDPShardedStrategy,
    DDPSpawnShardedStrategy,
    DDPSpawnStrategy,
    DDPStrategy,
    DeepSpeedStrategy,
    SingleDeviceStrategy,
    SingleTPUStrategy,
    XLAStrategy,
)
from lightning_lite.strategies.ddp_spawn import _DDP_FORK_ALIASES
from lightning_lite.utilities.exceptions import MisconfigurationException


def test_accelerator_choice_cpu():
    connector = _Connector()
    assert isinstance(connector.accelerator, CPUAccelerator)
    assert isinstance(connector.strategy, SingleDeviceStrategy)


@RunIf(tpu=True, standalone=True)
@pytest.mark.parametrize(
    ["accelerator", "devices"], [("tpu", None), ("tpu", 1), ("tpu", [1]), ("tpu", 8), ("auto", 1), ("auto", 8)]
)
@mock.patch.dict(os.environ, os.environ.copy(), clear=True)
def test_accelerator_choice_tpu(accelerator, devices):
    connector = _Connector(accelerator=accelerator, devices=devices)
    assert isinstance(connector.accelerator, TPUAccelerator)
    if devices is None or (isinstance(devices, int) and devices > 1):
        # accelerator=tpu, devices=None (default) maps to devices=auto (8) and then chooses XLAStrategy
        # This behavior may change in the future: https://github.com/Lightning-AI/lightning/issues/10606
        assert isinstance(connector.strategy, XLAStrategy)
    else:
        assert isinstance(connector.strategy, SingleTPUStrategy)


@RunIf(skip_windows=True, standalone=True)
def test_strategy_choice_ddp_on_cpu():
    """Test that selecting DDPStrategy on CPU works."""
    _test_strategy_choice_ddp_and_cpu(ddp_strategy_class=DDPStrategy)


@RunIf(skip_windows=True)
def test_strategy_choice_ddp_spawn_on_cpu():
    """Test that selecting DDPSpawnStrategy on CPU works."""
    _test_strategy_choice_ddp_and_cpu(ddp_strategy_class=DDPSpawnStrategy)


def _test_strategy_choice_ddp_and_cpu(ddp_strategy_class):
    connector = _Connector(
        strategy=ddp_strategy_class(find_unused_parameters=True),
        accelerator="cpu",
        devices=2,
    )
    assert isinstance(connector.strategy, ddp_strategy_class)
    assert isinstance(connector.accelerator, CPUAccelerator)
    assert connector.strategy.num_processes == 2
    assert connector.strategy.parallel_devices == [torch.device("cpu")] * 2


@mock.patch.dict(
    os.environ,
    {
        "SLURM_NTASKS": "2",
        "SLURM_JOB_NAME": "SOME_NAME",
        "SLURM_NODEID": "0",
        "LOCAL_RANK": "0",
        "SLURM_PROCID": "0",
        "SLURM_LOCALID": "0",
    },
)
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=0)
def test_custom_cluster_environment_in_slurm_environment(_):
    """Test that we choose the custom cluster even when SLURM or TE flags are around."""

    class CustomCluster(LightningEnvironment):
        @property
        def main_address(self):
            return "asdf"

        @property
        def creates_processes_externally(self) -> bool:
            return True

    connector = _Connector(
        plugins=[CustomCluster()],
        accelerator="cpu",
        strategy="ddp",
        devices=2,
    )
    assert isinstance(connector.accelerator, CPUAccelerator)
    assert isinstance(connector.strategy, DDPStrategy)
    assert isinstance(connector.strategy.cluster_environment, CustomCluster)


@mock.patch.dict(
    os.environ,
    {
        "SLURM_NTASKS": "2",
        "SLURM_NTASKS_PER_NODE": "1",
        "SLURM_JOB_NAME": "SOME_NAME",
        "SLURM_NODEID": "0",
        "LOCAL_RANK": "0",
        "SLURM_PROCID": "0",
        "SLURM_LOCALID": "0",
    },
)
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=0)
def test_custom_accelerator(*_):
    class Accel(Accelerator):
        def setup_device(self, device: torch.device) -> None:
            pass

        def get_device_stats(self, device: torch.device) -> Dict[str, Any]:
            pass

        def teardown(self) -> None:
            pass

        @staticmethod
        def parse_devices(devices):
            return devices

        @staticmethod
        def get_parallel_devices(devices):
            return [torch.device("cpu")] * devices

        @staticmethod
        def auto_device_count() -> int:
            return 1

        @staticmethod
        def is_available() -> bool:
            return True

        @staticmethod
        def name() -> str:
            return "custom_acc_name"

    class Prec(Precision):
        pass

    class Strat(SingleDeviceStrategy):
        pass

    strategy = Strat(device=torch.device("cpu"), accelerator=Accel(), precision=Prec())
    connector = _Connector(strategy=strategy, devices=2)
    assert isinstance(connector.accelerator, Accel)
    assert isinstance(connector.strategy, Strat)
    assert isinstance(connector.precision, Prec)
    assert connector.strategy is strategy

    class Strat(DDPStrategy):
        pass

    strategy = Strat(accelerator=Accel(), precision=Prec())
    connector = _Connector(strategy=strategy, devices=2)
    assert isinstance(connector.accelerator, Accel)
    assert isinstance(connector.strategy, Strat)
    assert isinstance(connector.precision, Prec)
    assert connector.strategy is strategy


@mock.patch.dict(
    os.environ,
    {
        "SLURM_NTASKS": "2",
        "SLURM_NTASKS_PER_NODE": "1",
        "SLURM_JOB_NAME": "SOME_NAME",
        "SLURM_NODEID": "0",
        "LOCAL_RANK": "0",
        "SLURM_PROCID": "0",
        "SLURM_LOCALID": "0",
    },
)
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=0)
def test_dist_backend_accelerator_mapping(*_):
    connector = _Connector(strategy="ddp_spawn", accelerator="cpu", devices=2)
    assert isinstance(connector.accelerator, CPUAccelerator)
    assert isinstance(connector.strategy, DDPStrategy)
    assert connector.strategy.local_rank == 0


@RunIf(mps=False)
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=2)
def test_interactive_incompatible_backend_error(_, monkeypatch):
    monkeypatch.setattr(lightning_lite.connector, "_IS_INTERACTIVE", True)
    with pytest.raises(RuntimeError, match=r"strategy='ddp'\)`.*is not compatible"):
        _Connector(strategy="ddp", accelerator="gpu", devices=2)

    with pytest.raises(RuntimeError, match=r"strategy='ddp_spawn'\)`.*is not compatible"):
        _Connector(strategy="ddp_spawn", accelerator="gpu", devices=2)

    with pytest.raises(RuntimeError, match=r"strategy='ddp_sharded_spawn'\)`.*is not compatible"):
        _Connector(strategy="ddp_sharded_spawn", accelerator="gpu", devices=2)

    with pytest.raises(RuntimeError, match=r"strategy='ddp'\)`.*is not compatible"):
        # Edge case: _Connector maps dp to ddp if accelerator != gpu
        _Connector(strategy="dp")


@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=2)
def test_interactive_compatible_dp_strategy_gpu(_, monkeypatch):
    monkeypatch.setattr(lightning_lite.utilities.imports, "_IS_INTERACTIVE", True)
    connector = _Connector(strategy="dp", accelerator="gpu")
    assert connector.strategy.launcher is None


@RunIf(skip_windows=True)
def test_interactive_compatible_strategy_tpu(tpu_available, monkeypatch):
    monkeypatch.setattr(lightning_lite.utilities.imports, "_IS_INTERACTIVE", True)
    connector = _Connector(accelerator="tpu")
    assert connector.strategy.launcher.is_interactive_compatible


@RunIf(skip_windows=True)
def test_interactive_compatible_strategy_ddp_fork(monkeypatch):
    monkeypatch.setattr(lightning_lite.utilities.imports, "_IS_INTERACTIVE", True)
    connector = _Connector(strategy="ddp_fork", accelerator="cpu")
    assert connector.strategy.launcher.is_interactive_compatible


@RunIf(mps=False)
@pytest.mark.parametrize(
    ["strategy", "strategy_class"],
    [
        ("ddp", DDPStrategy),
        ("ddp_spawn", DDPSpawnStrategy),
        ("ddp_sharded", DDPShardedStrategy),
        ("ddp_sharded_spawn", DDPSpawnShardedStrategy),
        pytest.param("deepspeed", DeepSpeedStrategy, marks=RunIf(deepspeed=True)),
    ],
)
@pytest.mark.parametrize("devices", [1, 2])
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=2)
def test_strategy_choice_multi_node_gpu(_, strategy, strategy_class, devices):
    connector = _Connector(num_nodes=2, accelerator="gpu", strategy=strategy, devices=devices)
    assert isinstance(connector.strategy, strategy_class)


@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=0)
def test_cuda_accelerator_can_not_run_on_system(_):
    connector = _Connector(accelerator="cpu")
    assert isinstance(connector.accelerator, CPUAccelerator)

    with pytest.raises(
        RuntimeError,
        match="CUDAAccelerator` can not run on your system since the accelerator is not available.",
    ):
        _Connector(accelerator="cuda", devices=1)


@pytest.mark.skipif(TPUAccelerator.is_available(), reason="test requires missing TPU")
@mock.patch("lightning_lite.accelerators.tpu._XLA_AVAILABLE", True)
def test_tpu_accelerator_can_not_run_on_system():
    with pytest.raises(RuntimeError, match="TPUAccelerator` can not run on your system"):
        _Connector(accelerator="tpu", devices=8)


@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=2)
@pytest.mark.parametrize("device_count", (["0"], [0, "1"], ["GPU"], [["0", "1"], [0, 1]], [False]))
def test_accelererator_invalid_type_devices(_, device_count):
    with pytest.raises(
        MisconfigurationException, match=r"must be an int, a string, a sequence of ints or None, but you"
    ):
        _ = _Connector(accelerator="gpu", devices=device_count)


@RunIf(min_cuda_gpus=1)
def test_accelerator_gpu():
    connector = _Connector(accelerator="gpu", devices=1)
    assert isinstance(connector.accelerator, CUDAAccelerator)

    connector = _Connector(accelerator="gpu")
    assert isinstance(connector.accelerator, CUDAAccelerator)

    connector = _Connector(accelerator="auto", devices=1)
    assert isinstance(connector.accelerator, CUDAAccelerator)


@pytest.mark.parametrize(["devices", "strategy_class"], [(1, SingleDeviceStrategy), (5, DDPSpawnStrategy)])
def test_accelerator_cpu_with_devices(devices, strategy_class):
    connector = _Connector(accelerator="cpu", devices=devices)
    assert connector._parallel_devices == [torch.device("cpu")] * devices
    assert isinstance(connector.strategy, strategy_class)
    assert isinstance(connector.accelerator, CPUAccelerator)


@RunIf(min_cuda_gpus=2)
@pytest.mark.parametrize(
    ["devices", "strategy_class"], [(1, SingleDeviceStrategy), ([1], SingleDeviceStrategy), (2, DDPSpawnStrategy)]
)
def test_accelerator_gpu_with_devices(devices, strategy_class):
    connector = _Connector(accelerator="gpu", devices=devices)
    assert len(connector._parallel_devices) == len(devices) if isinstance(devices, list) else devices
    assert isinstance(connector.strategy, strategy_class)
    assert isinstance(connector.accelerator, CUDAAccelerator)


@RunIf(min_cuda_gpus=1)
def test_accelerator_auto_with_devices_gpu():
    connector = _Connector(accelerator="auto", devices=1)
    assert isinstance(connector.accelerator, CUDAAccelerator)
    assert connector._parallel_devices == [torch.device("cuda", 0)]


def test_set_devices_if_none_cpu():
    connector = _Connector(accelerator="cpu", devices=3)
    assert connector._parallel_devices == [torch.device("cpu")] * 3


def test_unsupported_strategy_types_on_cpu_and_fallback():
    with pytest.warns(UserWarning, match="is not supported on CPUs, hence setting `strategy='ddp"):
        connector = _Connector(strategy="dp", devices=2)
    assert isinstance(connector.strategy, DDPStrategy)


def test_invalid_accelerator_choice():
    with pytest.raises(ValueError, match="You selected an invalid accelerator name: `accelerator='cocofruit'`"):
        _Connector(accelerator="cocofruit")


def test_invalid_strategy_choice():
    with pytest.raises(ValueError, match="You selected an invalid strategy name: `strategy='cocofruit'`"):
        _Connector(strategy="cocofruit")


@pytest.mark.parametrize(
    ["strategy", "strategy_class"],
    [
        ("ddp_spawn", DDPSpawnStrategy),
        ("ddp_spawn_find_unused_parameters_false", DDPSpawnStrategy),
        ("ddp", DDPStrategy),
        ("ddp_find_unused_parameters_false", DDPStrategy),
    ],
)
def test_strategy_choice_cpu_str(strategy, strategy_class):
    connector = _Connector(strategy=strategy, accelerator="cpu", devices=2)
    assert isinstance(connector.strategy, strategy_class)


@pytest.mark.parametrize("strategy_class", [DDPSpawnStrategy, DDPStrategy])
def test_strategy_choice_cpu_instance(strategy_class):
    connector = _Connector(strategy=strategy_class(), accelerator="cpu", devices=2)
    assert isinstance(connector.strategy, strategy_class)


@RunIf(min_cuda_gpus=2)
@pytest.mark.parametrize(
    ["strategy", "strategy_class"],
    [
        ("ddp_spawn", DDPSpawnStrategy),
        ("ddp_spawn_find_unused_parameters_false", DDPSpawnStrategy),
        ("ddp", DDPStrategy),
        ("ddp_find_unused_parameters_false", DDPStrategy),
        ("dp", DataParallelStrategy),
        ("ddp_sharded", DDPShardedStrategy),
        ("ddp_sharded_spawn", DDPSpawnShardedStrategy),
        pytest.param("deepspeed", DeepSpeedStrategy, marks=RunIf(deepspeed=True)),
    ],
)
def test_strategy_choice_gpu_str(strategy, strategy_class):
    connector = _Connector(strategy=strategy, accelerator="gpu", devices=2)
    assert isinstance(connector.strategy, strategy_class)


@RunIf(fairscale=True)
@pytest.mark.parametrize(
    "strategy,expected_strategy", [("ddp_sharded", DDPShardedStrategy), ("ddp_sharded_spawn", DDPSpawnShardedStrategy)]
)
@pytest.mark.parametrize(
    "precision,expected_precision", [(16, NativeMixedPrecision), (32, Precision), ("bf16", NativeMixedPrecision)]
)
def test_strategy_choice_sharded(strategy, expected_strategy, precision, expected_precision):
    connector = _Connector(strategy=strategy, devices=1, precision=precision)
    assert isinstance(connector.strategy, expected_strategy)
    assert isinstance(connector.precision, expected_precision)


@RunIf(min_cuda_gpus=2)
@pytest.mark.parametrize("strategy_class", [DDPSpawnStrategy, DDPStrategy])
def test_strategy_choice_gpu_instance(strategy_class):
    connector = _Connector(strategy=strategy_class(), accelerator="gpu", devices=2)
    assert isinstance(connector.strategy, strategy_class)


@RunIf(min_cuda_gpus=2)
@pytest.mark.parametrize("strategy_class", [DDPSpawnStrategy, DDPStrategy])
def test_device_type_when_strategy_instance_gpu_passed(strategy_class):
    connector = _Connector(strategy=strategy_class(), accelerator="gpu", devices=2)
    assert isinstance(connector.strategy, strategy_class)
    assert isinstance(connector.accelerator, CUDAAccelerator)


@pytest.mark.parametrize("precision", [1, 12, "invalid"])
def test_validate_precision_type(precision):
    with pytest.raises(ValueError, match=f"Precision {repr(precision)} is invalid"):
        _Connector(precision=precision)


def test_strategy_choice_ddp_spawn_cpu():
    connector = _Connector(strategy="ddp_spawn", accelerator="cpu", devices=2)
    assert isinstance(connector.accelerator, CPUAccelerator)
    assert isinstance(connector.strategy, DDPSpawnStrategy)
    assert isinstance(connector.strategy.cluster_environment, LightningEnvironment)
    assert connector.strategy.launcher._start_method == "spawn"


@RunIf(skip_windows=True)
@mock.patch("lightning_lite.connector._IS_INTERACTIVE", True)
def test_strategy_choice_ddp_fork_in_interactive():
    """Test that when accelerator and strategy are unspecified, the connector chooses DDP Fork in interactive
    environments by default."""
    connector = _Connector(devices=2)
    assert isinstance(connector.accelerator, CPUAccelerator)
    assert isinstance(connector.strategy, DDPSpawnStrategy)
    assert isinstance(connector.strategy.cluster_environment, LightningEnvironment)
    assert connector.strategy.launcher._start_method == "fork"


@RunIf(skip_windows=True)
def test_strategy_choice_ddp_fork_cpu():
    connector = _Connector(strategy="ddp_fork", accelerator="cpu", devices=2)
    assert isinstance(connector.accelerator, CPUAccelerator)
    assert isinstance(connector.strategy, DDPSpawnStrategy)
    assert isinstance(connector.strategy.cluster_environment, LightningEnvironment)
    assert connector.strategy.launcher._start_method == "fork"


@mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"})
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=2)
@mock.patch("lightning_lite.accelerators.mps.MPSAccelerator.is_available", return_value=False)
def test_strategy_choice_ddp(*_):
    connector = _Connector(strategy="ddp", accelerator="gpu", devices=1)
    assert isinstance(connector.accelerator, CUDAAccelerator)
    assert isinstance(connector.strategy, DDPStrategy)
    assert isinstance(connector.strategy.cluster_environment, LightningEnvironment)


@mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"})
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=2)
@mock.patch("lightning_lite.accelerators.mps.MPSAccelerator.is_available", return_value=False)
def test_strategy_choice_ddp_spawn(*_):
    connector = _Connector(strategy="ddp_spawn", accelerator="gpu", devices=1)
    assert isinstance(connector.accelerator, CUDAAccelerator)
    assert isinstance(connector.strategy, DDPSpawnStrategy)
    assert isinstance(connector.strategy.cluster_environment, LightningEnvironment)


@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=2)
@pytest.mark.parametrize("job_name,expected_env", [("some_name", SLURMEnvironment), ("bash", LightningEnvironment)])
@pytest.mark.parametrize("strategy", ["ddp", DDPStrategy])
def test_strategy_choice_ddp_slurm(_, strategy, job_name, expected_env):
    if not isinstance(strategy, str):
        strategy = strategy()

    with mock.patch.dict(
        os.environ,
        {
            "CUDA_VISIBLE_DEVICES": "0,1",
            "SLURM_NTASKS": "2",
            "SLURM_NTASKS_PER_NODE": "1",
            "SLURM_JOB_NAME": job_name,
            "SLURM_NODEID": "0",
            "SLURM_PROCID": "1",
            "SLURM_LOCALID": "1",
        },
    ):
        connector = _Connector(strategy=strategy, accelerator="cuda", devices=2)
        assert isinstance(connector.accelerator, CUDAAccelerator)
        assert isinstance(connector.strategy, DDPStrategy)
        assert isinstance(connector.strategy.cluster_environment, expected_env)


@mock.patch.dict(
    os.environ,
    {
        "CUDA_VISIBLE_DEVICES": "0,1",
        "WORLD_SIZE": "2",
        "LOCAL_WORLD_SIZE": "2",
        "RANK": "1",
        "LOCAL_RANK": "1",
        "GROUP_RANK": "0",
        "TORCHELASTIC_RUN_ID": "1",
    },
)
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=2)
@mock.patch("lightning_lite.accelerators.mps.MPSAccelerator.is_available", return_value=False)
def test_strategy_choice_ddp_te(*_):
    connector = _Connector(strategy="ddp", accelerator="gpu", devices=2)
    assert isinstance(connector.accelerator, CUDAAccelerator)
    assert isinstance(connector.strategy, DDPStrategy)
    assert isinstance(connector.strategy.cluster_environment, TorchElasticEnvironment)
    assert connector.strategy.cluster_environment.local_rank() == 1
    assert connector.strategy.local_rank == 1


@mock.patch.dict(
    os.environ,
    {
        "WORLD_SIZE": "2",
        "LOCAL_WORLD_SIZE": "2",
        "RANK": "1",
        "LOCAL_RANK": "1",
        "GROUP_RANK": "0",
        "TORCHELASTIC_RUN_ID": "1",
    },
)
def test_strategy_choice_ddp_cpu_te():
    connector = _Connector(strategy="ddp_spawn", accelerator="cpu", devices=2)
    assert isinstance(connector.accelerator, CPUAccelerator)
    assert isinstance(connector.strategy, DDPStrategy)
    assert isinstance(connector.strategy.cluster_environment, TorchElasticEnvironment)
    assert connector.strategy.cluster_environment.local_rank() == 1
    assert connector.strategy.local_rank == 1


@mock.patch.dict(
    os.environ,
    {
        "CUDA_VISIBLE_DEVICES": "0",
        "KUBERNETES_PORT": "tcp://127.0.0.1:443",
        "MASTER_ADDR": "1.2.3.4",
        "MASTER_PORT": "500",
        "WORLD_SIZE": "20",
        "RANK": "1",
    },
)
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=1)
@mock.patch("lightning_lite.accelerators.mps.MPSAccelerator.is_available", return_value=False)
def test_strategy_choice_ddp_kubeflow(*_):
    connector = _Connector(strategy="ddp", accelerator="gpu", devices=1)
    assert isinstance(connector.accelerator, CUDAAccelerator)
    assert isinstance(connector.strategy, DDPStrategy)
    assert isinstance(connector.strategy.cluster_environment, KubeflowEnvironment)
    assert connector.strategy.cluster_environment.local_rank() == 0
    assert connector.strategy.local_rank == 0


@mock.patch.dict(
    os.environ,
    {
        "KUBERNETES_PORT": "tcp://127.0.0.1:443",
        "MASTER_ADDR": "1.2.3.4",
        "MASTER_PORT": "500",
        "WORLD_SIZE": "20",
        "RANK": "1",
    },
)
def test_strategy_choice_ddp_cpu_kubeflow():
    connector = _Connector(strategy="ddp_spawn", accelerator="cpu", devices=2)
    assert isinstance(connector.accelerator, CPUAccelerator)
    assert isinstance(connector.strategy, DDPStrategy)
    assert isinstance(connector.strategy.cluster_environment, KubeflowEnvironment)
    assert connector.strategy.cluster_environment.local_rank() == 0
    assert connector.strategy.local_rank == 0


@mock.patch.dict(
    os.environ,
    {
        "SLURM_NTASKS": "2",
        "SLURM_NTASKS_PER_NODE": "1",
        "SLURM_JOB_NAME": "SOME_NAME",
        "SLURM_NODEID": "0",
        "LOCAL_RANK": "0",
        "SLURM_PROCID": "0",
        "SLURM_LOCALID": "0",
    },
)
@pytest.mark.parametrize("strategy", ["ddp", DDPStrategy()])
def test_strategy_choice_ddp_cpu_slurm(strategy):
    connector = _Connector(strategy=strategy, accelerator="cpu", devices=2)
    assert isinstance(connector.accelerator, CPUAccelerator)
    assert isinstance(connector.strategy, DDPStrategy)
    assert isinstance(connector.strategy.cluster_environment, SLURMEnvironment)
    assert connector.strategy.local_rank == 0


@mock.patch.dict(os.environ, {}, clear=True)
def test_unsupported_tpu_choice(tpu_available):
    with pytest.raises(NotImplementedError, match=r"accelerator='tpu', precision=64\)` is not implemented"):
        _Connector(accelerator="tpu", precision=64)

    # if user didn't set strategy, _Connector will choose the TPUSingleStrategy or XLAStrategy
    with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `SingleTPUStrategy`"), pytest.warns(
        UserWarning, match=r"accelerator='tpu', precision=16\)` but native AMP is not supported"
    ):
        _Connector(accelerator="tpu", precision=16, strategy="ddp")

    # wrong precision plugin type
    strategy = XLAStrategy(accelerator=TPUAccelerator(), precision=Precision())
    with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `TPUPrecision` plugin"):
        _Connector(strategy=strategy, devices=8)

    # wrong strategy type
    strategy = DDPStrategy(accelerator=TPUAccelerator(), precision=TPUPrecision())
    with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `SingleTPUStrategy`"):
        _Connector(strategy=strategy, devices=8)


@mock.patch("lightning_lite.accelerators.cuda.CUDAAccelerator.is_available", return_value=False)
@mock.patch("lightning_lite.accelerators.mps.MPSAccelerator.is_available", return_value=False)
def test_devices_auto_choice_cpu(tpu_available, *_):
    connector = _Connector(accelerator="auto", devices="auto")
    assert isinstance(connector.accelerator, CPUAccelerator)
    assert isinstance(connector.strategy, SingleDeviceStrategy)
    assert connector.strategy.root_device == torch.device("cpu")


@RunIf(mps=False)
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=2)
def test_devices_auto_choice_gpu(*_):
    connector = _Connector(accelerator="auto", devices="auto")
    assert isinstance(connector.accelerator, CUDAAccelerator)
    assert isinstance(connector.strategy, DDPSpawnStrategy)
    assert len(connector._parallel_devices) == 2


@RunIf(mps=True)
def test_devices_auto_choice_mps():
    connector = _Connector(accelerator="auto", devices="auto")
    assert isinstance(connector.accelerator, MPSAccelerator)
    assert isinstance(connector.strategy, SingleDeviceStrategy)
    assert connector.strategy.root_device == torch.device("mps", 0)
    assert connector._parallel_devices == [torch.device("mps", 0)]


@pytest.mark.parametrize(
    ["parallel_devices", "accelerator"],
    [([torch.device("cpu")], "cuda"), ([torch.device("cuda", i) for i in range(8)], "tpu")],
)
def test_parallel_devices_in_strategy_conflict_with_accelerator(parallel_devices, accelerator):
    with pytest.raises(ValueError, match=r"parallel_devices set through"):
        _Connector(strategy=DDPStrategy(parallel_devices=parallel_devices), accelerator=accelerator)


@pytest.mark.parametrize(
    ["plugins", "expected"],
    [
        ([LightningEnvironment(), SLURMEnvironment()], "ClusterEnvironment"),
        ([TorchCheckpointIO(), TorchCheckpointIO()], "CheckpointIO"),
        (
            [Precision(), DoublePrecision(), LightningEnvironment(), SLURMEnvironment()],
            "Precision, ClusterEnvironment",
        ),
    ],
)
def test_plugin_only_one_instance_for_one_type(plugins, expected):
    with pytest.raises(ValueError, match=f"Received multiple values for {expected}"):
        _Connector(plugins=plugins)


@pytest.mark.parametrize("accelerator", ("cpu", "cuda", "mps", "tpu"))
@pytest.mark.parametrize("devices", ("0", 0, []))
def test_passing_zero_and_empty_list_to_devices_flag(accelerator, devices):
    with pytest.raises(ValueError, match="value is not a valid input using"):
        _Connector(accelerator=accelerator, devices=devices)


@pytest.mark.parametrize(
    "expected_accelerator_flag,expected_accelerator_class",
    [
        pytest.param("cuda", CUDAAccelerator, marks=RunIf(min_cuda_gpus=1)),
        pytest.param("mps", MPSAccelerator, marks=RunIf(mps=True)),
    ],
)
def test_gpu_accelerator_backend_choice(expected_accelerator_flag, expected_accelerator_class):
    connector = _Connector(accelerator="gpu")
    assert connector._accelerator_flag == expected_accelerator_flag
    assert isinstance(connector.accelerator, expected_accelerator_class)


@mock.patch("lightning_lite.accelerators.mps.MPSAccelerator.is_available", return_value=False)
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=1)
def test_gpu_accelerator_backend_choice_cuda(*_):
    connector = _Connector(accelerator="gpu")
    assert connector._accelerator_flag == "cuda"
    assert isinstance(connector.accelerator, CUDAAccelerator)


@RunIf(min_torch="1.12")
@mock.patch("lightning_lite.accelerators.mps.MPSAccelerator.is_available", return_value=True)
@mock.patch("lightning_lite.accelerators.mps._get_all_available_mps_gpus", return_value=[0])
def test_gpu_accelerator_backend_choice_mps(*_):
    connector = _Connector(accelerator="gpu")
    assert connector._accelerator_flag == "mps"
    assert isinstance(connector.accelerator, MPSAccelerator)


@mock.patch("lightning_lite.accelerators.mps.MPSAccelerator.is_available", return_value=False)
@mock.patch("lightning_lite.accelerators.cuda.CUDAAccelerator.is_available", return_value=False)
def test_gpu_accelerator_no_gpu_backend_found_error(*_):
    with pytest.raises(RuntimeError, match="No supported gpu backend found!"):
        _Connector(accelerator="gpu")


@pytest.mark.parametrize("strategy", _DDP_FORK_ALIASES)
@mock.patch(
    "lightning_lite.connector.torch.multiprocessing.get_all_start_methods",
    return_value=[],
)
def test_ddp_fork_on_unsupported_platform(_, strategy):
    with pytest.raises(ValueError, match="process forking is not supported on this platform"):
        _Connector(strategy=strategy)


@mock.patch("lightning_lite.plugins.precision.native_amp._TORCH_GREATER_EQUAL_1_10", True)
def test_precision_selection_16_on_cpu_warns():
    with pytest.warns(
        UserWarning, match=r"precision=16\)` but native AMP is not supported on CPU. Using `precision='bf16"
    ):
        _Connector(precision=16)


@mock.patch("lightning_lite.plugins.precision.native_amp._TORCH_GREATER_EQUAL_1_10", False)
def test_precision_selection_16_raises_torch_version(monkeypatch):
    with pytest.raises(ImportError, match="must install torch greater or equal to 1.10"):
        _Connector(accelerator="cpu", precision=16)
    with pytest.raises(ImportError, match="must install torch greater or equal to 1.10"):
        _Connector(accelerator="cpu", precision="bf16")


class MyNativeAMP(NativeMixedPrecision):
    pass


@RunIf(mps=False)
@pytest.mark.parametrize("strategy,devices", [("ddp", 2), ("ddp_spawn", 2)])
@pytest.mark.parametrize(
    "is_custom_plugin,plugin_cls",
    [(False, NativeMixedPrecision), (True, MyNativeAMP)],
)
@mock.patch("lightning_lite.plugins.precision.native_amp._TORCH_GREATER_EQUAL_1_10", True)
def test_precision_selection_amp_ddp(strategy, devices, is_custom_plugin, plugin_cls):
    plugin = None
    if is_custom_plugin:
        plugin = plugin_cls(16, "cpu")
    connector = _Connector(
        precision=16,
        devices=devices,
        strategy=strategy,
        plugins=plugin,
    )
    assert isinstance(connector.precision, plugin_cls)


@pytest.mark.parametrize(
    ["strategy", "strategy_cls"], [("DDP", DDPStrategy), ("DDP_FIND_UNUSED_PARAMETERS_FALSE", DDPStrategy)]
)
def test_strategy_str_passed_being_case_insensitive(strategy, strategy_cls):
    connector = _Connector(strategy=strategy)
    assert isinstance(connector.strategy, strategy_cls)


@pytest.mark.parametrize("precision", ["64", "32", "16", pytest.param("bf16", marks=RunIf(min_torch="1.10"))])
@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=1)
def test_precision_from_environment(_, precision):
    """Test that the precision input can be set through the environment variable."""
    with mock.patch.dict(os.environ, {"LT_PRECISION": precision}):
        connector = _Connector(accelerator="cuda")  # need to use cuda, because AMP not available on CPU
    assert isinstance(connector.precision, Precision)


@pytest.mark.parametrize(
    "accelerator, strategy, expected_accelerator, expected_strategy",
    [
        ("cpu", None, CPUAccelerator, SingleDeviceStrategy),
        ("cpu", "ddp", CPUAccelerator, DDPStrategy),
        pytest.param("mps", None, MPSAccelerator, SingleDeviceStrategy, marks=RunIf(mps=True)),
        pytest.param("cuda", "dp", CUDAAccelerator, DataParallelStrategy, marks=RunIf(min_cuda_gpus=1)),
        pytest.param(
            "cuda", "deepspeed", CUDAAccelerator, DeepSpeedStrategy, marks=RunIf(min_cuda_gpus=1, deepspeed=True)
        ),
    ],
)
def test_accelerator_strategy_from_environment(accelerator, strategy, expected_accelerator, expected_strategy):
    """Test that the accelerator and strategy input can be set through the environment variables."""
    env_vars = {"LT_ACCELERATOR": accelerator}
    if strategy is not None:
        env_vars["LT_STRATEGY"] = strategy

    with mock.patch.dict(os.environ, env_vars):
        connector = _Connector()
        assert isinstance(connector.accelerator, expected_accelerator)
        assert isinstance(connector.strategy, expected_strategy)


@mock.patch("lightning_lite.accelerators.cuda.num_cuda_devices", return_value=8)
def test_devices_from_environment(*_):
    """Test that the devices and number of nodes can be set through the environment variables."""
    with mock.patch.dict(os.environ, {"LT_DEVICES": "2", "LT_NUM_NODES": "3"}):
        connector = _Connector(accelerator="cuda")
        assert isinstance(connector.accelerator, CUDAAccelerator)
        assert isinstance(connector.strategy, DDPStrategy)
        assert len(connector._parallel_devices) == 2
        assert connector._num_nodes_flag == 3


def test_arguments_from_environment_collision():
    """Test that the connector raises an error when the CLI settings conflict with settings in the code."""
    with mock.patch.dict(os.environ, {"LT_ACCELERATOR": "cpu"}):
        with pytest.raises(
            ValueError, match=escape("Your code has `LightningLite(accelerator='cuda', ...)` but it conflicts")
        ):
            _Connector(accelerator="cuda")

    with mock.patch.dict(os.environ, {"LT_STRATEGY": "ddp"}):
        with pytest.raises(
            ValueError, match=escape("Your code has `LightningLite(strategy='ddp_spawn', ...)` but it conflicts")
        ):
            _Connector(strategy="ddp_spawn")

    with mock.patch.dict(os.environ, {"LT_DEVICES": "2"}):
        with pytest.raises(ValueError, match=escape("Your code has `LightningLite(devices=3, ...)` but it conflicts")):
            _Connector(devices=3)

    with mock.patch.dict(os.environ, {"LT_NUM_NODES": "3"}):
        with pytest.raises(
            ValueError, match=escape("Your code has `LightningLite(num_nodes=2, ...)` but it conflicts")
        ):
            _Connector(num_nodes=2)

    with mock.patch.dict(os.environ, {"LT_PRECISION": "16"}):
        with pytest.raises(
            ValueError, match=escape("Your code has `LightningLite(precision=64, ...)` but it conflicts")
        ):
            _Connector(precision=64)