[tune] logger refactor part 1: move classes and utilities to own files (#11746)

* [tune] logger refactor part 1: move classes and utilities to own files

* Fix circular dependency

* Remove uneeded pretty print copy

* Apply suggestions from code review

python/ray/tune/__init__.py

@@ -8,7 +8,7 @@ from ray.tune.stopper import Stopper, EarlyStopping
 from ray.tune.registry import register_env, register_trainable
 from ray.tune.trainable import Trainable
 from ray.tune.durable_trainable import DurableTrainable
-from ray.tune.trial_runner import Callback
+from ray.tune.callback import Callback
 from ray.tune.suggest import grid_search
 from ray.tune.session import (
     report, get_trial_dir, get_trial_name, get_trial_id, make_checkpoint_dir,
@@ -22,7 +22,7 @@ from ray.tune.suggest import create_searcher
 from ray.tune.schedulers import create_scheduler
 
 __all__ = [
-    "Trainable", "DurableTrainable", "TuneError", "Callback", "grid_search",
+    "Trainable", "DurableTrainable", "Callback", "TuneError", "grid_search",
     "register_env", "register_trainable", "run", "run_experiments",
     "with_parameters", "Stopper", "EarlyStopping", "Experiment", "function",
     "sample_from", "track", "uniform", "quniform", "choice", "randint",

python/ray/tune/analysis/experiment_analysis.py

@@ -17,7 +17,7 @@ from ray.tune.error import TuneError
 from ray.tune.result import EXPR_PROGRESS_FILE, EXPR_PARAM_FILE,\
     CONFIG_PREFIX, TRAINING_ITERATION
 from ray.tune.trial import Trial
-from ray.tune.trainable import TrainableUtil
+from ray.tune.utils.trainable import TrainableUtil
 
 logger = logging.getLogger(__name__)
 

python/ray/tune/callback.py

@@ -0,0 +1,202 @@
+from typing import Dict, List
+
+from ray.tune.checkpoint_manager import Checkpoint
+from ray.tune.trial import Trial
+
+
+class Callback:
+    """Tune base callback that can be extended and passed to a ``TrialRunner``
+
+    Tune callbacks are called from within the ``TrialRunner`` class. There are
+    several hooks that can be used, all of which are found in the submethod
+    definitions of this base class.
+
+    The parameters passed to the ``**info`` dict vary between hooks. The
+    parameters passed are described in the docstrings of the methods.
+
+    This example will print a metric each time a result is received:
+
+    .. code-block:: python
+
+        from ray import tune
+        from ray.tune import Callback
+
+
+        class MyCallback(Callback):
+            def on_trial_result(self, iteration, trials, trial, result,
+                                **info):
+                print(f"Got result: {result['metric']}")
+
+
+        def train(config):
+            for i in range(10):
+                tune.report(metric=i)
+
+
+        tune.run(
+            train,
+            callbacks=[MyCallback()])
+
+    """
+
+    def on_step_begin(self, iteration: int, trials: List[Trial], **info):
+        """Called at the start of each tuning loop step.
+
+        Arguments:
+            iteration (int): Number of iterations of the tuning loop.
+            trials (List[Trial]): List of trials.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_step_end(self, iteration: int, trials: List[Trial], **info):
+        """Called at the end of each tuning loop step.
+
+        The iteration counter is increased before this hook is called.
+
+        Arguments:
+            iteration (int): Number of iterations of the tuning loop.
+            trials (List[Trial]): List of trials.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_trial_start(self, iteration: int, trials: List[Trial], trial: Trial,
+                       **info):
+        """Called after starting a trial instance.
+
+        Arguments:
+            iteration (int): Number of iterations of the tuning loop.
+            trials (List[Trial]): List of trials.
+            trial (Trial): Trial that just has been started.
+            **info: Kwargs dict for forward compatibility.
+
+        """
+        pass
+
+    def on_trial_restore(self, iteration: int, trials: List[Trial],
+                         trial: Trial, **info):
+        """Called after restoring a trial instance.
+
+        Arguments:
+            iteration (int): Number of iterations of the tuning loop.
+            trials (List[Trial]): List of trials.
+            trial (Trial): Trial that just has been restored.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_trial_save(self, iteration: int, trials: List[Trial], trial: Trial,
+                      **info):
+        """Called after receiving a checkpoint from a trial.
+
+        Arguments:
+            iteration (int): Number of iterations of the tuning loop.
+            trials (List[Trial]): List of trials.
+            trial (Trial): Trial that just saved a checkpoint.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_trial_result(self, iteration: int, trials: List[Trial],
+                        trial: Trial, result: Dict, **info):
+        """Called after receiving a result from a trial.
+
+        The search algorithm and scheduler are notified before this
+        hook is called.
+
+        Arguments:
+            iteration (int): Number of iterations of the tuning loop.
+            trials (List[Trial]): List of trials.
+            trial (Trial): Trial that just sent a result.
+            result (Dict): Result that the trial sent.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_trial_complete(self, iteration: int, trials: List[Trial],
+                          trial: Trial, **info):
+        """Called after a trial instance completed.
+
+        The search algorithm and scheduler are notified before this
+        hook is called.
+
+        Arguments:
+            iteration (int): Number of iterations of the tuning loop.
+            trials (List[Trial]): List of trials.
+            trial (Trial): Trial that just has been completed.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_trial_error(self, iteration: int, trials: List[Trial], trial: Trial,
+                       **info):
+        """Called after a trial instance failed (errored).
+
+        The search algorithm and scheduler are notified before this
+        hook is called.
+
+        Arguments:
+            iteration (int): Number of iterations of the tuning loop.
+            trials (List[Trial]): List of trials.
+            trial (Trial): Trial that just has errored.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+    def on_checkpoint(self, iteration: int, trials: List[Trial], trial: Trial,
+                      checkpoint: Checkpoint, **info):
+        """Called after a trial saved a checkpoint with Tune.
+
+        Arguments:
+            iteration (int): Number of iterations of the tuning loop.
+            trials (List[Trial]): List of trials.
+            trial (Trial): Trial that just has errored.
+            checkpoint (Checkpoint): Checkpoint object that has been saved
+                by the trial.
+            **info: Kwargs dict for forward compatibility.
+        """
+        pass
+
+
+class CallbackList:
+    """Call multiple callbacks at once."""
+
+    def __init__(self, callbacks: List[Callback]):
+        self._callbacks = callbacks
+
+    def on_step_begin(self, **info):
+        for callback in self._callbacks:
+            callback.on_step_begin(**info)
+
+    def on_step_end(self, **info):
+        for callback in self._callbacks:
+            callback.on_step_end(**info)
+
+    def on_trial_start(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_start(**info)
+
+    def on_trial_restore(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_restore(**info)
+
+    def on_trial_save(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_save(**info)
+
+    def on_trial_result(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_result(**info)
+
+    def on_trial_complete(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_complete(**info)
+
+    def on_trial_error(self, **info):
+        for callback in self._callbacks:
+            callback.on_trial_error(**info)
+
+    def on_checkpoint(self, **info):
+        for callback in self._callbacks:
+            callback.on_checkpoint(**info)

python/ray/tune/config_parser.py

@@ -8,7 +8,7 @@ from six import string_types
 from ray.tune import TuneError
 from ray.tune.trial import Trial
 from ray.tune.resources import json_to_resources
-from ray.tune.logger import _SafeFallbackEncoder
+from ray.tune.utils.util import SafeFallbackEncoder
 
 
 def make_parser(parser_creator=None, **kwargs):
@@ -143,7 +143,7 @@ def to_argv(config):
         elif isinstance(v, bool):
             pass
         else:
-            argv.append(json.dumps(v, cls=_SafeFallbackEncoder))
+            argv.append(json.dumps(v, cls=SafeFallbackEncoder))
     return argv
 
 

python/ray/tune/integration/horovod.py

@@ -7,7 +7,7 @@ from filelock import FileLock
 import ray
 from ray import tune
 from ray.tune.resources import Resources
-from ray.tune.trainable import TrainableUtil
+from ray.tune.utils.trainable import TrainableUtil
 from ray.tune.result import RESULT_DUPLICATE
 from ray.tune.logger import NoopLogger
 

python/ray/tune/integration/torch.py

@@ -16,7 +16,7 @@ from ray.tune.result import RESULT_DUPLICATE
 from ray.tune.logger import NoopLogger
 from ray.tune.function_runner import wrap_function
 from ray.tune.resources import Resources
-from ray.tune.trainable import TrainableUtil
+from ray.tune.utils.trainable import TrainableUtil
 from ray.tune.utils import detect_checkpoint_function
 from ray.util.sgd.torch.utils import setup_process_group, setup_address
 from ray.util.sgd.torch.constants import NCCL_TIMEOUT_S

python/ray/tune/logger.py

@@ -1,12 +1,15 @@
 import csv
 import json
 import logging
-import os
-import yaml
 import numbers
 import numpy as np
+import os
+import yaml
+
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Type, Union
 
 import ray.cloudpickle as cloudpickle
+from ray.tune.utils.util import SafeFallbackEncoder
 from ray.util.debug import log_once
 from ray.tune.result import (NODE_IP, TRAINING_ITERATION, TIME_TOTAL_S,
                              TIMESTEPS_TOTAL, EXPR_PARAM_FILE,
@@ -15,6 +18,9 @@ from ray.tune.result import (NODE_IP, TRAINING_ITERATION, TIME_TOTAL_S,
 from ray.tune.syncer import get_node_syncer
 from ray.tune.utils import flatten_dict
 
+if TYPE_CHECKING:
+    from ray.tune.trial import Trial  # noqa: F401
+
 logger = logging.getLogger(__name__)
 
 tf = None
@@ -34,7 +40,10 @@ class Logger:
         trial (Trial): Trial object for the logger to access.
     """
 
-    def __init__(self, config, logdir, trial=None):
+    def __init__(self,
+                 config: Dict,
+                 logdir: str,
+                 trial: Optional["Trial"] = None):
         self.config = config
         self.logdir = logdir
         self.trial = trial
@@ -89,7 +98,7 @@ class MLFLowLogger(Logger):
             client.log_param(self._run_id, key, value)
         self.client = client
 
-    def on_result(self, result):
+    def on_result(self, result: Dict):
         for key, value in result.items():
             if not isinstance(value, float):
                 continue
@@ -113,8 +122,8 @@ class JsonLogger(Logger):
         local_file = os.path.join(self.logdir, EXPR_RESULT_FILE)
         self.local_out = open(local_file, "a")
 
-    def on_result(self, result):
-        json.dump(result, self, cls=_SafeFallbackEncoder)
+    def on_result(self, result: Dict):
+        json.dump(result, self, cls=SafeFallbackEncoder)
         self.write("\n")
         self.local_out.flush()
 
@@ -127,7 +136,7 @@ class JsonLogger(Logger):
     def close(self):
         self.local_out.close()
 
-    def update_config(self, config):
+    def update_config(self, config: Dict):
         self.config = config
         config_out = os.path.join(self.logdir, EXPR_PARAM_FILE)
         with open(config_out, "w") as f:
@@ -136,7 +145,7 @@ class JsonLogger(Logger):
                 f,
                 indent=2,
                 sort_keys=True,
-                cls=_SafeFallbackEncoder)
+                cls=SafeFallbackEncoder)
         config_pkl = os.path.join(self.logdir, EXPR_PARAM_PICKLE_FILE)
         with open(config_pkl, "wb") as f:
             cloudpickle.dump(self.config, f)
@@ -159,7 +168,7 @@ class CSVLogger(Logger):
         self._file = open(progress_file, "a")
         self._csv_out = None
 
-    def on_result(self, result):
+    def on_result(self, result: Dict):
         tmp = result.copy()
         if "config" in tmp:
             del tmp["config"]
@@ -203,7 +212,7 @@ class TBXLogger(Logger):
         self._file_writer = SummaryWriter(self.logdir, flush_secs=30)
         self.last_result = None
 
-    def on_result(self, result):
+    def on_result(self, result: Dict):
         step = result.get(TIMESTEPS_TOTAL) or result[TRAINING_ITERATION]
 
         tmp = result.copy()
@@ -313,11 +322,11 @@ class UnifiedLogger(Logger):
     """
 
     def __init__(self,
-                 config,
-                 logdir,
-                 trial=None,
-                 loggers=None,
-                 sync_function=None):
+                 config: Dict,
+                 logdir: str,
+                 trial: Optional["Trial"] = None,
+                 loggers: Optional[List[Type[Logger]]] = None,
+                 sync_function: Union[None, Callable, str] = None):
         if loggers is None:
             self._logger_cls_list = DEFAULT_LOGGERS
         else:

python/ray/tune/ray_trial_executor.py

@@ -18,7 +18,7 @@ from ray.tune.function_runner import FunctionRunner
 from ray.tune.logger import NoopLogger
 from ray.tune.result import TRIAL_INFO, STDOUT_FILE, STDERR_FILE
 from ray.tune.resources import Resources
-from ray.tune.trainable import TrainableUtil
+from ray.tune.utils.trainable import TrainableUtil
 from ray.tune.trial import Trial, Checkpoint, Location, TrialInfo
 from ray.tune.trial_executor import TrialExecutor
 from ray.tune.utils import warn_if_slow

python/ray/tune/schedulers/pbt.py

@@ -11,7 +11,7 @@ from ray.tune import trial_runner
 from ray.tune import trial_executor
 from ray.tune.error import TuneError
 from ray.tune.result import TRAINING_ITERATION
-from ray.tune.logger import _SafeFallbackEncoder
+from ray.tune.utils.util import SafeFallbackEncoder
 from ray.tune.sample import Domain, Function
 from ray.tune.schedulers import FIFOScheduler, TrialScheduler
 from ray.tune.suggest.variant_generator import format_vars
@@ -503,13 +503,13 @@ class PopulationBasedTraining(FIFOScheduler):
         ]
         # Log to global file.
         with open(os.path.join(trial.local_dir, "pbt_global.txt"), "a+") as f:
-            print(json.dumps(policy, cls=_SafeFallbackEncoder), file=f)
+            print(json.dumps(policy, cls=SafeFallbackEncoder), file=f)
         # Overwrite state in target trial from trial_to_clone.
         if os.path.exists(trial_to_clone_path):
             shutil.copyfile(trial_to_clone_path, trial_path)
         # Log new exploit in target trial log.
         with open(trial_path, "a+") as f:
-            f.write(json.dumps(policy, cls=_SafeFallbackEncoder) + "\n")
+            f.write(json.dumps(policy, cls=SafeFallbackEncoder) + "\n")
 
     def _get_new_config(self, trial, trial_to_clone):
         """Gets new config for trial by exploring trial_to_clone's config."""

python/ray/tune/tests/test_cluster.py

@@ -19,7 +19,7 @@ from ray.tune.ray_trial_executor import RayTrialExecutor
 from ray.tune.resources import Resources
 from ray.tune.suggest import BasicVariantGenerator
 from ray.tune.syncer import CloudSyncer
-from ray.tune.trainable import TrainableUtil
+from ray.tune.utils.trainable import TrainableUtil
 from ray.tune.trial import Trial
 from ray.tune.trial_runner import TrialRunner
 from ray.tune.utils.mock import (MockDurableTrainer, MockRemoteTrainer,

python/ray/tune/tests/test_function_api.py

@@ -9,7 +9,7 @@ from ray.rllib import _register_all
 
 from ray import tune
 from ray.tune.logger import NoopLogger
-from ray.tune.trainable import TrainableUtil
+from ray.tune.utils.trainable import TrainableUtil
 from ray.tune.function_runner import with_parameters, wrap_function, \
     FuncCheckpointUtil
 from ray.tune.result import TRAINING_ITERATION

python/ray/tune/tests/test_trainable_util.py

@@ -5,7 +5,7 @@ import unittest
 
 import ray.utils
 
-from ray.tune.trainable import TrainableUtil
+from ray.tune.utils.trainable import TrainableUtil
 
 
 class TrainableUtilTest(unittest.TestCase):

python/ray/tune/tests/test_trial_runner_callbacks.py

@@ -13,7 +13,8 @@ from ray.tune.ray_trial_executor import RayTrialExecutor
 from ray.tune.result import TRAINING_ITERATION
 
 from ray.tune.trial import Trial
-from ray.tune.trial_runner import Callback, TrialRunner
+from ray.tune.callback import Callback
+from ray.tune.trial_runner import TrialRunner
 
 
 class TestCallback(Callback):
@@ -45,7 +46,7 @@ class TestCallback(Callback):
     def on_trial_complete(self, **info):
         self.state["trial_complete"] = info
 
-    def on_trial_fail(self, **info):
+    def on_trial_error(self, **info):
         self.state["trial_fail"] = info
 
 

python/ray/tune/trainable.py

@@ -3,16 +3,13 @@ from contextlib import redirect_stdout, redirect_stderr
 from datetime import datetime
 
 import copy
-import io
 import logging
-import glob
 import os
 import pickle
 import platform
 
-import pandas as pd
+from ray.tune.utils.trainable import TrainableUtil
 from ray.tune.utils.util import Tee
-from six import string_types
 import shutil
 import tempfile
 import time
@@ -20,7 +17,6 @@ import uuid
 
 import ray
 from ray.util.debug import log_once
-from ray.tune.logger import UnifiedLogger
 from ray.tune.result import (
     DEFAULT_RESULTS_DIR, TIME_THIS_ITER_S, TIMESTEPS_THIS_ITER, DONE,
     TIMESTEPS_TOTAL, EPISODES_THIS_ITER, EPISODES_TOTAL, TRAINING_ITERATION,
@@ -32,155 +28,6 @@ logger = logging.getLogger(__name__)
 SETUP_TIME_THRESHOLD = 10
 
 
-class TrainableUtil:
-    @staticmethod
-    def process_checkpoint(checkpoint, parent_dir, trainable_state):
-        saved_as_dict = False
-        if isinstance(checkpoint, string_types):
-            if not checkpoint.startswith(parent_dir):
-                raise ValueError(
-                    "The returned checkpoint path must be within the "
-                    "given checkpoint dir {}: {}".format(
-                        parent_dir, checkpoint))
-            checkpoint_path = checkpoint
-            if os.path.isdir(checkpoint_path):
-                # Add trailing slash to prevent tune metadata from
-                # being written outside the directory.
-                checkpoint_path = os.path.join(checkpoint_path, "")
-        elif isinstance(checkpoint, dict):
-            saved_as_dict = True
-            checkpoint_path = os.path.join(parent_dir, "checkpoint")
-            with open(checkpoint_path, "wb") as f:
-                pickle.dump(checkpoint, f)
-        else:
-            raise ValueError("Returned unexpected type {}. "
-                             "Expected str or dict.".format(type(checkpoint)))
-
-        with open(checkpoint_path + ".tune_metadata", "wb") as f:
-            trainable_state["saved_as_dict"] = saved_as_dict
-            pickle.dump(trainable_state, f)
-        return checkpoint_path
-
-    @staticmethod
-    def pickle_checkpoint(checkpoint_path):
-        """Pickles checkpoint data."""
-        checkpoint_dir = TrainableUtil.find_checkpoint_dir(checkpoint_path)
-        data = {}
-        for basedir, _, file_names in os.walk(checkpoint_dir):
-            for file_name in file_names:
-                path = os.path.join(basedir, file_name)
-                with open(path, "rb") as f:
-                    data[os.path.relpath(path, checkpoint_dir)] = f.read()
-        # Use normpath so that a directory path isn't mapped to empty string.
-        name = os.path.relpath(
-            os.path.normpath(checkpoint_path), checkpoint_dir)
-        name += os.path.sep if os.path.isdir(checkpoint_path) else ""
-        data_dict = pickle.dumps({
-            "checkpoint_name": name,
-            "data": data,
-        })
-        return data_dict
-
-    @staticmethod
-    def checkpoint_to_object(checkpoint_path):
-        data_dict = TrainableUtil.pickle_checkpoint(checkpoint_path)
-        out = io.BytesIO()
-        if len(data_dict) > 10e6:  # getting pretty large
-            logger.info("Checkpoint size is {} bytes".format(len(data_dict)))
-        out.write(data_dict)
-        return out.getvalue()
-
-    @staticmethod
-    def find_checkpoint_dir(checkpoint_path):
-        """Returns the directory containing the checkpoint path.
-
-        Raises:
-            FileNotFoundError if the directory is not found.
-        """
-        if not os.path.exists(checkpoint_path):
-            raise FileNotFoundError("Path does not exist", checkpoint_path)
-        if os.path.isdir(checkpoint_path):
-            checkpoint_dir = checkpoint_path
-        else:
-            checkpoint_dir = os.path.dirname(checkpoint_path)
-        while checkpoint_dir != os.path.dirname(checkpoint_dir):
-            if os.path.exists(os.path.join(checkpoint_dir, ".is_checkpoint")):
-                break
-            checkpoint_dir = os.path.dirname(checkpoint_dir)
-        else:
-            raise FileNotFoundError("Checkpoint directory not found for {}"
-                                    .format(checkpoint_path))
-        return checkpoint_dir
-
-    @staticmethod
-    def make_checkpoint_dir(checkpoint_dir, index, override=False):
-        """Creates a checkpoint directory within the provided path.
-
-        Args:
-            checkpoint_dir (str): Path to checkpoint directory.
-            index (str): A subdirectory will be created
-                at the checkpoint directory named 'checkpoint_{index}'.
-            override (bool): Deletes checkpoint_dir before creating
-                a new one.
-        """
-        suffix = "checkpoint"
-        if index is not None:
-            suffix += "_{}".format(index)
-        checkpoint_dir = os.path.join(checkpoint_dir, suffix)
-
-        if override and os.path.exists(checkpoint_dir):
-            shutil.rmtree(checkpoint_dir)
-        os.makedirs(checkpoint_dir, exist_ok=True)
-        # Drop marker in directory to identify it as a checkpoint dir.
-        open(os.path.join(checkpoint_dir, ".is_checkpoint"), "a").close()
-        return checkpoint_dir
-
-    @staticmethod
-    def create_from_pickle(obj, tmpdir):
-        info = pickle.loads(obj)
-        data = info["data"]
-        checkpoint_path = os.path.join(tmpdir, info["checkpoint_name"])
-
-        for relpath_name, file_contents in data.items():
-            path = os.path.join(tmpdir, relpath_name)
-
-            # This may be a subdirectory, hence not just using tmpdir
-            os.makedirs(os.path.dirname(path), exist_ok=True)
-            with open(path, "wb") as f:
-                f.write(file_contents)
-        return checkpoint_path
-
-    @staticmethod
-    def get_checkpoints_paths(logdir):
-        """ Finds the checkpoints within a specific folder.
-
-        Returns a pandas DataFrame of training iterations and checkpoint
-        paths within a specific folder.
-
-        Raises:
-            FileNotFoundError if the directory is not found.
-        """
-        marker_paths = glob.glob(
-            os.path.join(logdir, "checkpoint_*/.is_checkpoint"))
-        iter_chkpt_pairs = []
-        for marker_path in marker_paths:
-            chkpt_dir = os.path.dirname(marker_path)
-            metadata_file = glob.glob(
-                os.path.join(chkpt_dir, "*.tune_metadata"))
-            if len(metadata_file) != 1:
-                raise ValueError(
-                    "{} has zero or more than one tune_metadata.".format(
-                        chkpt_dir))
-
-            chkpt_path = metadata_file[0][:-len(".tune_metadata")]
-            chkpt_iter = int(chkpt_dir[chkpt_dir.rfind("_") + 1:])
-            iter_chkpt_pairs.append([chkpt_iter, chkpt_path])
-
-        chkpt_df = pd.DataFrame(
-            iter_chkpt_pairs, columns=["training_iteration", "chkpt_path"])
-        return chkpt_df
-
-
 class Trainable:
     """Abstract class for trainable models, functions, etc.
 
@@ -594,6 +441,8 @@ class Trainable:
             self._result_logger = logger_creator(config)
             self._logdir = self._result_logger.logdir
         else:
+            from ray.tune.logger import UnifiedLogger
+
             logdir_prefix = datetime.today().strftime("%Y-%m-%d_%H-%M-%S")
             ray.utils.try_to_create_directory(DEFAULT_RESULTS_DIR)
             self._logdir = tempfile.mkdtemp(

python/ray/tune/trial.py

@@ -20,7 +20,7 @@ from ray.tune.logger import pretty_print, UnifiedLogger
 from ray.tune.registry import get_trainable_cls, validate_trainable
 from ray.tune.result import DEFAULT_RESULTS_DIR, DONE, TRAINING_ITERATION
 from ray.tune.resources import Resources, json_to_resources, resources_to_json
-from ray.tune.trainable import TrainableUtil
+from ray.tune.utils.trainable import TrainableUtil
 from ray.tune.utils import date_str, flatten_dict
 from ray.utils import binary_to_hex, hex_to_binary
 

python/ray/tune/trial_runner.py

@@ -1,5 +1,3 @@
-from typing import Dict, List
-
 import click
 from datetime import datetime
 import json
@@ -12,8 +10,8 @@ import types
 import ray.cloudpickle as cloudpickle
 from ray.services import get_node_ip_address
 from ray.tune import TuneError
+from ray.tune.callback import CallbackList
 from ray.tune.stopper import NoopStopper
-from ray.tune.progress_reporter import trial_progress_str
 from ray.tune.ray_trial_executor import RayTrialExecutor
 from ray.tune.result import (TIME_THIS_ITER_S, RESULT_DUPLICATE,
                              SHOULD_CHECKPOINT)
@@ -71,186 +69,6 @@ class _TuneFunctionDecoder(json.JSONDecoder):
         return cloudpickle.loads(hex_to_binary(obj["value"]))
 
 
-class Callback:
-    """Tune base callback that can be extended and passed to a ``TrialRunner``
-
-    Tune callbacks are called from within the ``TrialRunner`` class. There are
-    several hooks that can be used, all of which are found in the submethod
-    definitions of this base class.
-
-    The parameters passed to the ``**info`` dict vary between hooks. The
-    parameters passed are described in the docstrings of the methods.
-
-    This example will print a metric each time a result is received:
-
-    .. code-block:: python
-
-        from ray import tune
-        from ray.tune import Callback
-
-
-        class MyCallback(Callback):
-            def on_trial_result(self, iteration, trials, trial, result,
-                                **info):
-                print(f"Got result: {result['metric']}")
-
-
-        def train(config):
-            for i in range(10):
-                tune.report(metric=i)
-
-
-        tune.run(
-            train,
-            callbacks=[MyCallback()])
-
-    """
-
-    def on_step_begin(self, iteration: int, trials: List[Trial], **info):
-        """Called at the start of each tuning loop step.
-
-        Arguments:
-            iteration (int): Number of iterations of the tuning loop.
-            trials (List[Trial]): List of trials.
-            **info: Kwargs dict for forward compatibility.
-        """
-        pass
-
-    def on_step_end(self, iteration: int, trials: List[Trial], **info):
-        """Called at the end of each tuning loop step.
-
-        The iteration counter is increased before this hook is called.
-
-        Arguments:
-            iteration (int): Number of iterations of the tuning loop.
-            trials (List[Trial]): List of trials.
-            **info: Kwargs dict for forward compatibility.
-        """
-        pass
-
-    def on_trial_start(self, iteration: int, trials: List[Trial], trial: Trial,
-                       **info):
-        """Called after starting a trial instance.
-
-        Arguments:
-            iteration (int): Number of iterations of the tuning loop.
-            trials (List[Trial]): List of trials.
-            trial (Trial): Trial that just has been started.
-            **info: Kwargs dict for forward compatibility.
-
-        """
-        pass
-
-    def on_trial_restore(self, iteration: int, trials: List[Trial],
-                         trial: Trial, **info):
-        """Called after restoring a trial instance.
-
-        Arguments:
-            iteration (int): Number of iterations of the tuning loop.
-            trials (List[Trial]): List of trials.
-            trial (Trial): Trial that just has been restored.
-            **info: Kwargs dict for forward compatibility.
-        """
-        pass
-
-    def on_trial_save(self, iteration: int, trials: List[Trial], trial: Trial,
-                      **info):
-        """Called after receiving a checkpoint from a trial.
-
-        Arguments:
-            iteration (int): Number of iterations of the tuning loop.
-            trials (List[Trial]): List of trials.
-            trial (Trial): Trial that just saved a checkpoint.
-            **info: Kwargs dict for forward compatibility.
-        """
-        pass
-
-    def on_trial_result(self, iteration: int, trials: List[Trial],
-                        trial: Trial, result: Dict, **info):
-        """Called after receiving a result from a trial.
-
-        The search algorithm and scheduler are notified before this
-        hook is called.
-
-        Arguments:
-            iteration (int): Number of iterations of the tuning loop.
-            trials (List[Trial]): List of trials.
-            trial (Trial): Trial that just sent a result.
-            result (Dict): Result that the trial sent.
-            **info: Kwargs dict for forward compatibility.
-        """
-        pass
-
-    def on_trial_complete(self, iteration: int, trials: List[Trial],
-                          trial: Trial, **info):
-        """Called after a trial instance completed.
-
-        The search algorithm and scheduler are notified before this
-        hook is called.
-
-        Arguments:
-            iteration (int): Number of iterations of the tuning loop.
-            trials (List[Trial]): List of trials.
-            trial (Trial): Trial that just has been completed.
-            **info: Kwargs dict for forward compatibility.
-        """
-        pass
-
-    def on_trial_fail(self, iteration: int, trials: List[Trial], trial: Trial,
-                      **info):
-        """Called after a trial instance failed (errored).
-
-        The search algorithm and scheduler are notified before this
-        hook is called.
-
-        Arguments:
-            iteration (int): Number of iterations of the tuning loop.
-            trials (List[Trial]): List of trials.
-            trial (Trial): Trial that just has errored.
-            **info: Kwargs dict for forward compatibility.
-        """
-        pass
-
-
-class _CallbackList:
-    """Call multiple callbacks at once."""
-
-    def __init__(self, callbacks: List[Callback]):
-        self._callbacks = callbacks
-
-    def on_step_begin(self, **info):
-        for callback in self._callbacks:
-            callback.on_step_begin(**info)
-
-    def on_step_end(self, **info):
-        for callback in self._callbacks:
-            callback.on_step_end(**info)
-
-    def on_trial_start(self, **info):
-        for callback in self._callbacks:
-            callback.on_trial_start(**info)
-
-    def on_trial_restore(self, **info):
-        for callback in self._callbacks:
-            callback.on_trial_restore(**info)
-
-    def on_trial_save(self, **info):
-        for callback in self._callbacks:
-            callback.on_trial_save(**info)
-
-    def on_trial_result(self, **info):
-        for callback in self._callbacks:
-            callback.on_trial_result(**info)
-
-    def on_trial_complete(self, **info):
-        for callback in self._callbacks:
-            callback.on_trial_complete(**info)
-
-    def on_trial_fail(self, **info):
-        for callback in self._callbacks:
-            callback.on_trial_fail(**info)
-
-
 class TrialRunner:
     """A TrialRunner implements the event loop for scheduling trials on Ray.
 
@@ -400,7 +218,7 @@ class TrialRunner:
                 self._local_checkpoint_dir,
                 TrialRunner.CKPT_FILE_TMPL.format(self._session_str))
 
-        self._callbacks = _CallbackList(callbacks or [])
+        self._callbacks = CallbackList(callbacks or [])
 
     @property
     def resumed(self):
@@ -618,6 +436,8 @@ class TrialRunner:
         self.trial_executor.try_checkpoint_metadata(trial)
 
     def debug_string(self, delim="\n"):
+        from ray.tune.progress_reporter import trial_progress_str
+
         result_keys = [
             list(t.last_result) for t in self.get_trials() if t.last_result
         ]
@@ -724,6 +544,7 @@ class TrialRunner:
         """
         try:
             result = self.trial_executor.fetch_result(trial)
+            result.update(trial_id=trial.trial_id)
             is_duplicate = RESULT_DUPLICATE in result
             force_checkpoint = result.get(SHOULD_CHECKPOINT, False)
             # TrialScheduler and SearchAlgorithm still receive a
@@ -740,10 +561,23 @@ class TrialRunner:
             flat_result = flatten_dict(result)
             if self._stopper(trial.trial_id,
                              result) or trial.should_stop(flat_result):
+                result.update(done=True)
+
                 # Hook into scheduler
                 self._scheduler_alg.on_trial_complete(self, trial, flat_result)
                 self._search_alg.on_trial_complete(
                     trial.trial_id, result=flat_result)
+
+                # If this is not a duplicate result, the callbacks should
+                # be informed about the result.
+                if not is_duplicate:
+                    with warn_if_slow("callbacks.on_trial_result"):
+                        self._callbacks.on_trial_result(
+                            iteration=self._iteration,
+                            trials=self._trials,
+                            trial=trial,
+                            result=result.copy())
+
                 self._callbacks.on_trial_complete(
                     iteration=self._iteration,
                     trials=self._trials,
@@ -771,6 +605,7 @@ class TrialRunner:
                             iteration=self._iteration,
                             trials=self._trials,
                             trial=trial)
+                    result.update(done=True)
 
             if not is_duplicate:
                 trial.update_last_result(
@@ -861,6 +696,11 @@ class TrialRunner:
         if checkpoint_value:
             try:
                 trial.saving_to.value = checkpoint_value
+                self._callbacks.on_checkpoint(
+                    iteration=self._iteration,
+                    trials=self._trials,
+                    trial=trial,
+                    checkpoint=trial.saving_to)
                 trial.on_checkpoint(trial.saving_to)
                 self.trial_executor.try_checkpoint_metadata(trial)
             except Exception:
@@ -909,7 +749,7 @@ class TrialRunner:
             else:
                 self._scheduler_alg.on_trial_error(self, trial)
                 self._search_alg.on_trial_complete(trial.trial_id, error=True)
-                self._callbacks.on_trial_fail(
+                self._callbacks.on_trial_error(
                     iteration=self._iteration,
                     trials=self._trials,
                     trial=trial)
@@ -969,6 +809,10 @@ class TrialRunner:
                     trial)
                 self._scheduler_alg.on_trial_error(self, trial)
                 self._search_alg.on_trial_complete(trial.trial_id, error=True)
+                self._callbacks.on_trial_error(
+                    iteration=self._iteration,
+                    trials=self._trials,
+                    trial=trial)
             else:
                 logger.debug("Trial %s: Restore dispatched correctly.", trial)
         else:
@@ -1051,6 +895,8 @@ class TrialRunner:
         elif trial.status in [Trial.PENDING, Trial.PAUSED]:
             self._scheduler_alg.on_trial_remove(self, trial)
             self._search_alg.on_trial_complete(trial.trial_id)
+            self._callbacks.on_trial_complete(
+                iteration=self._iteration, trials=self._trials, trial=trial)
         elif trial.status is Trial.RUNNING:
             try:
                 result = self.trial_executor.fetch_result(trial)
@@ -1058,11 +904,19 @@ class TrialRunner:
                 self._scheduler_alg.on_trial_complete(self, trial, result)
                 self._search_alg.on_trial_complete(
                     trial.trial_id, result=result)
+                self._callbacks.on_trial_complete(
+                    iteration=self._iteration,
+                    trials=self._trials,
+                    trial=trial)
             except Exception:
                 error_msg = traceback.format_exc()
                 logger.exception("Error processing event.")
                 self._scheduler_alg.on_trial_error(self, trial)
                 self._search_alg.on_trial_complete(trial.trial_id, error=True)
+                self._callbacks.on_trial_error(
+                    iteration=self._iteration,
+                    trials=self._trials,
+                    trial=trial)
                 error = True
 
         self.trial_executor.stop_trial(trial, error=error, error_msg=error_msg)

python/ray/tune/tune.py

@@ -457,7 +457,8 @@ def run_experiments(experiments,
                     reuse_actors=False,
                     trial_executor=None,
                     raise_on_failed_trial=True,
-                    concurrent=True):
+                    concurrent=True,
+                    callbacks=None):
     """Runs and blocks until all trials finish.
 
     Examples:
@@ -487,7 +488,8 @@ def run_experiments(experiments,
             reuse_actors=reuse_actors,
             trial_executor=trial_executor,
             raise_on_failed_trial=raise_on_failed_trial,
-            scheduler=scheduler).trials
+            scheduler=scheduler,
+            callbacks=callbacks).trials
     else:
         trials = []
         for exp in experiments:
@@ -501,5 +503,6 @@ def run_experiments(experiments,
                 reuse_actors=reuse_actors,
                 trial_executor=trial_executor,
                 raise_on_failed_trial=raise_on_failed_trial,
-                scheduler=scheduler).trials
+                scheduler=scheduler,
+                callbacks=callbacks).trials
         return trials

python/ray/tune/utils/trainable.py

@@ -0,0 +1,161 @@
+import glob
+import io
+import logging
+import shutil
+
+import pandas as pd
+import pickle
+import os
+
+from six import string_types
+
+logger = logging.getLogger(__name__)
+
+
+class TrainableUtil:
+    @staticmethod
+    def process_checkpoint(checkpoint, parent_dir, trainable_state):
+        saved_as_dict = False
+        if isinstance(checkpoint, string_types):
+            if not checkpoint.startswith(parent_dir):
+                raise ValueError(
+                    "The returned checkpoint path must be within the "
+                    "given checkpoint dir {}: {}".format(
+                        parent_dir, checkpoint))
+            checkpoint_path = checkpoint
+            if os.path.isdir(checkpoint_path):
+                # Add trailing slash to prevent tune metadata from
+                # being written outside the directory.
+                checkpoint_path = os.path.join(checkpoint_path, "")
+        elif isinstance(checkpoint, dict):
+            saved_as_dict = True
+            checkpoint_path = os.path.join(parent_dir, "checkpoint")
+            with open(checkpoint_path, "wb") as f:
+                pickle.dump(checkpoint, f)
+        else:
+            raise ValueError("Returned unexpected type {}. "
+                             "Expected str or dict.".format(type(checkpoint)))
+
+        with open(checkpoint_path + ".tune_metadata", "wb") as f:
+            trainable_state["saved_as_dict"] = saved_as_dict
+            pickle.dump(trainable_state, f)
+        return checkpoint_path
+
+    @staticmethod
+    def pickle_checkpoint(checkpoint_path):
+        """Pickles checkpoint data."""
+        checkpoint_dir = TrainableUtil.find_checkpoint_dir(checkpoint_path)
+        data = {}
+        for basedir, _, file_names in os.walk(checkpoint_dir):
+            for file_name in file_names:
+                path = os.path.join(basedir, file_name)
+                with open(path, "rb") as f:
+                    data[os.path.relpath(path, checkpoint_dir)] = f.read()
+        # Use normpath so that a directory path isn't mapped to empty string.
+        name = os.path.relpath(
+            os.path.normpath(checkpoint_path), checkpoint_dir)
+        name += os.path.sep if os.path.isdir(checkpoint_path) else ""
+        data_dict = pickle.dumps({
+            "checkpoint_name": name,
+            "data": data,
+        })
+        return data_dict
+
+    @staticmethod
+    def checkpoint_to_object(checkpoint_path):
+        data_dict = TrainableUtil.pickle_checkpoint(checkpoint_path)
+        out = io.BytesIO()
+        if len(data_dict) > 10e6:  # getting pretty large
+            logger.info("Checkpoint size is {} bytes".format(len(data_dict)))
+        out.write(data_dict)
+        return out.getvalue()
+
+    @staticmethod
+    def find_checkpoint_dir(checkpoint_path):
+        """Returns the directory containing the checkpoint path.
+
+        Raises:
+            FileNotFoundError if the directory is not found.
+        """
+        if not os.path.exists(checkpoint_path):
+            raise FileNotFoundError("Path does not exist", checkpoint_path)
+        if os.path.isdir(checkpoint_path):
+            checkpoint_dir = checkpoint_path
+        else:
+            checkpoint_dir = os.path.dirname(checkpoint_path)
+        while checkpoint_dir != os.path.dirname(checkpoint_dir):
+            if os.path.exists(os.path.join(checkpoint_dir, ".is_checkpoint")):
+                break
+            checkpoint_dir = os.path.dirname(checkpoint_dir)
+        else:
+            raise FileNotFoundError("Checkpoint directory not found for {}"
+                                    .format(checkpoint_path))
+        return checkpoint_dir
+
+    @staticmethod
+    def make_checkpoint_dir(checkpoint_dir, index, override=False):
+        """Creates a checkpoint directory within the provided path.
+
+        Args:
+            checkpoint_dir (str): Path to checkpoint directory.
+            index (str): A subdirectory will be created
+                at the checkpoint directory named 'checkpoint_{index}'.
+            override (bool): Deletes checkpoint_dir before creating
+                a new one.
+        """
+        suffix = "checkpoint"
+        if index is not None:
+            suffix += "_{}".format(index)
+        checkpoint_dir = os.path.join(checkpoint_dir, suffix)
+
+        if override and os.path.exists(checkpoint_dir):
+            shutil.rmtree(checkpoint_dir)
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        # Drop marker in directory to identify it as a checkpoint dir.
+        open(os.path.join(checkpoint_dir, ".is_checkpoint"), "a").close()
+        return checkpoint_dir
+
+    @staticmethod
+    def create_from_pickle(obj, tmpdir):
+        info = pickle.loads(obj)
+        data = info["data"]
+        checkpoint_path = os.path.join(tmpdir, info["checkpoint_name"])
+
+        for relpath_name, file_contents in data.items():
+            path = os.path.join(tmpdir, relpath_name)
+
+            # This may be a subdirectory, hence not just using tmpdir
+            os.makedirs(os.path.dirname(path), exist_ok=True)
+            with open(path, "wb") as f:
+                f.write(file_contents)
+        return checkpoint_path
+
+    @staticmethod
+    def get_checkpoints_paths(logdir):
+        """ Finds the checkpoints within a specific folder.
+
+        Returns a pandas DataFrame of training iterations and checkpoint
+        paths within a specific folder.
+
+        Raises:
+            FileNotFoundError if the directory is not found.
+        """
+        marker_paths = glob.glob(
+            os.path.join(logdir, "checkpoint_*/.is_checkpoint"))
+        iter_chkpt_pairs = []
+        for marker_path in marker_paths:
+            chkpt_dir = os.path.dirname(marker_path)
+            metadata_file = glob.glob(
+                os.path.join(chkpt_dir, "*.tune_metadata"))
+            if len(metadata_file) != 1:
+                raise ValueError(
+                    "{} has zero or more than one tune_metadata.".format(
+                        chkpt_dir))
+
+            chkpt_path = metadata_file[0][:-len(".tune_metadata")]
+            chkpt_iter = int(chkpt_dir[chkpt_dir.rfind("_") + 1:])
+            iter_chkpt_pairs.append([chkpt_iter, chkpt_path])
+
+        chkpt_df = pd.DataFrame(
+            iter_chkpt_pairs, columns=["training_iteration", "chkpt_path"])
+        return chkpt_df

python/ray/tune/utils/util.py

@@ -1,5 +1,7 @@
 import copy
+import json
 import logging
+import numbers
 import os
 import inspect
 import threading
@@ -538,6 +540,31 @@ def detect_config_single(func):
     return use_config_single
 
 
+class SafeFallbackEncoder(json.JSONEncoder):
+    def __init__(self, nan_str="null", **kwargs):
+        super(SafeFallbackEncoder, self).__init__(**kwargs)
+        self.nan_str = nan_str
+
+    def default(self, value):
+        try:
+            if np.isnan(value):
+                return self.nan_str
+
+            if (type(value).__module__ == np.__name__
+                    and isinstance(value, np.ndarray)):
+                return value.tolist()
+
+            if issubclass(type(value), numbers.Integral):
+                return int(value)
+            if issubclass(type(value), numbers.Number):
+                return float(value)
+
+            return super(SafeFallbackEncoder, self).default(value)
+
+        except Exception:
+            return str(value)  # give up, just stringify it (ok for logs)
+
+
 if __name__ == "__main__":
     ray.init()
     X = pin_in_object_store("hello")