fixing Win failed import (Lightning-AI#1163)

* version

* try fix distrib

* update try import

Author: Borda

.github/workflows/rebase.yml

@@ -2,8 +2,8 @@ name: Automatic Rebase
 # https://github.com/marketplace/actions/automatic-rebase
 
 on:
-  issue_comment:
-    types: [created]
+  - pull_request
+
 jobs:
   rebase:
     name: Rebase

pytorch_lightning/__init__.py

@@ -1,6 +1,6 @@
 """Root package info."""
 
-__version__ = '0.7.1'
+__version__ = '0.7.2-dev'
 __author__ = 'William Falcon et al.'
 __author_email__ = 'waf2107@columbia.edu'
 __license__ = 'Apache-2.0'

pytorch_lightning/core/hooks.py

@@ -22,10 +22,10 @@
 
 try:
     from apex import amp
-
-    APEX_AVAILABLE = True
 except ImportError:
     APEX_AVAILABLE = False
+else:
+    APEX_AVAILABLE = True
 
 
 class ModelHooks(torch.nn.Module):

pytorch_lightning/core/lightning.py

@@ -9,7 +9,7 @@
 
 import torch
 from torch import Tensor
-from torch.distributed import init_process_group
+import torch.distributed as torch_distrib
 from torch.nn.parallel import DistributedDataParallel
 from torch.optim import Adam
 from torch.optim.optimizer import Optimizer
@@ -24,10 +24,10 @@
 
 try:
     import torch_xla.core.xla_model as xm
-    XLA_AVAILABLE = True
-
 except ImportError:
     XLA_AVAILABLE = False
+else:
+    XLA_AVAILABLE = True
 
 
 class LightningModule(ABC, GradInformation, ModelIO, ModelHooks):
@@ -859,7 +859,7 @@ def init_ddp_connection(self):
 
         root_node = self.trainer.resolve_root_node_address(root_node)
         os.environ['MASTER_ADDR'] = root_node
-        init_process_group('nccl', rank=proc_rank, world_size=world_size)
+        torch_distrib.init_process_group('nccl', rank=proc_rank, world_size=world_size)
 
     def configure_apex(
             self,

pytorch_lightning/trainer/auto_mix_precision.py

@@ -1,13 +1,12 @@
-
+import logging as log
 from abc import ABC
 
 try:
     from apex import amp
-
-    APEX_AVAILABLE = True
 except ImportError:
     APEX_AVAILABLE = False
-import logging as log
+else:
+    APEX_AVAILABLE = True
 
 
 class TrainerAMPMixin(ABC):

pytorch_lightning/trainer/data_loading.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Union, List, Tuple, Callable
 
-import torch.distributed as dist
+import torch.distributed as torch_distrib
 from torch.utils.data import SequentialSampler, DataLoader
 from torch.utils.data.distributed import DistributedSampler
 
@@ -224,7 +224,7 @@ def request_dataloader(self, dataloader_fx: Callable) -> DataLoader:
         # get the function we'll use to get data
         if self.use_ddp or self.use_ddp2:
             # all processes wait until data download has happened
-            dist.barrier()
+            torch_distrib.barrier()
 
         # data download/load on TPU
         elif self.use_tpu and XLA_AVAILABLE:

pytorch_lightning/trainer/trainer.py

@@ -8,7 +8,7 @@
 
 import torch
 from torch import optim
-import torch.distributed as dist
+import torch.distributed as torch_distrib
 import torch.multiprocessing as mp
 from torch.optim.optimizer import Optimizer
 from torch.utils.data import DataLoader
@@ -748,7 +748,7 @@ def run_pretrain_routine(self, model: LightningModule):
             self.logger.save()
 
         if self.use_ddp or self.use_ddp2:
-            dist.barrier()
+            torch_distrib.barrier()
 
         # wait for all models to restore weights
         if self.on_tpu and XLA_AVAILABLE:

pytorch_lightning/trainer/training_io.py

@@ -100,7 +100,7 @@
 from typing import Union
 
 import torch
-import torch.distributed as dist
+import torch.distributed as torch_distrib
 
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.loggers import LightningLoggerBase
@@ -177,7 +177,7 @@ def restore_weights(self, model):
         # wait for all models to restore weights
         if self.use_ddp or self.use_ddp2:
             # wait for all processes to catch up
-            dist.barrier()
+            torch_distrib.barrier()
 
         # wait for all models to restore weights
         if self.on_tpu and XLA_AVAILABLE: