Add ConfigProto.isolate_session_state option for the distributed runtime.

Setting this option to true when creating a session ensures that no stateful resources (variables, queues, iterators, etc.) will be visible to any other session running on the same server, and those resources will be deleted when the session is closed. The default behavior, namely that all `tf.Variable` objects are shared by default and most other resources are shared when their `shared_name` attr is non-empty, is preserved. This change augments the semantics of the WorkerService.CreateWorkerSession RPC. Now, if the server_def in the request is empty, it implies that the worker should use its default ClusterSpec. Note that clusters created using ClusterSpec propagation always have isolated session state, and are unaffected by this change. PiperOrigin-RevId: 177173545

Author: mrry

tensorflow/core/common_runtime/device.h

@@ -131,7 +131,7 @@ class Device : public DeviceBase {
  OpSegment* op_segment() { return &op_seg_; }
 
  // Returns the resource manager associated w/ this device.
- ResourceMgr* resource_manager() { return rmgr_; }
+ virtual ResourceMgr* resource_manager() { return rmgr_; }
 
  // Summarizes the status of this Device, for debugging.
  string DebugString() const { return ProtoDebugString(device_attributes_); }

tensorflow/core/common_runtime/renamed_device.cc

@@ -21,7 +21,8 @@ namespace tensorflow {
 /* static */
 Device* RenamedDevice::NewRenamedDevice(const string& new_base,
  Device* underlying,
- bool owns_underlying) {
+ bool owns_underlying,
+ bool isolate_session_state) {
  DeviceNameUtils::ParsedName parsed_name;
  CHECK(DeviceNameUtils::ParseFullName(new_base, &parsed_name));
  DeviceNameUtils::ParsedName underlying_parsed_name =
@@ -35,15 +36,17 @@ Device* RenamedDevice::NewRenamedDevice(const string& new_base,
  parsed_name.id);
  DeviceAttributes attributes(underlying->attributes());
  attributes.set_name(name);
- return new RenamedDevice(underlying, attributes, owns_underlying);
+ return new RenamedDevice(underlying, attributes, owns_underlying,
+ isolate_session_state);
 }
 
 RenamedDevice::RenamedDevice(Device* underlying,
  const DeviceAttributes& attributes,
- bool owns_underlying)
+ bool owns_underlying, bool isolate_session_state)
  : Device(underlying->env(), attributes),
  underlying_(underlying),
- owns_underlying_(owns_underlying) {}
+ owns_underlying_(owns_underlying),
+ isolate_session_state_(isolate_session_state) {}
 
 RenamedDevice::~RenamedDevice() {
  if (owns_underlying_) {

tensorflow/core/common_runtime/renamed_device.h

@@ -29,7 +29,9 @@ namespace tensorflow {
 class RenamedDevice : public Device {
  public:
  static Device* NewRenamedDevice(const string& new_base, Device* underlying,
- bool owns_underlying);
+ bool owns_underlying,
+ bool isolate_session_state);
+
  ~RenamedDevice() override;
 
  // Below are virtual methods defined on DeviceBase
@@ -113,11 +115,21 @@ class RenamedDevice : public Device {
  return underlying_->FillContextMap(graph, device_context_map);
  }
 
+ // Returns the resource manager associated w/ this device.
+ ResourceMgr* resource_manager() override {
+ if (isolate_session_state_) {
+ return Device::resource_manager();
+ } else {
+ return underlying_->resource_manager();
+ }
+ }
+
  private:
  RenamedDevice(Device* underlying, const DeviceAttributes& attributes,
- bool owns_underlying);
+ bool owns_underlying, bool isolate_session_state);
  Device* const underlying_;
  const bool owns_underlying_;
+ const bool isolate_session_state_;
 };
 
 } // namespace tensorflow

tensorflow/core/distributed_runtime/BUILD

@@ -140,6 +140,7 @@ cc_library(
  hdrs = ["session_mgr.h"],
  deps = [
  ":graph_mgr",
+ ":worker_cache_wrapper",
  ":worker_session",
  "//tensorflow/core:core_cpu_internal",
  "//tensorflow/core:lib",
@@ -263,6 +264,17 @@ cc_library(
  ],
 )
 
+cc_library(
+ name = "worker_cache_wrapper",
+ hdrs = ["worker_cache_wrapper.h"],
+ deps = [
+ ":worker_cache",
+ ":worker_interface",
+ "//tensorflow/core:lib",
+ "//tensorflow/core:protos_all_cc",
+ ],
+)
+
 cc_library(
  name = "remote_device",
  srcs = ["remote_device.cc"],

tensorflow/core/distributed_runtime/master_session.cc

@@ -1049,7 +1049,10 @@ Status MasterSession::Create(GraphDef* graph_def,
  TF_RETURN_IF_ERROR(GraphExecutionState::MakeForBaseGraph(
  graph_def, execution_options, &execution_state_));
  }
- if (options.cluster_def != nullptr) {
+ // TODO(b/36574172): Remove these conditions when ClusterSpec
+ // propagation is supported in all servers.
+ if (options.cluster_def != nullptr ||
+ session_opts_.config.isolate_session_state()) {
  should_delete_worker_sessions_ = true;
  return CreateWorkerSessions(options);
  }
@@ -1058,10 +1061,9 @@ Status MasterSession::Create(GraphDef* graph_def,
 
 Status MasterSession::CreateWorkerSessions(
  const WorkerCacheFactoryOptions& options) {
- CHECK(worker_cache_) << "CreateWorkerSessions should be called only with "
- << "dynamic cluster membership.";
  std::vector<string> worker_names;
- worker_cache_->ListWorkers(&worker_names);
+ WorkerCacheInterface* worker_cache = get_worker_cache();
+ worker_cache->ListWorkers(&worker_names);
 
  struct WorkerGroup {
  // The worker name. (Not owned.)
@@ -1079,10 +1081,10 @@ Status MasterSession::CreateWorkerSessions(
  std::vector<WorkerGroup> workers(worker_names.size());
 
  // Release the workers.
- auto cleanup = gtl::MakeCleanup([this, &workers] {
+ auto cleanup = gtl::MakeCleanup([this, &workers, worker_cache] {
  for (auto&& worker_group : workers) {
  if (worker_group.worker != nullptr) {
- worker_cache_->ReleaseWorker(*worker_group.name, worker_group.worker);
+ worker_cache->ReleaseWorker(*worker_group.name, worker_group.worker);
  }
  }
  });
@@ -1091,11 +1093,19 @@ Status MasterSession::CreateWorkerSessions(
  // Create all the workers & kick off the computations.
  for (size_t i = 0; i < worker_names.size(); ++i) {
  workers[i].name = &worker_names[i];
- workers[i].worker = worker_cache_->CreateWorker(worker_names[i]);
+ workers[i].worker = worker_cache->CreateWorker(worker_names[i]);
  workers[i].request.set_session_handle(handle_);
- *workers[i].request.mutable_server_def()->mutable_cluster() =
- *options.cluster_def;
- workers[i].request.mutable_server_def()->set_protocol(*options.protocol);
+ if (options.cluster_def) {
+ *workers[i].request.mutable_server_def()->mutable_cluster() =
+ *options.cluster_def;
+ workers[i].request.mutable_server_def()->set_protocol(*options.protocol);
+ // Session state is always isolated when ClusterSpec propagation
+ // is in use.
+ workers[i].request.set_isolate_session_state(true);
+ } else {
+ workers[i].request.set_isolate_session_state(
+ session_opts_.config.isolate_session_state());
+ }
 
  DeviceNameUtils::ParsedName name;
  if (!DeviceNameUtils::ParseFullName(worker_names[i], &name)) {
@@ -1162,7 +1172,7 @@ Status MasterSession::DeleteWorkerSessions() {
  // Create all the workers & kick off the computations.
  for (size_t i = 0; i < worker_names.size(); ++i) {
  workers[i].name = &worker_names[i];
- workers[i].worker = worker_cache_->CreateWorker(worker_names[i]);
+ workers[i].worker = worker_cache->CreateWorker(worker_names[i]);
  workers[i].request.set_session_handle(handle_);
  }
 

tensorflow/core/distributed_runtime/session_mgr.cc

@@ -20,7 +20,10 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/device_mgr.h"
 #include "tensorflow/core/common_runtime/renamed_device.h"
 #include "tensorflow/core/distributed_runtime/graph_mgr.h"
+#include "tensorflow/core/distributed_runtime/worker_cache_wrapper.h"
 #include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/protobuf/cluster.pb.h"
+#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
 
 namespace tensorflow {
 
@@ -29,7 +32,10 @@ SessionMgr::SessionMgr(
  std::unique_ptr<WorkerCacheInterface> default_worker_cache,
  WorkerCacheFactory worker_cache_factory)
  : worker_env_(worker_env),
- legacy_session_("", default_worker_name, std::move(default_worker_cache),
+ default_worker_cache_(std::move(default_worker_cache)),
+ legacy_session_("", default_worker_name,
+ std::unique_ptr<WorkerCacheInterface>(
+ new WorkerCacheWrapper(default_worker_cache_.get())),
  std::unique_ptr<DeviceMgr>(worker_env->device_mgr),
  std::unique_ptr<GraphMgr>(
  new GraphMgr(worker_env, worker_env->device_mgr))),
@@ -41,7 +47,8 @@ string SessionMgr::WorkerNameFromServerDef(const ServerDef& server_def) {
 }
 
 Status SessionMgr::CreateSession(const string& session,
- const ServerDef& server_def) {
+ const ServerDef& server_def,
+ bool isolate_session_state) {
  mutex_lock l(mu_);
  if (session.empty()) {
  return errors::InvalidArgument("Session must be non-empty.");
@@ -50,12 +57,18 @@ Status SessionMgr::CreateSession(const string& session,
  const string worker_name = WorkerNameFromServerDef(server_def);
 
  WorkerCacheInterface* worker_cache = nullptr;
- TF_RETURN_IF_ERROR(worker_cache_factory_(server_def, &worker_cache));
+ if (server_def.cluster().job().empty()) {
+ worker_cache = new WorkerCacheWrapper(default_worker_cache_.get());
+ } else {
+ TF_RETURN_IF_ERROR(worker_cache_factory_(server_def, &worker_cache));
+ }
 
+ CHECK(!worker_env_->local_devices.empty())
+ << "The WorkerEnv must have at least one device in `local_devices`.";
  std::vector<Device*> renamed_devices;
  for (Device* d : worker_env_->local_devices) {
- renamed_devices.push_back(
- RenamedDevice::NewRenamedDevice(worker_name, d, false));
+ renamed_devices.push_back(RenamedDevice::NewRenamedDevice(
+ worker_name, d, false, isolate_session_state));
  }
  std::unique_ptr<DeviceMgr> device_mgr(new DeviceMgr(renamed_devices));
 

tensorflow/core/distributed_runtime/session_mgr.h

@@ -45,7 +45,8 @@ class SessionMgr {
  ~SessionMgr() {}
 
  // Allocates state for a new session.
- Status CreateSession(const string& session, const ServerDef& server_def);
+ Status CreateSession(const string& session, const ServerDef& server_def,
+ bool isolate_session_state);
 
  // Locates the worker session for a given session handle
  WorkerSession* WorkerSessionForSession(const string& session);
@@ -71,6 +72,7 @@ class SessionMgr {
  // legacy_session_ is deleted. Further, we must ensure that WorkerSession's
  // device_mgr is deleted after WorkerSession's graph_mgr.
 
+ std::unique_ptr<WorkerCacheInterface> default_worker_cache_;
  WorkerSession legacy_session_;
 
  const WorkerCacheFactory worker_cache_factory_;

tensorflow/core/distributed_runtime/session_mgr_test.cc

@@ -22,14 +22,36 @@ limitations under the License.
 
 namespace tensorflow {
 
+class FakeDevice : public Device {
+ private:
+ explicit FakeDevice(const DeviceAttributes& device_attributes)
+ : Device(nullptr, device_attributes) {}
+
+ public:
+ Status Sync() override { return errors::Unimplemented("FakeDevice::Sync()"); }
+
+ Allocator* GetAllocator(AllocatorAttributes attr) override { return nullptr; }
+
+ static std::unique_ptr<Device> MakeCPU(const string& name) {
+ DeviceAttributes device_attributes;
+ device_attributes.set_name(name);
+ device_attributes.set_device_type(DeviceType("FakeCPU").type());
+ return std::unique_ptr<Device>(new FakeDevice(device_attributes));
+ }
+};
+
 class SessionMgrTest : public ::testing::Test {
  protected:
  SessionMgrTest()
- : mgr_(&env_, "/job:mnist/replica:0/task:0",
- std::unique_ptr<WorkerCacheInterface>(),
- factory_),
- legacy_session_(mgr_.WorkerSessionForSession("novel_session_id")) {}
+ : device_(FakeDevice::MakeCPU(
+ "/job:mnist/replica:0/task:0/device:fakecpu:0")),
+ mgr_(&env_, "/job:mnist/replica:0/task:0",
+ std::unique_ptr<WorkerCacheInterface>(), factory_),
+ legacy_session_(mgr_.WorkerSessionForSession("novel_session_id")) {
+ env_.local_devices = {device_.get()};
+ }
 
+ std::unique_ptr<Device> device_;
  WorkerEnv env_;
  SessionMgr::WorkerCacheFactory factory_ =
  [](const ServerDef& server_def, WorkerCacheInterface** worker_cache) {
@@ -42,14 +64,48 @@ class SessionMgrTest : public ::testing::Test {
 
 TEST_F(SessionMgrTest, CreateSessionSimple) {
  ServerDef server_def;
+ server_def.set_job_name("worker");
+ server_def.set_task_index(3);
+
  string session_handle = "test_session_handle";
- TF_EXPECT_OK(mgr_.CreateSession(session_handle, server_def));
+ TF_EXPECT_OK(mgr_.CreateSession(session_handle, server_def, true));
  WorkerSession* session = mgr_.WorkerSessionForSession(session_handle);
  EXPECT_NE(nullptr, session) << "Session for " << session_handle << "was null";
  EXPECT_NE(mgr_.LegacySession(), session);
  TF_EXPECT_OK(mgr_.DeleteSession(session_handle));
 }
 
+TEST_F(SessionMgrTest, CreateSessionIsolateSessionState) {
+ ServerDef server_def;
+ server_def.set_job_name("worker");
+ server_def.set_task_index(3);
+
+ TF_EXPECT_OK(mgr_.CreateSession("handle_1", server_def, false));
+ WorkerSession* session_1 = mgr_.WorkerSessionForSession("handle_1");
+ std::vector<Device*> devices_1 = session_1->device_mgr->ListDevices();
+ EXPECT_EQ(1, devices_1.size());
+
+ TF_EXPECT_OK(mgr_.CreateSession("handle_2", server_def, false));
+ WorkerSession* session_2 = mgr_.WorkerSessionForSession("handle_2");
+ std::vector<Device*> devices_2 = session_2->device_mgr->ListDevices();
+ EXPECT_EQ(1, devices_2.size());
+
+ TF_EXPECT_OK(mgr_.CreateSession("handle_3", server_def, true));
+ WorkerSession* session_3 = mgr_.WorkerSessionForSession("handle_3");
+ std::vector<Device*> devices_3 = session_3->device_mgr->ListDevices();
+ EXPECT_EQ(1, devices_3.size());
+
+ TF_EXPECT_OK(mgr_.CreateSession("handle_4", server_def, true));
+ WorkerSession* session_4 = mgr_.WorkerSessionForSession("handle_4");
+ std::vector<Device*> devices_4 = session_4->device_mgr->ListDevices();
+ EXPECT_EQ(1, devices_4.size());
+
+ EXPECT_EQ(devices_1[0]->resource_manager(), devices_2[0]->resource_manager());
+ EXPECT_NE(devices_1[0]->resource_manager(), devices_3[0]->resource_manager());
+ EXPECT_NE(devices_1[0]->resource_manager(), devices_4[0]->resource_manager());
+ EXPECT_NE(devices_3[0]->resource_manager(), devices_4[0]->resource_manager());
+}
+
 TEST_F(SessionMgrTest, LegacySession) {
  ServerDef server_def;
  string session_handle = "";

tensorflow/core/distributed_runtime/worker.cc

@@ -44,7 +44,8 @@ void Worker::CreateWorkerSessionAsync(const CreateWorkerSessionRequest* request,
  CreateWorkerSessionResponse* response,
  StatusCallback done) {
  Status s = env_->session_mgr->CreateSession(request->session_handle(),
- request->server_def());
+ request->server_def(),
+ request->isolate_session_state());
  done(s);
 }