[tune] Avoid scheduler blocking, add reuse_actors optimization

python/ray/tune/ray_trial_executor.py

@@ -13,6 +13,7 @@
 from ray.tune.logger import NoopLogger
 from ray.tune.trial import Trial, Resources, Checkpoint
 from ray.tune.trial_executor import TrialExecutor
+from ray.tune.util import warn_if_slow
 
 logger = logging.getLogger(__name__)
 
@@ -114,9 +115,6 @@ def _stop_trial(self, trial, error=False, error_msg=None,
                 stop_tasks = []
                 stop_tasks.append(trial.runner.stop.remote())
                 stop_tasks.append(trial.runner.__ray_terminate__.remote())
-                # TODO(ekl)  seems like wait hangs when killing actors
-                _, unfinished = ray.wait(
-                    stop_tasks, num_returns=2, timeout=0.25)
         except Exception:
             logger.exception("Error stopping runner for Trial %s", str(trial))
             self.set_status(trial, Trial.ERROR)
@@ -207,7 +205,8 @@ def reset_trial(self, trial, new_config, new_experiment_tag):
         trial.experiment_tag = new_experiment_tag
         trial.config = new_config
         trainable = trial.runner
-        reset_val = ray.get(trainable.reset_config.remote(new_config))
+        with warn_if_slow("reset_config"):
+            reset_val = ray.get(trainable.reset_config.remote(new_config))
         return reset_val
 
     def get_running_trials(self):
@@ -228,7 +227,8 @@ def fetch_result(self, trial):
         if not trial_future:
             raise ValueError("Trial was not running.")
         self._running.pop(trial_future[0])
-        result = ray.get(trial_future[0])
+        with warn_if_slow("fetch_result"):
+            result = ray.get(trial_future[0])
         return result
 
     def _commit_resources(self, resources):
@@ -368,7 +368,8 @@ def save(self, trial, storage=Checkpoint.DISK):
         if storage == Checkpoint.MEMORY:
             trial._checkpoint.value = trial.runner.save_to_object.remote()
         else:
-            trial._checkpoint.value = ray.get(trial.runner.save.remote())
+            with warn_if_slow("save_to_disk"):
+                trial._checkpoint.value = ray.get(trial.runner.save.remote())
         return trial._checkpoint.value
 
     def restore(self, trial, checkpoint=None):
@@ -389,11 +390,12 @@ def restore(self, trial, checkpoint=None):
             value = checkpoint.value
             if checkpoint.storage == Checkpoint.MEMORY:
                 assert type(value) != Checkpoint, type(value)
-                ray.get(trial.runner.restore_from_object.remote(value))
+                trial.runner.restore_from_object.remote(value)
             else:
                 worker_ip = ray.get(trial.runner.current_ip.remote())
                 trial.sync_logger_to_new_location(worker_ip)
-                ray.get(trial.runner.restore.remote(value))
+                with warn_if_slow("restore_from_disk"):
+                    ray.get(trial.runner.restore.remote(value))
             trial.last_result = checkpoint.last_result
             return True
         except Exception:

python/ray/tune/schedulers/pbt.py

@@ -218,23 +218,29 @@ def _exploit(self, trial_executor, trial, trial_to_clone):
         trial_state = self._trial_state[trial]
         new_state = self._trial_state[trial_to_clone]
         if not new_state.last_checkpoint:
-            logger.warning("[pbt]: no checkpoint for trial."
-                           " Skip exploit for Trial {}".format(trial))
+            logger.info("[pbt]: no checkpoint for trial."
+                        " Skip exploit for Trial {}".format(trial))
             return
         new_config = explore(trial_to_clone.config, self._hyperparam_mutations,
                              self._resample_probability,
                              self._custom_explore_fn)
-        logger.warning("[exploit] transferring weights from trial "
-                       "{} (score {}) -> {} (score {})".format(
-                           trial_to_clone, new_state.last_score, trial,
-                           trial_state.last_score))
+        logger.info("[exploit] transferring weights from trial "
+                    "{} (score {}) -> {} (score {})".format(
+                        trial_to_clone, new_state.last_score, trial,
+                        trial_state.last_score))
         # TODO(ekl) restarting the trial is expensive. We should implement a
         # lighter way reset() method that can alter the trial config.
         new_tag = make_experiment_tag(trial_state.orig_tag, new_config,
                                       self._hyperparam_mutations)
         reset_successful = trial_executor.reset_trial(trial, new_config,
                                                       new_tag)
-        if not reset_successful:
+        if reset_successful:
+            trial_executor.restore(
+                trial, Checkpoint.from_object(new_state.last_checkpoint))
+        else:
+            logger.warning(
+                "Your Trainable does not support reset_config(); falling "
+                "back to slower path of stopping and restarting trial.")
             trial_executor.stop_trial(trial, stop_logger=False)
             trial.config = new_config
             trial.experiment_tag = new_tag

python/ray/tune/util.py

@@ -2,12 +2,16 @@
 from __future__ import division
 from __future__ import print_function
 
+import logging
 import base64
 import copy
 import numpy as np
+import time
 
 import ray
 
+logger = logging.getLogger(__name__)
+
 _pinned_objects = []
 PINNED_OBJECT_PREFIX = "ray.tune.PinnedObject:"
 
@@ -36,6 +40,28 @@ def get_pinned_object(pinned_id):
             ObjectID(base64.b64decode(pinned_id[len(PINNED_OBJECT_PREFIX):]))))
 
 
+class warn_if_slow(object):
+    """Prints a warning if a given operation is slower than 100ms.
+
+    Example:
+        >>> with expect_fast("some_operation"):
+        ...    ray.get(something)
+    """
+
+    def __init__(self, name):
+        self.name = name
+
+    def __enter__(self):
+        self.start = time.time()
+
+    def __exit__(self, type, value, traceback):
+        now = time.time()
+        if now - self.start > 0.1:
+            logger.warning("The `{}` operation took {} seconds to complete, ".
+                           format(self.name, now - self.start) +
+                           "which may be a performance bottleneck.")
+
+
 def merge_dicts(d1, d2):
     """Returns a new dict that is d1 and d2 deep merged."""
     merged = copy.deepcopy(d1)