Fix Trainer.plugins type declaration (#7288)

* Fix trainer.plugins type declaration

* Don't ClusterEnvironment(Plugin)

* fix import error, yapf formatter

* Add test

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

Author: leezu

pytorch_lightning/trainer/trainer.py

@@ -30,6 +30,7 @@
 from pytorch_lightning.core.step_result import Result
 from pytorch_lightning.loggers import LightningLoggerBase
 from pytorch_lightning.plugins import Plugin
+from pytorch_lightning.plugins.environments import ClusterEnvironment
 from pytorch_lightning.profiler import BaseProfiler
 from pytorch_lightning.trainer.callback_hook import TrainerCallbackHookMixin
 from pytorch_lightning.trainer.configuration_validator import ConfigValidator
@@ -137,7 +138,7 @@ def __init__(
         terminate_on_nan: bool = False,
         auto_scale_batch_size: Union[str, bool] = False,
         prepare_data_per_node: bool = True,
-        plugins: Optional[Union[Plugin, str, list]] = None,
+        plugins: Optional[Union[List[Union[Plugin, ClusterEnvironment, str]], Plugin, ClusterEnvironment, str]] = None,
         amp_backend: str = 'native',
         amp_level: str = 'O2',
         distributed_backend: Optional[str] = None,

tests/utilities/test_cli.py

@@ -27,6 +27,7 @@
 
 from pytorch_lightning import LightningDataModule, LightningModule, Trainer
 from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
+from pytorch_lightning.plugins.environments import SLURMEnvironment
 from pytorch_lightning.utilities import _TPU_AVAILABLE
 from pytorch_lightning.utilities.cli import LightningArgumentParser, LightningCLI, SaveConfigCallback
 from tests.helpers import BoringDataModule, BoringModel
@@ -280,6 +281,22 @@ def on_fit_start(self):
     assert cli.trainer.ran_asserts
 
 
+def test_lightning_cli_args_cluster_environments(tmpdir):
+    plugins = [dict(class_path='pytorch_lightning.plugins.environments.SLURMEnvironment')]
+
+    class TestModel(BoringModel):
+
+        def on_fit_start(self):
+            # Ensure SLURMEnvironment is set, instead of default LightningEnvironment
+            assert isinstance(self.trainer.accelerator_connector._cluster_environment, SLURMEnvironment)
+            self.trainer.ran_asserts = True
+
+    with mock.patch('sys.argv', ['any.py', f'--trainer.plugins={json.dumps(plugins)}']):
+        cli = LightningCLI(TestModel, trainer_defaults=dict(default_root_dir=str(tmpdir), fast_dev_run=True))
+
+    assert cli.trainer.ran_asserts
+
+
 def test_lightning_cli_args(tmpdir):
 
     cli_args = [
@@ -390,6 +407,7 @@ def test_lightning_cli_print_config():
 def test_lightning_cli_submodules(tmpdir):
 
     class MainModule(BoringModel):
+
         def __init__(
             self,
             submodule1: LightningModule,