Enable reuse of a DirectSession after a run times out.

mrry · tensorflower-gardener · commit f46fe646a26f · 2016-11-09T13:48:52.000-08:00
Fixes tensorflow#5115. Change: 138676008
diff --git a/tensorflow/core/common_runtime/direct_session.cc b/tensorflow/core/common_runtime/direct_session.cc
@@ -407,6 +407,7 @@ Status DirectSession::Run(const RunOptions& run_options,
   // Create a run state and start execution.
   RunState run_state(input_tensor_names, output_names);
   run_state.rendez = new IntraProcessRendezvous(device_mgr_.get());
+  CancellationManager step_cancellation_manager;
 
   // Send inputs.
   TF_RETURN_IF_ERROR(SendInputs(inputs, executors_and_keys, run_state.rendez));
@@ -425,7 +426,7 @@ Status DirectSession::Run(const RunOptions& run_options,
   Executor::Args args;
   args.step_id = step_id_counter_.fetch_add(1);
   args.rendezvous = run_state.rendez;
-  args.cancellation_manager = cancellation_manager_;
+  args.cancellation_manager = &step_cancellation_manager;
   args.runner = [this, pool](Executor::Args::Closure c) {
     SchedClosure(pool, std::move(c));
   };
@@ -464,13 +465,33 @@ Status DirectSession::Run(const RunOptions& run_options,
   }
 #endif  // GOOGLE_CUDA
 
+  // Register this step with session's cancellation manager, so that
+  // `Session::Close()` will cancel the step.
+  CancellationToken cancellation_token =
+      cancellation_manager_->get_cancellation_token();
+  bool already_cancelled = !cancellation_manager_->RegisterCallback(
+      cancellation_token, [&step_cancellation_manager]() {
+        step_cancellation_manager.StartCancel();
+      });
+  if (already_cancelled) {
+    // NOTE(mrry): If we don't explicitly notify
+    // `run_state.executors_done`, the RunState destructor would
+    // block on this notification.
+    run_state.executors_done.Notify();
+    delete barrier;
+    return errors::Cancelled("Run call was cancelled");
+  }
+
   for (const auto& item : executors_and_keys->items) {
     item.executor->RunAsync(args, barrier->Get());
   }
 
-  WaitForNotification(&run_state, run_options.timeout_in_ms() > 0
-                                      ? run_options.timeout_in_ms()
-                                      : operation_timeout_in_ms_);
+  WaitForNotification(&run_state, &step_cancellation_manager,
+                      run_options.timeout_in_ms() > 0
+                          ? run_options.timeout_in_ms()
+                          : operation_timeout_in_ms_);
+
+  cancellation_manager_->DeregisterCallback(cancellation_token);
 
 #if GOOGLE_CUDA
   if (tracer) {
@@ -687,7 +708,8 @@ Status DirectSession::PRun(const string& handle, const NamedTensorList& inputs,
               run_state->pending_outputs.size() == 0);
     }
     if (done) {
-      WaitForNotification(run_state, operation_timeout_in_ms_);
+      WaitForNotification(run_state, cancellation_manager_,
+                          operation_timeout_in_ms_);
       partial_runs_.erase(handle);
     }
   }
@@ -1158,6 +1180,7 @@ DirectSession::RunState::~RunState() {
 }
 
 void DirectSession::WaitForNotification(RunState* run_state,
+                                        CancellationManager* cm,
                                         int64 timeout_in_ms) {
   if (timeout_in_ms > 0) {
     bool notified = WaitForNotificationWithTimeout(&run_state->executors_done,
@@ -1168,12 +1191,7 @@ void DirectSession::WaitForNotification(RunState* run_state,
         run_state->status.Update(Status(error::DEADLINE_EXCEEDED,
                                         "Timed out waiting for notification"));
       }
-      // TODO(sherrym): This cancels all steps in the session, even ones that
-      // have not exceeded their deadline. An alternative would be to use a
-      // two-level cancellation manager with a Session-global one containing
-      // several step-local ones. Probably the RunState should have its own
-      // CancellationManager.
-      cancellation_manager_->StartCancel();
+      cm->StartCancel();
     }
   } else {
     run_state->executors_done.WaitForNotification();
diff --git a/tensorflow/core/common_runtime/direct_session.h b/tensorflow/core/common_runtime/direct_session.h
@@ -209,7 +209,10 @@ class DirectSession : public Session {
 
   // Use the appropriate WaitForNotification function based on whether
   // operation_timeout_in_ms is greater than 0.
-  void WaitForNotification(RunState* run_state, int64 timeout_in_ms);
+  //
+  // If the timeout expires, the `cm->StartCancel()` will be called.
+  void WaitForNotification(RunState* run_state, CancellationManager* cm,
+                           int64 timeout_in_ms);
 
   ::tensorflow::Status CheckNotClosed() {
     mutex_lock l(closed_lock_);
diff --git a/tensorflow/python/kernel_tests/fifo_queue_test.py b/tensorflow/python/kernel_tests/fifo_queue_test.py
@@ -1512,6 +1512,23 @@ def testDequeueWithTimeout(self):
                                    "Timed out waiting for notification"):
         sess.run(dequeued_t)
 
+  def testReusableAfterTimeout(self):
+    with self.test_session() as sess:
+      q = tf.FIFOQueue(10, tf.float32)
+      dequeued_t = q.dequeue()
+      enqueue_op = q.enqueue(37)
+
+      with self.assertRaisesRegexp(tf.errors.DeadlineExceededError,
+                                   "Timed out waiting for notification"):
+        sess.run(dequeued_t, options=tf.RunOptions(timeout_in_ms=10))
+
+      with self.assertRaisesRegexp(tf.errors.DeadlineExceededError,
+                                   "Timed out waiting for notification"):
+        sess.run(dequeued_t, options=tf.RunOptions(timeout_in_ms=10))
+
+      sess.run(enqueue_op)
+      self.assertEqual(37, sess.run(dequeued_t))
+
 
 class QueueContainerTest(tf.test.TestCase):
 
