chore: import upstream snapshot with attribution

python/ray/tune/logger/logger.py

@@ -0,0 +1,155 @@
+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)