wehub-resource-sync
156 lines
4.7 KiB
Python
156 lines
4.7 KiB
Python
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Dict, List, Optional, Set
|
|
|
|
import pyarrow
|
|
import yaml
|
|
|
|
from ray.air._internal.json import SafeFallbackEncoder
|
|
from ray.tune.callback import Callback
|
|
from ray.util.annotations import DeveloperAPI, PublicAPI
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.tune.experiment.trial import Trial # noqa: F401
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Apply flow style for sequences of this length
|
|
_SEQUENCE_LEN_FLOW_STYLE = 3
|
|
|
|
|
|
@PublicAPI
|
|
class LoggerCallback(Callback):
|
|
"""Base class for experiment-level logger callbacks
|
|
|
|
This base class defines a general interface for logging events,
|
|
like trial starts, restores, ends, checkpoint saves, and receiving
|
|
trial results.
|
|
|
|
Callbacks implementing this interface should make sure that logging
|
|
utilities are cleaned up properly on trial termination, i.e. when
|
|
``log_trial_end`` is received. This includes e.g. closing files.
|
|
"""
|
|
|
|
def log_trial_start(self, trial: "Trial"):
|
|
"""Handle logging when a trial starts.
|
|
|
|
Args:
|
|
trial: Trial object.
|
|
"""
|
|
pass
|
|
|
|
def log_trial_restore(self, trial: "Trial"):
|
|
"""Handle logging when a trial restores.
|
|
|
|
Args:
|
|
trial: Trial object.
|
|
"""
|
|
pass
|
|
|
|
def log_trial_save(self, trial: "Trial"):
|
|
"""Handle logging when a trial saves a checkpoint.
|
|
|
|
Args:
|
|
trial: Trial object.
|
|
"""
|
|
pass
|
|
|
|
def log_trial_result(self, iteration: int, trial: "Trial", result: Dict):
|
|
"""Handle logging when a trial reports a result.
|
|
|
|
Args:
|
|
iteration: Iteration of the experiment that this result belongs to.
|
|
trial: Trial object.
|
|
result: Result dictionary.
|
|
"""
|
|
pass
|
|
|
|
def log_trial_end(self, trial: "Trial", failed: bool = False):
|
|
"""Handle logging when a trial ends.
|
|
|
|
Args:
|
|
trial: Trial object.
|
|
failed: True if the Trial finished gracefully, False if
|
|
it failed (e.g. when it raised an exception).
|
|
"""
|
|
pass
|
|
|
|
def on_trial_result(
|
|
self,
|
|
iteration: int,
|
|
trials: List["Trial"],
|
|
trial: "Trial",
|
|
result: Dict,
|
|
**info,
|
|
):
|
|
self.log_trial_result(iteration, trial, result)
|
|
|
|
def on_trial_start(
|
|
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
|
):
|
|
self.log_trial_start(trial)
|
|
|
|
def on_trial_restore(
|
|
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
|
):
|
|
self.log_trial_restore(trial)
|
|
|
|
def on_trial_save(
|
|
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
|
):
|
|
self.log_trial_save(trial)
|
|
|
|
def on_trial_complete(
|
|
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
|
):
|
|
self.log_trial_end(trial, failed=False)
|
|
|
|
def on_trial_error(
|
|
self, iteration: int, trials: List["Trial"], trial: "Trial", **info
|
|
):
|
|
self.log_trial_end(trial, failed=True)
|
|
|
|
def _restore_from_remote(self, file_name: str, trial: "Trial") -> None:
|
|
if not trial.checkpoint:
|
|
# If there's no checkpoint, there's no logging artifacts to restore
|
|
# since we're starting from scratch.
|
|
return
|
|
|
|
local_file = Path(trial.local_path, file_name).as_posix()
|
|
remote_file = Path(trial.storage.trial_fs_path, file_name).as_posix()
|
|
|
|
try:
|
|
pyarrow.fs.copy_files(
|
|
remote_file,
|
|
local_file,
|
|
source_filesystem=trial.storage.storage_filesystem,
|
|
)
|
|
logger.debug(f"Copied {remote_file} to {local_file}")
|
|
except FileNotFoundError:
|
|
logger.warning(f"Remote file not found: {remote_file}")
|
|
except Exception:
|
|
logger.exception(f"Error downloading {remote_file}")
|
|
|
|
|
|
class _RayDumper(yaml.SafeDumper):
|
|
def represent_sequence(self, tag, sequence, flow_style=None):
|
|
if len(sequence) > _SEQUENCE_LEN_FLOW_STYLE:
|
|
return super().represent_sequence(tag, sequence, flow_style=True)
|
|
return super().represent_sequence(tag, sequence, flow_style=flow_style)
|
|
|
|
|
|
@DeveloperAPI
|
|
def pretty_print(result, exclude: Optional[Set[str]] = None):
|
|
result = result.copy()
|
|
result.update(config=None) # drop config from pretty print
|
|
result.update(hist_stats=None) # drop hist_stats from pretty print
|
|
out = {}
|
|
for k, v in result.items():
|
|
if v is not None and (exclude is None or k not in exclude):
|
|
out[k] = v
|
|
|
|
cleaned = json.dumps(out, cls=SafeFallbackEncoder)
|
|
return yaml.dump(json.loads(cleaned), Dumper=_RayDumper, default_flow_style=False)
|