Repository: Lightning-AI/pytorch-lightning
Pull request: Enable more shorthand strategy names in the Fabric CLI
## What does this PR do?

Fixes #16479

For example, this is now possible:

```
lightning run model --strategy=deepspeed_stage_3 train.py
```

## Before submitting
- [x] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/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?
- [x] Did you make sure to update the documentation with your changes? (if necessary)
- [x] 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?
- [x] Did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

 - [x] Is this pull request ready for review? (if not, please submit in draft mode)
 - [x] Check that all items from **Before submitting** are resolved
 - [x] Make sure the title is self-explanatory and the description concisely explains the PR
 - [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?
I made sure I had fun coding 🙃


cc @borda @carmocca @justusschock @awaelchli

Problem statement:
Expose available strategies from registry in Fabric CLI
### Description & Motivation

The Fabric CLI has the following options for strategy:

```
lightning run model --help

  --strategy [ddp|dp|deepspeed]   Strategy for how to run across multiple
                                  devices.
```

But special keywords like `deepspeed_stage_3` are not exposed. Why?

The initial implementation was restrictive and minimal with intention. The CLI can't support strategy options like "ddp-spawn", because it assumes the launcher is a script launcher.

### Pitch

Allow all entries from `StrategyRegistery.available_strategies()`, but filter out the ones that have the `MultiProcessingLauncher` (ddp-spawn, ddp-fork, etc.).

### Alternatives

1) Hardcode the subset of options without `StrategyRegistery.available_strategies()`
2) Keep as is

### Additional context

Raised by @justusschock 

cc @borda @carmocca @justusschock @awaelchli

Hints from issue discussion:

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/src/lightning_fabric/CHANGELOG.md b/src/lightning_fabric/CHANGELOG.md
index 4c5c406db5747..7e6feac648f43 100644
--- a/src/lightning_fabric/CHANGELOG.md
+++ b/src/lightning_fabric/CHANGELOG.md
@@ -22,6 +22,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
   * `Fabric.load` returns a dictionary of objects that weren't loaded into the state
   * `Strategy.save_checkpoint` and `Fabric.load_checkpoint` are now responsible for accessing the state of the model and optimizers
 
+- Enabled all shorthand strategy names that can be supported in the CLI ([#16485](https://github.com/Lightning-AI/lightning/pull/16485))
+
 
 ### Deprecated
 
diff --git a/src/lightning_fabric/cli.py b/src/lightning_fabric/cli.py
index a50c580c37ad8..f022743c5289f 100644
--- a/src/lightning_fabric/cli.py
+++ b/src/lightning_fabric/cli.py
@@ -13,12 +13,14 @@
 # limitations under the License.
 import logging
 import os
+import re
 from argparse import Namespace
 from typing import Any, List, Optional
 
 from lightning_utilities.core.imports import RequirementCache
 
 from lightning_fabric.accelerators import CPUAccelerator, CUDAAccelerator, MPSAccelerator
+from lightning_fabric.strategies import STRATEGY_REGISTRY
 from lightning_fabric.utilities.device_parser import _parse_gpu_ids
 
 _log = logging.getLogger(__name__)
@@ -26,9 +28,17 @@
 _CLICK_AVAILABLE = RequirementCache("click")
 
 _SUPPORTED_ACCELERATORS = ("cpu", "gpu", "cuda", "mps", "tpu")
-_SUPPORTED_STRATEGIES = ("ddp", "dp", "deepspeed")
 _SUPPORTED_PRECISION = ("64", "32", "16", "bf16")
 
+
+def _get_supported_strategies() -> List[str]:
+    """Returns strategy choices from the registry, with the ones removed that are incompatible to be launched from
+    the CLI or ones that require further configuration by the user."""
+    available_strategies = STRATEGY_REGISTRY.available_strategies()
+    excluded = r".*(spawn|fork|notebook|xla|tpu|offload).*"
+    return [strategy for strategy in available_strategies if not re.match(excluded, strategy)]
+
+
 if _CLICK_AVAILABLE:
     import click
 
@@ -50,7 +60,7 @@
     )
     @click.option(
         "--strategy",
-        type=click.Choice(_SUPPORTED_STRATEGIES),
+        type=click.Choice(_get_supported_strategies()),
         default=None,
         help="Strategy for how to run across multiple devices.",
     )
diff --git a/tests/tests_fabric/test_cli.py b/tests/tests_fabric/test_cli.py
index 4fc0c64a2f30b..effff5f276d64 100644
--- a/tests/tests_fabric/test_cli.py
+++ b/tests/tests_fabric/test_cli.py
@@ -11,7 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import contextlib
 import os
+from io import StringIO
 from unittest import mock
 from unittest.mock import Mock
 
@@ -54,7 +56,7 @@ def test_cli_env_vars_accelerator(_, accelerator, monkeypatch, fake_script):
     assert os.environ["LT_ACCELERATOR"] == accelerator
 
 
-@pytest.mark.parametrize("strategy", ["dp", "ddp", "deepspeed"])
+@pytest.mark.parametrize("strategy", ["dp", "ddp", "deepspeed", "deepspeed_stage_3"])
 @mock.patch.dict(os.environ, os.environ.copy(), clear=True)
 @mock.patch("lightning_fabric.accelerators.cuda.num_cuda_devices", return_value=2)
 def test_cli_env_vars_strategy(_, strategy, monkeypatch, fake_script):
@@ -65,6 +67,15 @@ def test_cli_env_vars_strategy(_, strategy, monkeypatch, fake_script):
     assert os.environ["LT_STRATEGY"] == strategy
 
 
+@pytest.mark.parametrize("strategy", ["ddp_spawn", "ddp_fork", "ddp_notebook", "deepspeed_stage_3_offload"])
+def test_cli_env_vars_unsupported_strategy(strategy, fake_script):
+    ioerr = StringIO()
+    with pytest.raises(SystemExit) as e, contextlib.redirect_stderr(ioerr):
+        _run_model.main([fake_script, "--strategy", strategy])
+    assert e.value.code == 2
+    assert f"Invalid value for '--strategy': '{strategy}'" in ioerr.getvalue()
+
+
 @pytest.mark.parametrize("devices", ["1", "2", "0,", "1,0", "-1"])
 @mock.patch.dict(os.environ, os.environ.copy(), clear=True)
 @mock.patch("lightning_fabric.accelerators.cuda.num_cuda_devices", return_value=2)


Code review comments:
diff_hunk: @@ -54,7 +56,7 @@ def test_cli_env_vars_accelerator(_, accelerator, monkeypatch, fake_script):
     assert os.environ["LT_ACCELERATOR"] == accelerator
 
 
-@pytest.mark.parametrize("strategy", ["dp", "ddp", "deepspeed"])
+@pytest.mark.parametrize("strategy", ["dp", "ddp", "deepspeed", "deepspeed_stage_3"])
path: tests/tests_fabric/test_cli.py
text: @user2:
```suggestion
@user1.mark.parametrize("strategy", _get_supported_strategies())
```