From 4116c64698f777275f47707225f04281e3bcf7ba Mon Sep 17 00:00:00 2001
From: Eric Liang <ekhliang@gmail.com>
Date: Thu, 29 Mar 2018 15:36:49 -0700
Subject: [PATCH] [tune] Remove rllib dep again, and add a test (#1792)

* tune should not depend on rllib

* fix dep test

* Tue Mar 27 16:55:41 PDT 2018

* f401
---
 .travis.yml                             |  1 +
 python/ray/tune/test/dependency_test.py | 26 +++++++++++++++++++++++++
 python/ray/tune/trial.py                | 13 ++++++++++---
 python/ray/tune/tune.py                 |  3 ---
 4 files changed, 37 insertions(+), 6 deletions(-)
 create mode 100644 python/ray/tune/test/dependency_test.py

diff --git a/.travis.yml b/.travis.yml
index fc3375dec..0168ba459 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -136,6 +136,7 @@ script:
   - python -m pytest python/ray/dataframe/test/test_series.py
 
   # ray tune tests
+  - python python/ray/tune/test/dependency_test.py
   - python -m pytest python/ray/tune/test/trial_runner_test.py
   - python -m pytest python/ray/tune/test/trial_scheduler_test.py
   - python -m pytest python/ray/tune/test/tune_server_test.py
diff --git a/python/ray/tune/test/dependency_test.py b/python/ray/tune/test/dependency_test.py
new file mode 100644
index 000000000..54c7cd2e3
--- /dev/null
+++ b/python/ray/tune/test/dependency_test.py
@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import sys
+
+import ray
+from ray.tune import register_trainable, run_experiments
+
+
+def f(config, reporter):
+    reporter(timesteps_total=1)
+
+
+if __name__ == "__main__":
+    ray.init()
+    register_trainable("my_class", f)
+    run_experiments({
+        "test": {
+            "run": "my_class",
+            "stop": {"training_iteration": 1}
+        }
+    })
+    assert 'ray.rllib' not in sys.modules, "RLlib should not be imported"
diff --git a/python/ray/tune/trial.py b/python/ray/tune/trial.py
index 80adbc0b5..908a23d72 100644
--- a/python/ray/tune/trial.py
+++ b/python/ray/tune/trial.py
@@ -56,6 +56,11 @@ class Resources(
         return self.gpu + self.extra_gpu
 
 
+def has_trainable(trainable_name):
+    return ray.tune.registry._default_registry.contains(
+        ray.tune.registry.TRAINABLE_CLASS, trainable_name)
+
+
 class Trial(object):
     """A trial object holds the state for one model training run.
 
@@ -83,9 +88,11 @@ class Trial(object):
         in ray.tune.config_parser.
         """
 
-        if not ray.tune.registry._default_registry.contains(
-                ray.tune.registry.TRAINABLE_CLASS, trainable_name):
-            raise TuneError("Unknown trainable: " + trainable_name)
+        if not has_trainable(trainable_name):
+            # Make sure rllib agents are registered
+            from ray import rllib  # noqa: F401
+            if not has_trainable(trainable_name):
+                raise TuneError("Unknown trainable: " + trainable_name)
 
         if stopping_criterion:
             for k in stopping_criterion:
diff --git a/python/ray/tune/tune.py b/python/ray/tune/tune.py
index 28e6fba9b..61e8e0722 100644
--- a/python/ray/tune/tune.py
+++ b/python/ray/tune/tune.py
@@ -49,9 +49,6 @@ def run_experiments(experiments, scheduler=None, with_server=False,
         verbose (bool): How much output should be printed for each trial.
     """
 
-    # Make sure rllib agents are registered
-    from ray import rllib  # noqa # pylint: disable=unused-import
-
     if scheduler is None:
         scheduler = FIFOScheduler()