Fix incorrect str plugin type (#19133)

Author: awaelchli

src/lightning/fabric/connector.py

@@ -65,8 +65,7 @@
 from lightning.fabric.utilities.device_parser import _determine_root_gpu_device
 from lightning.fabric.utilities.imports import _IS_INTERACTIVE
 
-_PLUGIN = Union[Precision, ClusterEnvironment, CheckpointIO]
-_PLUGIN_INPUT = Union[_PLUGIN, str]
+_PLUGIN_INPUT = Union[Precision, ClusterEnvironment, CheckpointIO]
 
 
 class _Connector:
@@ -84,15 +83,9 @@ class _Connector:
                backend (registed these too, and _strategy_type could be deprecated)
 
         C. plugins flag could be:
-            1. List of str, which could contain:
-                i. precision str (Not supported in the old accelerator_connector version)
-                ii. checkpoint_io str (Not supported in the old accelerator_connector version)
-                iii. cluster_environment str (Not supported in the old accelerator_connector version)
-            2. List of class, which could contains:
-                i. precision class (should be removed, and precision flag should allow user pass classes)
-                ii. checkpoint_io class
-                iii. cluster_environment class
-
+            1. precision class (should be removed, and precision flag should allow user pass classes)
+            2. checkpoint_io class
+            3. cluster_environment class
 
     priorities which to take when:
         A. Class > str

src/lightning/pytorch/plugins/__init__.py

@@ -13,8 +13,7 @@
 from lightning.pytorch.plugins.precision.transformer_engine import TransformerEnginePrecision
 from lightning.pytorch.plugins.precision.xla import XLAPrecision
 
-PLUGIN = Union[Precision, ClusterEnvironment, CheckpointIO, LayerSync]
-PLUGIN_INPUT = Union[PLUGIN, str]
+_PLUGIN_INPUT = Union[Precision, ClusterEnvironment, CheckpointIO, LayerSync]
 
 __all__ = [
     "AsyncCheckpointIO",

src/lightning/pytorch/trainer/connectors/accelerator_connector.py

@@ -36,7 +36,7 @@
 from lightning.pytorch.accelerators.mps import MPSAccelerator
 from lightning.pytorch.accelerators.xla import XLAAccelerator
 from lightning.pytorch.plugins import (
-    PLUGIN_INPUT,
+    _PLUGIN_INPUT,
     CheckpointIO,
     DeepSpeedPrecision,
     DoublePrecision,
@@ -81,7 +81,7 @@ def __init__(
         num_nodes: int = 1,
         accelerator: Union[str, Accelerator] = "auto",
         strategy: Union[str, Strategy] = "auto",
-        plugins: Optional[Union[PLUGIN_INPUT, List[PLUGIN_INPUT]]] = None,
+        plugins: Optional[Union[_PLUGIN_INPUT, List[_PLUGIN_INPUT]]] = None,
         precision: Optional[_PRECISION_INPUT] = None,
         sync_batchnorm: bool = False,
         benchmark: Optional[bool] = None,
@@ -96,20 +96,14 @@ def __init__(
                 2. accelerator str
                 3. accelerator auto
 
-            B. strategy flag could be :
+            B. strategy flag could be:
                 1. strategy class
                 2. strategy str registered with StrategyRegistry
 
             C. plugins flag could be:
-                1. List of str, which could contain:
-                    i. precision str (Not supported in the old accelerator_connector version)
-                    ii. checkpoint_io str (Not supported in the old accelerator_connector version)
-                    iii. cluster_environment str (Not supported in the old accelerator_connector version)
-                2. List of class, which could contains:
-                    i. precision class (should be removed, and precision flag should allow user pass classes)
-                    ii. checkpoint_io class
-                    iii. cluster_environment class
-
+                1. precision class (should be removed, and precision flag should allow user pass classes)
+                2. checkpoint_io class
+                3. cluster_environment class
 
         priorities which to take when:
             A. Class > str
@@ -175,7 +169,7 @@ def _check_config_and_set_final_flags(
         strategy: Union[str, Strategy],
         accelerator: Union[str, Accelerator],
         precision: Optional[_PRECISION_INPUT],
-        plugins: Optional[Union[PLUGIN_INPUT, List[PLUGIN_INPUT]]],
+        plugins: Optional[Union[_PLUGIN_INPUT, List[_PLUGIN_INPUT]]],
         sync_batchnorm: bool,
     ) -> None:
         """This method checks:

src/lightning/pytorch/trainer/trainer.py

@@ -47,7 +47,7 @@
 from lightning.pytorch.loops.evaluation_loop import _EvaluationLoop
 from lightning.pytorch.loops.fit_loop import _FitLoop
 from lightning.pytorch.loops.utilities import _parse_loop_limits, _reset_progress
-from lightning.pytorch.plugins import PLUGIN_INPUT, Precision
+from lightning.pytorch.plugins import _PLUGIN_INPUT, Precision
 from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.strategies import ParallelStrategy, Strategy
 from lightning.pytorch.trainer import call, setup
@@ -128,7 +128,7 @@ def __init__(
         profiler: Optional[Union[Profiler, str]] = None,
         detect_anomaly: bool = False,
         barebones: bool = False,
-        plugins: Optional[Union[PLUGIN_INPUT, List[PLUGIN_INPUT]]] = None,
+        plugins: Optional[Union[_PLUGIN_INPUT, List[_PLUGIN_INPUT]]] = None,
         sync_batchnorm: bool = False,
         reload_dataloaders_every_n_epochs: int = 0,
         default_root_dir: Optional[_PATH] = None,