Change DeviceFactory functions that create devices to propagate

Statuses, so that failures to initialize devices don't crash
the program.

Changes swig for device_lib to be a lot simpler, thanks to mrry@
and keveman@'s help.

Change allocation of eigen scratch memory to go through the allocator.

Re-enable test for local devices now that python3 issue is fixed.
Change: 129678132

8 files changed, 116 insertions, 71 deletions

diff --git a/tensorflow/core/common_runtime/constant_folding.cc b/tensorflow/core/common_runtime/constant_folding.cc
index 2689dcc276..3d31ca8ce7 100644
--- a/tensorflow/core/common_runtime/constant_folding.cc
+++ b/tensorflow/core/common_runtime/constant_folding.cc
@@ -173,9 +173,9 @@ Device* GetCPUDevice() {
   mutex_lock l(mu);
   if (!device) {
     std::vector<Device*> devices;
-    DeviceFactory::GetFactory(DEVICE_CPU)
-        ->CreateDevices(SessionOptions{}, "", &devices);
-    if (devices.size() > 0) {
+    Status s = DeviceFactory::GetFactory(DEVICE_CPU)
+                   ->CreateDevices(SessionOptions{}, "", &devices);
+    if (s.ok() && devices.size() > 0) {
       device = devices[0];
     }
   }
diff --git a/tensorflow/core/common_runtime/device_factory.cc b/tensorflow/core/common_runtime/device_factory.cc
index 7d0a2380cb..8104f44636 100644
--- a/tensorflow/core/common_runtime/device_factory.cc
+++ b/tensorflow/core/common_runtime/device_factory.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <unordered_map>
 #include <vector>
 
+#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/mutex.h"
@@ -74,25 +75,26 @@ DeviceFactory* DeviceFactory::GetFactory(const string& device_type) {
   return it->second.factory.get();
 }
 
-void DeviceFactory::AddDevices(const SessionOptions& options,
-                               const string& name_prefix,
-                               std::vector<Device*>* devices) {
+Status DeviceFactory::AddDevices(const SessionOptions& options,
+                                 const string& name_prefix,
+                                 std::vector<Device*>* devices) {
   // CPU first.
   auto cpu_factory = GetFactory("CPU");
   if (!cpu_factory) {
-    LOG(FATAL)
-        << "CPU Factory not registered.  Did you link in threadpool_device?";
+    return errors::NotFound(
+        "CPU Factory not registered.  Did you link in threadpool_device?");
   }
   size_t init_size = devices->size();
   cpu_factory->CreateDevices(options, name_prefix, devices);
   if (devices->size() == init_size) {
-    LOG(FATAL) << "No CPU devices are available in this process";
+    return errors::NotFound("No CPU devices are available in this process");
   }
 
   // Then GPU.
   auto gpu_factory = GetFactory("GPU");
   if (gpu_factory) {
-    gpu_factory->CreateDevices(options, name_prefix, devices);
+    TF_RETURN_IF_ERROR(
+        gpu_factory->CreateDevices(options, name_prefix, devices));
   }
 
   // Then the rest.
@@ -100,9 +102,11 @@ void DeviceFactory::AddDevices(const SessionOptions& options,
   for (auto& p : device_factories()) {
     auto factory = p.second.factory.get();
     if (factory != cpu_factory && factory != gpu_factory) {
-      factory->CreateDevices(options, name_prefix, devices);
+      TF_RETURN_IF_ERROR(factory->CreateDevices(options, name_prefix, devices));
     }
   }
+
+  return Status::OK();
 }
 
 Device* DeviceFactory::NewDevice(const string& type,
diff --git a/tensorflow/core/common_runtime/device_factory.h b/tensorflow/core/common_runtime/device_factory.h
index bb1f93623f..3852dd2223 100644
--- a/tensorflow/core/common_runtime/device_factory.h
+++ b/tensorflow/core/common_runtime/device_factory.h
@@ -18,6 +18,8 @@ limitations under the License.
 
 #include <string>
 #include <vector>
+
+#include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
@@ -36,9 +38,9 @@ class DeviceFactory {
   // any device type specific properties/counts listed in "options".
   //
   // CPU devices are added first.
-  static void AddDevices(const SessionOptions& options,
-                         const string& name_prefix,
-                         std::vector<Device*>* devices);
+  static Status AddDevices(const SessionOptions& options,
+                           const string& name_prefix,
+                           std::vector<Device*>* devices);
 
   // Helper for tests.  Create a single device of type "type".  The
   // returned device is always numbered zero, so if creating multiple
@@ -47,9 +49,9 @@ class DeviceFactory {
                            const string& name_prefix);
 
   // Most clients should call AddDevices() instead.
-  virtual void CreateDevices(const SessionOptions& options,
-                             const string& name_prefix,
-                             std::vector<Device*>* devices) = 0;
+  virtual Status CreateDevices(const SessionOptions& options,
+                               const string& name_prefix,
+                               std::vector<Device*>* devices) = 0;
 };
 
 namespace dfactory {
diff --git a/tensorflow/core/common_runtime/direct_session.cc b/tensorflow/core/common_runtime/direct_session.cc
index 9e20aee879..55e83155b9 100644
--- a/tensorflow/core/common_runtime/direct_session.cc
+++ b/tensorflow/core/common_runtime/direct_session.cc
@@ -1045,8 +1045,13 @@ class DirectSessionFactory : public SessionFactory {
       EnableCPUAllocatorFullStats(true);
     }
     std::vector<Device*> devices;
-    DeviceFactory::AddDevices(options, "/job:localhost/replica:0/task:0",
-                              &devices);
+    Status s = DeviceFactory::AddDevices(
+        options, "/job:localhost/replica:0/task:0", &devices);
+    if (!s.ok()) {
+      LOG(ERROR) << s;
+      return nullptr;
+    }
+
     return new DirectSession(options, new DeviceMgr(devices));
   }
 };
diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.cc b/tensorflow/core/common_runtime/gpu/gpu_device.cc
index c9163d6e08..65615891ef 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_device.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_device.cc
@@ -229,22 +229,38 @@ BaseGPUDevice::BaseGPUDevice(const SessionOptions& options, const string& name,
       gpu_allocator_(gpu_allocator),
       cpu_allocator_(cpu_allocator),
       gpu_id_(gpu_id),
-      sync_every_op_(sync_every_op) {
+      sync_every_op_(sync_every_op),
+      max_streams_(max_streams) {
   ProcessState::singleton()->EnableGPUDevice();
+}
+
+BaseGPUDevice::~BaseGPUDevice() {
+  delete gpu_device_info_;
+  for (auto ctx : device_contexts_) ctx->Unref();
+  for (auto& stream_group : streams_) {
+    delete stream_group.compute;
+    delete stream_group.host_to_device;
+    delete stream_group.device_to_host;
+    delete stream_group.device_to_device;
+  }
+}
 
-  executor_ = GPUMachineManager()->ExecutorForDevice(gpu_id_).ValueOrDie();
-  if (!executor_) {
-    LOG(ERROR) << "Failed to get StreamExecutor for device " << gpu_id_;
-    return;
+Status BaseGPUDevice::Init(const SessionOptions& options) {
+  auto executor_status = GPUMachineManager()->ExecutorForDevice(gpu_id_);
+  if (!executor_status.status().ok()) {
+    return errors::Internal("Failed to get StreamExecutor for device ",
+                            gpu_id_);
   }
+
+  executor_ = executor_status.ValueOrDie();
   em_.reset(new EventMgr(executor_, options.config.gpu_options()));
 
-  if (max_streams < 1) {
-    LOG(FATAL) << "Invalid value for max_streams.";
+  if (max_streams_ < 1) {
+    return errors::InvalidArgument("Invalid value for max_streams.");
   }
 
   // Create the specified number of GPU streams
-  for (int i = 0; i < max_streams; i++) {
+  for (int i = 0; i < max_streams_; i++) {
     auto stream = new gpu::Stream(executor_);
     stream->Init();
     VLOG(2) << "Created stream[" << i << "] = " << stream;
@@ -267,14 +283,24 @@ BaseGPUDevice::BaseGPUDevice(const SessionOptions& options, const string& name,
     streams_.push_back({stream, host_to_device_stream, device_to_host_stream,
                         device_to_device_stream});
 
-    perftools::gputools::DeviceMemory<char> mem =
-        executor_->AllocateArray<char>(Eigen::kCudaScratchSize +
-                                       sizeof(unsigned int));
-    scratch_.push_back(static_cast<char*>(mem.opaque()));
+    size_t scratch_buffer_size = Eigen::kCudaScratchSize + sizeof(unsigned int);
+    void* scratch_buffer = gpu_allocator_->AllocateRaw(
+        Allocator::kAllocatorAlignment, scratch_buffer_size);
+    if (scratch_buffer == nullptr) {
+      return errors::FailedPrecondition(
+          "Failed to allocate scratch buffer for device ", gpu_id_);
+    }
+    scratch_.push_back(static_cast<char*>(scratch_buffer));
+
+    perftools::gputools::DeviceMemory<char> mem(
+        perftools::gputools::DeviceMemoryBase(scratch_buffer,
+                                              scratch_buffer_size));
+
     bool ok = executor_->SynchronousMemZero(
         &mem, Eigen::kCudaScratchSize + sizeof(unsigned int));
     if (!ok) {
-      LOG(FATAL) << "Failed to initialize device " << gpu_id;
+      return errors::FailedPrecondition(
+          "Failed to memcopy into scratch buffer for device ", gpu_id_);
     }
 
     device_contexts_.push_back(
@@ -286,17 +312,8 @@ BaseGPUDevice::BaseGPUDevice(const SessionOptions& options, const string& name,
   gpu_device_info_->default_context = device_contexts_[0];
   gpu_device_info_->event_mgr = em_.get();
   set_tensorflow_gpu_device_info(gpu_device_info_);
-}
 
-BaseGPUDevice::~BaseGPUDevice() {
-  delete gpu_device_info_;
-  for (auto ctx : device_contexts_) ctx->Unref();
-  for (auto& stream_group : streams_) {
-    delete stream_group.compute;
-    delete stream_group.host_to_device;
-    delete stream_group.device_to_host;
-    delete stream_group.device_to_device;
-  }
+  return Status::OK();
 }
 
 bool BaseGPUDevice::RequiresRecordingAccessedTensors() const {
@@ -571,9 +588,9 @@ void BaseGPUDevice::ReinitializeGpuDevice(OpKernelContext* context,
   }
 }
 
-void BaseGPUDeviceFactory::CreateDevices(const SessionOptions& options,
-                                         const string& name_prefix,
-                                         std::vector<Device*>* devices) {
+Status BaseGPUDeviceFactory::CreateDevices(const SessionOptions& options,
+                                           const string& name_prefix,
+                                           std::vector<Device*>* devices) {
   int n = INT_MAX;
   auto iter = options.config.device_count().find("GPU");
   if (iter != options.config.device_count().end()) {
@@ -585,9 +602,15 @@ void BaseGPUDeviceFactory::CreateDevices(const SessionOptions& options,
     n = valid_gpu_ids.size();
   }
   for (int i = 0; i < n; i++) {
-    devices->push_back(CreateGPUDevice(
-        options, strings::StrCat(name_prefix, "/gpu:", i), valid_gpu_ids[i]));
+    BaseGPUDevice* gpu_device;
+    TF_RETURN_IF_ERROR(CreateGPUDevice(options,
+                                       strings::StrCat(name_prefix, "/gpu:", i),
+                                       valid_gpu_ids[i], &gpu_device));
+    TF_RETURN_IF_ERROR(gpu_device->Init(options));
+    devices->push_back(gpu_device);
   }
+
+  return Status::OK();
 }
 
 namespace {
@@ -615,8 +638,9 @@ static string GetShortDeviceDescription(int device_id,
                          ", pci bus id: ", desc.pci_bus_id());
 }
 
-LocalDevice* BaseGPUDeviceFactory::CreateGPUDevice(
-    const SessionOptions& options, const string& name, int gpu_id) {
+Status BaseGPUDeviceFactory::CreateGPUDevice(const SessionOptions& options,
+                                             const string& name, int gpu_id,
+                                             BaseGPUDevice** out_device) {
   CHECK_GE(gpu_id, 0);
 
   // Look up the device, to see its attributes.
@@ -675,12 +699,14 @@ LocalDevice* BaseGPUDeviceFactory::CreateGPUDevice(
           << " numa: " << numa_node << " pci: " << desc.pci_bus_id();
 
   ProcessState* process_state = ProcessState::singleton();
-  return CreateGPUDevice(
+  *out_device = CreateGPUDevice(
       options, name, allocated_bytes, bus_adjacency, gpu_id,
       GetShortDeviceDescription(gpu_id, desc),
       process_state->GetGPUAllocator(options.config.gpu_options(), gpu_id,
                                      allocated_memory),
       process_state->GetCPUAllocator(numa_node));
+
+  return Status::OK();
 }
 
 static int GetDefaultMinGPUMultiprocessorCount(gpu::Platform* gpu_manager) {
diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.h b/tensorflow/core/common_runtime/gpu/gpu_device.h
index 4ac9c4021d..03090aa537 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_device.h
+++ b/tensorflow/core/common_runtime/gpu/gpu_device.h
@@ -48,6 +48,9 @@ class BaseGPUDevice : public LocalDevice {
 
   ~BaseGPUDevice() override;
 
+  // Initialize the device and return the status of initialization.
+  Status Init(const SessionOptions& options);
+
   // GPU devices require the Op Compute method to save a reference to
   // any temporary tensors that are allocated until the Op execution
   // completes.
@@ -97,6 +100,7 @@ class BaseGPUDevice : public LocalDevice {
   mutex trace_mu_;
   int gpu_id_ = -1;
   const bool sync_every_op_ = false;
+  const int32 max_streams_;
   std::unique_ptr<EventMgr> em_;
 
   void ReinitializeDevice(OpKernelContext* context, PerOpGpuDevice* device,
@@ -105,19 +109,19 @@ class BaseGPUDevice : public LocalDevice {
 
 class BaseGPUDeviceFactory : public DeviceFactory {
  public:
-  void CreateDevices(const SessionOptions& options, const string& name_prefix,
-                     std::vector<Device*>* devices) override;
+  Status CreateDevices(const SessionOptions& options, const string& name_prefix,
+                       std::vector<Device*>* devices) override;
 
  private:
-  LocalDevice* CreateGPUDevice(const SessionOptions& options,
-                               const string& name, int gpu_id);
-
-  virtual LocalDevice* CreateGPUDevice(const SessionOptions& options,
-                                       const string& name, Bytes memory_limit,
-                                       BusAdjacency bus_adjacency, int gpu_id,
-                                       const string& physical_device_desc,
-                                       Allocator* gpu_allocator,
-                                       Allocator* cpu_allocator) = 0;
+  Status CreateGPUDevice(const SessionOptions& options, const string& name,
+                         int gpu_id, BaseGPUDevice** out_device);
+
+  virtual BaseGPUDevice* CreateGPUDevice(const SessionOptions& options,
+                                         const string& name, Bytes memory_limit,
+                                         BusAdjacency bus_adjacency, int gpu_id,
+                                         const string& physical_device_desc,
+                                         Allocator* gpu_allocator,
+                                         Allocator* cpu_allocator) = 0;
 
   void GetValidDeviceIds(std::vector<int>* ids);
 };
diff --git a/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc b/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc
index 96a4c06d2b..a7e3773ac8 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc
@@ -49,12 +49,12 @@ class GPUDevice : public BaseGPUDevice {
 
 class GPUDeviceFactory : public BaseGPUDeviceFactory {
  private:
-  LocalDevice* CreateGPUDevice(const SessionOptions& options,
-                               const string& name, Bytes memory_limit,
-                               BusAdjacency bus_adjacency, int gpu_id,
-                               const string& physical_device_desc,
-                               Allocator* gpu_allocator,
-                               Allocator* cpu_allocator) override {
+  BaseGPUDevice* CreateGPUDevice(const SessionOptions& options,
+                                 const string& name, Bytes memory_limit,
+                                 BusAdjacency bus_adjacency, int gpu_id,
+                                 const string& physical_device_desc,
+                                 Allocator* gpu_allocator,
+                                 Allocator* cpu_allocator) override {
     return new GPUDevice(options, name, memory_limit, bus_adjacency, gpu_id,
                          physical_device_desc, gpu_allocator, cpu_allocator);
   }
@@ -89,8 +89,8 @@ class GPUCompatibleCPUDevice : public ThreadPoolDevice {
 // The associated factory.
 class GPUCompatibleCPUDeviceFactory : public DeviceFactory {
  public:
-  void CreateDevices(const SessionOptions& options, const string& name_prefix,
-                     std::vector<Device*>* devices) override {
+  Status CreateDevices(const SessionOptions& options, const string& name_prefix,
+                       std::vector<Device*>* devices) override {
     int n = 1;
     auto iter = options.config.device_count().find("CPU");
     if (iter != options.config.device_count().end()) {
@@ -101,6 +101,8 @@ class GPUCompatibleCPUDeviceFactory : public DeviceFactory {
       devices->push_back(new GPUCompatibleCPUDevice(
           options, name, Bytes(256 << 20), BUS_ANY, cpu_allocator()));
     }
+
+    return Status::OK();
   }
 };
 REGISTER_LOCAL_DEVICE_FACTORY("CPU", GPUCompatibleCPUDeviceFactory, 50);
diff --git a/tensorflow/core/common_runtime/threadpool_device_factory.cc b/tensorflow/core/common_runtime/threadpool_device_factory.cc
index 5ae9ee0e36..93772a750e 100644
--- a/tensorflow/core/common_runtime/threadpool_device_factory.cc
+++ b/tensorflow/core/common_runtime/threadpool_device_factory.cc
@@ -26,8 +26,8 @@ namespace tensorflow {
 // TODO(zhifengc/tucker): Figure out the bytes of available RAM.
 class ThreadPoolDeviceFactory : public DeviceFactory {
  public:
-  void CreateDevices(const SessionOptions& options, const string& name_prefix,
-                     std::vector<Device*>* devices) override {
+  Status CreateDevices(const SessionOptions& options, const string& name_prefix,
+                       std::vector<Device*>* devices) override {
     // TODO(zhifengc/tucker): Figure out the number of available CPUs
     // and/or NUMA configuration.
     int n = 1;
@@ -40,6 +40,8 @@ class ThreadPoolDeviceFactory : public DeviceFactory {
       devices->push_back(new ThreadPoolDevice(options, name, Bytes(256 << 20),
                                               BUS_ANY, cpu_allocator()));
     }
+
+    return Status::OK();
   }
 };
 REGISTER_LOCAL_DEVICE_FACTORY("CPU", ThreadPoolDeviceFactory);