wehub-resource-sync
75 lines
2.6 KiB
Python
75 lines
2.6 KiB
Python
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
|