undo pause logic

Signed-off-by: Kai Fricke <kai@anyscale.com>

python/ray/tune/BUILD

@@ -359,31 +359,31 @@ py_test(
     size = "large",
     srcs = ["tests/test_trial_scheduler.py"],
     deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "medium_instance", "no_new_storage"],
+    tags = ["team:ml", "exclusive", "medium_instance"],
 )
 
 py_test(
     name = "test_trial_scheduler_pbt",
     size = "large",
     srcs = ["tests/test_trial_scheduler_pbt.py"],
     deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "medium_instance", "no_new_storage"],
+    tags = ["team:ml", "exclusive", "medium_instance"],
 )
 
 py_test(
     name = "test_trial_scheduler_resource_changing",
     size = "small",
     srcs = ["tests/test_trial_scheduler_resource_changing.py"],
     deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "no_new_storage"],
+    tags = ["team:ml", "exclusive"],
 )
 
 py_test(
     name = "test_tune_restore_warm_start",
     size = "large",
     srcs = ["tests/test_tune_restore_warm_start.py"],
     deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "no_new_storage"],
+    tags = ["team:ml", "exclusive"],
 )
 
 py_test(
@@ -565,7 +565,7 @@ py_test(
     size = "medium",
     srcs = ["examples/bohb_example.py"],
     deps = [":tune_lib"],
-    tags = ["team:ml", "exclusive", "example", "no_new_storage"]
+    tags = ["team:ml", "exclusive", "example"]
 )
 
 py_test(

python/ray/tune/execution/tune_controller.py

@@ -1518,19 +1518,18 @@ def _process_trial_failure(
             exception: Exception prior to invoking this method.
         """
         self._has_errored = True
-        if trial.status == Trial.RUNNING:
-            if trial.should_recover():
-                self._try_recover(trial, exc=exception)
-                self._callbacks.on_trial_recover(
-                    iteration=self._iteration, trials=self._trials, trial=trial
-                )
-            else:
-                self._scheduler_alg.on_trial_error(self, trial)
-                self._search_alg.on_trial_complete(trial.trial_id, error=True)
-                self._schedule_trial_stop(trial, exception=exception)
-                self._callbacks.on_trial_error(
-                    iteration=self._iteration, trials=self._trials, trial=trial
-                )
+        if trial.status == Trial.RUNNING and trial.should_recover():
+            self._try_recover(trial, exc=exception)
+            self._callbacks.on_trial_recover(
+                iteration=self._iteration, trials=self._trials, trial=trial
+            )
+        elif trial.status in {Trial.RUNNING, Trial.PENDING}:
+            self._scheduler_alg.on_trial_error(self, trial)
+            self._search_alg.on_trial_complete(trial.trial_id, error=True)
+            self._schedule_trial_stop(trial, exception=exception)
+            self._callbacks.on_trial_error(
+                iteration=self._iteration, trials=self._trials, trial=trial
+            )
 
     def _schedule_trial_stop(self, trial: Trial, exception: Optional[Exception] = None):
         if trial.status == Trial.ERROR:
@@ -1619,38 +1618,13 @@ def _schedule_graceful_trial_stop(self, trial: Trial):
             self._schedule_trial_stop(trial)
 
     def _schedule_trial_pause(self, trial: Trial, should_checkpoint: bool = True):
-        if _use_storage_context():
-            if trial not in self._trial_to_actor:
-                logger.debug(
-                    f"Trial PAUSE requested for trial {trial} but trial is already "
-                    f"stopping. Ignoring."
-                )
-                return
-
-            if should_checkpoint:
-                # We need to wait for the save to finish before stopping the trial.
-                def stop_after_save_result(*args, **kwargs):
-                    self._on_saving_result(*args, **kwargs)
-                    self._schedule_trial_stop(trial)
-                    self._set_trial_status(trial, Trial.PAUSED)
-
-                # NOTE: Ensure that the trial is PAUSED while it's saving a checkpoint.
-                self._set_trial_status(trial, Trial.PAUSED)
-                self._schedule_trial_task(
-                    trial=trial,
-                    method_name="save",
-                    on_result=stop_after_save_result,
-                    on_error=self._trial_task_failure,
-                )
-                trial.temporary_state.saving_to = True
-            else:
-                self._schedule_trial_stop(trial)
-                self._set_trial_status(trial, Trial.PAUSED)
-
-            return
-
         if should_checkpoint:
-            self._schedule_trial_save(trial, storage=CheckpointStorage.MEMORY)
+            self._schedule_trial_save(
+                trial,
+                storage=CheckpointStorage.PERSISTENT
+                if _use_storage_context()
+                else CheckpointStorage.MEMORY,
+            )
         self._schedule_trial_stop(trial)
         self._set_trial_status(trial, Trial.PAUSED)
 
@@ -1981,6 +1955,9 @@ def _process_trial_save(
 
         try:
             if _use_storage_context() and isinstance(checkpoint_value, _TrainingResult):
+                if not checkpoint_value.checkpoint:
+                    logger.debug(f"Got empty checkpoint for trial {trial}")
+                    return
                 try:
                     self._callbacks.on_checkpoint(
                         iteration=self._iteration,

python/ray/tune/schedulers/pbt.py

@@ -635,7 +635,7 @@ def _checkpoint_or_exploit(
                 state.last_checkpoint = trial.checkpoint
             else:
                 state.last_checkpoint = tune_controller._schedule_trial_save(
-                    trial, CheckpointStorage.MEMORY, result=state.last_result
+                    trial, CheckpointStorage.PERSISTENT, result=state.last_result
                 )
             self._num_checkpoints += 1
         else:
@@ -1089,7 +1089,7 @@ def on_trial_result(
         )
 
         checkpoint = tune_controller._schedule_trial_save(
-            trial, CheckpointStorage.MEMORY, result=result
+            trial, CheckpointStorage.PERSISTENT, result=result
         )
 
         new_tag = _make_experiment_tag(self.experiment_tag, new_config, new_config)

python/ray/tune/trainable/function_trainable.py

@@ -551,7 +551,6 @@ def save_checkpoint(self, checkpoint_dir: str = ""):
         if _use_storage_context():
             # TRAIN -> SAVE remote calls get processed sequentially,
             # so `_last_training_result.checkpoint` holds onto the latest ckpt.
-            assert self._last_training_result.checkpoint
             return self._last_training_result
 
         checkpoint = self._status_reporter.get_checkpoint()