Repository: Lightning-AI/pytorch-lightning

Issue:
Automatically pick available GPU
Thanks for this great library! 

## 🚀 Feature
I would like to change the behavior of this code:
``` python
    trainer = pl.Trainer(
        ... snip ...,
        gpus=1,
    )
```
Currently, when setting `gpus` to an integer `n`, the first `n` GPUs are automatically used. 

I would like to change the behavior such that when multiple GPUs are available, the trainer picks the first available GPU.

### Motivation
When running multiple jobs in parallel on a server with multiple available GPUs, I get the error: 
```
RuntimeError: CUDA error: all CUDA-capable devices are busy or unavailable
```
This is because all 4 running jobs are scheduled to GPU 0, even though I have 4 GPUs available. 

*Note*: the GPUs are configured to be in "exclusive mode", which means that only one process at a time can use them. 

### Pitch

I would like to change the behavior such that when multiple GPUs are available, the trainer picks the first available GPU.

### Alternatives

One could also fix this in client code, like this: 

``` python
def retry_jittered_backoff(f, num_retries=5):
    # Based on:
    # https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
    import time
    import random
    cap = 1.0                  # max sleep time is 1s
    base = 0.01                # initial sleep time is 10ms
    sleep = base               # initial sleep time is 10ms

    for i in range(num_retries):
        try:
            return f()
        except RuntimeError as e:
            if i == num_retries - 1:
                raise e
            else:
                continue
        time.sleep(sleep)
        sleep = min(cap, random.uniform(base, sleep * 3))


def pick_gpu():
    for i in range(torch.cuda.device_count()):
        torch.cuda.set_device(i)
        try:
            torch.ones(1).cuda()
        except RuntimeError:
            continue
        return i
    raise RuntimeError("No GPUs available.")


def main(.. snip ..):

    model = Model(.. snip ..)
    trainer = pl.Trainer(
        .. snip ..,
        gpus=[retry_jittered_backoff(pick_gpu)],
    )
    trainer.fit(model)

```

This is exactly the kind of boilerplate that I would hope `pytorch-lighting` can make redundant.  

## PR-able? 

Would you be open to consider accepting a PR? If so, could you give me some pointers where to add the above code?


Discussion:
Hey, thanks for your contribution! Great first issue!
you can just do 
```
gpus[2]
```
to pick the 3rd gpu no?

But I also like this.

Maybe:

```python
gpus=1
gpu_choice='manual|auto'
```

or

```
auto_gpu_choice=True
```
> you can just do
> 
> ```
> gpus[2]
> ```
> 
> to pick the 3rd gpu no?

You could, if the calling program assigns a gpu for you. You are out of luck, however, if you naively use `make -j 4` for instance. I think my point is that the program does not always know which gpu is available. 

> But I also like this.
> 
> Maybe:
> 
> ```python
> gpus=1
> gpu_choice='manual|auto'
> ```

That sounds good and pretty self-documenting. 

I will try to put a PR together :) 
@ahendriksen how is it going with drafting PR? :]

Resolving patch:
diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py
--- a/pytorch_lightning/trainer/distrib_parts.py
+++ b/pytorch_lightning/trainer/distrib_parts.py
@@ -339,7 +339,8 @@
 
 import os
 from abc import ABC, abstractmethod
-
+import time
+import random
 import torch
 
 from pytorch_lightning import _logger as log
@@ -648,3 +649,44 @@ def determine_root_gpu_device(gpus):
     root_gpu = gpus[0]
 
     return root_gpu
+
+
+def retry_jittered_backoff(f, num_retries=5):
+    # Based on:
+    # https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
+    cap = 1.0                  # max sleep time is 1s
+    base = 0.01                # initial sleep time is 10ms
+    sleep = base               # initial sleep time is 10ms
+
+    for i in range(num_retries):
+        try:
+            return f()
+        except RuntimeError as e:
+            if i == num_retries - 1:
+                raise e
+            else:
+                continue
+        time.sleep(sleep)
+        sleep = min(cap, random.uniform(base, sleep * 3))
+
+
+def pick_single_gpu(exclude_gpus=[]):
+    for i in range(torch.cuda.device_count()):
+        if i in exclude_gpus:
+            continue
+        # Try to allocate on device:
+        device = torch.device(f"cuda:{i}")
+        try:
+            torch.ones(1).to(device)
+        except RuntimeError:
+            continue
+        return i
+    raise RuntimeError("No GPUs available.")
+
+
+def pick_multiple_gpus(n):
+    picked = []
+    for _ in range(n):
+        picked.append(pick_single_gpu(exclude_gpus=picked))
+
+    return picked
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -22,7 +22,12 @@
 from pytorch_lightning.trainer.data_loading import TrainerDataLoadingMixin
 from pytorch_lightning.trainer.deprecated_api import TrainerDeprecatedAPITillVer0_8, TrainerDeprecatedAPITillVer0_9
 from pytorch_lightning.trainer.distrib_data_parallel import TrainerDDPMixin
-from pytorch_lightning.trainer.distrib_parts import TrainerDPMixin, parse_gpu_ids, determine_root_gpu_device
+from pytorch_lightning.trainer.distrib_parts import (
+    TrainerDPMixin,
+    parse_gpu_ids,
+    determine_root_gpu_device,
+    pick_multiple_gpus,
+)
 from pytorch_lightning.trainer.evaluation_loop import TrainerEvaluationLoopMixin
 from pytorch_lightning.trainer.logging import TrainerLoggingMixin
 from pytorch_lightning.trainer.model_hooks import TrainerModelHooksMixin
@@ -85,6 +90,7 @@ def __init__(
             process_position: int = 0,
             num_nodes: int = 1,
             gpus: Optional[Union[List[int], str, int]] = None,
+            auto_select_gpus: bool = False,
             num_tpu_cores: Optional[int] = None,
             log_gpu_memory: Optional[str] = None,
             progress_bar_refresh_rate: int = 1,
@@ -158,6 +164,13 @@ def __init__(
 
             gpus: Which GPUs to train on.
 
+            auto_select_gpus:
+
+                If enabled and `gpus` is an integer, pick available
+                gpus automatically. This is especially useful when
+                GPUs are configured to be in "exclusive mode", such
+                that only one process at a time can access them.
+
             num_tpu_cores: How many TPU cores to train on (1 or 8).
 
             log_gpu_memory: None, 'min_max', 'all'. Might slow performance
@@ -384,8 +397,12 @@ def __init__(
         self.accumulate_grad_batches = accumulate_grad_batches
         self.configure_accumulated_gradients(accumulate_grad_batches)
 
-        # allow int, string and gpu list
-        self.gpus = gpus
+        # for gpus allow int, string and gpu list
+        if auto_select_gpus and isinstance(gpus, int):
+            self.gpus = pick_multiple_gpus(gpus)
+        else:
+            self.gpus = gpus
+
         self.data_parallel_device_ids = parse_gpu_ids(self.gpus)
         self.root_gpu = determine_root_gpu_device(self.data_parallel_device_ids)
         self.root_device = torch.device("cpu")


Test patch:
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -685,3 +685,18 @@ def _optimizer_step(*args, **kwargs):
     model.prev_called_batch_idx = 0
 
     trainer.fit(model)
+
+
+def test_gpu_choice(tmpdir):
+    trainer_options = dict(
+        default_save_path=tmpdir,
+    )
+    # Only run if CUDA is available
+    if not torch.cuda.is_available():
+        return
+
+    num_gpus = torch.cuda.device_count()
+    Trainer(**trainer_options, gpus=num_gpus, auto_select_gpus=True)
+
+    with pytest.raises(RuntimeError, match=r'.*No GPUs available.*'):
+        Trainer(**trainer_options, gpus=num_gpus + 1, auto_select_gpus=True)