Add a test only method to reset ProcessState.

PiperOrigin-RevId: 186647005

diff --git a/tensorflow/core/common_runtime/gpu/gpu_device_test.cc b/tensorflow/core/common_runtime/gpu/gpu_device_test.cc
index b568232..f3935f6 100644 (file)
--- a/tensorflow/core/common_runtime/gpu/gpu_device_test.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_device_test.cc
@@ -18,42 +18,48 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/gpu/gpu_device.h"
 
 #include "tensorflow/core/common_runtime/gpu/gpu_init.h"
+#include "tensorflow/core/common_runtime/gpu/process_state.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/platform/test.h"
 
 namespace tensorflow {
-namespace {
 const char* kDeviceNamePrefix = "/job:localhost/replica:0/task:0";
 
-static SessionOptions MakeSessionOptions(
-    const string& visible_device_list = "",
-    double per_process_gpu_memory_fraction = 0, int gpu_device_count = 1,
-    const std::vector<std::vector<float>>& memory_limit_mb = {}) {
-  SessionOptions options;
-  ConfigProto* config = &options.config;
-  (*config->mutable_device_count())["GPU"] = gpu_device_count;
-  GPUOptions* gpu_options = config->mutable_gpu_options();
-  gpu_options->set_visible_device_list(visible_device_list);
-  gpu_options->set_per_process_gpu_memory_fraction(
-      per_process_gpu_memory_fraction);
-  for (const auto& v : memory_limit_mb) {
-    auto virtual_devices =
-        gpu_options->mutable_experimental()->add_virtual_devices();
-    for (float mb : v) {
-      virtual_devices->add_memory_limit_mb(mb);
+class GPUDeviceTest : public ::testing::Test {
+ public:
+  void TearDown() { ProcessState::singleton()->TestOnlyReset(); }
+
+ protected:
+  static SessionOptions MakeSessionOptions(
+      const string& visible_device_list = "",
+      double per_process_gpu_memory_fraction = 0, int gpu_device_count = 1,
+      const std::vector<std::vector<float>>& memory_limit_mb = {}) {
+    SessionOptions options;
+    ConfigProto* config = &options.config;
+    (*config->mutable_device_count())["GPU"] = gpu_device_count;
+    GPUOptions* gpu_options = config->mutable_gpu_options();
+    gpu_options->set_visible_device_list(visible_device_list);
+    gpu_options->set_per_process_gpu_memory_fraction(
+        per_process_gpu_memory_fraction);
+    for (const auto& v : memory_limit_mb) {
+      auto virtual_devices =
+          gpu_options->mutable_experimental()->add_virtual_devices();
+      for (float mb : v) {
+        virtual_devices->add_memory_limit_mb(mb);
+      }
     }
+    return options;
   }
-  return options;
-}
 
-static bool StartsWith(const string& lhs, const string& rhs) {
-  if (rhs.length() > lhs.length()) return false;
-  return lhs.substr(0, rhs.length()) == rhs;
-}
+  static bool StartsWith(const string& lhs, const string& rhs) {
+    if (rhs.length() > lhs.length()) return false;
+    return lhs.substr(0, rhs.length()) == rhs;
+  }
+};
 
-TEST(GPUDeviceTest, FailedToParseVisibleDeviceList) {
+TEST_F(GPUDeviceTest, FailedToParseVisibleDeviceList) {
   SessionOptions opts = MakeSessionOptions("0,abc");
   std::vector<tensorflow::Device*> devices;
   Status status = DeviceFactory::GetFactory("GPU")->CreateDevices(
@@ -63,7 +69,7 @@ TEST(GPUDeviceTest, FailedToParseVisibleDeviceList) {
       << status;
 }
 
-TEST(GPUDeviceTest, InvalidGpuId) {
+TEST_F(GPUDeviceTest, InvalidGpuId) {
   SessionOptions opts = MakeSessionOptions("100");
   std::vector<tensorflow::Device*> devices;
   Status status = DeviceFactory::GetFactory("GPU")->CreateDevices(
@@ -74,7 +80,7 @@ TEST(GPUDeviceTest, InvalidGpuId) {
       << status;
 }
 
-TEST(GPUDeviceTest, DuplicateEntryInVisibleDeviceList) {
+TEST_F(GPUDeviceTest, DuplicateEntryInVisibleDeviceList) {
   SessionOptions opts = MakeSessionOptions("0,0");
   std::vector<tensorflow::Device*> devices;
   Status status = DeviceFactory::GetFactory("GPU")->CreateDevices(
@@ -85,7 +91,7 @@ TEST(GPUDeviceTest, DuplicateEntryInVisibleDeviceList) {
       << status;
 }
 
-TEST(GPUDeviceTest, VirtualDeviceConfigConflictsWithMemoryFractionSettings) {
+TEST_F(GPUDeviceTest, VirtualDeviceConfigConflictsWithMemoryFractionSettings) {
   SessionOptions opts = MakeSessionOptions("0", 0.1, 1, {{}});
   std::vector<tensorflow::Device*> devices;
   Status status = DeviceFactory::GetFactory("GPU")->CreateDevices(
@@ -96,7 +102,7 @@ TEST(GPUDeviceTest, VirtualDeviceConfigConflictsWithMemoryFractionSettings) {
       << status;
 }
 
-TEST(GPUDeviceTest, GpuDeviceCountTooSmall) {
+TEST_F(GPUDeviceTest, GpuDeviceCountTooSmall) {
   // device_count is 0, but with one entry in visible_device_list and one
   // (empty) VirtualDevices messages.
   SessionOptions opts = MakeSessionOptions("0", 0, 0, {{}});
@@ -109,7 +115,7 @@ TEST(GPUDeviceTest, GpuDeviceCountTooSmall) {
       << status;
 }
 
-TEST(GPUDeviceTest, NotEnoughGpuInVisibleDeviceList) {
+TEST_F(GPUDeviceTest, NotEnoughGpuInVisibleDeviceList) {
   // Single entry in visible_device_list with two (empty) VirtualDevices
   // messages.
   SessionOptions opts = MakeSessionOptions("0", 0, 8, {{}, {}});
@@ -122,7 +128,7 @@ TEST(GPUDeviceTest, NotEnoughGpuInVisibleDeviceList) {
       << status;
 }
 
-TEST(GPUDeviceTest, VirtualDeviceConfigConflictsWithVisibleDeviceList) {
+TEST_F(GPUDeviceTest, VirtualDeviceConfigConflictsWithVisibleDeviceList) {
   // This test requires at least two visible GPU hardware.
   if (GPUMachineManager()->VisibleDeviceCount() < 2) return;
   // Three entries in visible_device_list with two (empty) VirtualDevices
@@ -139,7 +145,7 @@ TEST(GPUDeviceTest, VirtualDeviceConfigConflictsWithVisibleDeviceList) {
       << status;
 }
 
-TEST(GPUDeviceTest, EmptyVirtualDeviceConfig) {
+TEST_F(GPUDeviceTest, EmptyVirtualDeviceConfig) {
   // It'll create single virtual device when the virtual device config is empty.
   SessionOptions opts = MakeSessionOptions("0");
   std::vector<tensorflow::Device*> devices;
@@ -150,7 +156,7 @@ TEST(GPUDeviceTest, EmptyVirtualDeviceConfig) {
   for (auto d : devices) delete d;
 }
 
-TEST(GPUDeviceTest, SingleVirtualDeviceWithNoMemoryLimit) {
+TEST_F(GPUDeviceTest, SingleVirtualDeviceWithNoMemoryLimit) {
   // It'll create single virtual device for the gpu in question when
   // memory_limit_mb is unset.
   SessionOptions opts = MakeSessionOptions("0", 0, 1, {{}});
@@ -162,7 +168,7 @@ TEST(GPUDeviceTest, SingleVirtualDeviceWithNoMemoryLimit) {
   for (auto d : devices) delete d;
 }
 
-TEST(GPUDeviceTest, SingleVirtualDeviceWithMemoryLimit) {
+TEST_F(GPUDeviceTest, SingleVirtualDeviceWithMemoryLimit) {
   SessionOptions opts = MakeSessionOptions("0", 0, 1, {{123}});
   std::vector<tensorflow::Device*> devices;
   TF_CHECK_OK(DeviceFactory::GetFactory("GPU")->CreateDevices(
@@ -172,7 +178,7 @@ TEST(GPUDeviceTest, SingleVirtualDeviceWithMemoryLimit) {
   for (auto d : devices) delete d;
 }
 
-TEST(GPUDeviceTest, MultipleVirtualDevices) {
+TEST_F(GPUDeviceTest, MultipleVirtualDevices) {
   SessionOptions opts = MakeSessionOptions("0", 0, 1, {{123, 456}});
   std::vector<tensorflow::Device*> devices;
   TF_CHECK_OK(DeviceFactory::GetFactory("GPU")->CreateDevices(
@@ -195,7 +201,6 @@ TEST(GPUDeviceTest, MultipleVirtualDevices) {
   for (auto d : devices) delete d;
 }
 
-}  // namespace
 }  // namespace tensorflow
 
 #endif

diff --git a/tensorflow/core/common_runtime/gpu/process_state.cc b/tensorflow/core/common_runtime/gpu/process_state.cc
index 61013bd..866a03d 100644 (file)
--- a/tensorflow/core/common_runtime/gpu/process_state.cc
+++ b/tensorflow/core/common_runtime/gpu/process_state.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/framework/log_memory.h"
 #include "tensorflow/core/framework/tracking_allocator.h"
+#include "tensorflow/core/lib/gtl/stl_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/mutex.h"
@@ -318,4 +319,17 @@ void ProcessState::AddGPUAllocVisitor(int bus_id, AllocVisitor visitor) {
 #endif  // GOOGLE_CUDA
 }
 
+void ProcessState::TestOnlyReset() {
+  mutex_lock lock(mu_);
+  gpu_device_enabled_ = false;
+  gpu_visitors_.clear();
+  mem_desc_map_.clear();
+  gtl::STLDeleteElements(&cpu_allocators_);
+  gtl::STLDeleteElements(&gpu_allocators_);
+  gtl::STLDeleteElements(&cuda_host_allocators_);
+  gtl::STLDeleteElements(&cpu_al_);
+  gtl::STLDeleteElements(&gpu_al_);
+  gtl::STLDeleteElements(&cuda_al_);
+}
+
 }  // namespace tensorflow

diff --git a/tensorflow/core/common_runtime/gpu/process_state.h b/tensorflow/core/common_runtime/gpu/process_state.h
index f6e2349..bc2c418 100644 (file)
--- a/tensorflow/core/common_runtime/gpu/process_state.h
+++ b/tensorflow/core/common_runtime/gpu/process_state.h
@@ -114,6 +114,10 @@ class ProcessState {
  protected:
   ProcessState();
 
+  // Helper method for unit tests to reset the ProcessState singleton by
+  // cleaning up everything. Never use in production.
+  virtual void TestOnlyReset();
+
   static ProcessState* instance_;
   bool gpu_device_enabled_;
 
@@ -132,6 +136,8 @@ class ProcessState {
   std::vector<Allocator*> cpu_al_ GUARDED_BY(mu_);
   std::vector<Allocator*> gpu_al_ GUARDED_BY(mu_);
   std::vector<Allocator*> cuda_al_ GUARDED_BY(mu_);
+
+  friend class GPUDeviceTest;
 };
 
 namespace internal {