Repository: ray-project/ray_lightning (language: python)

Issue:
Upgrade to work with pytorch lightning 1.5 (and 1.6?)
There are a lot of backward incompatibilities in 1.5

https://github.com/PyTorchLightning/pytorch-lightning/releases/tag/1.5.0



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index b6b9f3c0..0d923dfc 100644
--- a/README.md
+++ b/README.md
@@ -34,6 +34,7 @@ Here are the supported PyTorch Lightning versions:
 | Ray Lightning | PyTorch Lightning |
 |---|---|
 | 0.1 | 1.4 |
+| 0.2 | 1.5 |
 
 
 ## PyTorch Distributed Data Parallel Plugin on Ray
diff --git a/ray_lightning/ray_ddp.py b/ray_lightning/ray_ddp.py
index cfa76795..25affd53 100644
--- a/ray_lightning/ray_ddp.py
+++ b/ray_lightning/ray_ddp.py
@@ -1,16 +1,21 @@
-import socket
-from contextlib import closing
-from typing import Callable, Dict, List, Union, Any, Optional
+from typing import Callable, Dict, List, Union, Any, Tuple, Optional
 
-import os
 from collections import defaultdict
+from contextlib import closing
+import os
+import socket
 
+import numpy as np
 import torch
 
+import pytorch_lightning as pl
 from pytorch_lightning.accelerators import CPUAccelerator
 from pytorch_lightning.plugins import DDPSpawnPlugin
 from pytorch_lightning import _logger as log, LightningModule
-from pytorch_lightning.utilities import rank_zero_only
+from pytorch_lightning.trainer.states import TrainerFn
+from pytorch_lightning.utilities import rank_zero_only, rank_zero_info
+from pytorch_lightning.utilities.apply_func import apply_to_collection
+from pytorch_lightning.utilities.seed import reset_seed
 
 import ray
 from ray.util import PublicAPI
@@ -20,7 +25,6 @@
 from ray_lightning.util import process_results, to_state_stream, \
     load_state_stream
 from ray_lightning.tune import TUNE_INSTALLED, is_session_enabled
-from ray_lightning.ray_environment import RayEnvironment
 
 
 def find_free_port():
@@ -116,11 +120,7 @@ def __init__(self,
                  **ddp_kwargs: Union[Any, Dict[str, Any]]):
         if not ray.is_initialized():
             ray.init()
-        super().__init__(
-            sync_batchnorm=None,
-            parallel_devices=[],
-            cluster_environment=RayEnvironment(world_size=num_workers),
-            **ddp_kwargs)
+        self._is_remote = False
         resources_per_worker = resources_per_worker if resources_per_worker \
             else {}
         self.nickname = "ddp_ray"
@@ -136,7 +136,23 @@ def __init__(self,
         self.additional_resources_per_worker = resources_per_worker
         self.workers = []
         self.init_hook = init_hook
+
         self._local_rank = 0
+        self._global_rank = 0
+        self._node_rank = 0
+
+        super().__init__(
+            parallel_devices=[], cluster_environment=None, **ddp_kwargs)
+
+    def __getstate__(self):
+        d = self.__dict__.copy()
+        # Don't serialize the workers.
+        del d["workers"]
+        return d
+
+    def __setstate__(self, d):
+        d["workers"] = []
+        self.__dict__.update(d)
 
     def _create_worker(self):
         """Creates Ray actor."""
@@ -146,35 +162,29 @@ def _create_worker(self):
             resources=self.additional_resources_per_worker).remote()
         return worker
 
-    def setup(self, model: LightningModule):
+    def setup(self):
         """Sets up PTL Trainer and creates the Ray actors."""
-        # Check that trainer attribute has been set when this method is called.
-        self._model = model
         self.workers = [self._create_worker() for _ in range(self.num_workers)]
         if self.init_hook:
             ray.get([w.execute.remote(self.init_hook) for w in self.workers])
 
-    def __getstate__(self):
-        d = self.__dict__.copy()
-        del d["workers"]
-        return d
-
-    def __setstate__(self, d):
-        d["workers"] = []
-        self.__dict__.update(d)
-
-    def get_local_ranks(self) -> Dict[int, int]:
-        """Creates a mapping of global ranks to local ranks."""
-        # Get the local ranks for all the workers and store as a dict.
-        # First get the IP address of each remote worker.
-        node_ips = ray.get([w.get_node_ip.remote() for w in self.workers])
-        rank_counter_dict = defaultdict(int)
-        global_to_local = [None] * self.num_workers
-        for global_rank in range(self.num_workers):
-            ip = node_ips[global_rank]
-            global_to_local[global_rank] = rank_counter_dict[ip]
-            rank_counter_dict[ip] += 1
-        return global_to_local
+    def setup_environment(self) -> None:
+        # Swap out the accelerator if necessary.
+        # This is needed to support CPU head with GPU workers or Ray Client.
+        current_accelerator = self.lightning_module.trainer.accelerator
+        if self.use_gpu and isinstance(current_accelerator, CPUAccelerator):
+            from weakref import proxy
+            from ray_lightning.util import DelayedGPUAccelerator
+            precision_plugin = current_accelerator.precision_plugin
+            new_accelerator = DelayedGPUAccelerator(
+                precision_plugin=precision_plugin, training_type_plugin=self)
+            self.lightning_module.trainer.accelerator_connector \
+                ._training_type_plugin = \
+                proxy(new_accelerator.training_type_plugin)
+            self.lightning_module.trainer.accelerator_connector \
+                ._precision_plugin = proxy(new_accelerator.precision_plugin)
+            self.lightning_module.trainer.accelerator_connector.accelerator \
+                = new_accelerator
 
     def _setup_env_vars(self):
         # Get rank 0 worker address and port for DDP connection.
@@ -191,7 +201,48 @@ def _setup_env_vars(self):
         values = [os.getenv(k) for k in keys]
         ray.get([w.set_env_vars.remote(keys, values) for w in self.workers])
 
-    def execution_loop(self, trainer, tune_enabled: bool = True):
+    def start_training(self, trainer):
+        results = self.execution_loop(tune_enabled=True)
+        # reset optimizers, since main process is never used for training and
+        # thus does not have a valid optim state.
+        trainer.optimizers = []
+        return results
+
+    def start_evaluating(self, trainer):
+        results = self.execution_loop(tune_enabled=False)
+        return results
+
+    def start_predicting(self, trainer):
+        results = self.execution_loop(tune_enabled=False)
+        return results
+
+    def get_local_ranks(self) -> Dict[int, Tuple[int, int]]:
+        """Creates a mapping of global ranks to local ranks/node ranks."""
+        # Get the local ranks for all the workers and store as a dict.
+        # First get the IP address of each remote worker.
+        node_ips = ray.get([w.get_node_ip.remote() for w in self.workers])
+
+        node_rank_map = {}
+        counter = 0
+        for ip in node_ips:
+            # If this is a new IP address, then increment counter.
+            if ip not in node_rank_map:
+                node_rank_map[ip] = counter
+                counter += 1
+
+        rank_counter_dict = defaultdict(int)
+        global_to_local = [None] * self.num_workers
+
+        for global_rank in range(self.num_workers):
+            ip = node_ips[global_rank]
+            global_to_local[global_rank] = (
+                rank_counter_dict[ip],  # local rank
+                node_rank_map[ip])  # node rank
+            rank_counter_dict[ip] += 1
+
+        return global_to_local
+
+    def execution_loop(self, tune_enabled: bool = True):
         """Main execution loop for training, testing, & prediction.
 
         Sets up the torch.distributed process group for each
@@ -199,11 +250,14 @@ def execution_loop(self, trainer, tune_enabled: bool = True):
         ``train_remote`` on each worker. If using with Ray Tune, create a
         communication queue to retrieve intermediate results, and process
         those results. Finally retrieve the training results from the rank 0
-        worker and return."""
+        worker and return.
+        """
 
         # Sets environment variables for all workers.
+        # This will set the MASTER_ADDR and MASTER_PORT on each Ray actor.
         self._setup_env_vars()
 
+        # Get the mapping from global ranks to the respective local ranks.
         self.global_to_local = self.get_local_ranks()
 
         model = self._model
@@ -216,65 +270,48 @@ def execution_loop(self, trainer, tune_enabled: bool = True):
             # Create communication queue and send to all the workers.
             queue = Queue(actor_options={"num_cpus": 0})
 
-        futures = [
+        self._futures = [
             self.workers[i].execute.remote(self.execute_remote, model_ref, i,
                                            queue)
             for i in range(self.num_workers)
         ]
 
-        results = process_results(futures, queue)
-        # Get the results, checkpoint path, and model weights from worker 0.
-        results, best_path, state_stream = results[0]
-        state_dict = load_state_stream(state_stream, to_gpu=self.use_gpu)
-        # Set the state for PTL using the output from remote training.
-        self._results = results
-        self._model = model
-        self._model.load_state_dict(state_dict)
-        if self.lightning_module.trainer.checkpoint_callback:
-            self.lightning_module.trainer.checkpoint_callback \
-                .best_model_path = best_path
+        process_results(self._futures, queue)
 
+        self._model = model
         if queue:
             # Shutdown the queue.
             queue.shutdown()
 
-        return results
+    def post_dispatch(self, trainer: "pl.Trainer"):
+        """Finalized fitting.
 
-    def setup_environment(self) -> None:
-        # Swap out the accelerator if necessary.
-        # This is needed to support CPU head with GPU workers or Ray Client.
-        current_accelerator = self.lightning_module.trainer.accelerator
-        if self.use_gpu and isinstance(current_accelerator, CPUAccelerator):
-            from weakref import proxy
-            from ray_lightning.util import DelayedGPUAccelerator
-            precision_plugin = current_accelerator.precision_plugin
-            new_accelerator = DelayedGPUAccelerator(
-                precision_plugin=precision_plugin, training_type_plugin=self)
-            self.lightning_module.trainer.accelerator_connector\
-                ._training_type_plugin = \
-                proxy(new_accelerator.training_type_plugin)
-            self.lightning_module.trainer.accelerator_connector\
-                ._precision_plugin = proxy(new_accelerator.precision_plugin)
-            self.lightning_module.trainer.accelerator_connector.accelerator \
-                = new_accelerator
+        1. Load model weights on driver.
+        2. Shutdown DDP process group and all Ray actors.
 
-    def start_training(self, trainer):
-        results = self.execution_loop(trainer, tune_enabled=True)
-        # reset optimizers, since main process is never used for training and
-        # thus does not have a valid optim state.
-        trainer.optimizers = []
-        return results
+        Shutdown the DDP process group and all the Ray
+        actors."""
 
-    def start_evaluating(self, trainer):
-        results = self.execution_loop(trainer, tune_enabled=False)
-        return results
+        results = ray.get(self._futures)
+        # Get the results, checkpoint path, and model weights from worker 0.
+        results, best_path, state_stream, callback_metrics = results[0]
+        self._results = results
 
-    def start_predicting(self, trainer):
-        results = self.execution_loop(trainer, tune_enabled=False)
-        return results
+        # From DDPSpawn.get_queue
+        self.lightning_module.trainer.callback_metrics.update(
+            apply_to_collection(callback_metrics,
+                                np.ndarray, lambda x: torch.tensor(x)))
 
-    def post_dispatch(self):
-        """Shutdown the DDP process group and all the Ray actors. """
+        # DDPSpawnPlugin.__recover_child_process_weights begin
+        # Difference here is that instead of writing the model weights to a
+        # file and loading it, we use the state dict of the model directly.
+        state_dict = load_state_stream(state_stream, to_gpu=self.use_gpu)
+        # Set the state for PTL using the output from remote training.
+        self.lightning_module.load_state_dict(state_dict)
+        if self.lightning_module.trainer.checkpoint_callback:
+            self.lightning_module.trainer.checkpoint_callback \
+                .best_model_path = best_path
+        # DDPSpawnPlugin.__recover_child_process_weights_end
 
         def shutdown_remote():
             torch.distributed.destroy_process_group()
@@ -289,87 +326,145 @@ def shutdown_remote():
 
     # All methods below are only executed in remote Ray workers.
 
+    def set_world_ranks(self, process_idx: int = 0):
+        """Set the appropriate rank attributes for the trainer."""
+        # Ranks should only be set once all the actors are created and
+        # training has begun (otherwise self.global_to_local has not been
+        # initialized).
+        # If this method is called on the driver (i.e. self._is_remote is
+        # False, then do a no-op).
+        if self._is_remote:
+            self._global_rank = process_idx
+            self._local_rank, self._node_rank = self.global_to_local[
+                self.global_rank]
+
+    def _worker_setup(self, process_idx: int):
+        reset_seed()
+        self.set_world_ranks(process_idx)
+        rank_zero_only.rank = self.global_rank
+
+        # Taken from pytorch_lightning.utilities.distributed
+        # .init_dist_connection
+        # Modified to not use cluster environment.
+        if torch.distributed.is_available(
+        ) and not torch.distributed.is_initialized():
+            log.info(
+                f"initializing distributed: GLOBAL_RANK: {self.global_rank}, "
+                f"MEMBER: {self.global_rank + 1}/{self.world_size}")
+            torch.distributed.init_process_group(
+                self.torch_distributed_backend,
+                rank=self.global_rank,
+                world_size=self.world_size)
+
+            # on rank=0 let everyone know training is starting
+            rank_zero_info(
+                f"{'-' * 100}\n"
+                f"distributed_backend={self.torch_distributed_backend}\n"
+                f"All distributed processes registered. Starting with "
+                f"{self.world_size} processes\n"
+                f"{'-' * 100}\n")
+
     def execute_remote(self,
                        model: LightningModule,
                        global_rank: int,
                        queue: Queue = None):
-        """Train/test/eval function to be executed on each remote worker."""
+        """Train/test/eval function to be executed on each remote worker.
+
+        Modified from DDPSpawn._wrapped_function and DDPSpawn.new_process
+
+        """
         assert isinstance(self, RayPlugin)
         # This method should be executed remotely in each worker.
         self._model = model
-        self.lightning_module.trainer.accelerator_connector\
+        self.lightning_module.trainer._accelerator_connector \
             ._training_type_plugin = self
-        self.lightning_module.trainer.accelerator.training_type_plugin = self
-        self.cluster_environment.set_global_rank(global_rank)
-        self.cluster_environment.set_remote_execution(True)
+        self.lightning_module.trainer._accelerator_connector.accelerator \
+            .training_type_plugin = self
+
+        # TODO: See if this is necessary.
+        self.lightning_module.trainer._data_connector.prepare_data()
+
+        # Set _is_remote to True so that self.set_world_ranks will
+        # properly set the ranks.
+        self._is_remote = True
 
         if queue is not None:
             # Initialize session.
             init_session(rank=global_rank, queue=queue)
 
-        # Calling new_process will call
-        # transfer_distrib_spawn_state_on_fit_end.
-        # We override that method and have it just set attributes.
-        # Then we can just return those attributes here.
-        super(RayPlugin, self).new_process(
-            process_idx=global_rank,
-            trainer=self.lightning_module.trainer,
-            mp_queue=None)
-        # Only need results from worker 0.
+        self._worker_setup(process_idx=global_rank)
+
+        # Below is modified from DDPSpawn.new_process
+
+        # Move the model to the correct device.
+        self.model_to_device()
+
+        # TODO: Support syncbatchnorm.
+        # skip wrapping the model if we are not fitting as no gradients
+        # need to be exchanged.
+        trainer_fn = self.lightning_module.trainer.state.fn
+        if trainer_fn == TrainerFn.FITTING:
+            self.configure_ddp()
+
+        self.barrier()
+
+        results = self.lightning_module.trainer.run_stage()
+
+        # __transfer_distrib_spawn_state_on_fit_end start
         if self.global_rank == 0:
-            return self.results, self.best_model_path, self.model_state_stream
+            checkpoint_callback = \
+                self.lightning_module.trainer.checkpoint_callback
+            best_model_path = checkpoint_callback.best_model_path if \
+                checkpoint_callback else None
+
+            # PyTorch Lightning saves the model weights in a temp file and
+            # loads it back on the driver.
+            # This won't work in a multi-node setup though, so we return the
+            # model state stream directly.
+            model_state_stream = to_state_stream(
+                self.lightning_module.state_dict())
+
+            # From DDPSpawn.add_to_queue
+            callback_metrics: dict = apply_to_collection(
+                self.lightning_module.trainer.callback_metrics,
+                torch.Tensor, lambda x: x.cpu().numpy(
+                ))  # send as numpy to avoid issues with memory sharing
+
+            return_val = results, best_model_path, model_state_stream, \
+                callback_metrics
         else:
-            return None
-
-    def init_ddp_connection(self,
-                            global_rank: int,
-                            world_size: int,
-                            is_slurm_managing_tasks: bool = False) -> None:
-        """Process group creation to be executed on each remote worker."""
-        torch_backend = os.getenv("PL_TORCH_DISTRIBUTED_BACKEND")
-        if torch_backend is None:
-            torch_backend = "nccl" if self.use_gpu else "gloo"
-
-        if not torch.distributed.is_initialized():
-            log.info(f"initializing ddp: GLOBAL_RANK: {global_rank}, MEMBER:"
-                     f" {global_rank + 1}/{world_size}")
-            torch.distributed.init_process_group(
-                backend=torch_backend,
-                rank=global_rank,
-                world_size=world_size,
-            )
+            return_val = None
+        # __transfer_distrib_spawn_state_on_fit_end end
 
-    def set_world_ranks(self, process_idx: int = 0):
-        """Set the appropriate rank attribues for the trainer."""
-        assert self.cluster_environment is not None
-        if self.cluster_environment.is_remote():
-            self._local_rank = self.global_to_local[self.global_rank]
-            self.cluster_environment.set_global_rank(self.global_rank)
-            self.cluster_environment.set_world_size(self.num_workers)
-            rank_zero_only.rank = self.cluster_environment.global_rank()
+        self.lightning_module.trainer._call_teardown_hook()
+
+        return return_val
+
+    @property
+    def world_size(self) -> int:
+        return self.num_workers
+
+    @property
+    def local_rank(self) -> int:
+        return self._local_rank
+
+    @property
+    def global_rank(self) -> int:
+        return self._global_rank
+
+    @property
+    def node_rank(self) -> int:
+        return self._node_rank
 
     @property
     def root_device(self):
-        # Ray already sets CUDA_VISIBLE_DEVICES for each process.
         if self.use_gpu and torch.cuda.is_available():
+            # Ray already sets CUDA_VISIBLE_DEVICES for each process.
+            # So the device is the 0th index in CUDA_VISIBLE_DEVICES
             return torch.device("cuda", 0)
         else:
             return torch.device("cpu")
 
-    def transfer_distrib_spawn_state_on_fit_end(self, results):
-        """Sets the training output as attributes so it can be retrieved."""
-        if self.global_rank == 0:
-            # Save training results as attributes.
-            self._results = results
-            self.model_state_stream = \
-                to_state_stream(self.lightning_module.state_dict())
-            best_model_path = None
-            if self.lightning_module.trainer.checkpoint_callback is not None:
-                best_model_path = \
-                    self.lightning_module.trainer.checkpoint_callback\
-                        .best_model_path
-            self.best_model_path = best_model_path
-
     @property
     def distributed_sampler_kwargs(self):
         """Returns the args to use for torch.data.DistributedSampler."""
@@ -378,10 +473,5 @@ def distributed_sampler_kwargs(self):
         return distributed_sampler_kwargs
 
     @property
-    def require_distributed_sampler(self):
-        """This plugin requires a distributed sampler."""
-        return True
-
-    @property
-    def is_distributed(self):
+    def _is_single_process_single_device(self):
         return True
diff --git a/ray_lightning/ray_ddp_sharded.py b/ray_lightning/ray_ddp_sharded.py
index 60cbeb6d..b63ab1a1 100644
--- a/ray_lightning/ray_ddp_sharded.py
+++ b/ray_lightning/ray_ddp_sharded.py
@@ -1,10 +1,34 @@
 from pytorch_lightning.plugins import DDPSpawnShardedPlugin
+from pytorch_lightning.plugins.precision.sharded_native_amp import \
+    ShardedNativeMixedPrecisionPlugin
+from pytorch_lightning.utilities import _FAIRSCALE_AVAILABLE
+
+if _FAIRSCALE_AVAILABLE:
+    from fairscale.optim.grad_scaler import ShardedGradScaler
 
 from ray.util import PublicAPI
 
 from ray_lightning import RayPlugin
 
 
+# C3 linearization of parent classes will do breadth first since both
+# RayPlugin and DDPSpawnShardedPlugin share a common parent of DDPSpawnPlugin
 @PublicAPI(stability="beta")
 class RayShardedPlugin(RayPlugin, DDPSpawnShardedPlugin):
-    pass
+    def execute_remote(self, model, global_rank, queue):
+        # Need to set self._model here otherwise self.lightning_module will
+        # return None.
+        self._model = model
+
+        # This is copied from `DDPSpawnShardedPlugin.new_process`.
+        # As of PTL 1.5, this is the only difference between
+        # `DDPSpawnShardedPlugin` and `DDPSpawnPlugin`.
+        precision_plugin = self.lightning_module.trainer.accelerator\
+            .precision_plugin
+        if isinstance(precision_plugin, ShardedNativeMixedPrecisionPlugin):
+            precision_plugin.scaler = ShardedGradScaler()
+
+        # After setting the grad scaler, we can now call the default
+        # `RayPlugin.execute_remote`.
+        return super().execute_remote(
+            model=self._model, global_rank=global_rank, queue=queue)
diff --git a/ray_lightning/ray_environment.py b/ray_lightning/ray_environment.py
deleted file mode 100644
index 79ac7a98..00000000
--- a/ray_lightning/ray_environment.py
+++ /dev/null
@@ -1,45 +0,0 @@
-from pytorch_lightning.plugins.environments import ClusterEnvironment
-from pytorch_lightning.utilities import rank_zero_only
-
-
-class RayEnvironment(ClusterEnvironment):
-    """Environment for PTL training on a Ray cluster."""
-
-    def __init__(self, world_size):
-        self.set_world_size(world_size)
-        self._global_rank = 0
-        self._is_remote = False
-
-    def creates_children(self) -> bool:
-        return False
-
-    def master_address(self) -> str:
-        raise NotImplementedError
-
-    def master_port(self) -> int:
-        raise NotImplementedError
-
-    def world_size(self) -> int:
-        return self._world_size
-
-    def set_world_size(self, size: int) -> None:
-        self._world_size = size
-
-    def global_rank(self) -> int:
-        return self._global_rank
-
-    def set_global_rank(self, rank: int) -> None:
-        self._global_rank = rank
-        rank_zero_only.rank = rank
-
-    def set_remote_execution(self, is_remote: bool) -> None:
-        self._is_remote = is_remote
-
-    def is_remote(self) -> bool:
-        return self._is_remote
-
-    def local_rank(self) -> int:
-        raise NotImplementedError
-
-    def node_rank(self) -> int:
-        raise NotImplementedError
diff --git a/ray_lightning/ray_horovod.py b/ray_lightning/ray_horovod.py
index a31405bf..8fb5a08e 100644
--- a/ray_lightning/ray_horovod.py
+++ b/ray_lightning/ray_horovod.py
@@ -1,5 +1,5 @@
 import torch
-from pytorch_lightning import LightningModule
+import pytorch_lightning as pl
 from pytorch_lightning.plugins import HorovodPlugin
 from pytorch_lightning.utilities import rank_zero_only
 
@@ -114,9 +114,8 @@ def world_size(self) -> int:
             return self.num_workers
         return hvd.size()
 
-    def setup(self, model: LightningModule):
+    def setup(self):
         """Creates the RayExecutor object."""
-        self._model = model
         settings = RayExecutor.create_settings(timeout_s=30)
         self.executor = RayExecutor(
             settings,
@@ -170,9 +169,10 @@ def start_training(self, trainer):
     def train_remote(self, model: ObjectRef, queue: Queue = None, **kwargs):
         """Training function to be executed on each remote worker."""
         self._model = ray.get(model)
-        self.lightning_module.trainer.accelerator_connector\
+        self.lightning_module.trainer._accelerator_connector\
             ._training_type_plugin = self
-        self.lightning_module.trainer.accelerator.training_type_plugin = self
+        self.lightning_module.trainer._accelerator_connector.accelerator\
+            .training_type_plugin = self
 
         hvd.init()
         rank_zero_only.rank = self.global_rank
@@ -201,7 +201,7 @@ def train_remote(self, model: ObjectRef, queue: Queue = None, **kwargs):
         return results, to_state_stream(self.lightning_module.state_dict()), \
             best_model_path
 
-    def post_dispatch(self):
+    def post_dispatch(self, trainer: "pl.Trainer"):
         """Shuts down the RayExecutor."""
         self.executor.shutdown()
 
diff --git a/ray_lightning/tune.py b/ray_lightning/tune.py
index 0584c93b..8fc7642b 100644
--- a/ray_lightning/tune.py
+++ b/ray_lightning/tune.py
@@ -110,7 +110,7 @@ def __init__(
         def _get_report_dict(self, trainer: Trainer,
                              pl_module: LightningModule):
             # Don't report if just doing initial validation sanity checks.
-            if trainer.running_sanity_check:
+            if trainer.sanity_checking:
                 return
             if not self._metrics:
                 report_dict = {
@@ -167,7 +167,7 @@ def _create_checkpoint(checkpoint_stream, global_step: int,
                     f.write(checkpoint_stream)
 
         def _handle(self, trainer: Trainer, pl_module: LightningModule):
-            if trainer.running_sanity_check:
+            if trainer.sanity_checking:
                 return
             checkpoint_dict = trainer.checkpoint_connector.dump_checkpoint()
             # Convert to a state stream first.
diff --git a/ray_lightning/util.py b/ray_lightning/util.py
index 034a5f26..f696e1ef 100644
--- a/ray_lightning/util.py
+++ b/ray_lightning/util.py
@@ -39,7 +39,7 @@ def _handle_queue(queue):
             item()
 
 
-def process_results(training_result_futures, queue):
+def process_results(training_result_futures, queue=None):
     """Process results from the queue, and return results from the futures."""
     not_ready = training_result_futures
     while not_ready:
diff --git a/setup.py b/setup.py
index b2b1d9aa..3a70bd2d 100644
--- a/setup.py
+++ b/setup.py
@@ -9,4 +9,4 @@
     long_description="Custom Pytorch Lightning distributed plugins "
     "built on top of distributed computing framework Ray.",
     url="https://github.com/ray-project/ray_lightning_accelerators",
-    install_requires=["pytorch-lightning", "ray"])
+    install_requires=["pytorch-lightning==1.5.*", "ray"])


Test patch:
diff --git a/ray_lightning/tests/test_ddp.py b/ray_lightning/tests/test_ddp.py
index f2d066ae..dce87d39 100644
--- a/ray_lightning/tests/test_ddp.py
+++ b/ray_lightning/tests/test_ddp.py
@@ -8,6 +8,7 @@
 from pytorch_lightning.callbacks import EarlyStopping
 
 import ray
+from ray.cluster_utils import Cluster
 
 from ray_lightning import RayPlugin
 from ray_lightning.tests.utils import get_trainer, train_test, \
@@ -21,6 +22,13 @@ def ray_start_2_cpus():
     ray.shutdown()
 
 
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
 @pytest.fixture
 def ray_start_4_cpus_4_extra():
     address_info = ray.init(num_cpus=4, resources={"extra": 4})
@@ -41,6 +49,17 @@ def seed():
     pl.seed_everything(0)
 
 
+@pytest.fixture
+def ray_start_cluster_2_node_2_cpu():
+    cluster = Cluster()
+    cluster.add_node(num_cpus=2)
+    cluster.add_node(num_cpus=2)
+    address_info = ray.init(cluster.address)
+    yield address_info
+    ray.shutdown()
+    cluster.shutdown()
+
+
 @pytest.mark.parametrize("num_workers", [1, 2])
 def test_actor_creation(tmpdir, ray_start_2_cpus, num_workers):
     """Tests whether the appropriate number of training actors are created."""
@@ -50,11 +69,49 @@ def check_num_actor():
         assert len(ray.state.actors()) == num_workers
 
     model.on_epoch_end = check_num_actor
+
     plugin = RayPlugin(num_workers=num_workers)
     trainer = get_trainer(tmpdir, plugins=[plugin])
     trainer.fit(model)
 
 
+def test_global_local_ranks(ray_start_4_cpus):
+    """Tests local rank and node rank map is correct."""
+
+    @ray.remote
+    class Node1Actor:
+        def get_node_ip(self):
+            return "1"
+
+    @ray.remote
+    class Node2Actor:
+        def get_node_ip(self):
+            return "2"
+
+    plugin = RayPlugin(num_workers=4, use_gpu=False)
+
+    # 2 workers on "Node 1", 2 workers on "Node 2"
+    plugin.workers = [
+        Node1Actor.remote(),
+        Node1Actor.remote(),
+        Node2Actor.remote(),
+        Node2Actor.remote()
+    ]
+
+    global_to_local = plugin.get_local_ranks()
+
+    assert len(global_to_local) == 4
+    local_ranks = {ranks[0] for ranks in global_to_local}
+    node_ranks = {ranks[1] for ranks in global_to_local}
+
+    assert local_ranks == set(range(2))
+    assert node_ranks == set(range(2))
+
+    # Make sure the rank 0 worker has local rank and node rank of 0.
+    assert global_to_local[0][0] == 0
+    assert global_to_local[0][1] == 0
+
+
 @pytest.mark.parametrize("num_workers", [1, 2])
 @pytest.mark.parametrize("extra_resource_per_worker", [1, 2])
 @pytest.mark.parametrize("num_cpus_per_worker", [1, 2])
diff --git a/ray_lightning/tests/utils.py b/ray_lightning/tests/utils.py
index 7175722d..0cc2a1c5 100644
--- a/ray_lightning/tests/utils.py
+++ b/ray_lightning/tests/utils.py
@@ -3,13 +3,15 @@
 
 import torch
 import torch.nn.functional as F
-from pytorch_lightning.plugins import Plugin
 from torch.utils.data import Dataset
 
 import pytorch_lightning as pl
+from pytorch_lightning.plugins import PLUGIN_INPUT
 from pytorch_lightning import LightningModule, Callback, Trainer, \
     LightningDataModule
 
+import torchmetrics
+
 
 class RandomDataset(Dataset):
     def __init__(self, size: int, length: int):
@@ -106,7 +108,7 @@ def __init__(self, config, data_dir=None):
         self.layer_1 = torch.nn.Linear(28 * 28, layer_1)
         self.layer_2 = torch.nn.Linear(layer_1, layer_2)
         self.layer_3 = torch.nn.Linear(layer_2, 10)
-        self.accuracy = pl.metrics.Accuracy()
+        self.accuracy = torchmetrics.Accuracy()
 
     def forward(self, x):
         batch_size, channels, width, height = x.size()
@@ -146,12 +148,11 @@ def validation_epoch_end(self, outputs):
 
 
 def get_trainer(dir,
-                plugins: List[Plugin],
+                plugins: List[PLUGIN_INPUT],
                 use_gpu: bool = False,
                 max_epochs: int = 1,
                 limit_train_batches: int = 10,
                 limit_val_batches: int = 10,
-                progress_bar_refresh_rate: int = 0,
                 callbacks: Optional[List[Callback]] = None,
                 checkpoint_callback: bool = True,
                 **trainer_kwargs) -> Trainer:
@@ -165,7 +166,7 @@ def get_trainer(dir,
         max_epochs=max_epochs,
         limit_train_batches=limit_train_batches,
         limit_val_batches=limit_val_batches,
-        progress_bar_refresh_rate=progress_bar_refresh_rate,
+        enable_progress_bar=False,
         checkpoint_callback=checkpoint_callback,
         **trainer_kwargs)
     return trainer
@@ -198,7 +199,7 @@ def predict_test(trainer: Trainer, model: LightningModule,
     model = trainer.lightning_module
     dm.setup(stage="test")
     test_loader = dm.test_dataloader()
-    acc = pl.metrics.Accuracy()
+    acc = torchmetrics.Accuracy()
     for batch in test_loader:
         x, y = batch
         with torch.no_grad():
diff --git a/requirements-test.txt b/requirements-test.txt
index 493dc4a6..dc34f13f 100644
--- a/requirements-test.txt
+++ b/requirements-test.txt
@@ -4,8 +4,9 @@ flake8-comprehensions
 flake8-quotes
 yapf==0.23.0
 pytest
-pytorch-lightning==1.4.7
+pytorch-lightning==1.5.8
 lightning-bolts==0.3.3
 ray[tune]
 torch==1.8.1
+torchmetrics
 torchvision