[tune] Initial track integration

python/ray/tune/__init__.py

@@ -14,5 +14,5 @@
 __all__ = [
     "Trainable", "TuneError", "grid_search", "register_env",
     "register_trainable", "run", "run_experiments", "Experiment", "function",
-    "sample_from", "uniform", "choice", "randint", "randn"
+    "sample_from", "track", "uniform", "choice", "randint", "randn"
 ]

python/ray/tune/automlboard/backend/collector.py

@@ -14,7 +14,7 @@
 from ray.tune.automlboard.models.models import JobRecord, \
     TrialRecord, ResultRecord
 from ray.tune.result import DEFAULT_RESULTS_DIR, JOB_META_FILE, \
-    EXPR_PARARM_FILE, EXPR_RESULT_FILE, EXPR_META_FILE
+    EXPR_PARAM_FILE, EXPR_RESULT_FILE, EXPR_META_FILE
 
 
 class CollectorService(object):
@@ -327,7 +327,7 @@ def _build_trial_meta(cls, expr_dir):
         if not meta:
             job_id = expr_dir.split("/")[-2]
             trial_id = expr_dir[-8:]
-            params = parse_json(os.path.join(expr_dir, EXPR_PARARM_FILE))
+            params = parse_json(os.path.join(expr_dir, EXPR_PARAM_FILE))
             meta = {
                 "trial_id": trial_id,
                 "job_id": job_id,
@@ -349,7 +349,7 @@ def _build_trial_meta(cls, expr_dir):
         if meta.get("end_time", None):
             meta["end_time"] = timestamp2date(meta["end_time"])
 
-        meta["params"] = parse_json(os.path.join(expr_dir, EXPR_PARARM_FILE))
+        meta["params"] = parse_json(os.path.join(expr_dir, EXPR_PARAM_FILE))
 
         return meta
 

python/ray/tune/examples/track_example.py

@@ -0,0 +1,71 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import keras
+from keras.datasets import mnist
+from keras.models import Sequential
+from keras.layers import (Dense, Dropout, Flatten, Conv2D, MaxPooling2D)
+
+from ray.tune import track
+from ray.tune.examples.utils import TuneKerasCallback, get_mnist_data
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--smoke-test", action="store_true", help="Finish quickly for testing")
+parser.add_argument(
+    "--lr",
+    type=float,
+    default=0.01,
+    metavar="LR",
+    help="learning rate (default: 0.01)")
+parser.add_argument(
+    "--momentum",
+    type=float,
+    default=0.5,
+    metavar="M",
+    help="SGD momentum (default: 0.5)")
+parser.add_argument(
+    "--hidden", type=int, default=64, help="Size of hidden layer.")
+args, _ = parser.parse_known_args()
+
+
+def train_mnist(args):
+    track.init(trial_name="track-example", trial_config=vars(args))
+    batch_size = 128
+    num_classes = 10
+    epochs = 1 if args.smoke_test else 12
+    mnist.load()
+    x_train, y_train, x_test, y_test, input_shape = get_mnist_data()
+
+    model = Sequential()
+    model.add(
+        Conv2D(
+            32, kernel_size=(3, 3), activation="relu",
+            input_shape=input_shape))
+    model.add(Conv2D(64, (3, 3), activation="relu"))
+    model.add(MaxPooling2D(pool_size=(2, 2)))
+    model.add(Dropout(0.5))
+    model.add(Flatten())
+    model.add(Dense(args.hidden, activation="relu"))
+    model.add(Dropout(0.5))
+    model.add(Dense(num_classes, activation="softmax"))
+
+    model.compile(
+        loss="categorical_crossentropy",
+        optimizer=keras.optimizers.SGD(lr=args.lr, momentum=args.momentum),
+        metrics=["accuracy"])
+
+    model.fit(
+        x_train,
+        y_train,
+        batch_size=batch_size,
+        epochs=epochs,
+        validation_data=(x_test, y_test),
+        callbacks=[TuneKerasCallback(track.metric)])
+    track.shutdown()
+
+
+if __name__ == "__main__":
+    train_mnist(args)

python/ray/tune/examples/utils.py

@@ -15,7 +15,9 @@ def __init__(self, reporter, logs={}):
 
     def on_train_end(self, epoch, logs={}):
         self.reporter(
-            timesteps_total=self.iteration, done=1, mean_accuracy=logs["acc"])
+            timesteps_total=self.iteration,
+            done=1,
+            mean_accuracy=logs.get("acc"))
 
     def on_batch_end(self, batch, logs={}):
         self.iteration += 1

python/ray/tune/function_runner.py

@@ -5,9 +5,11 @@
 import logging
 import sys
 import time
+import inspect
 import threading
 from six.moves import queue
 
+from ray.tune import track
 from ray.tune import TuneError
 from ray.tune.trainable import Trainable
 from ray.tune.result import TIME_THIS_ITER_S, RESULT_DUPLICATE
@@ -244,6 +246,17 @@ def _report_thread_runner_error(self, block=False):
 
 
 def wrap_function(train_func):
+
+    use_track = False
+    try:
+        func_args = inspect.getargspec(train_func).args
+        use_track = ("reporter" not in func_args and len(func_args) == 1)
+        if use_track:
+            logger.info("tune.track signature detected.")
+    except Exception:
+        logger.info(
+            "Function inspection failed - assuming reporter signature.")
+
     class WrappedFunc(FunctionRunner):
         def _trainable_func(self, config, reporter):
             output = train_func(config, reporter)
@@ -253,4 +266,12 @@ def _trainable_func(self, config, reporter):
             reporter(**{RESULT_DUPLICATE: True})
             return output
 
-    return WrappedFunc
+    class WrappedTrackFunc(FunctionRunner):
+        def _trainable_func(self, config, reporter):
+            track.init(_tune_reporter=reporter)
+            output = train_func(config)
+            reporter(**{RESULT_DUPLICATE: True})
+            track.shutdown()
+            return output
+
+    return WrappedTrackFunc if use_track else WrappedFunc

python/ray/tune/logger.py

@@ -50,6 +50,11 @@ def on_result(self, result):
 
         raise NotImplementedError
 
+    def update_config(self, config):
+        """Updates the config for all loggers."""
+
+        pass
+
     def close(self):
         """Releases all resources used by this logger."""
 
@@ -68,17 +73,7 @@ def on_result(self, result):
 
 class JsonLogger(Logger):
     def _init(self):
-        config_out = os.path.join(self.logdir, "params.json")
-        with open(config_out, "w") as f:
-            json.dump(
-                self.config,
-                f,
-                indent=2,
-                sort_keys=True,
-                cls=_SafeFallbackEncoder)
-        config_pkl = os.path.join(self.logdir, "params.pkl")
-        with open(config_pkl, "wb") as f:
-            cloudpickle.dump(self.config, f)
+        self.update_config(self.config)
         local_file = os.path.join(self.logdir, "result.json")
         self.local_out = open(local_file, "a")
 
@@ -96,6 +91,15 @@ def flush(self):
     def close(self):
         self.local_out.close()
 
+    def update_config(self, config):
+        self.config = config
+        config_out = os.path.join(self.logdir, "params.json")
+        with open(config_out, "w") as f:
+            json.dump(self.config, f, cls=_SafeFallbackEncoder)
+        config_pkl = os.path.join(self.logdir, "params.pkl")
+        with open(config_pkl, "wb") as f:
+            cloudpickle.dump(self.config, f)
+
 
 def to_tf_values(result, path):
     values = []
@@ -231,6 +235,10 @@ def on_result(self, result):
         self._log_syncer.set_worker_ip(result.get(NODE_IP))
         self._log_syncer.sync_if_needed()
 
+    def update_config(self, config):
+        for _logger in self._loggers:
+            _logger.update_config(config)
+
     def close(self):
         for _logger in self._loggers:
             _logger.close()

python/ray/tune/result.py

@@ -68,7 +68,7 @@
 EXPR_META_FILE = "trial_status.json"
 
 # File that stores parameters of the trial.
-EXPR_PARARM_FILE = "params.json"
+EXPR_PARAM_FILE = "params.json"
 
 # File that stores the progress of the trial.
 EXPR_PROGRESS_FILE = "progress.csv"

python/ray/tune/tests/test_track.py

@@ -0,0 +1,84 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import pandas as pd
+import unittest
+
+import ray
+from ray import tune
+from ray.tune import track
+from ray.tune.result import EXPR_PARAM_FILE, EXPR_RESULT_FILE
+
+
+def _check_json_val(fname, key, val):
+    with open(fname, "r") as f:
+        df = pd.read_json(f, typ="frame", lines=True)
+        return key in df.columns and (df[key].tail(n=1) == val).all()
+
+
+class TrackApiTest(unittest.TestCase):
+    def tearDown(self):
+        track.shutdown()
+        ray.shutdown()
+
+    def testSessionInitShutdown(self):
+        self.assertTrue(track._session is None)
+
+        # Checks that the singleton _session is created/destroyed
+        # by track.init() and track.shutdown()
+        for _ in range(2):
+            # do it twice to see that we can reopen the session
+            track.init(trial_name="test_init")
+            self.assertTrue(track._session is not None)
+            track.shutdown()
+            self.assertTrue(track._session is None)
+
+    def testLogCreation(self):
+        """Checks that track.init() starts logger and creates log files."""
+        track.init(trial_name="test_init")
+        session = track.get_session()
+        self.assertTrue(session is not None)
+
+        self.assertTrue(os.path.isdir(session.logdir))
+
+        params_path = os.path.join(session.logdir, EXPR_PARAM_FILE)
+        result_path = os.path.join(session.logdir, EXPR_RESULT_FILE)
+
+        self.assertTrue(os.path.exists(params_path))
+        self.assertTrue(os.path.exists(result_path))
+        self.assertTrue(session.logdir == track.trial_dir())
+
+    def testMetric(self):
+        track.init(trial_name="test_log")
+        session = track.get_session()
+        for i in range(5):
+            track.log(test=i)
+        result_path = os.path.join(session.logdir, EXPR_RESULT_FILE)
+        self.assertTrue(_check_json_val(result_path, "test", i))
+
+    def testRayOutput(self):
+        """Checks that local and remote log format are the same."""
+        ray.init()
+
+        def testme(config):
+            for i in range(config["iters"]):
+                track.log(iteration=i, hi="test")
+
+        trials = tune.run(testme, config={"iters": 5})
+        trial_res = trials[0].last_result
+        self.assertTrue(trial_res["hi"], "test")
+        self.assertTrue(trial_res["training_iteration"], 5)
+
+    def testLocalMetrics(self):
+        """Checks that metric state is updated correctly."""
+        track.init(trial_name="test_logs")
+        session = track.get_session()
+        self.assertEqual(set(session.trial_config.keys()), {"trial_id"})
+
+        result_path = os.path.join(session.logdir, EXPR_RESULT_FILE)
+        track.log(test=1)
+        self.assertTrue(_check_json_val(result_path, "test", 1))
+        track.log(iteration=1, test=2)
+        self.assertTrue(_check_json_val(result_path, "test", 2))

python/ray/tune/track/__init__.py

@@ -0,0 +1,71 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+
+from ray.tune.track.session import TrackSession
+
+logger = logging.getLogger(__name__)
+
+_session = None
+
+
+def get_session():
+    global _session
+    if not _session:
+        raise ValueError("Session not detected. Try `track.init()`?")
+    return _session
+
+
+def init(ignore_reinit_error=True, **session_kwargs):
+    """Initializes the global trial context for this process.
+
+    This creates a TrackSession object and the corresponding hooks for logging.
+
+    Examples:
+        >>> from ray.tune import track
+        >>> track.init()
+    """
+    global _session
+
+    if _session:
+        # TODO(ng): would be nice to stack crawl at creation time to report
+        # where that initial trial was created, and that creation line
+        # info is helpful to keep around anyway.
+        reinit_msg = "A session already exists in the current context."
+        if ignore_reinit_error:
+            if not _session.is_tune_session:
+                logger.warning(reinit_msg)
+            return
+        else:
+            raise ValueError(reinit_msg)
+
+    _session = TrackSession(**session_kwargs)
+
+
+def shutdown():
+    """Cleans up the trial and removes it from the global context."""
+
+    global _session
+    if _session:
+        _session.close()
+    _session = None
+
+
+def log(**kwargs):
+    """Applies TrackSession.log to the trial in the current context."""
+    _session = get_session()
+    return _session.log(**kwargs)
+
+
+def trial_dir():
+    """Returns the directory where trial results are saved.
+
+    This includes json data containing the session's parameters and metrics.
+    """
+    _session = get_session()
+    return _session.logdir
+
+
+__all__ = ["TrackSession", "session", "log", "trial_dir", "init", "shutdown"]