Repository: ray-project/ray

Issue:
[tune] PBT (Memory Checkpointing) using TF-Saver is broken
Another issue on the same topic is that [TensorFlow examples](https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/tune_mnist_ray_hyperband.py) does not work correctly for any training scheduler which requires checkpointing. The issue is related (as far as my understanding goes) to the fact that after running trainable.restore(...) ray deletes the temporary directory with the TensorFlow checkpoint. The TensorFlow method self.saver.restore(self.sess, path) is just a declaration of restoring process, and it does not restore model in the place where declared. The restoration of the model variables happens when session.run(...) is performed, and this is the moment when variables are loaded from saved checkpoint to the tf.Graph().  Since checkpoint directory is deleted (in trainable.restore(...) ), the actual process fails.

Why ray.tune delete checkpoint directory in trainable.restore(...)?

_Originally posted by @agniszczotka in https://github.com/ray-project/ray/issues/2856#issuecomment-459011993_
[tune] Checkpointing with tensorflow no longer works
<!--
General questions should be asked on the mailing list ray-dev@googlegroups.com.
Questions about how to use Ray should be asked on
[StackOverflow](https://stackoverflow.com/questions/tagged/ray).

Before submitting an issue, please fill out the following form.
-->

### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
- **Ray installed from (source or binary)**: binary
- **Ray version**: 0.6.2
- **Python version**: 3.6.7
- **Exact command to reproduce**:
https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/tune_mnist_ray_hyperband.py


<!--
You can obtain the Ray version with

python -c "import ray; print(ray.__version__)"
-->

### Describe the problem

Saving and restoring only works with single files in the newest version for ray. Tensorflow stores multiple files during checkpointing. Thus, the tensorflow example for storing and saving is broken.

```
    def _save(self, checkpoint_dir):
        return self.saver.save(
            self.sess, checkpoint_dir + "/save", global_step=self.iterations)

    def _restore(self, path):
        return self.saver.restore(self.sess, path)
```
https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/tune_mnist_ray_hyperband.py

### Source code / logs
<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->
ValueError: The returned checkpoint path does not exist: ray_results/2019-02-06_15-56-40yekao5r5/checkpoint_0/save



Discussion:

I believe this commit broke the functionality
https://github.com/ray-project/ray/commit/f9b58d7b0252f312e7b0983ffefabf68d3e7a939#diff-0ba1385a99926375f91eca677328f64c
Replacing the save and restore functions with the old variants fixes the problem
```
class ExtendedTrainable(Trainable, ABC):

    def save(self, checkpoint_dir=None):
        """Saves the current model state to a checkpoint.

        Subclasses should override ``_save()`` instead to save state.
        This method dumps additional metadata alongside the saved path.

        Args:
            checkpoint_dir (str): Optional dir to place the checkpoint.

        Returns:
            Checkpoint path that may be passed to restore().
        """
        checkpoint_path = self._save(checkpoint_dir or self.logdir)
        with open(checkpoint_path + ".tune_metadata", "wb") as file:
            pickle.dump({
                "experiment_id": self._experiment_id,
                "iteration": self._iteration,
                "timesteps_total": self._timesteps_total,
                "time_total": self._time_total,
                "episodes_total": self._episodes_total,
            }, file)

        return checkpoint_path

    def restore(self, checkpoint_path):
        with open(checkpoint_path + ".tune_metadata", "rb") as file:
            metadata = pickle.load(file)
        self._experiment_id = metadata["experiment_id"]
        self._iteration = metadata["iteration"]
        self._timesteps_total = metadata["timesteps_total"]
        self._time_total = metadata["time_total"]
        self._episodes_total = metadata["episodes_total"]
        self._restore(checkpoint_path)
        self._restored = True
```

Resolving patch:
diff --git a/python/ray/tune/examples/tune_mnist_ray_hyperband.py b/python/ray/tune/examples/tune_mnist_ray_hyperband.py
--- a/python/ray/tune/examples/tune_mnist_ray_hyperband.py
+++ b/python/ray/tune/examples/tune_mnist_ray_hyperband.py
@@ -199,11 +199,13 @@ def _train(self):
         return {"mean_accuracy": train_accuracy}
 
     def _save(self, checkpoint_dir):
-        return self.saver.save(
+        prefix = self.saver.save(
             self.sess, checkpoint_dir + "/save", global_step=self.iterations)
+        return {"prefix": prefix}
 
-    def _restore(self, path):
-        return self.saver.restore(self.sess, path)
+    def _restore(self, ckpt_data):
+        prefix = ckpt_data["prefix"]
+        return self.saver.restore(self.sess, prefix)
 
 
 # !!! Example of using the ray.tune Python API !!!
@@ -229,7 +231,7 @@ def _restore(self, path):
     }
 
     if args.smoke_test:
-        mnist_spec['stop']['training_iteration'] = 2
+        mnist_spec['stop']['training_iteration'] = 20
         mnist_spec['num_samples'] = 2
 
     ray.init()


Test patch: