Repository: Lightning-AI/pytorch-lightning
Pull request: Fabric: Enable "auto" for `devices` and `accelerator` as cli arguments
## What does this PR do?

Fixes #20451 

This PR adds support for the `auto` variable for `accelerator` and `devices`  in the Fabric CLI. 

-  When `auto` is passed to `--accelerator` we now detect whether the machine has either `mps`, `cuda`, or `cpu` and dynamically runs on the available device, this is the same case when no accelerator is specified. 

- For `auto` in the `--devices` argument the devices default to `1` as specified in the docstring of the argument here https://github.com/Lightning-AI/pytorch-lightning/blob/c03660a8f84c7c1a09f3da230061b88e84430111/src/lightning/fabric/cli.py#L80

<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

No breaking changes were introduced.

<details>
  <summary><b>Before submitting</b></summary>

- Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/Lightning-AI/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- Did you make sure to **update the documentation** with your changes? (if necessary)
- Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- Did you list all the **breaking changes** introduced by this pull request?
- Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

</details>

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

<details>
  <summary>Reviewer checklist</summary>

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

</details>

<!--

Did you have fun?

Make sure you had fun coding 🙃

-->


<!-- readthedocs-preview pytorch-lightning start -->
----
📚 Documentation preview 📚: https://pytorch-lightning--20913.org.readthedocs.build/en/20913/

<!-- readthedocs-preview pytorch-lightning end -->

Problem statement:
Error: Invalid value for '--accelerator': 'auto' is not one of 'cpu', 'gpu', 'cuda', 'mps', 'tpu'.
### Bug description

when run this codes
```
fabric run ./path/to/train.py \
    --devices=auto \
    --accelerator=auto
```
get problem below
```
Error: Invalid value for '--accelerator': 'auto' is not one of 'cpu', 'gpu', 'cuda', 'mps', 'tpu'.
```

### What version are you seeing the problem on?

v2.4

### How to reproduce the bug

_No response_

### Error messages and logs

```
# Error messages and logs here please
```


### Environment

<details>
  <summary>Current environment</summary>

```
#- PyTorch Lightning Version (e.g., 2.4.0):
#- PyTorch Version (e.g., 2.4):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
```

</details>


### More info

_No response_

cc @borda @justusschock @awaelchli

Hints from issue discussion:
hey @wkd88 that's correct, would you want to give it a shot in a quick PR?

here's the file that needs to be updated: https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/cli.py
> hey @wkd88 that's correct, would you want to give it a shot in a quick PR?
> 
> here's the file that needs to be updated: https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/cli.py

Hi @lantiga I want to try a PR on this issue. 

I’ve updated the file to include the "auto" option for the accelerator. You can find the PR here: #20473 . 
Hello @Borda I would like to take on this issue.

Thanks

Commit message: use a provate method to avoid dependency errors

Patch to review:
diff --git a/src/lightning/fabric/cli.py b/src/lightning/fabric/cli.py
index 2268614abb97b..0e419c8c03bcd 100644
--- a/src/lightning/fabric/cli.py
+++ b/src/lightning/fabric/cli.py
@@ -21,7 +21,7 @@
 from lightning_utilities.core.imports import RequirementCache
 from typing_extensions import get_args
 
-from lightning.fabric.accelerators import CPUAccelerator, CUDAAccelerator, MPSAccelerator
+from lightning.fabric.accelerators import CPUAccelerator, CUDAAccelerator, MPSAccelerator, XLAAccelerator
 from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT_STR, _PRECISION_INPUT_STR_ALIAS
 from lightning.fabric.strategies import STRATEGY_REGISTRY
 from lightning.fabric.utilities.consolidate_checkpoint import _process_cli_args
@@ -34,7 +34,18 @@
 _CLICK_AVAILABLE = RequirementCache("click")
 _LIGHTNING_SDK_AVAILABLE = RequirementCache("lightning_sdk")
 
-_SUPPORTED_ACCELERATORS = ("cpu", "gpu", "cuda", "mps", "tpu")
+_SUPPORTED_ACCELERATORS = ("cpu", "gpu", "cuda", "mps", "tpu", "auto")
+
+
+def _choose_auto_accelerator() -> str:
+    """Choose the best available accelerator for the current environment."""
+    if CUDAAccelerator.is_available():
+        return "cuda"
+    if MPSAccelerator.is_available():
+        return "mps"
+    if XLAAccelerator.is_available():
+        return "tpu"
+    return "cpu"
 
 
 def _get_supported_strategies() -> list[str]:
@@ -187,6 +198,14 @@ def _set_env_variables(args: Namespace) -> None:
 
 def _get_num_processes(accelerator: str, devices: str) -> int:
     """Parse the `devices` argument to determine how many processes need to be launched on the current machine."""
+
+    if accelerator == "auto" or accelerator is None:
+        accelerator = _choose_auto_accelerator()
+    if devices == "auto":
+        if accelerator == "cuda" or accelerator == "mps" or accelerator == "cpu":
+            devices = "1"
+        else:
+            raise ValueError(f"Cannot default to '1' device for accelerator='{accelerator}'")
     if accelerator == "gpu":
         parsed_devices = _parse_gpu_ids(devices, include_cuda=True, include_mps=True)
     elif accelerator == "cuda":
diff --git a/tests/tests_fabric/test_cli.py b/tests/tests_fabric/test_cli.py
index 944584114184b..e71c42bb46e13 100644
--- a/tests/tests_fabric/test_cli.py
+++ b/tests/tests_fabric/test_cli.py
@@ -46,7 +46,7 @@ def test_run_env_vars_defaults(monkeypatch, fake_script):
     assert "LT_PRECISION" not in os.environ
 
 
-@pytest.mark.parametrize("accelerator", ["cpu", "gpu", "cuda", pytest.param("mps", marks=RunIf(mps=True))])
+@pytest.mark.parametrize("accelerator", ["cpu", "gpu", "cuda", "auto", pytest.param("mps", marks=RunIf(mps=True))])
 @mock.patch.dict(os.environ, os.environ.copy(), clear=True)
 @mock.patch("lightning.fabric.accelerators.cuda.num_cuda_devices", return_value=2)
 def test_run_env_vars_accelerator(_, accelerator, monkeypatch, fake_script):
@@ -85,7 +85,7 @@ def test_run_env_vars_unsupported_strategy(strategy, fake_script):
     assert f"Invalid value for '--strategy': '{strategy}'" in ioerr.getvalue()
 
 
-@pytest.mark.parametrize("devices", ["1", "2", "0,", "1,0", "-1"])
+@pytest.mark.parametrize("devices", ["1", "2", "0,", "1,0", "-1", "auto"])
 @mock.patch.dict(os.environ, os.environ.copy(), clear=True)
 @mock.patch("lightning.fabric.accelerators.cuda.num_cuda_devices", return_value=2)
 def test_run_env_vars_devices_cuda(_, devices, monkeypatch, fake_script):
@@ -97,7 +97,7 @@ def test_run_env_vars_devices_cuda(_, devices, monkeypatch, fake_script):
 
 
 @RunIf(mps=True)
-@pytest.mark.parametrize("accelerator", ["mps", "gpu"])
+@pytest.mark.parametrize("accelerator", ["mps", "gpu", "auto"])
 @mock.patch.dict(os.environ, os.environ.copy(), clear=True)
 def test_run_env_vars_devices_mps(accelerator, monkeypatch, fake_script):
     monkeypatch.setitem(sys.modules, "torch.distributed.run", Mock())


Code review comments:
diff_hunk: @@ -34,7 +34,18 @@
 _CLICK_AVAILABLE = RequirementCache("click")
 _LIGHTNING_SDK_AVAILABLE = RequirementCache("lightning_sdk")
 
-_SUPPORTED_ACCELERATORS = ("cpu", "gpu", "cuda", "mps", "tpu")
+_SUPPORTED_ACCELERATORS = ("cpu", "gpu", "cuda", "mps", "tpu", "auto")
+
+
+def _choose_auto_accelerator() -> str:
path: src/lightning/fabric/cli.py
text: @user1:
so why do we have now the same function?

@author:
I tried to import from https://github.com/Lightning-AI/pytorch-lightning/blob/c85660c0da7777c667309a3ef937276353aacc27/src/lightning/pytorch/trainer/connectors/accelerator_connector.py#L333
But there were some dependencies errors, one for `torchmetrics` on some runners.

@author:
https://github.com/Lightning-AI/pytorch-lightning/actions/runs/15755765823/job/44410806075

@user1:
I see so them move the function to fabric and use it later in Connector
https://github.com/Lightning-AI/pytorch-lightning/tree/master/src/lightning/fabric/utilities