diff --git a/python/ray/tune/registry.py b/python/ray/tune/registry.py index 204d79c93..212ac031b 100644 --- a/python/ray/tune/registry.py +++ b/python/ray/tune/registry.py @@ -2,6 +2,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import logging from types import FunctionType import ray @@ -17,6 +18,8 @@ KNOWN_CATEGORIES = [ TRAINABLE_CLASS, ENV_CREATOR, RLLIB_MODEL, RLLIB_PREPROCESSOR ] +logger = logging.getLogger(__name__) + def register_trainable(name, trainable): """Register a trainable function or class. @@ -30,8 +33,16 @@ def register_trainable(name, trainable): from ray.tune.trainable import Trainable, wrap_function - if isinstance(trainable, FunctionType): + if isinstance(trainable, type): + logger.debug("Detected class for trainable.") + elif isinstance(trainable, FunctionType): + logger.debug("Detected function for trainable.") trainable = wrap_function(trainable) + elif callable(trainable): + logger.warning( + "Detected unknown callable for trainable. Converting to class.") + trainable = wrap_function(trainable) + if not issubclass(trainable, Trainable): raise TypeError("Second argument must be convertable to Trainable", trainable) diff --git a/python/ray/tune/test/trial_runner_test.py b/python/ray/tune/test/trial_runner_test.py index a92c3b0f1..ab63f1286 100644 --- a/python/ray/tune/test/trial_runner_test.py +++ b/python/ray/tune/test/trial_runner_test.py @@ -112,6 +112,24 @@ class TrainableFunctionApiTest(unittest.TestCase): self.assertRaises(TypeError, lambda: register_trainable("foo", B())) self.assertRaises(TypeError, lambda: register_trainable("foo", A)) + def testRegisterTrainableCallable(self): + def dummy_fn(config, reporter, steps): + reporter(timesteps_total=steps, done=True) + + from functools import partial + steps = 500 + register_trainable("test", partial(dummy_fn, steps=steps)) + [trial] = run_experiments({ + "foo": { + "run": "test", + "config": { + "script_min_iter_time_s": 0, + }, + } + }) + self.assertEqual(trial.status, Trial.TERMINATED) + self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], steps) + def testBuiltInTrainableResources(self): class B(Trainable): @classmethod