Replace CHECK with returning an InternalError on failing to create python tuple

swachhandl · pak-laura · commit 6fc67e408f23 · 2022-11-14T14:54:38.000-08:00
Returns InternalError if a non UTF-8 string is passed in as token, instead of crashing.

PiperOrigin-RevId: 478123498
diff --git a/tensorflow/python/lib/core/py_func.cc b/tensorflow/python/lib/core/py_func.cc
@@ -83,8 +83,8 @@ bool IsCPUDevice(const Device* d) {
   return d == nullptr || d->tensorflow_accelerator_device_info() == nullptr;
 }
 
-// Givens the 'call', prepares the token and inputs as a python tuple
-// that is appropriate for calling the trampoline.
+// Given the 'call', prepares the token and inputs as a python tuple that is
+// appropriate for calling the trampoline.
 Status MakeArgTuple(const PyCall* call, TFE_Context* ctx, PyObject** tuple) {
   int64_t n = call->ins.size();
   PyObject* lst = PyList_New(n);
@@ -119,8 +119,12 @@ Status MakeArgTuple(const PyCall* call, TFE_Context* ctx, PyObject** tuple) {
     PyList_SetItem(lst, i, arg);
   }
   *tuple = Py_BuildValue("(ssN)", call->token.c_str(), device_name, lst);
-  CHECK(*tuple);
-  return Status::OK();
+  if (*tuple == nullptr) {
+    return errors::Internal(
+        "Failed to create python tuple. Please make sure `token` is a "
+        "well-formed UTF-8 string.");
+  }
+  return OkStatus();
 }
 
 bool IsSingleNone(PyObject* obj) {
diff --git a/tensorflow/python/ops/script_ops_test.py b/tensorflow/python/ops/script_ops_test.py
@@ -16,8 +16,11 @@
 
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import gen_script_ops
+from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.ops import script_ops
 from tensorflow.python.ops.script_ops import numpy_function
 from tensorflow.python.platform import test
@@ -87,5 +90,30 @@ def func_stateful(a, b):
                      2)  # as stateful, func is guaranteed to execute twice
 
 
+class PyFunctionTest(test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_variable_arguments(self):
+
+    def plus(a, b):
+      return a + b
+
+    v1 = resource_variable_ops.ResourceVariable(1)
+    self.evaluate(v1.initializer)
+
+    actual_result = script_ops.eager_py_func(plus, [v1, 2], dtypes.int32)
+    expect_result = constant_op.constant(3, dtypes.int32)
+    self.assertAllEqual(actual_result, expect_result)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_fail_on_non_utf8_token(self):
+    value = constant_op.constant(value=[1, 2])
+    token = b"\xb0"
+    data_type = [dtypes.int32]
+    with self.assertRaises((errors.InternalError, UnicodeDecodeError)):
+      self.evaluate(
+          gen_script_ops.py_func(input=[value], token=token, Tout=data_type))
+
+
 if __name__ == "__main__":
   test.main()
