From f7b19c41e3c0e4c131a78c632e3e5ce380194fe4 Mon Sep 17 00:00:00 2001 From: Kai Fricke Date: Tue, 3 Nov 2020 16:48:09 +0100 Subject: [PATCH] [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 --- doc/source/tune/api_docs/internals.rst | 2 +- python/ray/tune/__init__.py | 4 +- .../ray/tune/analysis/experiment_analysis.py | 2 +- python/ray/tune/callback.py | 202 ++++++++++++++++ python/ray/tune/config_parser.py | 4 +- python/ray/tune/integration/horovod.py | 2 +- python/ray/tune/integration/torch.py | 2 +- python/ray/tune/logger.py | 39 +-- python/ray/tune/ray_trial_executor.py | 2 +- python/ray/tune/schedulers/pbt.py | 6 +- python/ray/tune/tests/test_cluster.py | 2 +- python/ray/tune/tests/test_function_api.py | 2 +- python/ray/tune/tests/test_trainable_util.py | 2 +- .../tune/tests/test_trial_runner_callbacks.py | 5 +- python/ray/tune/trainable.py | 157 +----------- python/ray/tune/trial.py | 2 +- python/ray/tune/trial_runner.py | 224 +++--------------- python/ray/tune/tune.py | 9 +- python/ray/tune/utils/trainable.py | 161 +++++++++++++ python/ray/tune/utils/util.py | 27 +++ 20 files changed, 481 insertions(+), 375 deletions(-) create mode 100644 python/ray/tune/callback.py create mode 100644 python/ray/tune/utils/trainable.py diff --git a/doc/source/tune/api_docs/internals.rst b/doc/source/tune/api_docs/internals.rst index 1ac7b4852..d2d7e001b 100644 --- a/doc/source/tune/api_docs/internals.rst +++ b/doc/source/tune/api_docs/internals.rst @@ -116,7 +116,7 @@ Trial Callbacks --------- -.. autoclass:: ray.tune.trial_runner.Callback +.. autoclass:: ray.tune.callback.Callback :members: diff --git a/python/ray/tune/__init__.py b/python/ray/tune/__init__.py index f3babecbb..8171b0756 100644 --- a/python/ray/tune/__init__.py +++ b/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", diff --git a/python/ray/tune/analysis/experiment_analysis.py b/python/ray/tune/analysis/experiment_analysis.py index f5fb830c3..bb151812d 100644 --- a/python/ray/tune/analysis/experiment_analysis.py +++ b/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__) diff --git a/python/ray/tune/callback.py b/python/ray/tune/callback.py new file mode 100644 index 000000000..c20070cfa --- /dev/null +++ b/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) diff --git a/python/ray/tune/config_parser.py b/python/ray/tune/config_parser.py index bb1efe836..4da1860b3 100644 --- a/python/ray/tune/config_parser.py +++ b/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 diff --git a/python/ray/tune/integration/horovod.py b/python/ray/tune/integration/horovod.py index d1eda4070..b89412e0a 100644 --- a/python/ray/tune/integration/horovod.py +++ b/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 diff --git a/python/ray/tune/integration/torch.py b/python/ray/tune/integration/torch.py index 96305d9a5..d71e90cb8 100644 --- a/python/ray/tune/integration/torch.py +++ b/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 diff --git a/python/ray/tune/logger.py b/python/ray/tune/logger.py index 94e20c5d5..5df2a0d31 100644 --- a/python/ray/tune/logger.py +++ b/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: diff --git a/python/ray/tune/ray_trial_executor.py b/python/ray/tune/ray_trial_executor.py index 1ed639678..5ad050d25 100644 --- a/python/ray/tune/ray_trial_executor.py +++ b/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 diff --git a/python/ray/tune/schedulers/pbt.py b/python/ray/tune/schedulers/pbt.py index 4df3dbbb3..24a0d97cc 100644 --- a/python/ray/tune/schedulers/pbt.py +++ b/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.""" diff --git a/python/ray/tune/tests/test_cluster.py b/python/ray/tune/tests/test_cluster.py index c4c3ba261..c7ea4faf4 100644 --- a/python/ray/tune/tests/test_cluster.py +++ b/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, diff --git a/python/ray/tune/tests/test_function_api.py b/python/ray/tune/tests/test_function_api.py index 200938120..4bb27ebd9 100644 --- a/python/ray/tune/tests/test_function_api.py +++ b/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 diff --git a/python/ray/tune/tests/test_trainable_util.py b/python/ray/tune/tests/test_trainable_util.py index b88948afa..25860eb1c 100644 --- a/python/ray/tune/tests/test_trainable_util.py +++ b/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): diff --git a/python/ray/tune/tests/test_trial_runner_callbacks.py b/python/ray/tune/tests/test_trial_runner_callbacks.py index 05cf507d2..45551f801 100644 --- a/python/ray/tune/tests/test_trial_runner_callbacks.py +++ b/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 diff --git a/python/ray/tune/trainable.py b/python/ray/tune/trainable.py index 3b17267a1..ccae6d287 100644 --- a/python/ray/tune/trainable.py +++ b/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( diff --git a/python/ray/tune/trial.py b/python/ray/tune/trial.py index 77037c6b8..c69e3b362 100644 --- a/python/ray/tune/trial.py +++ b/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 diff --git a/python/ray/tune/trial_runner.py b/python/ray/tune/trial_runner.py index 57badc2fc..b06da19bd 100644 --- a/python/ray/tune/trial_runner.py +++ b/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) diff --git a/python/ray/tune/tune.py b/python/ray/tune/tune.py index 3d189f409..4a0486e7e 100644 --- a/python/ray/tune/tune.py +++ b/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 diff --git a/python/ray/tune/utils/trainable.py b/python/ray/tune/utils/trainable.py new file mode 100644 index 000000000..b7509f028 --- /dev/null +++ b/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 diff --git a/python/ray/tune/utils/util.py b/python/ray/tune/utils/util.py index e77449a76..c6328adfe 100644 --- a/python/ray/tune/utils/util.py +++ b/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")