[tune] Avoid scheduler blocking, add reuse_actors optimization (#4218)

python/ray/tune/ray_trial_executor.py

@@ -10,10 +10,11 @@ import time
 import traceback
 
 import ray
-from ray.tune.error import TuneError
+from ray.tune.error import TuneError, AbortTrialExecution
 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__)
 
@@ -30,7 +31,7 @@ class _LocalWrapper(object):
 class RayTrialExecutor(TrialExecutor):
     """An implemention of TrialExecutor based on Ray."""
 
-    def __init__(self, queue_trials=False):
+    def __init__(self, queue_trials=False, reuse_actors=False):
         super(RayTrialExecutor, self).__init__(queue_trials)
         self._running = {}
         # Since trial resume after paused should not run
@@ -40,21 +41,46 @@ class RayTrialExecutor(TrialExecutor):
         self._avail_resources = Resources(cpu=0, gpu=0)
         self._committed_resources = Resources(cpu=0, gpu=0)
         self._resources_initialized = False
+        self._reuse_actors = reuse_actors
+        self._cached_actor = None
         if ray.is_initialized():
             self._update_avail_resources()
 
-    def _setup_runner(self, trial):
-        cls = ray.remote(
-            num_cpus=trial.resources.cpu,
-            num_gpus=trial.resources.gpu,
-            resources=trial.resources.custom_resources)(
-                trial._get_trainable_cls())
+    def _setup_runner(self, trial, reuse_allowed):
+        if (self._reuse_actors and reuse_allowed
+                and self._cached_actor is not None):
+            logger.debug("Reusing cached runner {} for {}".format(
+                self._cached_actor, trial.trial_id))
+            existing_runner = self._cached_actor
+            self._cached_actor = None
+        else:
+            if self._cached_actor:
+                logger.debug(
+                    "Cannot reuse cached runner {} for new trial".format(
+                        self._cached_actor))
+                self._cached_actor.stop.remote()
+                self._cached_actor.__ray_terminate__.remote()
+                self._cached_actor = None
+            existing_runner = None
+            cls = ray.remote(
+                num_cpus=trial.resources.cpu,
+                num_gpus=trial.resources.gpu,
+                resources=trial.resources.custom_resources)(
+                    trial._get_trainable_cls())
 
         trial.init_logger()
         # We checkpoint metadata here to try mitigating logdir duplication
         self.try_checkpoint_metadata(trial)
         remote_logdir = trial.logdir
 
+        if existing_runner:
+            trial.runner = existing_runner
+            if not self.reset_trial(trial, trial.config, trial.experiment_tag):
+                raise AbortTrialExecution(
+                    "Trial runner reuse requires reset_trial() to be "
+                    "implemented and return True.")
+            return existing_runner
+
         def logger_creator(config):
             # Set the working dir in the remote process, for user file writes
             if not os.path.exists(remote_logdir):
@@ -86,7 +112,10 @@ class RayTrialExecutor(TrialExecutor):
         """
         prior_status = trial.status
         self.set_status(trial, Trial.RUNNING)
-        trial.runner = self._setup_runner(trial)
+        trial.runner = self._setup_runner(
+            trial,
+            reuse_allowed=checkpoint is not None
+            or trial._checkpoint.value is not None)
         if not self.restore(trial, checkpoint):
             if trial.status == Trial.ERROR:
                 raise RuntimeError(
@@ -126,12 +155,18 @@ class RayTrialExecutor(TrialExecutor):
         try:
             trial.write_error_log(error_msg)
             if hasattr(trial, 'runner') and trial.runner:
-                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)
+                if (not error and self._reuse_actors
+                        and self._cached_actor is None):
+                    logger.debug("Reusing actor for {}".format(trial.runner))
+                    self._cached_actor = trial.runner
+                else:
+                    logger.info(
+                        "Destroying actor for trial {}. If your trainable is "
+                        "slow to initialize, consider setting "
+                        "reuse_actors=True to reduce actor creation "
+                        "overheads.".format(trial))
+                    trial.runner.stop.remote()
+                    trial.runner.__ray_terminate__.remote()
         except Exception:
             logger.exception("Error stopping runner for Trial %s", str(trial))
             self.set_status(trial, Trial.ERROR)
@@ -152,11 +187,13 @@ class RayTrialExecutor(TrialExecutor):
         self._commit_resources(trial.resources)
         try:
             self._start_trial(trial, checkpoint)
-        except Exception:
+        except Exception as e:
             logger.exception("Error starting runner for Trial %s", str(trial))
             error_msg = traceback.format_exc()
             time.sleep(2)
             self._stop_trial(trial, error=True, error_msg=error_msg)
+            if isinstance(e, AbortTrialExecution):
+                return  # don't retry fatal Tune errors
             try:
                 # This forces the trial to not start from checkpoint.
                 trial.clear_checkpoint()
@@ -222,7 +259,8 @@ class RayTrialExecutor(TrialExecutor):
         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):
@@ -249,7 +287,8 @@ class RayTrialExecutor(TrialExecutor):
         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])
 
         # For local mode
         if isinstance(result, _LocalWrapper):
@@ -400,7 +439,8 @@ class RayTrialExecutor(TrialExecutor):
         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):
@@ -421,11 +461,12 @@ class RayTrialExecutor(TrialExecutor):
             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: