Add forced checkpoint logic for PBT

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

python/ray/air/_internal/checkpoint_manager.py

@@ -72,6 +72,8 @@ def __init__(
         self.metrics = flatten_dict(metrics) if metrics else {}
         self.node_ip = node_ip or self.metrics.get(NODE_IP, None)
 
+        self.forced = False
+
         if (
             dir_or_data is not None
             and storage_mode == CheckpointStorage.MEMORY
@@ -82,6 +84,9 @@ def __init__(
                 f"as their data. Got: {dir_or_data}"
             )
 
+    def set_forced_bit(self, val: bool):
+        self.forced = val
+
     def commit(self, path: Optional[Path] = None) -> None:
         """Commit checkpoint to disk, if needed.
 

python/ray/tune/execution/checkpoint_manager.py

@@ -72,8 +72,10 @@ def handle_checkpoint(self, checkpoint: _TrackedCheckpoint):
             )
             self._process_persistent_checkpoint(checkpoint)
 
-    def on_checkpoint(self, checkpoint: _TrackedCheckpoint):
+    def on_checkpoint(self, checkpoint: _TrackedCheckpoint, force: bool = False):
         """Ray Tune's entrypoint"""
+        if force:
+            checkpoint.set_forced_bit(True)
         # Todo (krfricke): Replace with handle_checkpoint.
         self.handle_checkpoint(checkpoint)
 
@@ -100,17 +102,13 @@ def newest_persistent_checkpoint(self):
     @property
     def newest_checkpoint(self):
         """Returns the newest checkpoint (based on training iteration)."""
-        # NOTE: For PBT Debugging purposes only.
-        # TODO(justinvyu): Remove this.
-        # print("[DEBUGGING] newest_memory_checkpoint.id = ", self.newest_memory_checkpoint.id)
-        # print("[DEBUGGING] newest_persistent_checkpoint.id = ", self.newest_persistent_checkpoint.id)
-
-        # newest_checkpoint = max(
-        #     [self.newest_memory_checkpoint, self.newest_persistent_checkpoint],
-        #     key=lambda c: c.id,
-        # )
-        # return newest_checkpoint
-        return self.newest_memory_checkpoint
+        checkpoints = [self.newest_memory_checkpoint, self.newest_persistent_checkpoint]
+        # Always prefer forced checkpoints
+        # If multiple are forced, take the one with the highest checkpoint id
+        checkpoints.sort(key=lambda c: c.id)
+        checkpoints.sort(key=lambda c: c.forced)
+        newest_checkpoint = checkpoints[-1]
+        return newest_checkpoint
 
     @property
     def newest_memory_checkpoint(self):

python/ray/tune/experiment/trial.py

@@ -683,13 +683,13 @@ def clear_checkpoint(self):
         self.restoring_from = None
         self.invalidate_json_state()
 
-    def on_checkpoint(self, checkpoint: _TrackedCheckpoint):
+    def on_checkpoint(self, checkpoint: _TrackedCheckpoint, force: bool = False):
         """Hook for handling checkpoints taken by the Trainable.
 
         Args:
             checkpoint: Checkpoint taken.
         """
-        self.checkpoint_manager.on_checkpoint(checkpoint)
+        self.checkpoint_manager.on_checkpoint(checkpoint, force=force)
         self.invalidate_json_state()
 
     def on_restore(self):

python/ray/tune/schedulers/pbt.py

@@ -692,7 +692,10 @@ def _exploit(
             trial_executor.set_status(trial, Trial.PAUSED)
         trial.set_experiment_tag(new_tag)
         trial.set_config(new_config)
-        trial.on_checkpoint(new_state.last_checkpoint)
+        # Create a shallow copy to keep forced bit modifications isolated, while
+        # avoiding copying the underlying `dir_or_data`
+        exploit_checkpoint = copy.copy(new_state.last_checkpoint)
+        trial.on_checkpoint(exploit_checkpoint, force=True)
 
         self._num_perturbations += 1
         # Transfer over the last perturbation time as well