From d81605e9e70306737c30b726085a2228ee251f3b Mon Sep 17 00:00:00 2001
From: Eric Liang <ekhliang@gmail.com>
Date: Wed, 5 Sep 2018 17:45:09 -0700
Subject: [PATCH] [tune] Add a time/timesteps since last restore metric (#2819)

* rsm

* always log to avoid changing schema for csv writer

* add iter since restore

* update

* criteria warn
---
 python/ray/tune/logger.py                 |  1 +
 python/ray/tune/test/trial_runner_test.py | 32 +++++++++++++++++++++++
 python/ray/tune/trainable.py              | 13 ++++++++-
 python/ray/tune/trial.py                  |  4 ++-
 4 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/python/ray/tune/logger.py b/python/ray/tune/logger.py
index c1ffbcafe..f9b76240c 100644
--- a/python/ray/tune/logger.py
+++ b/python/ray/tune/logger.py
@@ -140,6 +140,7 @@ class _TFLogger(Logger):
         }, ["ray", "tune"])
         iteration_stats = tf.Summary(value=iteration_value)
         self._file_writer.add_summary(iteration_stats, t)
+        self._file_writer.flush()
 
     def flush(self):
         self._file_writer.flush()
diff --git a/python/ray/tune/test/trial_runner_test.py b/python/ray/tune/test/trial_runner_test.py
index e34137a1d..f3926e71f 100644
--- a/python/ray/tune/test/trial_runner_test.py
+++ b/python/ray/tune/test/trial_runner_test.py
@@ -1009,6 +1009,38 @@ class TrialRunnerTest(unittest.TestCase):
         self.assertEqual(ray.get(trials[1].runner.get_info.remote()), 1)
         self.addCleanup(os.remove, path)
 
+    def testRestoreMetricsAfterCheckpointing(self):
+        ray.init(num_cpus=1, num_gpus=1)
+        runner = TrialRunner(BasicVariantGenerator())
+        kwargs = {
+            "resources": Resources(cpu=1, gpu=1),
+        }
+        runner.add_trial(Trial("__fake", **kwargs))
+        trials = runner.get_trials()
+
+        runner.step()
+        self.assertEqual(trials[0].status, Trial.RUNNING)
+        self.assertEqual(ray.get(trials[0].runner.set_info.remote(1)), 1)
+        path = runner.trial_executor.save(trials[0])
+        runner.trial_executor.stop_trial(trials[0])
+        kwargs["restore_path"] = path
+
+        runner.add_trial(Trial("__fake", **kwargs))
+        trials = runner.get_trials()
+
+        runner.step()
+        self.assertEqual(trials[0].status, Trial.TERMINATED)
+        self.assertEqual(trials[1].status, Trial.RUNNING)
+        runner.step()
+        self.assertEqual(trials[1].last_result["timesteps_since_restore"], 10)
+        self.assertEqual(trials[1].last_result["iterations_since_restore"], 1)
+        self.assertGreater(trials[1].last_result["time_since_restore"], 0)
+        runner.step()
+        self.assertEqual(trials[1].last_result["timesteps_since_restore"], 20)
+        self.assertEqual(trials[1].last_result["iterations_since_restore"], 2)
+        self.assertGreater(trials[1].last_result["time_since_restore"], 0)
+        self.addCleanup(os.remove, path)
+
     def testCheckpointingAtEnd(self):
         ray.init(num_cpus=1, num_gpus=1)
         runner = TrialRunner(BasicVariantGenerator())
diff --git a/python/ray/tune/trainable.py b/python/ray/tune/trainable.py
index 5ff0b5187..898938d5e 100644
--- a/python/ray/tune/trainable.py
+++ b/python/ray/tune/trainable.py
@@ -75,6 +75,10 @@ class Trainable(object):
         self._iteration = 0
         self._time_total = 0.0
         self._timesteps_total = None
+        self._time_since_restore = 0.0
+        self._timesteps_since_restore = 0
+        self._iterations_since_restore = 0
+        self._restored = False
         self._setup()
         self._initialize_ok = True
         self._local_ip = ray.services.get_node_ip_address()
@@ -143,12 +147,14 @@ class Trainable(object):
         result = result.copy()
 
         self._iteration += 1
+        self._iterations_since_restore += 1
 
         if result.get(TIME_THIS_ITER_S) is not None:
             time_this_iter = result[TIME_THIS_ITER_S]
         else:
             time_this_iter = time.time() - start
         self._time_total += time_this_iter
+        self._time_since_restore += time_this_iter
 
         result.setdefault(DONE, False)
 
@@ -157,6 +163,7 @@ class Trainable(object):
             if self._timesteps_total is None:
                 self._timesteps_total = 0
             self._timesteps_total += result[TIMESTEPS_THIS_ITER]
+            self._timesteps_since_restore += result[TIMESTEPS_THIS_ITER]
 
         # self._timesteps_total should not override user-provided total
         result.setdefault(TIMESTEPS_TOTAL, self._timesteps_total)
@@ -176,7 +183,10 @@ class Trainable(object):
             pid=os.getpid(),
             hostname=os.uname()[1],
             node_ip=self._local_ip,
-            config=self.config)
+            config=self.config,
+            time_since_restore=self._time_since_restore,
+            timesteps_since_restore=self._timesteps_since_restore,
+            iterations_since_restore=self._iterations_since_restore)
 
         self._result_logger.on_result(result)
 
@@ -248,6 +258,7 @@ class Trainable(object):
         self._iteration = metadata[1]
         self._timesteps_total = metadata[2]
         self._time_total = metadata[3]
+        self._restored = True
 
     def restore_from_object(self, obj):
         """Restores training state from a checkpoint object.
diff --git a/python/ray/tune/trial.py b/python/ray/tune/trial.py
index 4ee4d9171..0008f2a2d 100644
--- a/python/ray/tune/trial.py
+++ b/python/ray/tune/trial.py
@@ -199,7 +199,9 @@ class Trial(object):
 
         for criteria, stop_value in self.stopping_criterion.items():
             if criteria not in result:
-                raise TuneError("Stopping Criteria not provided in result.")
+                raise TuneError(
+                    "Stopping criteria {} not provided in result {}.".format(
+                        criteria, result))
             if result[criteria] >= stop_value:
                 return True