Tune improvements and Cleanup (#69)

* wip

* update

* updates

* add test

* try

* add back

* fix

Author: amogkam

.github/workflows/test.yaml

@@ -37,7 +37,6 @@ jobs:
           python -m pip install codecov
           python -m pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-2.0.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl
           if [ -f requirements-test.txt ]; then python -m pip install -r requirements-test.txt; fi
-          HOROVOD_WITH_GLOO=1 HOROVOD_WITHOUT_MPI=1 HOROVOD_WITHOUT_MXNET=1 pip install git+https://github.com/horovod/horovod.git
       - name: Install package
         run: |
           python -m pip install -e .
@@ -124,7 +123,6 @@ jobs:
           python -m pip install codecov
           python -m pip install -U ray
           if [ -f requirements-test.txt ]; then python -m pip install -r requirements-test.txt; fi
-          HOROVOD_WITH_GLOO=1 HOROVOD_WITHOUT_MPI=1 HOROVOD_WITHOUT_MXNET=1 pip install -U git+https://github.com/horovod/horovod.git
       - name: Install package
         run: |
           python -m pip install -e .
@@ -193,4 +191,35 @@ jobs:
           pushd ray_lightning/tests
           echo "running examples with Ray Client 1" && python -m pytest -v --durations=0 -x test_client.py
           echo "running examples with Ray Client 2" && python -m pytest -v --durations=0 -x test_client_2.py
-          echo "running examples with Ray Client 3" && python -m pytest -v --durations=0 -x test_client_3.py
+          echo "running examples with Ray Client 3" && python -m pytest -v --durations=0 -x test_client_3.py
+
+  test_linux_compat:
+    # Test compatibility when optional libraries are not installed.
+    runs-on: ubuntu-latest
+    timeout-minutes: 40
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python 3.7
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.7
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install --upgrade setuptools
+          python -m pip install codecov
+          python -m pip install -U ray
+          if [ -f requirements-test.txt ]; then python -m pip install -r requirements-test.txt; fi
+          HOROVOD_WITH_GLOO=1 HOROVOD_WITHOUT_MPI=1 HOROVOD_WITHOUT_MXNET=1 pip install -U git+https://github.com/horovod/horovod.git
+      - name: Uninstall unavailable dependencies
+        run: |
+          # Uninstall Tune
+          pip uninstall -y tabulate
+      - name: Install package
+        run: |
+          python -m pip install -e .
+      - name: Test with Pytest
+        run: |
+          pushd ray_lightning/tests
+          python -m pytest -v --durations=0 -x test_ddp.py
+          python -m pytest -v --durations=0 -x test_horovod.py

README.md

@@ -86,8 +86,10 @@ You no longer have to set environment variables or configurations and run your t
 Example using `ray_lightning` with Tune:
 
 ```python
+from ray import tune
+
 from ray_lightning import RayPlugin
-from ray_lightning.tune import TuneReportCallback
+from ray_lightning.tune import TuneReportCallback, get_tune_ddp_resources
 
 def train_mnist(config):
 
@@ -111,17 +113,14 @@ config = {
     "batch_size": tune.choice([32, 64, 128]),
 }
 
-# Make sure to specify how many actors each training run will create via the "extra_cpu" field.
+# Make sure to pass in ``resources_per_trial`` using the ``get_tune_ddp_resources`` utility.
 analysis = tune.run(
         train_mnist,
         metric="loss",
         mode="min",
         config=config,
         num_samples=num_samples,
-        resources_per_trial={
-            "cpu": 1,
-            "extra_cpu": 4
-        },
+        resources_per_trial=get_tune_ddp_resources(num_workers=4),
         name="tune_mnist")
 
 print("Best hyperparameters found were: ", analysis.best_config)

ray_lightning/examples/ray_ddp_example.py

@@ -10,7 +10,7 @@
 
 import ray
 from ray import tune
-from ray_lightning.tune import TuneReportCallback
+from ray_lightning.tune import TuneReportCallback, get_tune_ddp_resources
 from ray_lightning import RayPlugin
 from ray_lightning.tests.utils import LightningMNISTClassifier
 
@@ -108,11 +108,8 @@ def tune_mnist(data_dir,
         mode="min",
         config=config,
         num_samples=num_samples,
-        resources_per_trial={
-            "cpu": 1,
-            "extra_cpu": num_workers,
-            "extra_gpu": num_workers * int(use_gpu)
-        },
+        resources_per_trial=get_tune_ddp_resources(
+            num_workers=num_workers, use_gpu=use_gpu),
         name="tune_mnist")
 
     print("Best hyperparameters found were: ", analysis.best_config)

ray_lightning/examples/ray_ddp_tune.py

@@ -7,7 +7,7 @@
 import pytorch_lightning as pl
 import ray
 from ray import tune
-from ray_lightning.tune import TuneReportCallback
+from ray_lightning.tune import TuneReportCallback, get_tune_ddp_resources
 from ray_lightning import RayPlugin
 from ray_lightning.tests.utils import LightningMNISTClassifier
 
@@ -71,11 +71,8 @@ def tune_mnist(data_dir,
         mode="min",
         config=config,
         num_samples=num_samples,
-        resources_per_trial={
-            "cpu": 1,
-            "extra_cpu": num_workers,
-            "extra_gpu": num_workers * int(use_gpu)
-        },
+        resources_per_trial=get_tune_ddp_resources(
+            num_workers=num_workers, use_gpu=use_gpu),
         name="tune_mnist")
 
     print("Best hyperparameters found were: ", analysis.best_config)

ray_lightning/ray_ddp.py

@@ -1,4 +1,6 @@
 import io
+import socket
+from contextlib import closing
 from typing import Callable, Dict, List, Union, Any
 
 import os
@@ -12,7 +14,6 @@
 from pytorch_lightning.utilities import rank_zero_only
 
 import ray
-from ray.util.sgd.utils import find_free_port
 from ray.util.queue import Queue
 
 from ray_lightning.session import init_session
@@ -21,6 +22,13 @@
 from ray_lightning.ray_environment import RayEnvironment
 
 
+def find_free_port():
+    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
+
+
 @ray.remote
 class RayExecutor:
     """A class to execute any arbitrary function remotely."""
@@ -39,7 +47,7 @@ def set_env_vars(self, keys: List[str], values: List[str]):
 
     def get_node_ip(self):
         """Returns the IP address of the node that this Ray actor is on."""
-        return ray.services.get_node_ip_address()
+        return ray.util.get_node_ip_address()
 
     def execute(self, fn: Callable, *args, **kwargs):
         """Execute the provided function and return the result."""

ray_lightning/ray_horovod.py

@@ -74,6 +74,7 @@ def __init__(self,
                  num_hosts: int = 1,
                  num_slots: int = 1,
                  use_gpu: bool = False):
+
         if not HOROVOD_AVAILABLE:
             raise RuntimeError("Please intall Horovod to use this plugin.")
         if not ray.is_initialized():

ray_lightning/tune.py

@@ -6,6 +6,7 @@
 from pytorch_lightning import Trainer, LightningModule
 
 from ray_lightning.session import put_queue, get_actor_rank
+from ray_lightning.util import Unavailable
 
 try:
     from ray import tune
@@ -14,15 +15,31 @@
     TUNE_INSTALLED = True
 except ImportError:
     tune = None
-    TuneCallback = object
+    TuneCallback = Unavailable
 
     def is_session_enabled():
         return False
 
+    get_tune_ddp_resources = Unavailable
+
     TUNE_INSTALLED = False
 
 if TUNE_INSTALLED:
 
+    def get_tune_ddp_resources(num_workers: int = 1,
+                               cpus_per_worker: int = 1,
+                               use_gpu: bool = False) -> Dict[str, int]:
+        """Returns the PlacementGroupFactory to use for Ray Tune."""
+        from ray.tune import PlacementGroupFactory
+
+        head_bundle = {"CPU": 1}
+        child_bundle = {"CPU": cpus_per_worker, "GPU": int(use_gpu)}
+        child_bundles = [child_bundle.copy() for _ in range(num_workers)]
+        bundles = [head_bundle] + child_bundles
+        placement_group_factory = PlacementGroupFactory(
+            bundles, strategy="PACK")
+        return placement_group_factory
+
     class TuneReportCallback(TuneCallback):
         """Distributed PyTorch Lightning to Ray Tune reporting callback
 
@@ -197,3 +214,8 @@ def __init__(
         def _handle(self, trainer: Trainer, pl_module: LightningModule):
             self._checkpoint._handle(trainer, pl_module)
             self._report._handle(trainer, pl_module)
+
+else:
+    # If Tune is not installed.
+    TuneReportCallback = Unavailable
+    TuneReportCheckpointCallback = Unavailable

ray_lightning/util.py

@@ -17,6 +17,8 @@ def setup(self, trainer: Trainer, model: LightningModule) -> None:
         return super(GPUAccelerator, self).setup(trainer, model)
 
     def on_train_start(self) -> None:
+        if "cuda" not in str(self.root_device):
+            raise RuntimeError("GPUs were requested but are not available.")
         super(DelayedGPUAccelerator, self).on_train_start()