[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>
2026-08-17 11:25:34 +08:00 · 2020-01-02 20:40:53 -08:00
parent 57061a15cf
commit ca651af1d7
30 changed files with 1006 additions and 349 deletions
@@ -0,0 +1,98 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from ray.tune.trainable import Trainable, TrainableUtil
+from ray.tune.syncer import get_cloud_sync_client
+
+
+class DurableTrainable(Trainable):
+    """Abstract class for a remote-storage backed fault-tolerant Trainable.
+
+    Supports checkpointing to and restoring from remote storage. To use this
+    class, implement the same private methods as ray.tune.Trainable (`_save`,
+    `_train`, `_restore`, `reset_config`, `_setup`, `_stop`).
+
+    .. warning:: This class is currently **experimental** and may
+        be subject to change.
+
+    Run this with Tune as follows. Setting `sync_to_driver=False` disables
+    syncing to the driver to avoid keeping redundant checkpoints around, as
+    well as preventing the driver from syncing up the same checkpoint.
+
+    See ``tune/trainable.py``.
+
+    Attributes:
+        remote_checkpoint_dir (str):  Upload directory (S3 or GS path).
+        storage_client: Tune-internal interface for interacting with external
+            storage.
+
+    >>> tune.run(MyDurableTrainable, sync_to_driver=False)
+    """
+
+    def __init__(self, remote_checkpoint_dir, *args, **kwargs):
+        """Initializes a DurableTrainable.
+
+        Args:
+            remote_checkpoint_dir (str): Upload directory (S3 or GS path).
+        """
+        super(DurableTrainable, self).__init__(*args, **kwargs)
+        self.remote_checkpoint_dir = remote_checkpoint_dir
+        self.storage_client = self._create_storage_client()
+
+    def save(self, checkpoint_dir=None):
+        """Saves the current model state to a checkpoint, persisted remotely.
+
+        The storage client must provide durability for
+        restoration to work. That is, once ``storage.client.wait()``
+        returns after a checkpoint `sync up`, the checkpoint is considered
+        committed and can be used to restore the trainable.
+
+        Args:
+            checkpoint_dir (Optional[str]): Optional dir to place the
+                checkpoint. Must be ``logdir`` or a sub-directory.
+
+        Returns:
+            Checkpoint path or prefix that may be passed to restore().
+        """
+        if checkpoint_dir:
+            if checkpoint_dir.starts_with(os.path.abspath(self.logdir)):
+                raise ValueError("`checkpoint_dir` must be `self.logdir`, or "
+                                 "a sub-directory.")
+
+        checkpoint_path = super(DurableTrainable, self).save(checkpoint_dir)
+        self.storage_client.sync_up(self.logdir, self.remote_checkpoint_dir)
+        self.storage_client.wait()
+        return checkpoint_path
+
+    def restore(self, checkpoint_path):
+        """Restores training state from a given checkpoint persisted remotely.
+
+        These checkpoints are returned from calls to save().
+
+        Args:
+            checkpoint_path (str): Local path to checkpoint.
+        """
+        self.storage_client.sync_down(self.remote_checkpoint_dir, self.logdir)
+        self.storage_client.wait()
+        super(DurableTrainable, self).restore(checkpoint_path)
+
+    def delete_checkpoint(self, checkpoint_path):
+        """Deletes checkpoint from both local and remote storage.
+
+        Args:
+            checkpoint_path (str): Local path to checkpoint.
+        """
+        super(DurableTrainable, self).delete_checkpoint(checkpoint_path)
+        local_dirpath = TrainableUtil.find_checkpoint_dir(checkpoint_path)
+        self.storage_client.delete(self._storage_path(local_dirpath))
+
+    def _create_storage_client(self):
+        """Returns a storage client."""
+        return get_cloud_sync_client(self.remote_checkpoint_dir)
+
+    def _storage_path(self, local_path):
+        rel_local_path = os.path.relpath(local_path, self.logdir)
+        return os.path.join(self.remote_checkpoint_dir, rel_local_path)