From 7eaf063f29442798506e57ff23001f35a981d6e0 Mon Sep 17 00:00:00 2001 From: Kai Fricke Date: Fri, 11 Sep 2020 01:39:44 +0100 Subject: [PATCH] [tune] wrapper function to pass arbitrary objects through the object store to trainables (#10679) --- doc/source/tune/_tutorials/_faq.rst | 31 +++++++++++ doc/source/tune/api_docs/execution.rst | 5 ++ doc/source/tune/user-guide.rst | 21 ++++--- python/ray/tune/__init__.py | 17 +++--- python/ray/tune/function_runner.py | 64 ++++++++++++++++++++++ python/ray/tune/registry.py | 24 ++++++++ python/ray/tune/tests/test_function_api.py | 38 ++++++++++++- 7 files changed, 180 insertions(+), 20 deletions(-) diff --git a/doc/source/tune/_tutorials/_faq.rst b/doc/source/tune/_tutorials/_faq.rst index 96ff28590..50a5f0015 100644 --- a/doc/source/tune/_tutorials/_faq.rst +++ b/doc/source/tune/_tutorials/_faq.rst @@ -220,6 +220,37 @@ of 4 machines with 1 GPU each, the trial will never be scheduled. In other words, you will have to make sure that your Ray cluster has machines that can actually fulfill your resource requests. +How can I pass further parameter values to my trainable function? +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Ray Tune expects your trainable functions to accept only up to two parameters, +``config`` and ``checkpoint_dir``. But sometimes there are cases where +you want to pass constant arguments, like the number of epochs to run, +or a dataset to train on. Ray Tune offers a wrapper function to achieve +just that, called ``tune.with_parameters()``: + +.. code-block:: python + + from ray import tune + + import numpy as np + + def train(config, checkpoint_dir=None, num_epochs=10, data=None): + for i in range(num_epochs): + for sample in data: + # ... train on sample + + # Some huge dataset + data = np.random.random(size=100000000) + + tune.run( + tune.with_parameters(train, num_epochs=10, data=data)) + + +This function works similarly to ``functools.partial``, but it stores +the parameters directly in the Ray object store. This means that you +can pass even huge objects like datasets, and Ray makes sure that these +are efficiently stored and retrieved on your cluster machines. + Further Questions or Issues? ---------------------------- diff --git a/doc/source/tune/api_docs/execution.rst b/doc/source/tune/api_docs/execution.rst index 74df5eff7..c984d3894 100644 --- a/doc/source/tune/api_docs/execution.rst +++ b/doc/source/tune/api_docs/execution.rst @@ -18,6 +18,11 @@ tune.Experiment .. autofunction:: ray.tune.Experiment +tune.with_parameters +-------------------- + +.. autofunction:: ray.tune.with_parameters + .. _tune-stop-ref: Stopper (tune.Stopper) diff --git a/doc/source/tune/user-guide.rst b/doc/source/tune/user-guide.rst index 2d5b4b55b..6c22c9c72 100644 --- a/doc/source/tune/user-guide.rst +++ b/doc/source/tune/user-guide.rst @@ -209,26 +209,25 @@ disable cross-node syncing: Handling Large Datasets ----------------------- -You often will want to compute a large object (e.g., training data, model weights) on the driver and use that object within each trial. Tune provides a ``pin_in_object_store`` utility function that can be used to broadcast such large objects. Objects pinned in this way will never be evicted from the Ray object store while the driver process is running, and can be efficiently retrieved from any task via ``get_pinned_object``. +You often will want to compute a large object (e.g., training data, model weights) on the driver and use that object within each trial. + +Tune provides a wrapper function ``tune.with_parameters()`` that allows you to broadcast large objects to your trainable. +Objects passed with this wrapper will be stored on the Ray object store and will be automatically fetched +and passed to your trainable as a parameter. .. code-block:: python - import ray from ray import tune - from ray.tune.utils import pin_in_object_store, get_pinned_object import numpy as np - ray.init() + def f(config, data=None): + pass + # use data - # X_id can be referenced in closures - X_id = pin_in_object_store(np.random.random(size=100000000)) + data = np.random.random(size=100000000) - def f(config): - X = get_pinned_object(X_id) - # use X - - tune.run(f) + tune.run(tune.with_parameters(f, data=data)) .. _tune-stopping: diff --git a/python/ray/tune/__init__.py b/python/ray/tune/__init__.py index a7659dc14..ef898d04a 100644 --- a/python/ray/tune/__init__.py +++ b/python/ray/tune/__init__.py @@ -1,5 +1,6 @@ from ray.tune.error import TuneError from ray.tune.tune import run_experiments, run +from ray.tune.function_runner import with_parameters from ray.tune.syncer import SyncConfig from ray.tune.experiment import Experiment from ray.tune.analysis import ExperimentAnalysis, Analysis @@ -21,12 +22,12 @@ from ray.tune.schedulers import create_scheduler __all__ = [ "Trainable", "DurableTrainable", "TuneError", "grid_search", - "register_env", "register_trainable", "run", "run_experiments", "Stopper", - "EarlyStopping", "Experiment", "function", "sample_from", "track", - "uniform", "quniform", "choice", "randint", "qrandint", "randn", "qrandn", - "loguniform", "qloguniform", "ExperimentAnalysis", "Analysis", - "CLIReporter", "JupyterNotebookReporter", "ProgressReporter", "report", - "get_trial_dir", "get_trial_name", "get_trial_id", "make_checkpoint_dir", - "save_checkpoint", "checkpoint_dir", "SyncConfig", "create_searcher", - "create_scheduler" + "register_env", "register_trainable", "run", "run_experiments", + "with_parameters", "Stopper", "EarlyStopping", "Experiment", "function", + "sample_from", "track", "uniform", "quniform", "choice", "randint", + "qrandint", "randn", "qrandn", "loguniform", "qloguniform", + "ExperimentAnalysis", "Analysis", "CLIReporter", "JupyterNotebookReporter", + "ProgressReporter", "report", "get_trial_dir", "get_trial_name", + "get_trial_id", "make_checkpoint_dir", "save_checkpoint", "checkpoint_dir", + "SyncConfig", "create_searcher", "create_scheduler" ] diff --git a/python/ray/tune/function_runner.py b/python/ray/tune/function_runner.py index e722bacf6..84bc82772 100644 --- a/python/ray/tune/function_runner.py +++ b/python/ray/tune/function_runner.py @@ -7,6 +7,7 @@ import threading import traceback import uuid +from ray.tune.registry import parameter_registry from six.moves import queue from ray.util.debug import log_once @@ -507,3 +508,66 @@ def wrap_function(train_func, warn=True): return output return ImplicitFunc + + +def with_parameters(fn, **kwargs): + """Wrapper for function trainables to pass arbitrary large data objects. + + This wrapper function will store all passed parameters in the Ray + object store and retrieve them when calling the function. It can thus + be used to pass arbitrary data, even datasets, to Tune trainable functions. + + This can also be used as an alternative to `functools.partial` to pass + default arguments to trainables. + + Args: + fn: function to wrap + **kwargs: parameters to store in object store. + + + .. code-block:: python + + from ray import tune + + def train(config, data=None): + for sample in data: + # ... + tune.report(loss=loss) + + data = HugeDataset(download=True) + + tune.run( + tune.with_parameters(train, data=data), + #... + ) + + """ + prefix = f"{str(fn)}_" + for k, v in kwargs.items(): + parameter_registry.put(prefix + k, v) + + use_checkpoint = detect_checkpoint_function(fn) + + def inner(config, checkpoint_dir=None): + fn_kwargs = {} + if use_checkpoint: + default = checkpoint_dir + sig = inspect.signature(fn) + if "checkpoint_dir" in sig.parameters: + default = sig.parameters["checkpoint_dir"].default \ + or default + fn_kwargs["checkpoint_dir"] = default + + for k in kwargs: + fn_kwargs[k] = parameter_registry.get(prefix + k) + fn(config, **fn_kwargs) + + # Use correct function signature if no `checkpoint_dir` parameter is set + if not use_checkpoint: + + def _inner(config): + inner(config, checkpoint_dir=None) + + return _inner + + return inner diff --git a/python/ray/tune/registry.py b/python/ray/tune/registry.py index 1cb78f43e..27935e641 100644 --- a/python/ray/tune/registry.py +++ b/python/ray/tune/registry.py @@ -149,3 +149,27 @@ class _Registry: _global_registry = _Registry() ray.worker._post_init_hooks.append(_global_registry.flush_values) + + +class _ParameterRegistry: + def __init__(self): + self.to_flush = {} + self.references = {} + + def put(self, k, v): + self.to_flush[k] = v + if ray.is_initialized(): + self.flush() + + def get(self, k): + if not ray.is_initialized(): + return self.to_flush[k] + return ray.get(self.references[k]) + + def flush(self): + for k, v in self.to_flush.items(): + self.references[k] = ray.put(v) + + +parameter_registry = _ParameterRegistry() +ray.worker._post_init_hooks.append(parameter_registry.flush) diff --git a/python/ray/tune/tests/test_function_api.py b/python/ray/tune/tests/test_function_api.py index 04a70ae9c..b3714326b 100644 --- a/python/ray/tune/tests/test_function_api.py +++ b/python/ray/tune/tests/test_function_api.py @@ -10,7 +10,8 @@ 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.function_runner import wrap_function, FuncCheckpointUtil +from ray.tune.function_runner import with_parameters, wrap_function, \ + FuncCheckpointUtil from ray.tune.result import TRAINING_ITERATION @@ -419,3 +420,38 @@ class FunctionApiTest(unittest.TestCase): analysis = tune.run(train, max_failures=3) trial_dfs = list(analysis.trial_dataframes.values()) assert len(trial_dfs[0]["training_iteration"]) == 10 + + def testWithParameters(self): + class Data: + def __init__(self): + self.data = [0] * 500_000 + + data = Data() + data.data[100] = 1 + + def train(config, data=None): + data.data[101] = 2 # Changes are local + tune.report(metric=len(data.data), hundred=data.data[100]) + + trial_1, trial_2 = tune.run( + with_parameters(train, data=data), num_samples=2).trials + + self.assertEquals(data.data[101], 0) + self.assertEquals(trial_1.last_result["metric"], 500_000) + self.assertEquals(trial_1.last_result["hundred"], 1) + self.assertEquals(trial_2.last_result["metric"], 500_000) + self.assertEquals(trial_2.last_result["hundred"], 1) + + # With checkpoint dir parameter + def train(config, checkpoint_dir="DIR", data=None): + data.data[101] = 2 # Changes are local + tune.report(metric=len(data.data), cp=checkpoint_dir) + + trial_1, trial_2 = tune.run( + with_parameters(train, data=data), num_samples=2).trials + + self.assertEquals(data.data[101], 0) + self.assertEquals(trial_1.last_result["metric"], 500_000) + self.assertEquals(trial_1.last_result["cp"], "DIR") + self.assertEquals(trial_2.last_result["metric"], 500_000) + self.assertEquals(trial_2.last_result["cp"], "DIR")