Set a more descriptive error when a Run() call is cancelled due to session closure.

mrry · tensorflower-gardener · commit 75110127affb · 2017-02-16T12:47:28.000-08:00
Also fixes a potential memory leak, where the worker-side state of a failed Run() call
would not be cleaned up.
Change: 147752067
diff --git a/tensorflow/core/distributed_runtime/master.cc b/tensorflow/core/distributed_runtime/master.cc
@@ -101,8 +101,7 @@ void Master::GC() {
                        << "Note that if you are starting multiple replicas "
                        << "on a staggered delay, session_gc_seconds may need "
                        << "to be raised.";
-          sess->Close().IgnoreError();
-          sess->Unref();
+          sess->GarbageCollect();
         });
       }
     }
diff --git a/tensorflow/core/distributed_runtime/master_session.cc b/tensorflow/core/distributed_runtime/master_session.cc
@@ -1387,41 +1387,67 @@ Status MasterSession::DoRunWithLocalExecution(
     pss.collect_rpcs = ph->should_collect_rpcs();
   }
 
-  TF_RETURN_IF_ERROR(rcg->RunPartitions(env_, step_id, count,
-                                        execution_state_.get(), &pss, opts, req,
-                                        resp, cancellation_manager_, false));
+  Status s =
+      rcg->RunPartitions(env_, step_id, count, execution_state_.get(), &pss,
+                         opts, req, resp, cancellation_manager_, false);
+  if (s.ok()) {
+    pss.end_micros = Env::Default()->NowMicros();
 
-  pss.end_micros = Env::Default()->NowMicros();
-
-  // Schedule post-processing and cleanup to be done asynchronously.
+    // Schedule post-processing and cleanup to be done asynchronously.
+    rcg->ProcessStats(step_id, &pss, execution_state_.get(), ph.get(),
+                      req.options(), resp->mutable_metadata());
+  } else if (errors::IsCancelled(s)) {
+    mutex_lock l(mu_);
+    if (closed_) {
+      if (garbage_collected_) {
+        s = errors::Cancelled(
+            "Step was cancelled because the session was garbage collected due "
+            "to inactivity.");
+      } else {
+        s = errors::Cancelled(
+            "Step was cancelled by an explicit call to `Session::Close()`.");
+      }
+    }
+  }
   rcg->Ref();
-  rcg->ProcessStats(step_id, &pss, execution_state_.get(), ph.get(),
-                    req.options(), resp->mutable_metadata());
   rcg->CleanupPartitionsAsync(step_id, [rcg](const Status& s) {
     if (!s.ok()) {
       LOG(ERROR) << "Cleanup partition error: " << s;
     }
     rcg->Unref();
   });
-  return Status::OK();
+  return s;
 }
 
 Status MasterSession::Close() {
+  {
+    mutex_lock l(mu_);
+    closed_ = true;  // All subsequent calls to Run() or Extend() will fail.
+  }
   cancellation_manager_->StartCancel();
   std::vector<ReffedClientGraph*> to_unref;
   {
     mutex_lock l(mu_);
     while (num_running_ != 0) {
       num_running_is_zero_.wait(l);
     }
-    closed_ = true;  // All subsequent calls to Run() or Extend() will fail.
     ClearRunsTable(&to_unref, &run_graphs_);
     ClearRunsTable(&to_unref, &partial_run_graphs_);
   }
   for (ReffedClientGraph* rcg : to_unref) rcg->Unref();
   return Status::OK();
 }
 
+void MasterSession::GarbageCollect() {
+  {
+    mutex_lock l(mu_);
+    closed_ = true;
+    garbage_collected_ = true;
+  }
+  cancellation_manager_->StartCancel();
+  Unref();
+}
+
 MasterSession::RunState::RunState(const std::vector<string>& input_names,
                                   const std::vector<string>& output_names,
                                   ReffedClientGraph* rcg, const uint64 step_id,
diff --git a/tensorflow/core/distributed_runtime/master_session.h b/tensorflow/core/distributed_runtime/master_session.h
@@ -88,6 +88,12 @@ class MasterSession : public core::RefCounted {
   // Close() may block the caller thread for a long time.
   Status Close();
 
+  // Close this session and release a reference on "*this".
+  //
+  // Note that, unlike Close(), this method does not block on the
+  // completion of all work.
+  void GarbageCollect();
+
  private:
   SessionOptions session_opts_;
 
@@ -158,6 +164,7 @@ class MasterSession : public core::RefCounted {
   int32 num_running_ GUARDED_BY(mu_) = 0;
 
   bool closed_ GUARDED_BY(mu_) = false;
+  bool garbage_collected_ GUARDED_BY(mu_) = false;
 
   std::unordered_map<uint64, int64> subgraph_execution_counts_ GUARDED_BY(mu_);
 
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc
@@ -162,7 +162,7 @@ Status GrpcServer::Init() {
   builder.SetMaxMessageSize(std::numeric_limits<int32>::max());
   builder.SetOption(
       std::unique_ptr<::grpc::ServerBuilderOption>(new NoReusePortOption));
-  master_impl_.reset(new Master(&master_env_, 0.0));
+  master_impl_ = CreateMaster(&master_env_);
   master_service_ = NewGrpcMasterService(master_impl_.get(), &builder);
   worker_impl_.reset(NewGrpcWorker(&worker_env_));
   worker_service_ = NewGrpcWorkerService(worker_impl_.get(), &builder);
@@ -297,6 +297,10 @@ ChannelCreationFunction GrpcServer::GetChannelCreationFunction(
   return NewHostPortGrpcChannel;
 }
 
+std::unique_ptr<Master> GrpcServer::CreateMaster(MasterEnv* master_env) {
+  return std::unique_ptr<Master>(new Master(master_env, 0.0));
+}
+
 /* static */
 Status GrpcServer::Create(const ServerDef& server_def, Env* env,
                           std::unique_ptr<ServerInterface>* out_server) {
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h
@@ -63,6 +63,8 @@ class GrpcServer : public ServerInterface {
   virtual ChannelCreationFunction GetChannelCreationFunction(
       const ServerDef& server_def) const;
 
+  virtual std::unique_ptr<Master> CreateMaster(MasterEnv* master_env);
+
   // Returns the port to which this server is bound.
   // This method may only be called after `this->Init()` returns successfully.
   int bound_port() const { return bound_port_; }
diff --git a/tensorflow/python/training/server_lib_test.py b/tensorflow/python/training/server_lib_test.py
@@ -156,7 +156,8 @@ def testCloseCancelsBlockingOperation(self):
     sess.run(dequeue_t)
 
     def blocking_dequeue():
-      with self.assertRaises(errors_impl.CancelledError):
+      with self.assertRaisesRegexp(errors_impl.CancelledError,
+                                   "Session::Close"):
         sess.run(dequeue_t)
 
     blocking_thread = self.checkedThread(blocking_dequeue)

Original file line number	Diff line number	Diff line change
`@@ -101,8 +101,7 @@ void Master::GC() {`
`101`	`101`	`<< "Note that if you are starting multiple replicas "`
`102`	`102`	`<< "on a staggered delay, session_gc_seconds may need "`
`103`	`103`	`<< "to be raised.";`
`104`		`- sess->Close().IgnoreError();`
`105`		`- sess->Unref();`
	`104`	`+ sess->GarbageCollect();`
`106`	`105`	`});`
`107`	`106`	`}`
`108`	`107`	`}`