fix bucket_batch ci error

This commit is contained in:
yanghaitao1 2020-09-18 02:53:46 -04:00
parent 016a656b33
commit da15495d7b
3 changed files with 6 additions and 8 deletions

View File

@ -50,14 +50,8 @@ Status PyFuncOp::Compute(const TensorRow &input, TensorRow *output) {
} else {
ret_py_obj = this->py_func_ptr_();
}
// Object is none if pyfunc timeout
if (ret_py_obj.is_none()) {
MS_LOG(INFO) << "Pyfunc execute time out";
goto TimeoutError;
}
if (output_type_ != DataType::DE_UNKNOWN) {
RETURN_IF_NOT_OK(CastOutput(ret_py_obj, output));
} else {
if (py::isinstance<py::tuple>(ret_py_obj)) {
// In case of a n-m mapping, the return value will be a tuple of numpy arrays
@ -65,6 +59,10 @@ Status PyFuncOp::Compute(const TensorRow &input, TensorRow *output) {
// Iterate over two containers simultaneously for memory copy
for (size_t i = 0; i < ret_py_tuple.size(); i++) {
py::object ret_py_ele = ret_py_tuple[i];
// Object is none if pyfunc timeout
if (ret_py_ele.is_none()) {
goto TimeoutError;
}
if (!py::isinstance<py::array>(ret_py_ele)) {
goto ShapeMisMatch;
}

View File

@ -1979,7 +1979,7 @@ class _PythonCallable:
return result.get(60)
except multiprocessing.TimeoutError:
# Ensure c++ pyfunc threads exit normally if python sub-process is killed unnormally.
return None
return (None,)
except KeyboardInterrupt:
self.pool.terminate()
self.pool.join()

View File

@ -504,7 +504,7 @@ def test_celeba_padded():
count = 0
for _ in data.create_dict_iterator(num_epochs=1, output_numpy=True):
count = count + 1
assert count == 2
assert count == 4
if __name__ == '__main__':