From 18671bd8d61f5cda6e4a99efb9d20bb71faf3f50 Mon Sep 17 00:00:00 2001
From: Li Hongzhang <lihongzhang1@huawei.com>
Date: Wed, 22 Jul 2020 14:56:34 +0800
Subject: [PATCH] pretrigger the npu queryings and set timeout

---
 .../sysmetric/collector/_collect_npu.py       | 43 ++++++++++++++++---
 1 file changed, 37 insertions(+), 6 deletions(-)

diff --git a/mindinsight/sysmetric/collector/_collect_npu.py b/mindinsight/sysmetric/collector/_collect_npu.py
index 9f34c8a..62071f3 100644
--- a/mindinsight/sysmetric/collector/_collect_npu.py
+++ b/mindinsight/sysmetric/collector/_collect_npu.py
@@ -16,15 +16,37 @@
 
 import inspect
 from ctypes import CDLL, Structure, byref, c_char, c_int, c_uint, c_ulong, c_ushort
-from functools import lru_cache
+from functools import lru_cache, wraps
+from threading import Thread
 
 from mindinsight.sysmetric.common.log import logger
 
-try:
-    libsmi = CDLL('libdrvdsmi_host.so')
-except OSError:
-    logger.info('Failed to load libdrvdsmi_host.so.')
-    libsmi = None
+
+def _timeout(seconds, default):
+    """
+    The timeout decorator wait for specified seconds or return the default value.
+
+    Args:
+        seconds (float): The specified seconds.
+        default (Any): The default value.
+    """
+
+    def outer(fn):
+
+        def target(*args, **kwargs):
+            nonlocal default
+            default = fn(*args, **kwargs)
+
+        @wraps(fn)
+        def inner(*args, **kwargs):
+            thread = Thread(target=target, args=args, kwargs=kwargs)
+            thread.start()
+            thread.join(seconds)
+            return default
+
+        return inner
+
+    return outer
 
 
 def libsmicall(*args, **kwargs):
@@ -188,6 +210,7 @@ def dsmi_get_hbm_info(device_id):
     }
 
 
+@_timeout(0.2, 0)
 def dsmi_get_device_utilization_rate(device_id, device_type):
     """
     Get device utilization rate, %.
@@ -279,3 +302,11 @@ def collect_npu():
             'temperature': dsmi_get_device_temperature(device_id)
         })
     return npus
+
+
+try:
+    libsmi = CDLL('libdrvdsmi_host.so')
+    Thread(target=collect_npu).start()
+except OSError:
+    logger.info('Failed to load libdrvdsmi_host.so.')
+    libsmi = None
-- 
GitLab