[tune] Cluster Fault Tolerance (#3309)

This PR introduces cluster-level fault tolerance for Tune by checkpointing global state. This occurs with relatively high frequency and allows users to easily resume experiments when the cluster crashes.

Note that this PR may affect automated workflows due to auto-prompting, but this is resolvable.

python/ray/tune/trial.py

@@ -3,15 +3,22 @@ from __future__ import division
 from __future__ import print_function
 
 from collections import namedtuple
+import ray.cloudpickle as cloudpickle
+import copy
 from datetime import datetime
 import logging
+import json
 import time
 import tempfile
 import os
+
+# For compatibility under py2 to consider unicode as str
+from six import string_types
 from numbers import Number
 
 import ray
 from ray.tune import TuneError
+from ray.tune.log_sync import validate_sync_function
 from ray.tune.logger import pretty_print, UnifiedLogger
 # NOTE(rkn): We import ray.tune.registry here instead of importing the names we
 # need because there are cyclic imports that may cause specific names to not
@@ -19,7 +26,7 @@ from ray.tune.logger import pretty_print, UnifiedLogger
 import ray.tune.registry
 from ray.tune.result import (DEFAULT_RESULTS_DIR, DONE, HOSTNAME, PID,
                              TIME_TOTAL_S, TRAINING_ITERATION, TIMESTEPS_TOTAL)
-from ray.utils import random_string, binary_to_hex
+from ray.utils import random_string, binary_to_hex, hex_to_binary
 
 DEBUG_PRINT_INTERVAL = 5
 MAX_LEN_IDENTIFIER = 130
@@ -66,6 +73,36 @@ class Resources(
         return self.gpu + self.extra_gpu
 
 
+def json_to_resources(data):
+    if data is None or data == "null":
+        return None
+    if isinstance(data, string_types):
+        data = json.loads(data)
+    for k in data:
+        if k in ["driver_cpu_limit", "driver_gpu_limit"]:
+            raise TuneError(
+                "The field `{}` is no longer supported. Use `extra_cpu` "
+                "or `extra_gpu` instead.".format(k))
+        if k not in Resources._fields:
+            raise TuneError(
+                "Unknown resource type {}, must be one of {}".format(
+                    k, Resources._fields))
+    return Resources(
+        data.get("cpu", 1), data.get("gpu", 0), data.get("extra_cpu", 0),
+        data.get("extra_gpu", 0))
+
+
+def resources_to_json(resources):
+    if resources is None:
+        return None
+    return {
+        "cpu": resources.cpu,
+        "gpu": resources.gpu,
+        "extra_cpu": resources.extra_cpu,
+        "extra_gpu": resources.extra_gpu,
+    }
+
+
 def has_trainable(trainable_name):
     return ray.tune.registry._global_registry.contains(
         ray.tune.registry.TRAINABLE_CLASS, trainable_name)
@@ -133,12 +170,8 @@ class Trial(object):
         The args here take the same meaning as the command line flags defined
         in ray.tune.config_parser.
         """
-        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)
 
+        Trial._registration_check(trainable_name)
         # Trial config
         self.trainable_name = trainable_name
         self.config = config or {}
@@ -149,14 +182,15 @@ class Trial(object):
             or self._get_trainable_cls().default_resource_request(self.config))
         self.stopping_criterion = stopping_criterion or {}
         self.upload_dir = upload_dir
-        self.trial_name_creator = trial_name_creator
         self.custom_loggers = custom_loggers
         self.sync_function = sync_function
+        validate_sync_function(sync_function)
         self.verbose = True
         self.max_failures = max_failures
 
         # Local trial state that is updated during the run
         self.last_result = None
+        self.last_update_time = -float("inf")
         self.checkpoint_freq = checkpoint_freq
         self.checkpoint_at_end = checkpoint_at_end
         self._checkpoint = Checkpoint(
@@ -170,6 +204,18 @@ class Trial(object):
         self.error_file = None
         self.num_failures = 0
 
+        self.trial_name = None
+        if trial_name_creator:
+            self.trial_name = trial_name_creator(self)
+
+    @classmethod
+    def _registration_check(cls, 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)
+
     @classmethod
     def generate_id(cls):
         return binary_to_hex(random_string())[:8]
@@ -180,10 +226,14 @@ class Trial(object):
         if not self.result_logger:
             if not os.path.exists(self.local_dir):
                 os.makedirs(self.local_dir)
-            self.logdir = tempfile.mkdtemp(
-                prefix="{}_{}".format(
-                    str(self)[:MAX_LEN_IDENTIFIER], date_str()),
-                dir=self.local_dir)
+            if not self.logdir:
+                self.logdir = tempfile.mkdtemp(
+                    prefix="{}_{}".format(
+                        str(self)[:MAX_LEN_IDENTIFIER], date_str()),
+                    dir=self.local_dir)
+            elif not os.path.exists(self.logdir):
+                os.makedirs(self.logdir)
+
             self.result_logger = UnifiedLogger(
                 self.config,
                 self.logdir,
@@ -307,6 +357,7 @@ class Trial(object):
                 pretty_print(result).replace("\n", "\n  ")))
             self.last_debug = time.time()
         self.last_result = result
+        self.last_update_time = time.time()
         self.result_logger.on_result(self.last_result)
 
     def _get_trainable_cls(self):
@@ -327,8 +378,8 @@ class Trial(object):
 
         Can be overriden with a custom string creator.
         """
-        if self.trial_name_creator:
-            return self.trial_name_creator(self)
+        if self.trial_name:
+            return self.trial_name
 
         if "env" in self.config:
             env = self.config["env"]
@@ -340,3 +391,48 @@ class Trial(object):
         if self.experiment_tag:
             identifier += "_" + self.experiment_tag
         return identifier.replace("/", "_")
+
+    def __getstate__(self):
+        """Memento generator for Trial.
+
+        Sets RUNNING trials to PENDING, and flushes the result logger.
+        Note this can only occur if the trial holds a DISK checkpoint.
+        """
+        assert self._checkpoint.storage == Checkpoint.DISK, (
+            "Checkpoint must not be in-memory.")
+        state = self.__dict__.copy()
+        state["resources"] = resources_to_json(self.resources)
+
+        pickle_data = {
+            "_checkpoint": self._checkpoint,
+            "config": self.config,
+            "custom_loggers": self.custom_loggers,
+            "sync_function": self.sync_function
+        }
+
+        for key, value in pickle_data.items():
+            state[key] = binary_to_hex(cloudpickle.dumps(value))
+
+        state["runner"] = None
+        state["result_logger"] = None
+        if self.status == Trial.RUNNING:
+            state["status"] = Trial.PENDING
+        if self.result_logger:
+            self.result_logger.flush()
+            state["__logger_started__"] = True
+        else:
+            state["__logger_started__"] = False
+        return copy.deepcopy(state)
+
+    def __setstate__(self, state):
+        logger_started = state.pop("__logger_started__")
+        state["resources"] = json_to_resources(state["resources"])
+        for key in [
+                "_checkpoint", "config", "custom_loggers", "sync_function"
+        ]:
+            state[key] = cloudpickle.loads(hex_to_binary(state[key]))
+
+        self.__dict__.update(state)
+        Trial._registration_check(self.trainable_name)
+        if logger_started:
+            self.init_logger()