Fix support for passing -1 to find_usable_cuda_devices function (#16866)

Co-authored-by: Yi Heng Lim <[email protected]>
Co-authored-by: Adrian Wälchli <[email protected]>

Author: awaelchli

src/lightning_fabric/CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed edge cases in parsing device ids using NVML ([#16795](https://github.com/Lightning-AI/lightning/pull/16795))
 - Fixed DDP spawn hang on TPU Pods ([#16844](https://github.com/Lightning-AI/lightning/pull/16844))
+- Fixed an error when passing `find_usable_cuda_devices(num_devices=-1)` ([#16866](https://github.com/Lightning-AI/lightning/pull/16866))
 
 
 ## [1.9.3] - 2023-02-21

src/lightning_fabric/accelerators/cuda.py

@@ -114,7 +114,7 @@ def find_usable_cuda_devices(num_devices: int = -1) -> List[int]:
             # exit early if we found the right number of GPUs
             break
 
-    if len(available_devices) != num_devices:
+    if num_devices != -1 and len(available_devices) != num_devices:
         raise RuntimeError(
             f"You requested to find {num_devices} devices but only {len(available_devices)} are currently available."
             f" The devices {unavailable_devices} are occupied by other processes and can't be used at the moment."

tests/tests_fabric/accelerators/test_cuda.py

@@ -143,3 +143,9 @@ def test_find_usable_cuda_devices_error_handling():
         "lightning_fabric.accelerators.cuda.torch.tensor", tensor_mock
     ), pytest.raises(RuntimeError, match=escape("The devices [0, 1] are occupied by other processes")):
         find_usable_cuda_devices(2)
+
+    # Request for as many GPUs as there are, no error should be raised
+    with mock.patch("lightning_fabric.accelerators.cuda.num_cuda_devices", return_value=5), mock.patch(
+        "lightning_fabric.accelerators.cuda.torch.tensor"
+    ):
+        assert find_usable_cuda_devices(-1) == [0, 1, 2, 3, 4]