diff --git a/tests/tests_lite/utilities/test_device_parser.py b/tests/tests_lite/utilities/test_device_parser.py index bb6e1665ef..0f005d5ce3 100644 --- a/tests/tests_lite/utilities/test_device_parser.py +++ b/tests/tests_lite/utilities/test_device_parser.py @@ -16,7 +16,90 @@ from unittest import mock 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 @@ import lightning_lite.utilities.device_parser 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 diff --git a/tests/tests_pytorch/models/test_gpu.py b/tests/tests_pytorch/models/test_gpu.py index 1f15f2a596..3951a500b8 100644 --- a/tests/tests_pytorch/models/test_gpu.py +++ b/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, {