[tune] Test example checkpointing (#4728)

python/ray/tune/trainable.py

@@ -238,7 +238,7 @@ class Trainable(object):
             checkpoint_dir (str): Optional dir to place the checkpoint.
 
         Returns:
-            Checkpoint path that may be passed to restore().
+            Checkpoint path or prefix that may be passed to restore().
         """
 
         checkpoint_dir = os.path.join(checkpoint_dir or self.logdir,
@@ -248,16 +248,11 @@ class Trainable(object):
         checkpoint = self._save(checkpoint_dir)
         saved_as_dict = False
         if isinstance(checkpoint, string_types):
-            if (not checkpoint.startswith(checkpoint_dir)
-                    or checkpoint == checkpoint_dir):
+            if not checkpoint.startswith(checkpoint_dir):
                 raise ValueError(
                     "The returned checkpoint path must be within the "
                     "given checkpoint dir {}: {}".format(
                         checkpoint_dir, checkpoint))
-            if not os.path.exists(checkpoint):
-                raise ValueError(
-                    "The returned checkpoint path does not exist: {}".format(
-                        checkpoint))
             checkpoint_path = checkpoint
         elif isinstance(checkpoint, dict):
             saved_as_dict = True
@@ -265,9 +260,9 @@ class Trainable(object):
             with open(checkpoint_path, "wb") as f:
                 pickle.dump(checkpoint, f)
         else:
-            raise ValueError(
-                "`_save` must return a dict or string type: {}".format(
-                    str(type(checkpoint))))
+            raise ValueError("Returned unexpected type {}. "
+                             "Expected str or dict.".format(type(checkpoint)))
+
         with open(checkpoint_path + ".tune_metadata", "wb") as f:
             pickle.dump({
                 "experiment_id": self._experiment_id,
@@ -288,25 +283,25 @@ class Trainable(object):
         """
 
         tmpdir = tempfile.mkdtemp("save_to_object", dir=self.logdir)
-        checkpoint_prefix = self.save(tmpdir)
+        checkpoint_path = self.save(tmpdir)
 
+        # Save all files in subtree.
         data = {}
-        base_dir = os.path.dirname(checkpoint_prefix)
-        for path in os.listdir(base_dir):
-            path = os.path.join(base_dir, path)
-            if path.startswith(checkpoint_prefix):
+        for basedir, _, file_names in os.walk(tmpdir):
+            for file_name in file_names:
+                path = os.path.join(basedir, file_name)
+
                 with open(path, "rb") as f:
-                    data[os.path.basename(path)] = f.read()
+                    data[os.path.relpath(path, tmpdir)] = f.read()
 
         out = io.BytesIO()
         data_dict = pickle.dumps({
-            "checkpoint_name": os.path.basename(checkpoint_prefix),
+            "checkpoint_name": os.path.relpath(checkpoint_path, tmpdir),
             "data": data,
         })
         if len(data_dict) > 10e6:  # getting pretty large
             logger.info("Checkpoint size is {} bytes".format(len(data_dict)))
         out.write(data_dict)
-
         shutil.rmtree(tmpdir)
         return out.getvalue()
 
@@ -318,7 +313,6 @@ class Trainable(object):
         Subclasses should override ``_restore()`` instead to restore state.
         This method restores additional metadata saved with the checkpoint.
         """
-
         with open(checkpoint_path + ".tune_metadata", "rb") as f:
             metadata = pickle.load(f)
         self._experiment_id = metadata["experiment_id"]
@@ -330,6 +324,7 @@ class Trainable(object):
         if saved_as_dict:
             with open(checkpoint_path, "rb") as loaded_state:
                 checkpoint_dict = pickle.load(loaded_state)
+            checkpoint_dict.update(tune_checkpoint_path=checkpoint_path)
             self._restore(checkpoint_dict)
         else:
             self._restore(checkpoint_path)
@@ -343,14 +338,18 @@ class Trainable(object):
 
         These checkpoints are returned from calls to save_to_object().
         """
-
         info = pickle.loads(obj)
         data = info["data"]
         tmpdir = tempfile.mkdtemp("restore_from_object", dir=self.logdir)
         checkpoint_path = os.path.join(tmpdir, info["checkpoint_name"])
 
-        for file_name, file_contents in data.items():
-            with open(os.path.join(tmpdir, file_name), "wb") as f:
+        for relpath_name, file_contents in data.items():
+            path = os.path.join(tmpdir, relpath_name)
+
+            # This may be a subdirectory, hence not just using tmpdir
+            if not os.path.exists(os.path.dirname(path)):
+                os.makedirs(os.path.dirname(path))
+            with open(path, "wb") as f:
                 f.write(file_contents)
 
         self.restore(checkpoint_path)
@@ -412,7 +411,7 @@ class Trainable(object):
 
         Returns:
             checkpoint (str | dict): If string, the return value is
-                expected to be the checkpoint path that will be passed to
+                expected to be the checkpoint path or prefix to be passed to
                 `_restore()`. If dict, the return value will be automatically
                 serialized by Tune and passed to `_restore()`.