[data] Make checkpoint restore pluggable via CheckpointConfig checkpoint_filter_cls / checkpoint_manager_cls

python/ray/data/_internal/planner/checkpoint/plan_read_op.py

@@ -59,7 +59,8 @@ def create_checkpoint_filter_op(
     if info.type == fs.FileType.NotFound:
         return physical_input_op
 
-    checkpoint_manager = IdColumnCheckpointManager(
+    manager_cls = checkpoint_config.checkpoint_manager_cls or IdColumnCheckpointManager
+    checkpoint_manager = manager_cls(
         checkpoint_config=checkpoint_config,
         data_context=data_context,
     )
@@ -129,7 +130,10 @@ def __init__(
 
     def init_checkpoint_filter(self):
         """Called once per actor worker to materialize the filter."""
-        self._filter = NumpyArrayBasedCheckpointFilter(self._config, self._ref)
+        filter_cls = (
+            self._config.checkpoint_filter_cls or NumpyArrayBasedCheckpointFilter
+        )
+        self._filter = filter_cls(self._config, self._ref)
 
     def __call__(self, blocks: Iterable[Block], ctx: TaskContext) -> Iterable[Block]:
         assert self._filter is not None, "checkpoint filter was not initialized!"

python/ray/data/checkpoint/__init__.py

@@ -1,3 +1,27 @@
 from .interfaces import CheckpointBackend, CheckpointConfig
 
-__all__ = ["CheckpointConfig", "CheckpointBackend"]
+__all__ = [
+    "CheckpointConfig",
+    "CheckpointBackend",
+    "CheckpointFilter",
+    "CheckpointManager",
+    "IdColumnCheckpointManager",
+    "NumpyArrayBasedCheckpointFilter",
+]
+
+_LAZY_EXPORTS = (
+    "CheckpointFilter",
+    "CheckpointManager",
+    "IdColumnCheckpointManager",
+    "NumpyArrayBasedCheckpointFilter",
+)
+
+
+def __getattr__(name: str) -> type:
+    # Lazily import filter/manager classes to avoid a circular import:
+    # checkpoint_filter -> ray.data.context -> this package.
+    if name in _LAZY_EXPORTS:
+        from ray.data.checkpoint import checkpoint_filter
+
+        return getattr(checkpoint_filter, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

python/ray/data/checkpoint/checkpoint_filter.py

@@ -22,6 +22,7 @@
 from ray.data.context import DataContext
 from ray.data.datasource.path_util import _unwrap_protocol
 from ray.types import ObjectRef
+from ray.util.annotations import DeveloperAPI
 
 logger = logging.getLogger(__name__)
 
@@ -166,8 +167,19 @@ def convert_and_sort_checkpointed_ids(
     return checkpointed_ids_ndarray, checkpoint_size
 
 
+@DeveloperAPI
 class CheckpointManager(abc.ABC):
-    """Manage checkpoint data."""
+    """Manage checkpoint data.
+
+    Subclasses passed as ``CheckpointConfig.checkpoint_manager_cls`` must have
+    a constructor accepting ``(checkpoint_config=..., data_context=...)``
+    keyword arguments, and their ``load_checkpoint`` must return an
+    ``(ObjectRef, int)`` tuple: the ref is passed opaquely to the configured
+    ``CheckpointFilter`` class's constructor, and the int (size in bytes) is
+    used for the per-actor memory reservation of the filter actors. Returning
+    ``(None, 0)`` means there is no checkpoint data to restore from, and the
+    checkpoint filter operator is not added to the plan.
+    """
 
     def __init__(
         self,
@@ -365,13 +377,21 @@ def _validate_loaded_checkpoint(
         pass
 
 
+@DeveloperAPI
 class IdColumnCheckpointManager(CheckpointManager):
     """Manager for regular ID columns."""
 
 
+@DeveloperAPI
 class CheckpointFilter(abc.ABC):
     """Abstract class which defines the interface for filtering checkpointed rows
     based on varying backends.
+
+    Subclasses passed as ``CheckpointConfig.checkpoint_filter_cls`` must have a
+    constructor accepting ``(checkpoint_config, checkpointed_ids_ref)``, where
+    ``checkpointed_ids_ref`` is an ``ObjectRef`` to the sorted NumPy array of
+    checkpointed IDs. The class is instantiated once per checkpoint filter
+    actor on a remote worker, so it must be serializable (or importable) there.
     """
 
     def __init__(self, config: CheckpointConfig):
@@ -391,6 +411,7 @@ def filter_rows_for_block(self, block: Block) -> Block:
         raise NotImplementedError
 
 
+@DeveloperAPI
 class NumpyArrayBasedCheckpointFilter(CheckpointFilter):
     """CheckpointFilter for batch-based backends.
 

python/ray/data/checkpoint/interfaces.py

@@ -1,13 +1,18 @@
+import inspect
 import os
 import warnings
 from enum import Enum
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, Optional, Tuple, Type
 
 import pyarrow
 
 from ray.util.annotations import DeveloperAPI, PublicAPI
 
 if TYPE_CHECKING:
+    from ray.data.checkpoint.checkpoint_filter import (
+        CheckpointFilter,
+        CheckpointManager,
+    )
     from ray.data.datasource import PathPartitionFilter
 
 
@@ -42,6 +47,32 @@ class CheckpointBackend(Enum):
     """
 
 
+def _validate_checkpoint_cls(cls: type, base_cls: type, param_name: str) -> None:
+    """Validate that ``cls`` is a concrete subclass of ``base_cls``.
+
+    Args:
+        cls: The user-provided class to validate.
+        base_cls: The required base class.
+        param_name: The ``CheckpointConfig`` parameter name, for error messages.
+
+    Raises:
+        InvalidCheckpointingConfig: if ``cls`` is not a subclass of
+            ``base_cls``, or is abstract (so instantiation would fail later
+            inside a remote worker instead of at config construction).
+    """
+    if not (isinstance(cls, type) and issubclass(cls, base_cls)):
+        raise InvalidCheckpointingConfig(
+            f"`{param_name}` must be a subclass of `{base_cls.__name__}`, "
+            f"but got {cls}"
+        )
+    if inspect.isabstract(cls):
+        raise InvalidCheckpointingConfig(
+            f"`{param_name}` must be a concrete class, but {cls} is abstract "
+            "(it does not implement all abstract methods of "
+            f"`{base_cls.__name__}`)"
+        )
+
+
 @PublicAPI(stability="beta")
 class CheckpointConfig:
     """Configuration for checkpointing.
@@ -66,6 +97,24 @@ class CheckpointConfig:
             completed rows.
         checkpoint_path_partition_filter: Filter for checkpoint files to load during
             restoration when reading from `checkpoint_path`.
+        checkpoint_filter_cls: Override the :class:`~ray.data.checkpoint.CheckpointFilter`
+            subclass used to filter out already-checkpointed rows during
+            restoration. The class is instantiated once per checkpoint filter
+            actor with ``(checkpoint_config, checkpointed_ids_ref)``, where
+            ``checkpointed_ids_ref`` is the ``ObjectRef`` returned by the
+            checkpoint manager's ``load_checkpoint`` (by default, a sorted
+            NumPy array of checkpointed IDs). Defaults to
+            :class:`~ray.data.checkpoint.NumpyArrayBasedCheckpointFilter`.
+        checkpoint_manager_cls: Override the
+            :class:`~ray.data.checkpoint.CheckpointManager` subclass used to
+            load checkpoint data during restoration. The class is instantiated
+            on the driver with ``(checkpoint_config=..., data_context=...)``
+            and its ``load_checkpoint`` must return an ``(ObjectRef, int)``
+            tuple: the ref is passed opaquely to ``checkpoint_filter_cls``,
+            and the int (size in bytes) feeds the per-actor memory
+            reservation. Typically customized together with
+            ``checkpoint_filter_cls``. Defaults to
+            :class:`~ray.data.checkpoint.IdColumnCheckpointManager`.
     """
 
     DEFAULT_CHECKPOINT_PATH_BUCKET_ENV_VAR = "RAY_DATA_CHECKPOINT_PATH_BUCKET"
@@ -84,6 +133,8 @@ def __init__(
         override_backend: Optional[CheckpointBackend] = None,
         write_num_threads: int = 3,
         checkpoint_path_partition_filter: Optional["PathPartitionFilter"] = None,
+        checkpoint_filter_cls: Optional[Type["CheckpointFilter"]] = None,
+        checkpoint_manager_cls: Optional[Type["CheckpointManager"]] = None,
     ):
         self.id_column: Optional[str] = id_column
 
@@ -93,6 +144,20 @@ def __init__(
                 f"but got {self.id_column}"
             )
 
+        if checkpoint_filter_cls is not None:
+            from ray.data.checkpoint.checkpoint_filter import CheckpointFilter
+
+            _validate_checkpoint_cls(
+                checkpoint_filter_cls, CheckpointFilter, "checkpoint_filter_cls"
+            )
+
+        if checkpoint_manager_cls is not None:
+            from ray.data.checkpoint.checkpoint_filter import CheckpointManager
+
+            _validate_checkpoint_cls(
+                checkpoint_manager_cls, CheckpointManager, "checkpoint_manager_cls"
+            )
+
         if override_backend is not None:
             warnings.warn(
                 "`override_backend` is deprecated and will be removed in August 2025.",
@@ -113,6 +178,8 @@ def __init__(
         self.delete_checkpoint_on_success: bool = delete_checkpoint_on_success
         self.write_num_threads: int = write_num_threads
         self.checkpoint_path_partition_filter = checkpoint_path_partition_filter
+        self.checkpoint_filter_cls = checkpoint_filter_cls
+        self.checkpoint_manager_cls = checkpoint_manager_cls
         self.checkpoint_actor_pool_min_size = self.CHECKPOINT_ACTOR_POOL_MIN_SIZE
         self.checkpoint_actor_pool_max_size = self.CHECKPOINT_ACTOR_POOL_MAX_SIZE
         self.checkpoint_actor_memory_bytes = self.CHECKPOINT_ACTOR_MEMORY_BYTES

python/ray/data/tests/test_checkpoint.py

@@ -278,6 +278,53 @@ def test_skip_inference_with_overrides(self):
         assert config.filesystem is fs
         assert config.backend is CheckpointBackend.CLOUD_OBJECT_STORAGE
 
+    @pytest.mark.parametrize("filter_cls", [1, "not_a_class", int])
+    def test_invalid_checkpoint_filter_cls(self, filter_cls, local_path):
+        with pytest.raises(
+            InvalidCheckpointingConfig,
+            match="`checkpoint_filter_cls` must be a subclass of `CheckpointFilter`",
+        ):
+            CheckpointConfig(ID_COL, local_path, checkpoint_filter_cls=filter_cls)
+
+    def test_valid_checkpoint_filter_cls(self, local_path):
+        assert CheckpointConfig(ID_COL, local_path).checkpoint_filter_cls is None
+
+        config = CheckpointConfig(
+            ID_COL, local_path, checkpoint_filter_cls=NumpyArrayBasedCheckpointFilter
+        )
+        assert config.checkpoint_filter_cls is NumpyArrayBasedCheckpointFilter
+
+    @pytest.mark.parametrize("manager_cls", [1, "not_a_class", int])
+    def test_invalid_checkpoint_manager_cls(self, manager_cls, local_path):
+        with pytest.raises(
+            InvalidCheckpointingConfig,
+            match="`checkpoint_manager_cls` must be a subclass of `CheckpointManager`",
+        ):
+            CheckpointConfig(ID_COL, local_path, checkpoint_manager_cls=manager_cls)
+
+    def test_valid_checkpoint_manager_cls(self, local_path):
+        assert CheckpointConfig(ID_COL, local_path).checkpoint_manager_cls is None
+
+        config = CheckpointConfig(
+            ID_COL, local_path, checkpoint_manager_cls=IdColumnCheckpointManager
+        )
+        assert config.checkpoint_manager_cls is IdColumnCheckpointManager
+
+    def test_abstract_checkpoint_filter_cls(self, local_path):
+        """The abstract base class (or a still-abstract subclass) is rejected
+        at config construction instead of failing inside a filter actor."""
+        from ray.data.checkpoint.checkpoint_filter import CheckpointFilter
+
+        class StillAbstractFilter(CheckpointFilter):
+            pass
+
+        for cls in [CheckpointFilter, StillAbstractFilter]:
+            with pytest.raises(
+                InvalidCheckpointingConfig,
+                match="`checkpoint_filter_cls` must be a concrete class",
+            ):
+                CheckpointConfig(ID_COL, local_path, checkpoint_filter_cls=cls)
+
 
 @pytest.mark.parametrize(
     "backend,fs,data_path",
@@ -378,6 +425,87 @@ def __call__(self, batch):
     )
 
 
+def test_custom_checkpoint_filter_cls(
+    ray_start_10_cpus_shared, generate_sample_data_csv, tmp_path
+):
+    """A custom `checkpoint_filter_cls` replaces the default filter during restore."""
+
+    class NoOpCheckpointFilter(NumpyArrayBasedCheckpointFilter):
+        def filter_rows_for_block(self, block):
+            # Keep every row, including already-checkpointed ones.
+            return block
+
+    ctx = ray.data.DataContext.get_current()
+    ckpt_path = os.path.join(tmp_path, "ckpt")
+    ctx.checkpoint_config = CheckpointConfig(
+        id_column=ID_COL,
+        checkpoint_path=ckpt_path,
+        checkpoint_filter_cls=NoOpCheckpointFilter,
+    )
+
+    csv_file = generate_sample_data_csv()
+
+    # Pre-populate the checkpoint dir. The default filter would drop these IDs
+    # on restore; the no-op filter must keep them.
+    checkpointed_ids = list(range(SAMPLE_DATA_NUM_ROWS // 2))
+    os.makedirs(ckpt_path, exist_ok=True)
+    pq.write_table(
+        pa.table({ID_COL: checkpointed_ids}),
+        os.path.join(ckpt_path, "pre_checkpoint.parquet"),
+    )
+
+    output_path = os.path.join(tmp_path, "output")
+    ds = ray.data.read_csv(csv_file)
+    ds.write_parquet(output_path)
+
+    # Disable checkpointing before reading back to avoid filtering.
+    ctx.checkpoint_config = None
+    ds_readback = ray.data.read_parquet(output_path)
+    actual_output = sorted([row[ID_COL] for row in ds_readback.iter_rows()])
+    assert actual_output == list(range(SAMPLE_DATA_NUM_ROWS))
+
+
+def test_custom_checkpoint_manager_cls(
+    ray_start_10_cpus_shared, generate_sample_data_csv, tmp_path
+):
+    """A custom `checkpoint_manager_cls` replaces the default manager during restore."""
+
+    class EmptyCheckpointManager(IdColumnCheckpointManager):
+        def load_checkpoint(self, data_file_dir=None, data_file_filesystem=None):
+            # Report no checkpoint data, so no filter operator is added.
+            return None, 0
+
+    ctx = ray.data.DataContext.get_current()
+    ckpt_path = os.path.join(tmp_path, "ckpt")
+    ctx.checkpoint_config = CheckpointConfig(
+        id_column=ID_COL,
+        checkpoint_path=ckpt_path,
+        checkpoint_manager_cls=EmptyCheckpointManager,
+    )
+
+    csv_file = generate_sample_data_csv()
+
+    # Pre-populate the checkpoint dir. The default manager would load these
+    # IDs and filter them out on restore; the custom manager reports no
+    # checkpoint data, so every row must be written.
+    checkpointed_ids = list(range(SAMPLE_DATA_NUM_ROWS // 2))
+    os.makedirs(ckpt_path, exist_ok=True)
+    pq.write_table(
+        pa.table({ID_COL: checkpointed_ids}),
+        os.path.join(ckpt_path, "pre_checkpoint.parquet"),
+    )
+
+    output_path = os.path.join(tmp_path, "output")
+    ds = ray.data.read_csv(csv_file)
+    ds.write_parquet(output_path)
+
+    # Disable checkpointing before reading back to avoid filtering.
+    ctx.checkpoint_config = None
+    ds_readback = ray.data.read_parquet(output_path)
+    actual_output = sorted([row[ID_COL] for row in ds_readback.iter_rows()])
+    assert actual_output == list(range(SAMPLE_DATA_NUM_ROWS))
+
+
 @pytest.mark.parametrize(
     "backend,fs,data_path",
     [