wehub-resource-sync
54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
from typing import Any, Dict, List, Optional
|
|
|
|
from ray.train import Checkpoint
|
|
from ray.train.v2._internal.execution.context import TrainRunContext
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
|
|
@DeveloperAPI
|
|
class RayTrainCallback:
|
|
"""Base Ray Train callback interface."""
|
|
|
|
pass
|
|
|
|
|
|
@DeveloperAPI
|
|
class UserCallback(RayTrainCallback):
|
|
"""Callback interface for custom user-defined callbacks to handling events
|
|
during training.
|
|
|
|
This callback is called on the Ray Train controller process, not on the
|
|
worker processes.
|
|
"""
|
|
|
|
def after_report(
|
|
self,
|
|
run_context: TrainRunContext,
|
|
metrics: List[Dict[str, Any]],
|
|
checkpoint: Optional[Checkpoint],
|
|
):
|
|
"""Called after all workers have reported a metric + checkpoint
|
|
via `ray.train.report`.
|
|
|
|
Args:
|
|
run_context: The `TrainRunContext` for the current training run.
|
|
metrics: A list of metric dictionaries reported by workers,
|
|
where metrics[i] is the metrics dict reported by worker i.
|
|
checkpoint: A Checkpoint object that has been persisted to
|
|
storage. This is None if no workers reported a checkpoint
|
|
(e.g. `ray.train.report(metrics, checkpoint=None)`).
|
|
"""
|
|
pass
|
|
|
|
def after_exception(
|
|
self, run_context: TrainRunContext, worker_exceptions: Dict[int, Exception]
|
|
):
|
|
"""Called after one or more workers have raised an exception.
|
|
|
|
Args:
|
|
run_context: The `TrainRunContext` for the current training run.
|
|
worker_exceptions: A dict from worker world rank to the exception
|
|
raised by that worker.
|
|
"""
|
|
pass
|