update data process

dubhe_data_process/.gitignore

@@ -1,8 +1,7 @@
 # my
-log/*
 *.log
 *.swp
-.nfs* 
+.nfs*
 core
 *.jpg
 *.JPEG

dubhe_data_process/algorithm-annotation.py

@@ -26,12 +26,14 @@ import common.config as config
 import annotation as annotation
 import luascript.starttaskscript as start_script
 import logging
+import common.select_gpu as gpu
 
 logging.basicConfig(format='%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s',
                     level=logging.DEBUG)
 
 if __name__ == '__main__':
     """Automatic annotation algorithm entry."""
+    gpu.select_gpu()
     jsonData = config.loadJsonData(config.configPath)
     redisClient = f.getRedisConnection(jsonData["ip"], jsonData["port"], jsonData["database"], jsonData["password"])
     logging.info('init redis client %s', redisClient)

dubhe_data_process/algorithm-imagenet.py

@@ -26,12 +26,14 @@ import common.RedisUtil as f
 import common.config as config
 import luascript.starttaskscript as start_script
 import logging
+import common.select_gpu as gpu
 
 logging.basicConfig(format='%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s',
                     level=logging.DEBUG)
 
 if __name__ == '__main__':
     """Imagenet algorithm entry."""
+    gpu.select_gpu()
     jsonData = config.loadJsonData(config.configPath)
     redisClient = f.getRedisConnection(jsonData["ip"], jsonData["port"], jsonData["database"], jsonData["password"])
     logging.info('init redis client %s', redisClient)

dubhe_data_process/common/select_gpu.py

@@ -0,0 +1,41 @@
+"""
+/**
+* Copyright 2020 Zhejiang Lab. All Rights Reserved.
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*     http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+* =============================================================
+*/
+"""
+import os
+import sys
+import pynvml
+import logging
+
+pid = os.getpid()
+pynvml.nvmlInit()
+
+
+def select_gpu():
+    deviceCount = pynvml.nvmlDeviceGetCount()
+    for i in range(deviceCount):
+        logging.info('-------------get GPU information--------------')
+        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+        logging.info("Device:%s %s", i, pynvml.nvmlDeviceGetName(handle))
+        gpu_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+        logging.info('free:%s MB', gpu_info.free / (1000 * 1000))
+        if gpu_info.free / (1000 * 1000) > 3072:
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(i)
+            logging.info('use GPU:%s %s', i, pynvml.nvmlDeviceGetName(handle))
+            return
+    logging.info('No GPU is currently available')
+    sys.exit()