When an invalid op is registered, ensure that tf.load_op_library returns

Suharsh Sivakumar · tensorflower-gardener · commit e40bbecdfd11 · 2016-07-11T12:35:07.000-07:00
a nice error rather than silently generating invalid python code.

Fixes:
- On linux, a fatal log that should be triggered in C++ seems to be ignored
  causing the tf.load_op_library call to fail and exit.
  We fix this by propagating the status from op registration to raise an
  exception from python rather than print a fatal log.
  The reason why the fatal log is being ignored on linux is unclear, my
  theory is some sort of SWIG issue, but online research provided no results.
- Op registrations that fail are still added to the OpDef registration.
  This is only an issue when the Fatal log is ignored, but we fix this
  anyways.
Change: 127114085
diff --git a/tensorflow/core/framework/load_library.cc b/tensorflow/core/framework/load_library.cc
@@ -51,7 +51,10 @@ Status LoadLibrary(const char* library_filename, void** result,
   std::unordered_set<string> seen_op_names;
   {
     mutex_lock lock(mu);
-    OpRegistry::Global()->ProcessRegistrations();
+    Status s = OpRegistry::Global()->ProcessRegistrations();
+    if (!s.ok()) {
+      return s;
+    }
     TF_RETURN_IF_ERROR(OpRegistry::Global()->SetWatcher(
         [&op_list, &seen_op_names](const Status& s,
                                    const OpDef& opdef) -> Status {
@@ -62,18 +65,22 @@ Status LoadLibrary(const char* library_filename, void** result,
               return Status::OK();
             }
           }
-          *op_list.add_op() = opdef;
-          seen_op_names.insert(opdef.name());
+          if (s.ok()) {
+            *op_list.add_op() = opdef;
+            seen_op_names.insert(opdef.name());
+          }
           return s;
         }));
     OpRegistry::Global()->DeferRegistrations();
-    Status s = env->LoadLibrary(library_filename, &lib);
+    s = env->LoadLibrary(library_filename, &lib);
+    if (s.ok()) {
+      s = OpRegistry::Global()->ProcessRegistrations();
+    }
     if (!s.ok()) {
       OpRegistry::Global()->ClearDeferredRegistrations();
       TF_RETURN_IF_ERROR(OpRegistry::Global()->SetWatcher(nullptr));
       return s;
     }
-    OpRegistry::Global()->ProcessRegistrations();
     TF_RETURN_IF_ERROR(OpRegistry::Global()->SetWatcher(nullptr));
   }
   string str;
diff --git a/tensorflow/core/framework/op.cc b/tensorflow/core/framework/op.cc
@@ -64,7 +64,7 @@ Status OpRegistry::LookUp(const string& op_type_name,
   bool first_call = false;
   {  // Scope for lock.
     mutex_lock lock(mu_);
-    first_call = CallDeferred();
+    first_call = MustCallDeferred();
     res = gtl::FindWithDefault(registry_, op_type_name, nullptr);
     // Note: Can't hold mu_ while calling Export() below.
   }
@@ -93,7 +93,7 @@ Status OpRegistry::LookUp(const string& op_type_name,
 
 void OpRegistry::GetRegisteredOps(std::vector<OpDef>* op_defs) {
   mutex_lock lock(mu_);
-  CallDeferred();
+  MustCallDeferred();
   for (const auto& p : registry_) {
     op_defs->push_back(p.second->op_def);
   }
@@ -111,7 +111,7 @@ Status OpRegistry::SetWatcher(const Watcher& watcher) {
 
 void OpRegistry::Export(bool include_internal, OpList* ops) const {
   mutex_lock lock(mu_);
-  CallDeferred();
+  MustCallDeferred();
 
   std::vector<std::pair<string, const OpRegistrationData*>> sorted(
       registry_.begin(), registry_.end());
@@ -138,9 +138,9 @@ void OpRegistry::ClearDeferredRegistrations() {
   deferred_.clear();
 }
 
-void OpRegistry::ProcessRegistrations() const {
+Status OpRegistry::ProcessRegistrations() const {
   mutex_lock lock(mu_);
-  CallDeferred();
+  return CallDeferred();
 }
 
 string OpRegistry::DebugString(bool include_internal) const {
@@ -153,7 +153,7 @@ string OpRegistry::DebugString(bool include_internal) const {
   return ret;
 }
 
-bool OpRegistry::CallDeferred() const {
+bool OpRegistry::MustCallDeferred() const {
   if (initialized_) return false;
   initialized_ = true;
   for (int i = 0; i < deferred_.size(); ++i) {
@@ -163,6 +163,19 @@ bool OpRegistry::CallDeferred() const {
   return true;
 }
 
+Status OpRegistry::CallDeferred() const {
+  if (initialized_) return Status::OK();
+  initialized_ = true;
+  for (int i = 0; i < deferred_.size(); ++i) {
+    Status s = RegisterAlreadyLocked(deferred_[i]);
+    if (!s.ok()) {
+      return s;
+    }
+  }
+  deferred_.clear();
+  return Status::OK();
+}
+
 Status OpRegistry::RegisterAlreadyLocked(
     OpRegistrationDataFactory op_data_factory) const {
   std::unique_ptr<OpRegistrationData> op_reg_data(new OpRegistrationData);
diff --git a/tensorflow/core/framework/op.h b/tensorflow/core/framework/op.h
@@ -111,8 +111,9 @@ class OpRegistry : public OpRegistryInterface {
 
   // Process the current list of deferred registrations. Note that calls to
   // Export, LookUp and DebugString would also implicitly process the deferred
-  // registrations.
-  void ProcessRegistrations() const;
+  // registrations. Returns the status of the first failed op registration or
+  // Status::OK() otherwise.
+  Status ProcessRegistrations() const;
 
   // Defer the registrations until a later call to a function that processes
   // deferred registrations are made. Normally, registrations that happen after
@@ -126,8 +127,13 @@ class OpRegistry : public OpRegistryInterface {
  private:
   // Ensures that all the functions in deferred_ get called, their OpDef's
   // registered, and returns with deferred_ empty.  Returns true the first
-  // time it is called.
-  bool CallDeferred() const EXCLUSIVE_LOCKS_REQUIRED(mu_);
+  // time it is called. Prints a fatal log if any op registration fails.
+  bool MustCallDeferred() const EXCLUSIVE_LOCKS_REQUIRED(mu_);
+
+  // Calls the functions in deferred_ and registers their OpDef's
+  // It returns the Status of the first failed op registration or Status::OK()
+  // otherwise.
+  Status CallDeferred() const EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
   // Add 'def' to the registry with additional data 'data'. On failure, or if
   // there is already an OpDef with that name registered, returns a non-okay
diff --git a/tensorflow/core/framework/op_registration_test.cc b/tensorflow/core/framework/op_registration_test.cc
@@ -44,14 +44,16 @@ TEST(OpRegistrationTest, TestBasic) {
 TEST(OpRegistrationTest, TestDuplicate) {
   std::unique_ptr<OpRegistry> registry(new OpRegistry);
   Register("Foo", registry.get());
-  registry->ProcessRegistrations();
+  Status s = registry->ProcessRegistrations();
+  EXPECT_TRUE(s.ok());
 
   registry->SetWatcher([](const Status& s, const OpDef& op_def) -> Status {
     EXPECT_TRUE(errors::IsAlreadyExists(s));
     return Status::OK();
   });
   Register("Foo", registry.get());
-  registry->ProcessRegistrations();
+  s = registry->ProcessRegistrations();
+  EXPECT_TRUE(s.ok());
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/user_ops/BUILD b/tensorflow/user_ops/BUILD
@@ -44,6 +44,18 @@ py_tests(
     data = [":duplicate_op.so"],
 )
 
+tf_custom_op_library(
+    name = "invalid_op.so",
+    srcs = ["invalid_op.cc"],
+)
+
+py_tests(
+    name = "invalid_op_test",
+    size = "small",
+    srcs = ["invalid_op_test.py"],
+    data = [":invalid_op.so"],
+)
+
 filegroup(
     name = "all_files",
     srcs = glob(
diff --git a/tensorflow/user_ops/invalid_op.cc b/tensorflow/user_ops/invalid_op.cc
@@ -0,0 +1,27 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tensorflow {
+
+REGISTER_OP("Invalid")
+    .Attr("invalid attr: int32")  // invalid since the name has a space.
+    .Doc(R"doc(
+An op to test that invalid ops do not successfully generate invalid python code.
+)doc");
+
+}  // namespace tensorflow
diff --git a/tensorflow/user_ops/invalid_op_test.py b/tensorflow/user_ops/invalid_op_test.py
@@ -0,0 +1,35 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for custom user ops."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os.path
+
+import tensorflow as tf
+
+
+class InvalidOpTest(tf.test.TestCase):
+
+  def testBasic(self):
+    library_filename = os.path.join(tf.resource_loader.get_data_files_path(),
+                                    'invalid_op.so')
+    with self.assertRaises(tf.errors.InvalidArgumentError):
+      tf.load_op_library(library_filename)
+
+
+if __name__ == '__main__':
+  tf.test.main()
