[tune] Async restores and S3/GCP-capable trial FT (#6376)

* Initial commit for asynchronous save/restore

* Set stage for cloud checkpointable trainable.

* Refactor log_sync and sync_client.

* Add durable trainable impl.

* Support delete in cmd based client

* Fix some tests and such

* Cleanup, comments.

* Use upload_dir instead.

* Revert files belonging to other PR in split.

* Pass upload_dir into trainable init.

* Pickle checkpoint at driver, more robust checkpoint_dir discovery.

* Cleanup trainable helper functions, fix tests.

* Addressed comments.

* Fix bugs from cluster testing, add parameterized cluster tests.

* Add trainable util test

* package_ref

* pbt_address

* Fix bug after running pbt example (_save returning dir).

* get cluster tests running, other bug fixes.

* raise_errors

* Fix deleter bug, add durable trainable example.

* Fix cluster test bugs.

* filelock

* save/restore bug fixes

* .

* Working cluster tests.

* Lint, revert to tracking memory checkpoints.

* Documentation, cleanup

* fixinitialsync

* fix_one_test

* Fix cluster test bug

* nit

* lint

* Revert tune md change

* Fix basename bug for directories.

* lint

* fix_tests

* nit_fix

* Add __init__ file.

* Move to utils package

Co-authored-by: Richard Liaw <rliaw@berkeley.edu>

python/ray/tune/checkpoint_manager.py

@@ -5,8 +5,6 @@ from __future__ import print_function
 
 import heapq
 import logging
-import os
-import shutil
 
 try:
     FileNotFoundError
@@ -23,31 +21,18 @@ class Checkpoint:
 
     Attributes:
         storage (str): Storage type.
-        value (str): If storage==MEMORY, value is a Python object.
-            If storage==DISK, value is a path points to the checkpoint in disk.
+        value (str): If storage==MEMORY, it is a Python object.
+            If storage==PERSISTENT, it is a path to persistent storage.
     """
 
     MEMORY = "memory"
-    DISK = "disk"
+    PERSISTENT = "persistent"
 
     def __init__(self, storage, value, result=None):
         self.storage = storage
         self.value = value
         self.result = result or {}
 
-    def delete(self):
-        """Deletes checkpoint data if disk checkpoint."""
-        if self.storage == Checkpoint.DISK and self.value:
-            checkpoint_dir = self.value
-            if not os.path.exists(checkpoint_dir):
-                raise FileNotFoundError(
-                    "Attempted to delete checkpoint at {} but "
-                    "path was not found.".format(checkpoint_dir))
-            elif os.path.isfile(checkpoint_dir):
-                shutil.rmtree(os.path.dirname(checkpoint_dir))
-            else:
-                shutil.rmtree(checkpoint_dir)
-
     @staticmethod
     def from_object(value=None):
         """Creates a checkpoint from a Python object."""
@@ -72,13 +57,15 @@ class QueueItem:
 class CheckpointManager:
     """Manages checkpoints on the driver for a trial."""
 
-    def __init__(self, keep_checkpoints_num, checkpoint_score_attr):
+    def __init__(self, keep_checkpoints_num, checkpoint_score_attr, delete_fn):
         """Initializes a new CheckpointManager.
 
         Args:
             keep_checkpoints_num (int): Keep at least this many checkpoints.
             checkpoint_score_attr (str): Attribute to use to determine which
                 checkpoints to keep.
+            delete_fn (function): Function that deletes checkpoints. Must be
+                idempotent.
         """
         self.keep_checkpoints_num = keep_checkpoints_num or float("inf")
         assert self.keep_checkpoints_num > 0, (
@@ -88,7 +75,7 @@ class CheckpointManager:
             self._checkpoint_score_attr = checkpoint_score_attr[4:]
         else:
             self._checkpoint_score_attr = checkpoint_score_attr
-
+        self.delete = delete_fn
         self.newest_checkpoint = Checkpoint(Checkpoint.MEMORY, None)
         self._best_checkpoints = []
         self._membership = set()
@@ -101,9 +88,6 @@ class CheckpointManager:
 
         Args:
             checkpoint (Checkpoint): Trial state checkpoint.
-
-        Raises:
-            KeyError if checkpoint_score_attr not in result of checkpoint.
         """
         old_checkpoint = self.newest_checkpoint
         self.newest_checkpoint = checkpoint
@@ -112,7 +96,7 @@ class CheckpointManager:
             queue_item = QueueItem(self._priority(checkpoint), checkpoint)
         except KeyError:
             if old_checkpoint not in self._membership:
-                old_checkpoint.delete()
+                self.delete(old_checkpoint)
             logger.error("Result dict has no key: {}. "
                          "checkpoint_score_attr must be set to a key in the "
                          "result dict.".format(self._checkpoint_score_attr))
@@ -126,11 +110,11 @@ class CheckpointManager:
             self._membership.add(checkpoint)
             if worst in self._membership:
                 self._membership.remove(worst)
-            worst.delete()
+            self.delete(worst)
 
         # Remove the old checkpoint if it isn't one of the best ones.
-        if old_checkpoint not in self._membership:
-            old_checkpoint.delete()
+        if old_checkpoint.value and old_checkpoint not in self._membership:
+            self.delete(old_checkpoint)
 
     def best_checkpoints(self):
         """Returns best checkpoints, sorted by score."""
@@ -140,3 +124,13 @@ class CheckpointManager:
     def _priority(self, checkpoint):
         priority = checkpoint.result[self._checkpoint_score_attr]
         return -priority if self._checkpoint_score_desc else priority
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        # Avoid serializing lambda since it may capture cyclical dependencies.
+        state.pop("delete")
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self.delete = None