Overwrite trial restore path

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

python/ray/tune/execution/tune_controller.py

@@ -2079,7 +2079,7 @@ def _schedule_trial_restore(self, trial: Trial) -> bool:
             )
             return True
 
-        checkpoint = trial.checkpoint
+        checkpoint = trial.temporary_state.next_restore or trial.checkpoint
 
         if checkpoint.dir_or_data is None:
             logger.debug(f"Not restoring trial {trial}: No checkpoint found.")

python/ray/tune/experiment/trial.py

@@ -212,6 +212,7 @@ def __init__(self):
 
         self.saving_to = None
         self.restoring_from = None
+        self.next_restore = None
 
         self.num_restore_failures = 0
 
@@ -1091,6 +1092,9 @@ def on_checkpoint(self, checkpoint: Union[_TrackedCheckpoint, _TrainingResult]):
             # This index will get restored when the trial is restored and will
             # be passed to the Trainable as the starting checkpoint index.
             self.storage.current_checkpoint_index += 1
+            # Remove any next restore overrides - instead we should now restore
+            # from trial.checkpoint
+            self.temporary_state.next_restore = None
         else:
             self.run_metadata.checkpoint_manager.on_checkpoint(checkpoint)
         self.run_metadata.invalidate_cache()

python/ray/tune/schedulers/pbt.py

@@ -943,7 +943,8 @@ def _exploit(
         else:
             exploit = checkpoint_to_exploit
 
-        trial.on_checkpoint(exploit)
+        # Next trial restore should use this checkpoint
+        trial.temporary_state.next_restore = exploit
 
         self._num_perturbations += 1
         # Transfer over the last perturbation time as well

python/ray/tune/tests/test_trial_scheduler_pbt.py

@@ -177,7 +177,7 @@ def get_virt_mem(cls):
         )
 
         checkpoint_config = CheckpointConfig(
-            num_to_keep=2,
+            num_to_keep=3,
             checkpoint_frequency=2,
         )