fix: process is not exit when oom / memory err in feed mode

pull/9975/head
jonyguo 4 years ago
parent 772d0ef2a3
commit 742ec73fa3

@ -15,6 +15,8 @@
"""Built-in iterators.
"""
from abc import abstractmethod
import os
import signal
import weakref
import numpy as np
@ -160,7 +162,16 @@ class DictIterator(Iterator):
Returns:
Dict, the next record in the dataset.
"""
return {k: self._transform_tensor(t) for k, t in self._iterator.GetNextAsMap().items()}
try:
return {k: self._transform_tensor(t) for k, t in self._iterator.GetNextAsMap().items()}
except RuntimeError as err:
## maybe "Out of memory" / "MemoryError" error
logger.error("Got runtime err: {}.".format(err))
err_info = str(err)
if err_info.find("Out of memory") >= 0 or err_info.find("MemoryError") >= 0:
logger.error("Memory error occurred, process will exit.")
os.kill(os.getpid(), signal.SIGKILL)
raise err
class TupleIterator(Iterator):

Loading…
Cancel
Save