diff --git a/dubhe_data_process/.gitignore b/dubhe_data_process/.gitignore index 13c8453..aa56c11 100644 --- a/dubhe_data_process/.gitignore +++ b/dubhe_data_process/.gitignore @@ -1,8 +1,7 @@ # my -log/* *.log *.swp -.nfs* +.nfs* core *.jpg *.JPEG diff --git a/dubhe_data_process/algorithm-annotation.py b/dubhe_data_process/algorithm-annotation.py index 3df56f2..065283b 100644 --- a/dubhe_data_process/algorithm-annotation.py +++ b/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) diff --git a/dubhe_data_process/algorithm-imagenet.py b/dubhe_data_process/algorithm-imagenet.py index 9c6d703..a1f261f 100644 --- a/dubhe_data_process/algorithm-imagenet.py +++ b/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) diff --git a/dubhe_data_process/common/select_gpu.py b/dubhe_data_process/common/select_gpu.py new file mode 100644 index 0000000..6143884 --- /dev/null +++ b/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() diff --git a/dubhe_data_process/log/dev/.gitkeep b/dubhe_data_process/log/dev/.gitkeep new file mode 100644 index 0000000..e69de29