utils.py 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753
  1. import psutil
  2. import shutil
  3. import os
  4. import os.path as osp
  5. from enum import Enum
  6. import multiprocessing as mp
  7. from queue import Queue
  8. import time
  9. import threading
  10. from ctypes import CDLL, c_char, c_uint, c_ulonglong
  11. from _ctypes import byref, Structure, POINTER
  12. import platform
  13. import string
  14. import logging
  15. import socket
  16. import logging.handlers
  17. import requests
  18. import json
  19. from json import JSONEncoder
  20. class CustomEncoder(JSONEncoder):
  21. def default(self, o):
  22. return o.__dict__
  23. class ShareData():
  24. workspace = None
  25. workspace_dir = ""
  26. has_gpu = True
  27. monitored_processes = mp.Queue(4096)
  28. current_port = 8000
  29. running_boards = {}
  30. machine_info = dict()
  31. load_demo_proc_dict = {}
  32. load_demo_proj_data_dict = {}
  33. DatasetStatus = Enum(
  34. 'DatasetStatus', ('XEMPTY', 'XCHECKING', 'XCHECKFAIL', 'XCOPYING',
  35. 'XCOPYDONE', 'XCOPYFAIL', 'XSPLITED'),
  36. start=0)
  37. TaskStatus = Enum(
  38. 'TaskStatus', ('XUNINIT', 'XINIT', 'XDOWNLOADING', 'XTRAINING',
  39. 'XTRAINDONE', 'XEVALUATED', 'XEXPORTING', 'XEXPORTED',
  40. 'XTRAINEXIT', 'XDOWNLOADFAIL', 'XTRAINFAIL', 'XEVALUATING',
  41. 'XEVALUATEFAIL', 'XEXPORTFAIL', 'XPRUNEING', 'XPRUNETRAIN'),
  42. start=0)
  43. ProjectType = Enum(
  44. 'ProjectType', ('classification', 'detection', 'segmentation',
  45. 'instance_segmentation', 'remote_segmentation'),
  46. start=0)
  47. DownloadStatus = Enum(
  48. 'DownloadStatus',
  49. ('XDDOWNLOADING', 'XDDOWNLOADFAIL', 'XDDOWNLOADDONE', 'XDDECOMPRESSED'),
  50. start=0)
  51. PredictStatus = Enum(
  52. 'PredictStatus', ('XPRESTART', 'XPREDONE', 'XPREFAIL'), start=0)
  53. PruneStatus = Enum(
  54. 'PruneStatus', ('XSPRUNESTART', 'XSPRUNEING', 'XSPRUNEDONE', 'XSPRUNEFAIL',
  55. 'XSPRUNEEXIT'),
  56. start=0)
  57. PretrainedModelStatus = Enum(
  58. 'PretrainedModelStatus',
  59. ('XPINIT', 'XPSAVING', 'XPSAVEFAIL', 'XPSAVEDONE'),
  60. start=0)
  61. ExportedModelType = Enum(
  62. 'ExportedModelType', ('XQUANTMOBILE', 'XPRUNEMOBILE', 'XTRAINMOBILE',
  63. 'XQUANTSERVER', 'XPRUNESERVER', 'XTRAINSERVER'),
  64. start=0)
  65. process_pool = Queue(1000)
  66. def get_ip():
  67. try:
  68. s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  69. s.connect(('8.8.8.8', 80))
  70. ip = s.getsockname()[0]
  71. finally:
  72. s.close()
  73. return ip
  74. def get_logger(filename):
  75. flask_logger = logging.getLogger()
  76. flask_logger.setLevel(level=logging.INFO)
  77. fmt = '%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s:%(message)s'
  78. format_str = logging.Formatter(fmt)
  79. ch = logging.StreamHandler()
  80. ch.setLevel(level=logging.INFO)
  81. ch.setFormatter(format_str)
  82. th = logging.handlers.TimedRotatingFileHandler(
  83. filename=filename, when='D', backupCount=5, encoding='utf-8')
  84. th.setFormatter(format_str)
  85. flask_logger.addHandler(th)
  86. flask_logger.addHandler(ch)
  87. return flask_logger
  88. def start_process(target, args):
  89. global process_pool
  90. p = mp.Process(target=target, args=args)
  91. p.start()
  92. process_pool.put(p)
  93. def pkill(pid):
  94. """结束进程pid,和与其相关的子进程
  95. Args:
  96. pid(int): 进程id
  97. """
  98. try:
  99. parent = psutil.Process(pid)
  100. for child in parent.children(recursive=True):
  101. child.kill()
  102. parent.kill()
  103. except:
  104. print("Try to kill process {} failed.".format(pid))
  105. def set_folder_status(dirname, status, message=""):
  106. """设置目录状态
  107. Args:
  108. dirname(str): 目录路径
  109. status(DatasetStatus): 状态
  110. message(str): 需要写到状态文件里的信息
  111. """
  112. if not osp.isdir(dirname):
  113. raise Exception("目录路径{}不存在".format(dirname))
  114. tmp_file = osp.join(dirname, status.name + '.tmp')
  115. with open(tmp_file, 'w', encoding='utf-8') as f:
  116. f.write("{}\n".format(message))
  117. shutil.move(tmp_file, osp.join(dirname, status.name))
  118. for status_type in [
  119. DatasetStatus, TaskStatus, PredictStatus, PruneStatus,
  120. DownloadStatus, PretrainedModelStatus
  121. ]:
  122. for s in status_type:
  123. if s == status:
  124. continue
  125. if osp.exists(osp.join(dirname, s.name)):
  126. os.remove(osp.join(dirname, s.name))
  127. def get_folder_status(dirname, with_message=False):
  128. """获取目录状态
  129. Args:
  130. dirname(str): 目录路径
  131. with_message(bool): 是否需要返回状态文件内的信息
  132. """
  133. status = None
  134. closest_time = 0
  135. message = ''
  136. for status_type in [
  137. DatasetStatus, TaskStatus, PredictStatus, PruneStatus,
  138. DownloadStatus, PretrainedModelStatus
  139. ]:
  140. for s in status_type:
  141. if osp.exists(osp.join(dirname, s.name)):
  142. modify_time = os.stat(osp.join(dirname, s.name)).st_mtime
  143. if modify_time > closest_time:
  144. closest_time = modify_time
  145. status = getattr(status_type, s.name)
  146. if with_message:
  147. encoding = 'utf-8'
  148. try:
  149. f = open(
  150. osp.join(dirname, s.name),
  151. 'r',
  152. encoding=encoding)
  153. message = f.read()
  154. f.close()
  155. except:
  156. try:
  157. import chardet
  158. f = open(filename, 'rb')
  159. data = f.read()
  160. f.close()
  161. encoding = chardet.detect(data).get('encoding')
  162. f = open(
  163. osp.join(dirname, s.name),
  164. 'r',
  165. encoding=encoding)
  166. message = f.read()
  167. f.close()
  168. except:
  169. pass
  170. if with_message:
  171. return status, message
  172. return status
  173. def _machine_check_proc(queue, path):
  174. info = dict()
  175. p = PyNvml()
  176. gpu_num = 0
  177. try:
  178. # import paddle.fluid.core as core
  179. # gpu_num = core.get_cuda_device_count()
  180. p.nvml_init(path)
  181. gpu_num = p.nvml_device_get_count()
  182. driver_version = bytes.decode(p.nvml_system_get_driver_version())
  183. except:
  184. driver_version = "N/A"
  185. info['gpu_num'] = gpu_num
  186. info['gpu_free_mem'] = list()
  187. try:
  188. for i in range(gpu_num):
  189. handle = p.nvml_device_get_handle_by_index(i)
  190. meminfo = p.nvml_device_get_memory_info(handle)
  191. free_mem = meminfo.free / 1024 / 1024
  192. info['gpu_free_mem'].append(free_mem)
  193. except:
  194. pass
  195. info['cpu_num'] = os.environ.get('CPU_NUM', 1)
  196. info['driver_version'] = driver_version
  197. info['path'] = p.nvml_lib_path
  198. queue.put(info, timeout=2)
  199. def get_machine_info(path=None):
  200. queue = mp.Queue(1)
  201. p = mp.Process(target=_machine_check_proc, args=(queue, path))
  202. p.start()
  203. p.join()
  204. return queue.get(timeout=2)
  205. def download(url, target_path):
  206. if not osp.exists(target_path):
  207. os.makedirs(target_path)
  208. fname = osp.split(url)[-1]
  209. fullname = osp.join(target_path, fname)
  210. retry_cnt = 0
  211. DOWNLOAD_RETRY_LIMIT = 3
  212. while not (osp.exists(fullname)):
  213. if retry_cnt < DOWNLOAD_RETRY_LIMIT:
  214. retry_cnt += 1
  215. else:
  216. # 设置下载失败
  217. msg = "Download from {} failed. Retry limit reached".format(url)
  218. set_folder_status(target_path, DownloadStatus.XDDOWNLOADFAIL, msg)
  219. raise RuntimeError(msg)
  220. req = requests.get(url, stream=True)
  221. if req.status_code != 200:
  222. msg = "Downloading from {} failed with code {}!".format(
  223. url, req.status_code)
  224. set_folder_status(target_path, DownloadStatus.XDDOWNLOADFAIL, msg)
  225. raise RuntimeError(msg)
  226. # For protecting download interupted, download to
  227. # tmp_fullname firstly, move tmp_fullname to fullname
  228. # after download finished
  229. tmp_fullname = fullname + "_tmp"
  230. total_size = req.headers.get('content-length')
  231. set_folder_status(target_path, DownloadStatus.XDDOWNLOADING,
  232. total_size)
  233. with open(tmp_fullname, 'wb') as f:
  234. if total_size:
  235. download_size = 0
  236. for chunk in req.iter_content(chunk_size=1024):
  237. f.write(chunk)
  238. download_size += 1024
  239. else:
  240. for chunk in req.iter_content(chunk_size=1024):
  241. if chunk:
  242. f.write(chunk)
  243. shutil.move(tmp_fullname, fullname)
  244. set_folder_status(target_path, DownloadStatus.XDDOWNLOADDONE)
  245. return fullname
  246. def is_pic(filename):
  247. suffixes = {'JPEG', 'jpeg', 'JPG', 'jpg', 'BMP', 'bmp', 'PNG', 'png'}
  248. suffix = filename.strip().split('.')[-1]
  249. if suffix not in suffixes:
  250. return False
  251. return True
  252. def is_available(ip, port):
  253. s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  254. try:
  255. s.connect((ip, int(port)))
  256. s.shutdown(2)
  257. return False
  258. except:
  259. return True
  260. def list_files(dirname):
  261. """ 列出目录下所有文件(包括所属的一级子目录下文件)
  262. Args:
  263. dirname: 目录路径
  264. """
  265. def filter_file(f):
  266. if f.startswith('.'):
  267. return True
  268. if hasattr(PretrainedModelStatus, f):
  269. return True
  270. return False
  271. all_files = list()
  272. dirs = list()
  273. for f in os.listdir(dirname):
  274. if filter_file(f):
  275. continue
  276. if osp.isdir(osp.join(dirname, f)):
  277. dirs.append(f)
  278. else:
  279. all_files.append(f)
  280. for d in dirs:
  281. for f in os.listdir(osp.join(dirname, d)):
  282. if filter_file(f):
  283. continue
  284. if osp.isdir(osp.join(dirname, d, f)):
  285. continue
  286. all_files.append(osp.join(d, f))
  287. return all_files
  288. def copy_model_directory(src, dst, files=None, filter_files=[]):
  289. """从src目录copy文件至dst目录,
  290. 注意:拷贝前会先清空dst中的所有文件
  291. Args:
  292. src: 源目录路径
  293. dst: 目标目录路径
  294. files: 需要拷贝的文件列表(src的相对路径)
  295. """
  296. set_folder_status(dst, PretrainedModelStatus.XPSAVING, os.getpid())
  297. if files is None:
  298. files = list_files(src)
  299. try:
  300. message = '{} {}'.format(os.getpid(), len(files))
  301. set_folder_status(dst, PretrainedModelStatus.XPSAVING, message)
  302. if not osp.samefile(src, dst):
  303. for i, f in enumerate(files):
  304. items = osp.split(f)
  305. if len(items) > 2:
  306. continue
  307. if len(items) == 2:
  308. if not osp.isdir(osp.join(dst, items[0])):
  309. if osp.exists(osp.join(dst, items[0])):
  310. os.remove(osp.join(dst, items[0]))
  311. os.makedirs(osp.join(dst, items[0]))
  312. if f not in filter_files:
  313. shutil.copy(osp.join(src, f), osp.join(dst, f))
  314. set_folder_status(dst, PretrainedModelStatus.XPSAVEDONE)
  315. except Exception as e:
  316. import traceback
  317. error_info = traceback.format_exc()
  318. set_folder_status(dst, PretrainedModelStatus.XPSAVEFAIL, error_info)
  319. def copy_pretrained_model(src, dst):
  320. p = mp.Process(
  321. target=copy_model_directory, args=(src, dst, None, ['model.pdopt']))
  322. p.start()
  323. return p
  324. def _get_gpu_info(queue):
  325. gpu_info = dict()
  326. mem_free = list()
  327. mem_used = list()
  328. mem_total = list()
  329. import pycuda.driver as drv
  330. from pycuda.tools import clear_context_caches
  331. drv.init()
  332. driver_version = drv.get_driver_version()
  333. gpu_num = drv.Device.count()
  334. for gpu_id in range(gpu_num):
  335. dev = drv.Device(gpu_id)
  336. try:
  337. context = dev.make_context()
  338. free, total = drv.mem_get_info()
  339. context.pop()
  340. free = free // 1024 // 1024
  341. total = total // 1024 // 1024
  342. used = total - free
  343. except:
  344. free = 0
  345. total = 0
  346. used = 0
  347. mem_free.append(free)
  348. mem_used.append(used)
  349. mem_total.append(total)
  350. gpu_info['mem_free'] = mem_free
  351. gpu_info['mem_used'] = mem_used
  352. gpu_info['mem_total'] = mem_total
  353. gpu_info['driver_version'] = driver_version
  354. gpu_info['gpu_num'] = gpu_num
  355. queue.put(gpu_info)
  356. def get_gpu_info():
  357. try:
  358. import pycuda
  359. except:
  360. gpu_info = dict()
  361. message = "未检测到GPU \n 若存在GPU请确保安装pycuda \n 若未安装pycuda请使用'pip install pycuda'来安装"
  362. gpu_info['gpu_num'] = 0
  363. return gpu_info, message
  364. queue = mp.Queue(1)
  365. p = mp.Process(target=_get_gpu_info, args=(queue, ))
  366. p.start()
  367. p.join()
  368. gpu_info = queue.get(timeout=2)
  369. if gpu_info['gpu_num'] == 0:
  370. message = "未检测到GPU"
  371. else:
  372. message = "检测到GPU"
  373. return gpu_info, message
  374. class TrainLogReader(object):
  375. def __init__(self, log_file):
  376. self.log_file = log_file
  377. self.eta = None
  378. self.train_metrics = None
  379. self.eval_metrics = None
  380. self.download_status = None
  381. self.eval_done = False
  382. self.train_error = None
  383. self.train_stage = None
  384. self.running_duration = None
  385. def update(self):
  386. if not osp.exists(self.log_file):
  387. return
  388. if self.train_stage == "Train Error":
  389. return
  390. if self.download_status == "Failed":
  391. return
  392. if self.train_stage == "Train Complete":
  393. return
  394. logs = open(self.log_file, encoding='utf-8').read().strip().split('\n')
  395. self.eta = None
  396. self.train_metrics = None
  397. self.eval_metrics = None
  398. if self.download_status != "Done":
  399. self.download_status = None
  400. start_time_timestamp = osp.getctime(self.log_file)
  401. for line in logs[::1]:
  402. try:
  403. start_time_str = " ".join(line.split()[0:2])
  404. start_time_array = time.strptime(start_time_str,
  405. "%Y-%m-%d %H:%M:%S")
  406. start_time_timestamp = time.mktime(start_time_array)
  407. break
  408. except Exception as e:
  409. pass
  410. for line in logs[::-1]:
  411. if line.count('Train Complete!'):
  412. self.train_stage = "Train Complete"
  413. if line.count('Training stop with error!'):
  414. self.train_error = line
  415. if self.train_metrics is not None \
  416. and self.eval_metrics is not None and self.eval_done and self.eta is not None:
  417. break
  418. items = line.strip().split()
  419. if line.count('Model saved in'):
  420. self.eval_done = True
  421. if line.count('download completed'):
  422. self.download_status = 'Done'
  423. break
  424. if line.count('download failed'):
  425. self.download_status = 'Failed'
  426. break
  427. if self.download_status != 'Done':
  428. if line.count('[DEBUG]\tDownloading'
  429. ) and self.download_status is None:
  430. self.download_status = dict()
  431. if not line.endswith('KB/s'):
  432. continue
  433. speed = items[-1].strip('KB/s').split('=')[-1]
  434. download = items[-2].strip('M, ').split('=')[-1]
  435. total = items[-3].strip('M, ').split('=')[-1]
  436. self.download_status['speed'] = speed
  437. self.download_status['download'] = float(download)
  438. self.download_status['total'] = float(total)
  439. if self.eta is None:
  440. if line.count('eta') > 0 and (line[-3] == ':' or
  441. line[-4] == ':'):
  442. eta = items[-1].strip().split('=')[1]
  443. h, m, s = [int(x) for x in eta.split(':')]
  444. self.eta = h * 3600 + m * 60 + s
  445. if self.train_metrics is None:
  446. if line.count('[INFO]\t[TRAIN]') > 0 and line.count(
  447. 'Step') > 0:
  448. if not items[-1].startswith('eta'):
  449. continue
  450. self.train_metrics = dict()
  451. metrics = items[4:]
  452. for metric in metrics:
  453. try:
  454. name, value = metric.strip(', ').split('=')
  455. value = value.split('/')[0]
  456. if value.count('.') > 0:
  457. value = float(value)
  458. elif value == 'nan':
  459. value = 'nan'
  460. else:
  461. value = int(value)
  462. self.train_metrics[name] = value
  463. except:
  464. pass
  465. if self.eval_metrics is None:
  466. if line.count('[INFO]\t[EVAL]') > 0 and line.count(
  467. 'Finished') > 0:
  468. if not line.strip().endswith(' .'):
  469. continue
  470. self.eval_metrics = dict()
  471. metrics = items[5:]
  472. for metric in metrics:
  473. try:
  474. name, value = metric.strip(', ').split('=')
  475. value = value.split('/')[0]
  476. if value.count('.') > 0:
  477. value = float(value)
  478. else:
  479. value = int(value)
  480. self.eval_metrics[name] = value
  481. except:
  482. pass
  483. end_time_timestamp = osp.getmtime(self.log_file)
  484. t_diff = time.gmtime(end_time_timestamp - start_time_timestamp)
  485. self.running_duration = "{}小时{}分{}秒".format(
  486. t_diff.tm_hour, t_diff.tm_min, t_diff.tm_sec)
  487. class PruneLogReader(object):
  488. def init_attr(self):
  489. self.eta = None
  490. self.iters = None
  491. self.current = None
  492. self.progress = None
  493. def __init__(self, log_file):
  494. self.log_file = log_file
  495. self.init_attr()
  496. def update(self):
  497. if not osp.exists(self.log_file):
  498. return
  499. logs = open(self.log_file, encoding='utf-8').read().strip().split('\n')
  500. self.init_attr()
  501. for line in logs[::-1]:
  502. metric_loaded = True
  503. for k, v in self.__dict__.items():
  504. if v is None:
  505. metric_loaded = False
  506. break
  507. if metric_loaded:
  508. break
  509. if line.count("Total evaluate iters") > 0:
  510. items = line.split(',')
  511. for item in items:
  512. kv_list = item.strip().split()[-1].split('=')
  513. kv_list = [v.strip() for v in kv_list]
  514. setattr(self, kv_list[0], kv_list[1])
  515. class QuantLogReader:
  516. def __init__(self, log_file):
  517. self.log_file = log_file
  518. self.stage = None
  519. self.running_duration = None
  520. def update(self):
  521. if not osp.exists(self.log_file):
  522. return
  523. logs = open(self.log_file, encoding='utf-8').read().strip().split('\n')
  524. for line in logs[::-1]:
  525. items = line.strip().split(' ')
  526. if line.count('[Run batch data]'):
  527. info = items[-3][:-1].split('=')[1]
  528. batch_id = float(info.split('/')[0])
  529. batch_all = float(info.split('/')[1])
  530. self.running_duration = \
  531. batch_id / batch_all * (10.0 / 30.0)
  532. self.stage = 'Batch'
  533. break
  534. elif line.count('[Calculate weight]'):
  535. info = items[-3][:-1].split('=')[1]
  536. weight_id = float(info.split('/')[0])
  537. weight_all = float(info.split('/')[1])
  538. self.running_duration = \
  539. weight_id / weight_all * (3.0 / 30.0) + (10.0 / 30.0)
  540. self.stage = 'Weight'
  541. break
  542. elif line.count('[Calculate activation]'):
  543. info = items[-3][:-1].split('=')[1]
  544. activation_id = float(info.split('/')[0])
  545. activation_all = float(info.split('/')[1])
  546. self.running_duration = \
  547. activation_id / activation_all * (16.0 / 30.0) + (13.0 / 30.0)
  548. self.stage = 'Activation'
  549. break
  550. elif line.count('Finish quant!'):
  551. self.stage = 'Finish'
  552. break
  553. class PyNvml(object):
  554. """ Nvidia GPU驱动检测类,可检测当前GPU驱动版本"""
  555. class PrintableStructure(Structure):
  556. _fmt_ = {}
  557. def __str__(self):
  558. result = []
  559. for x in self._fields_:
  560. key = x[0]
  561. value = getattr(self, key)
  562. fmt = "%s"
  563. if key in self._fmt_:
  564. fmt = self._fmt_[key]
  565. elif "<default>" in self._fmt_:
  566. fmt = self._fmt_["<default>"]
  567. result.append(("%s: " + fmt) % (key, value))
  568. return self.__class__.__name__ + "(" + string.join(result,
  569. ", ") + ")"
  570. class c_nvmlMemory_t(PrintableStructure):
  571. _fields_ = [
  572. ('total', c_ulonglong),
  573. ('free', c_ulonglong),
  574. ('used', c_ulonglong),
  575. ]
  576. _fmt_ = {'<default>': "%d B"}
  577. ## Device structures
  578. class struct_c_nvmlDevice_t(Structure):
  579. pass # opaque handle
  580. c_nvmlDevice_t = POINTER(struct_c_nvmlDevice_t)
  581. def __init__(self):
  582. self.nvml_lib = None
  583. self.nvml_lib_refcount = 0
  584. self.lib_load_lock = threading.Lock()
  585. self.nvml_lib_path = None
  586. def nvml_init(self, nvml_lib_path=None):
  587. self.lib_load_lock.acquire()
  588. sysstr = platform.system()
  589. if nvml_lib_path is None or nvml_lib_path.strip() == "":
  590. if sysstr == "Windows":
  591. nvml_lib_path = osp.join(
  592. os.getenv("ProgramFiles", "C:/Program Files"),
  593. "NVIDIA Corporation/NVSMI")
  594. if not osp.exists(osp.join(nvml_lib_path, "nvml.dll")):
  595. nvml_lib_path = "C:\\Windows\\System32"
  596. elif sysstr == "Linux":
  597. p1 = "/usr/lib/x86_64-linux-gnu"
  598. p2 = "/usr/lib/i386-linux-gnu"
  599. if osp.exists(osp.join(p1, "libnvidia-ml.so.1")):
  600. nvml_lib_path = p1
  601. elif osp.exists(osp.join(p2, "libnvidia-ml.so.1")):
  602. nvml_lib_path = p2
  603. else:
  604. nvml_lib_path = ""
  605. else:
  606. nvml_lib_path = "N/A"
  607. nvml_lib_dir = nvml_lib_path
  608. if sysstr == "Windows":
  609. nvml_lib_path = osp.join(nvml_lib_dir, "nvml.dll")
  610. else:
  611. nvml_lib_path = osp.join(nvml_lib_dir, "libnvidia-ml.so.1")
  612. self.nvml_lib_path = nvml_lib_path
  613. try:
  614. self.nvml_lib = CDLL(nvml_lib_path)
  615. fn = self._get_fn_ptr("nvmlInit_v2")
  616. fn()
  617. if sysstr == "Windows":
  618. driver_version = bytes.decode(
  619. self.nvml_system_get_driver_version())
  620. if driver_version.strip() == "":
  621. nvml_lib_path = osp.join(nvml_lib_dir, "nvml9.dll")
  622. self.nvml_lib = CDLL(nvml_lib_path)
  623. fn = self._get_fn_ptr("nvmlInit_v2")
  624. fn()
  625. except Exception as e:
  626. raise e
  627. finally:
  628. self.lib_load_lock.release()
  629. self.lib_load_lock.acquire()
  630. self.nvml_lib_refcount += 1
  631. self.lib_load_lock.release()
  632. def create_string_buffer(self, init, size=None):
  633. if isinstance(init, bytes):
  634. if size is None:
  635. size = len(init) + 1
  636. buftype = c_char * size
  637. buf = buftype()
  638. buf.value = init
  639. return buf
  640. elif isinstance(init, int):
  641. buftype = c_char * init
  642. buf = buftype()
  643. return buf
  644. raise TypeError(init)
  645. def _get_fn_ptr(self, name):
  646. return getattr(self.nvml_lib, name)
  647. def nvml_system_get_driver_version(self):
  648. c_version = self.create_string_buffer(81)
  649. fn = self._get_fn_ptr("nvmlSystemGetDriverVersion")
  650. ret = fn(c_version, c_uint(81))
  651. return c_version.value
  652. def nvml_device_get_count(self):
  653. c_count = c_uint()
  654. fn = self._get_fn_ptr("nvmlDeviceGetCount_v2")
  655. ret = fn(byref(c_count))
  656. return c_count.value
  657. def nvml_device_get_handle_by_index(self, index):
  658. c_index = c_uint(index)
  659. device = PyNvml.c_nvmlDevice_t()
  660. fn = self._get_fn_ptr("nvmlDeviceGetHandleByIndex_v2")
  661. ret = fn(c_index, byref(device))
  662. return device
  663. def nvml_device_get_memory_info(self, handle):
  664. c_memory = PyNvml.c_nvmlMemory_t()
  665. fn = self._get_fn_ptr("nvmlDeviceGetMemoryInfo")
  666. ret = fn(handle, byref(c_memory))
  667. return c_memory