提交 6b439593 编写于 作者: L Li Hongzhang

fix npu timeout mechanism

上级 c35b45ff
...@@ -15,9 +15,10 @@ ...@@ -15,9 +15,10 @@
"""The npu collector.""" """The npu collector."""
import inspect import inspect
from collections import defaultdict
from ctypes import CDLL, Structure, byref, c_char, c_int, c_uint, c_ulong, c_ushort from ctypes import CDLL, Structure, byref, c_char, c_int, c_uint, c_ulong, c_ushort
from functools import lru_cache, wraps from functools import lru_cache, wraps
from threading import Thread from threading import Lock, Thread
from mindinsight.sysmetric.common.log import logger from mindinsight.sysmetric.common.log import logger
...@@ -32,17 +33,26 @@ def _timeout(seconds, default): ...@@ -32,17 +33,26 @@ def _timeout(seconds, default):
""" """
def outer(fn): def outer(fn):
cached, lockdict = {}, defaultdict(Lock)
def target(*args, **kwargs):
nonlocal default def target(*args):
default = fn(*args, **kwargs) lock = lockdict[args]
if lock.acquire(blocking=False):
try:
cached[args] = fn(*args)
finally:
lock.release()
else:
logger.debug('%s%r skipped.', fn.__name__, args)
@wraps(fn) @wraps(fn)
def inner(*args, **kwargs): def inner(*args):
thread = Thread(target=target, args=args, kwargs=kwargs) thread = Thread(target=target, args=args, daemon=True)
thread.start() thread.start()
thread.join(seconds) thread.join(seconds)
return default if thread.is_alive():
logger.debug('%s%r timeouted.', fn.__name__, args)
return cached.get(args, default)
return inner return inner
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册