Better PythonFunction error message (#883)

* Better PythonFunction error message Signed-off-by: Albert Wolant <[email protected]>
NVIDIA · May 16, 2019 · 527e2fd · 527e2fd
1 parent 783a3e9
commit 527e2fd
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/dali/pipeline/operators/python_function/python_function.cc b/dali/pipeline/operators/python_function/python_function.cc
@@ -120,8 +120,12 @@ void PythonFunctionImpl<CPUBackend>::RunImpl(SampleWorkspace *ws, const int idx)
   const auto &input = ws->Input<CPUBackend>(idx);
   auto &output = ws->Output<CPUBackend>(idx);
   py::gil_scoped_acquire guard{};
-  py::array output_array = python_function(TensorToNumpyArray(input));
-  CopyNumpyArrayToTensor(output, output_array);
+  try {
+    py::array output_array = python_function(TensorToNumpyArray(input));
+    CopyNumpyArrayToTensor(output, output_array);
+  } catch(const py::error_already_set & e) {
+    throw std::runtime_error(to_string("PythonFunction error: ") + to_string(e.what()));
+  }
 }
 
 DALI_REGISTER_OPERATOR(PythonFunctionImpl, PythonFunctionImpl<CPUBackend>, CPU);

diff --git a/dali/test/python/test_python_function_operator.py b/dali/test/python/test_python_function_operator.py
@@ -130,3 +130,14 @@ def test_python_operator_flip():
         dali_output, = dali_flip.run()
         for i in range(len(numpy_output)):
             assert numpy.array_equal(numpy_output.at(i), dali_output.at(i))
+
+def invalid_function(image):
+    return img
+
+def test_python_operator_invalid_function():
+    invalid_pipe = PythonOperatorPipeline(BATCH_SIZE, NUM_WORKERS, DEVICE_ID, SEED, images_dir, invalid_function)
+    invalid_pipe.build()
+    try:
+        invalid_pipe.run()
+    except Exception as e:
+        print(e)