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

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