From 742ec73fa305c8d7629dc9ce70dbb94b8925e06d Mon Sep 17 00:00:00 2001 From: jonyguo Date: Tue, 15 Dec 2020 15:06:47 +0800 Subject: [PATCH] fix: process is not exit when oom / memory err in feed mode --- mindspore/dataset/engine/iterators.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/mindspore/dataset/engine/iterators.py b/mindspore/dataset/engine/iterators.py index 343fe8dd07..c71f3c3b4e 100644 --- a/mindspore/dataset/engine/iterators.py +++ b/mindspore/dataset/engine/iterators.py @@ -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):