fix import error for DDPPlugin in oss

Summary: DDPPlugin has been renamed to DDPStrategy (as part of https://github.com/PyTorchLightning/pytorch-lightning/issues/10549), causing oss CI to fail. Simply skipping the import to unblock CI since DDP feature is not used in test.

Reviewed By: kazhang

Differential Revision: D33351636

fbshipit-source-id: 7a1881c8cd48d9ff17edd41137d27a976103fdde

tests/runner/test_runner_lightning_quantization.py

@@ -20,7 +20,6 @@ from d2go.utils.testing.helper import tempdir
 from d2go.utils.testing.lightning_test_module import TestModule
 from pytorch_lightning import Trainer, seed_everything
 from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
-from pytorch_lightning.plugins import DDPPlugin
 from torch.ao.quantization import (  # @manual; @manual
     default_dynamic_qconfig,
     get_default_qconfig,
@@ -197,7 +196,8 @@ class TestQuantizationAwareTraining(unittest.TestCase):
             callbacks=[qat],
             max_epochs=num_epochs,
             logger=False,
-            plugins=[DDPPlugin(find_unused_parameters=False)],
+            strategy="ddp_find_unused_parameters_false",
+            plugins=[],
         )
         trainer.fit(model)
 

tools/lightning_train_net.py

@@ -22,9 +22,16 @@ from pytorch_lightning.callbacks import Callback
 from pytorch_lightning.callbacks import LearningRateMonitor
 from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
 from pytorch_lightning.loggers import TensorBoardLogger
-from pytorch_lightning.plugins import DDPPlugin
 from torch.distributed import get_rank
 
+try:
+    from pytorch_lightning.plugins import DDPPlugin
+except ImportError:
+    assert os.getenv("OSSRUN") == "1"
+    # FIXME: DDPPlugin has been renamed to DDPStrategy, however internal version is
+    # not updated yet, temporally skipping the import in oss env in order to unblock
+    # CI where DPP is not used.
+
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("detectron2go.lightning.train_net")