chore: import upstream snapshot with attribution

python/ray/data/checkpoint/load_checkpoint_callback.py

@@ -0,0 +1,47 @@
+import logging
+
+from ray.data._internal.execution.execution_callback import (
+    ExecutionCallback,
+)
+from ray.data._internal.execution.streaming_executor import StreamingExecutor
+from ray.data.checkpoint import CheckpointConfig
+from ray.data.datasource.path_util import _unwrap_protocol
+
+logger = logging.getLogger(__name__)
+
+
+class LoadCheckpointCallback(ExecutionCallback):
+    """
+    ExecutionCallback that handles checkpoints. This Callback is responsible for
+    deleting the checkpoint directory when these conditions are met:
+    1. `delete_checkpoint_on_success` is True.
+    2. The job finishes successfully.
+    """
+
+    def __init__(
+        self,
+        config: CheckpointConfig,
+    ):
+        assert config is not None
+        self._config = config
+
+    def before_execution_starts(self, executor: StreamingExecutor):
+        assert self._config is executor._data_context.checkpoint_config
+
+    def _delete_checkpoint(self):
+        checkpoint_path_unwrapped = _unwrap_protocol(self._config.checkpoint_path)
+        filesystem = self._config.filesystem
+        filesystem.delete_dir(checkpoint_path_unwrapped)
+
+    def after_execution_succeeds(self, executor: StreamingExecutor):
+        assert self._config is executor._data_context.checkpoint_config
+
+        # Delete checkpoint data.
+        try:
+            if self._config.delete_checkpoint_on_success:
+                self._delete_checkpoint()
+        except Exception:
+            logger.warning("Failed to delete checkpoint data.", exc_info=True)
+
+    def after_execution_fails(self, executor: StreamingExecutor, error: Exception):
+        assert self._config is executor._data_context.checkpoint_config