[tune] Asynchronous saves (#6912)

* Support asynchronous saves

* Fix merge issues

* Add test, fix existing tests

* More informative warning

* Lint, remove print statements

* Address comments, add checkpoint.is_resolved fn

* Add more detailed comments

python/ray/tune/ray_trial_executor.py

@@ -554,7 +554,7 @@ class RayTrialExecutor(TrialExecutor):
         self._update_avail_resources()
 
     def save(self, trial, storage=Checkpoint.PERSISTENT, result=None):
-        """Saves the trial's state to a checkpoint.
+        """Saves the trial's state to a checkpoint asynchronously.
 
         Args:
             trial (Trial): The trial to be saved.
@@ -567,29 +567,16 @@ class RayTrialExecutor(TrialExecutor):
              Checkpoint object, or None if an Exception occurs.
         """
         result = result or trial.last_result
-
         with self._change_working_directory(trial):
             if storage == Checkpoint.MEMORY:
                 value = trial.runner.save_to_object.remote()
                 checkpoint = Checkpoint(storage, value, result)
-            else:
-                with warn_if_slow("save_checkpoint_to_storage"):
-                    # TODO(ujvl): Make this asynchronous.
-                    value = ray.get(trial.runner.save.remote())
-                    checkpoint = Checkpoint(storage, value, result)
-        with warn_if_slow("on_checkpoint", DEFAULT_GET_TIMEOUT) as profile:
-            try:
                 trial.on_checkpoint(checkpoint)
-            except Exception:
-                logger.exception("Trial %s: Error handling checkpoint %s",
-                                 trial, checkpoint.value)
-                return None
-        if profile.too_slow and trial.sync_on_checkpoint:
-            logger.warning(
-                "Consider turning off forced head-worker trial checkpoint "
-                "syncs by setting sync_on_checkpoint=False. Note that this "
-                "might result in faulty trial restoration for some worker "
-                "failure modes.")
+            else:
+                value = trial.runner.save.remote()
+                checkpoint = Checkpoint(storage, value, result)
+                trial.saving_to = checkpoint
+                self._running[value] = trial
         return checkpoint
 
     def restore(self, trial, checkpoint=None):