chore: import upstream snapshot with attribution

python/ray/train/_internal/framework_checkpoint.py

@@ -0,0 +1,57 @@
+from typing import TYPE_CHECKING, Optional
+
+import ray.cloudpickle as ray_pickle
+from ray._common.utils import binary_to_hex, hex_to_binary
+from ray.train._checkpoint import Checkpoint
+
+if TYPE_CHECKING:
+    from ray.data.preprocessor import Preprocessor
+
+PREPROCESSOR_KEY = "preprocessor_pkl"
+
+
+class FrameworkCheckpoint(Checkpoint):
+    """A checkpoint to preserve the functionality of legacy
+    framework-specific checkpoints.
+
+    Example:
+
+        >>> import tempfile
+        >>> from ray.data.preprocessor import Preprocessor
+        >>> checkpoint = FrameworkCheckpoint(tempfile.mkdtemp())
+        >>> checkpoint.get_preprocessor() is None
+        True
+        >>> preprocessor = Preprocessor()
+        >>> preprocessor._attr = 1234
+        >>> checkpoint.set_preprocessor(preprocessor)
+        >>> checkpoint.get_preprocessor()._attr
+        1234
+    """
+
+    def get_preprocessor(self) -> Optional["Preprocessor"]:
+        """Return the preprocessor stored in the checkpoint.
+
+        .. warning::
+
+            The checkpoint path must point to a **trusted** source.
+            The preprocessor is stored as a pickle blob inside the checkpoint
+            metadata. Loading a checkpoint from an untrusted path (shared
+            storage, downloaded artifact, checkpoint produced by a different
+            party) is equivalent to executing arbitrary Python code. Never
+            call this method on a checkpoint you do not fully control.
+
+        Returns:
+            The preprocessor stored in the checkpoint, or ``None`` if no
+            preprocessor was stored.
+        """
+        metadata = self.get_metadata()
+        preprocessor_bytes = metadata.get(PREPROCESSOR_KEY)
+        if preprocessor_bytes is None:
+            return None
+        return ray_pickle.loads(hex_to_binary(preprocessor_bytes))
+
+    def set_preprocessor(self, preprocessor: "Preprocessor"):
+        """Store a preprocessor with the checkpoint."""
+        self.update_metadata(
+            {PREPROCESSOR_KEY: binary_to_hex(ray_pickle.dumps(preprocessor))}
+        )