chore: import upstream snapshot with attribution

python/ray/train/v2/api/validation_config.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol
+
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.train import Checkpoint
+
+
+@PublicAPI(stability="alpha")
+class ValidationFn(Protocol):
+    """Protocol for a function that validates a checkpoint."""
+
+    def __call__(self, checkpoint: "Checkpoint", **kwargs: Any) -> Dict:
+        ...
+
+
+@dataclass
+@PublicAPI(stability="alpha")
+class ValidationTaskConfig:
+    """Configuration for a specific validation task, passed to report().
+
+    Args:
+        fn_kwargs: json-serializable keyword arguments to pass to the validation function.
+            Note that we always pass `checkpoint` as the first argument to the
+            validation function.
+        timeout_s: Timeout in seconds for this validation task.
+            ``None`` is no timeout.
+    """
+
+    fn_kwargs: Optional[Dict[str, Any]] = None
+    timeout_s: Optional[float] = None
+
+    def __post_init__(self):
+        if self.fn_kwargs is None:
+            self.fn_kwargs = {}
+        assert (
+            self.timeout_s is None or self.timeout_s > 0
+        ), f"The `timeout_s` should be None or greater than zero, actual value: {self.timeout_s}"
+
+
+@PublicAPI(stability="alpha")
+class ValidationConfig:
+    """Configuration for validation, passed to the trainer.
+
+    Args:
+        fn: The validation function to run on checkpoints.
+            This function should accept a checkpoint as the first argument
+            and return a dictionary of metrics.
+        task_config: Default configuration for validation tasks.
+            The fn_kwargs in this config can be overridden by
+            ValidationTaskConfig passed to report().
+        ray_remote_kwargs: Keyword arguments to pass to `ray.remote()` for the validation task.
+            This can be used to specify resource requirements, number of retries, etc.
+    """
+
+    def __init__(
+        self,
+        fn: ValidationFn,
+        task_config: Optional[ValidationTaskConfig] = None,
+        ray_remote_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        self.fn = fn
+        if task_config is None:
+            self.task_config = ValidationTaskConfig()
+        else:
+            self.task_config = task_config
+        # TODO: ray_remote_kwargs is not json-serializable because retry_exceptions
+        # can be a list of exception types. If ray core makes ray_remote_kwargs json-serializable
+        # we can move this to ValidationTaskConfig.
+        if ray_remote_kwargs is None:
+            self.ray_remote_kwargs = {}
+        else:
+            self.ray_remote_kwargs = ray_remote_kwargs