Move device parser tests inside Lite (#14586)

Author: awaelchli

tests/tests_lite/utilities/test_device_parser.py

@@ -16,7 +16,90 @@
 import pytest
 import torch
 
-import lightning_lite.utilities.device_parser
+from lightning_lite.utilities import device_parser
+from lightning_lite.utilities.exceptions import MisconfigurationException
+
+_PRETEND_N_OF_GPUS = 16
+
+
+@pytest.fixture
+def mocked_device_count(monkeypatch):
+    def device_count():
+        return _PRETEND_N_OF_GPUS
+
+    def is_available():
+        return True
+
+    monkeypatch.setattr(device_parser, "is_cuda_available", is_available)
+    monkeypatch.setattr(device_parser, "num_cuda_devices", device_count)
+
+
+@pytest.fixture
+def mocked_device_count_0(monkeypatch):
+    def device_count():
+        return 0
+
+    monkeypatch.setattr(device_parser, "num_cuda_devices", device_count)
+
+
+@pytest.mark.parametrize(
+    ["devices", "expected_root_gpu"],
+    [
+        pytest.param(None, None, id="No gpus, expect gpu root device to be None"),
+        pytest.param([0], 0, id="Oth gpu, expect gpu root device to be 0."),
+        pytest.param([1], 1, id="1st gpu, expect gpu root device to be 1."),
+        pytest.param([3], 3, id="3rd gpu, expect gpu root device to be 3."),
+        pytest.param([1, 2], 1, id="[1, 2] gpus, expect gpu root device to be 1."),
+    ],
+)
+def test_determine_root_gpu_device(devices, expected_root_gpu):
+    assert device_parser.determine_root_gpu_device(devices) == expected_root_gpu
+
+
+@pytest.mark.parametrize(
+    ["devices", "expected_gpu_ids"],
+    [
+        (None, None),
+        (0, None),
+        ([], None),
+        (1, [0]),
+        (3, [0, 1, 2]),
+        pytest.param(-1, list(range(_PRETEND_N_OF_GPUS)), id="-1 - use all gpus"),
+        ([0], [0]),
+        ([1, 3], [1, 3]),
+        ((1, 3), [1, 3]),
+        ("0", None),
+        ("3", [0, 1, 2]),
+        ("1, 3", [1, 3]),
+        ("2,", [2]),
+        pytest.param("-1", list(range(_PRETEND_N_OF_GPUS)), id="'-1' - use all gpus"),
+    ],
+)
+def test_parse_gpu_ids(mocked_device_count, devices, expected_gpu_ids):
+    assert device_parser.parse_gpu_ids(devices, include_cuda=True) == expected_gpu_ids
+
+
+@pytest.mark.parametrize("devices", [0.1, -2, False, [-1], [None], ["0"], [0, 0]])
+def test_parse_gpu_fail_on_unsupported_inputs(mocked_device_count, devices):
+    with pytest.raises(MisconfigurationException):
+        device_parser.parse_gpu_ids(devices, include_cuda=True)
+
+
+@pytest.mark.parametrize("devices", [[1, 2, 19], -1, "-1"])
+def test_parse_gpu_fail_on_non_existent_id(mocked_device_count_0, devices):
+    with pytest.raises(MisconfigurationException):
+        device_parser.parse_gpu_ids(devices, include_cuda=True)
+
+
+def test_parse_gpu_fail_on_non_existent_id_2(mocked_device_count):
+    with pytest.raises(MisconfigurationException):
+        device_parser.parse_gpu_ids([1, 2, 19], include_cuda=True)
+
+
+@pytest.mark.parametrize("devices", [-1, "-1"])
+def test_parse_gpu_returns_none_when_no_devices_are_available(mocked_device_count_0, devices):
+    with pytest.raises(MisconfigurationException):
+        device_parser.parse_gpu_ids(devices, include_cuda=True)
 
 
 @pytest.mark.skipif(
@@ -27,5 +110,5 @@
 def test_num_cuda_devices_without_forking(*_):
     """This merely tests that on platforms without fork support our helper functions fall back to the default
     implementation for determining cuda availability."""
-    assert lightning_lite.utilities.device_parser.is_cuda_available()
-    assert lightning_lite.utilities.device_parser.num_cuda_devices() == 2
+    assert device_parser.is_cuda_available()
+    assert device_parser.num_cuda_devices() == 2

tests/tests_pytorch/models/test_gpu.py

@@ -109,66 +109,6 @@ def test_root_gpu_property_0_raising(mocked_device_count_0, devices, expected_ro
         Trainer(accelerator="gpu", devices=devices, strategy=strategy)
 
 
-@pytest.mark.parametrize(
-    ["devices", "expected_root_gpu"],
-    [
-        pytest.param(None, None, id="No gpus, expect gpu root device to be None"),
-        pytest.param([0], 0, id="Oth gpu, expect gpu root device to be 0."),
-        pytest.param([1], 1, id="1st gpu, expect gpu root device to be 1."),
-        pytest.param([3], 3, id="3rd gpu, expect gpu root device to be 3."),
-        pytest.param([1, 2], 1, id="[1, 2] gpus, expect gpu root device to be 1."),
-    ],
-)
-def test_determine_root_gpu_device(devices, expected_root_gpu):
-    assert device_parser.determine_root_gpu_device(devices) == expected_root_gpu
-
-
-@pytest.mark.parametrize(
-    ["devices", "expected_gpu_ids"],
-    [
-        (None, None),
-        (0, None),
-        ([], None),
-        (1, [0]),
-        (3, [0, 1, 2]),
-        pytest.param(-1, list(range(PRETEND_N_OF_GPUS)), id="-1 - use all gpus"),
-        ([0], [0]),
-        ([1, 3], [1, 3]),
-        ((1, 3), [1, 3]),
-        ("0", None),
-        ("3", [0, 1, 2]),
-        ("1, 3", [1, 3]),
-        ("2,", [2]),
-        pytest.param("-1", list(range(PRETEND_N_OF_GPUS)), id="'-1' - use all gpus"),
-    ],
-)
-def test_parse_gpu_ids(mocked_device_count, devices, expected_gpu_ids):
-    assert device_parser.parse_gpu_ids(devices, include_cuda=True) == expected_gpu_ids
-
-
-@pytest.mark.parametrize("devices", [0.1, -2, False, [-1], [None], ["0"], [0, 0]])
-def test_parse_gpu_fail_on_unsupported_inputs(mocked_device_count, devices):
-    with pytest.raises(MisconfigurationException):
-        device_parser.parse_gpu_ids(devices, include_cuda=True)
-
-
-@pytest.mark.parametrize("devices", [[1, 2, 19], -1, "-1"])
-def test_parse_gpu_fail_on_non_existent_id(mocked_device_count_0, devices):
-    with pytest.raises(MisconfigurationException):
-        device_parser.parse_gpu_ids(devices, include_cuda=True)
-
-
-def test_parse_gpu_fail_on_non_existent_id_2(mocked_device_count):
-    with pytest.raises(MisconfigurationException):
-        device_parser.parse_gpu_ids([1, 2, 19], include_cuda=True)
-
-
-@pytest.mark.parametrize("devices", [-1, "-1"])
-def test_parse_gpu_returns_none_when_no_devices_are_available(mocked_device_count_0, devices):
-    with pytest.raises(MisconfigurationException):
-        device_parser.parse_gpu_ids(devices, include_cuda=True)
-
-
 @mock.patch.dict(
     os.environ,
     {