wehub-resource-sync
48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
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
|