Fix exception raised when using ThreadPool (#790)

This commit is contained in:
allegroai 2022-10-14 10:16:11 +03:00
parent d497869a1f
commit 57b7a31097

View File

@ -349,7 +349,7 @@ class SafeQueue(object):
if BackgroundMonitor.get_at_exit_state(): if BackgroundMonitor.get_at_exit_state():
self._q_put(obj) self._q_put(obj)
return return
self.__thread_pool.get().apply_async(self._q_put, args=(obj, )) self.__thread_pool.get().apply_async(self._q_put, args=(obj, False))
except: # noqa except: # noqa
pid = os.getpid() pid = os.getpid()
p = None p = None
@ -360,13 +360,16 @@ class SafeQueue(object):
pid = pid or os.getpid() pid = pid or os.getpid()
return len([p for p in self._q_size if p == pid]) return len([p for p in self._q_size if p == pid])
def _q_put(self, obj): def _q_put(self, obj, allow_raise=True):
# noinspection PyBroadException
try: try:
self._q.put(obj) self._q.put(obj)
except BaseException: except BaseException:
# make sure we zero the _q_size of the process dies (i.e. queue put fails) # make sure we zero the _q_size of the process dies (i.e. queue put fails)
self._q_size.clear() self._q_size.clear()
if allow_raise:
raise raise
return
pid = os.getpid() pid = os.getpid()
# GIL will make sure it is atomic # GIL will make sure it is atomic
# pop the First "counter" that is ours (i.e. pid == os.getpid()) # pop the First "counter" that is ours (i.e. pid == os.getpid())