[tune] wrapper function to pass arbitrary objects through the object store to trainables (#10679)

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?
 ----------------------------

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)

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:
 

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"
 ]

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

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)

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")