wehub-resource-sync
58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
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))}
|
|
)
|