Update tf.contrib.slim.evaluation.wait_for_new_checkpoint to continue waiting if

no checkpoint exists. Also add an optional timeout parameter.
Change: 127874450

2 files changed, 20 insertions, 3 deletions

diff --git a/tensorflow/contrib/slim/python/slim/evaluation.py b/tensorflow/contrib/slim/python/slim/evaluation.py
index 4668b030a7..e6314a9ce9 100644
--- a/tensorflow/contrib/slim/python/slim/evaluation.py
+++ b/tensorflow/contrib/slim/python/slim/evaluation.py
@@ -141,7 +141,8 @@ __all__ = ['evaluation', 'evaluation_loop', 'wait_for_new_checkpoint']
 
 def wait_for_new_checkpoint(checkpoint_dir,
                             last_checkpoint,
-                            seconds_to_sleep=1):
+                            seconds_to_sleep=1,
+                            timeout=None):
   """Waits until a new checkpoint file is found.
 
   Args:
@@ -149,13 +150,18 @@ def wait_for_new_checkpoint(checkpoint_dir,
     last_checkpoint: The last checkpoint path used.
     seconds_to_sleep: The number of seconds to sleep for before looking for a
       new checkpoint.
+    timeout: The maximum amount of time to wait. If left as `None`, then the
+      process will wait indefinitely.
 
   Returns:
-    a new checkpoint path.
+    a new checkpoint path, or None if the timeout was reached.
   """
+  stop_time = time.time() + timeout if timeout is not None else None
   while True:
     checkpoint_path = tf_saver.latest_checkpoint(checkpoint_dir)
-    if checkpoint_path == last_checkpoint:
+    if checkpoint_path is None or checkpoint_path == last_checkpoint:
+      if stop_time is not None and time.time() + seconds_to_sleep > stop_time:
+        return None
       time.sleep(seconds_to_sleep)
     else:
       return checkpoint_path
diff --git a/tensorflow/contrib/slim/python/slim/evaluation_test.py b/tensorflow/contrib/slim/python/slim/evaluation_test.py
index 077daefe83..e956742bcb 100644
--- a/tensorflow/contrib/slim/python/slim/evaluation_test.py
+++ b/tensorflow/contrib/slim/python/slim/evaluation_test.py
@@ -21,6 +21,7 @@ from __future__ import print_function
 
 import glob
 import os
+import time
 
 import numpy as np
 import tensorflow as tf
@@ -199,6 +200,16 @@ class EvaluationTest(tf.test.TestCase):
           sess, init_op=init_op, eval_op=update_op)
       self.assertAlmostEqual(accuracy.eval(), self._expected_accuracy)
 
+  def testLatestCheckpointReturnsNoneAfterTimeout(self):
+    start = time.time()
+    ret = slim.evaluation.wait_for_new_checkpoint(
+        '/non-existent-dir', 'foo', timeout=1.0, seconds_to_sleep=0.5)
+    end = time.time()
+    self.assertIsNone(ret)
+    # We've waited one time.
+    self.assertGreater(end, start + 0.5)
+    # The timeout kicked in.
+    self.assertLess(end, start + 1.1)
 
 if __name__ == '__main__':
   tf.test.main()