system.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. import sys
  2. import os
  3. import psutil
  4. import platform
  5. def pkill(pid):
  6. try:
  7. parent = psutil.Process(pid)
  8. for child in parent.children(recursive=True):
  9. child.kill()
  10. parent.kill()
  11. except:
  12. print("Try to kill process {} failed.".format(pid))
  13. def get_system_info(machine_info={}):
  14. if machine_info:
  15. return {'status': 1, 'info': machine_info}
  16. from .utils import get_gpu_info
  17. gpu_info, message = get_gpu_info()
  18. cpu_num = os.environ.get('CPU_NUM', 1)
  19. sysstr = platform.system()
  20. machine_info['message'] = message
  21. machine_info['cpu_num'] = cpu_num
  22. machine_info['gpu_num'] = gpu_info['gpu_num']
  23. machine_info['sysstr'] = sysstr
  24. if gpu_info['gpu_num'] > 0:
  25. machine_info['driver_version'] = gpu_info['driver_version']
  26. machine_info['gpu_free_mem'] = gpu_info['mem_free']
  27. return {'status': 1, 'info': machine_info}
  28. def get_gpu_memory_info(machine_info):
  29. gpu_mem_infos = list()
  30. if machine_info['gpu_num'] == 0:
  31. pass
  32. else:
  33. from .utils import get_gpu_info
  34. gpu_info, message = get_gpu_info()
  35. for i in range(gpu_info['gpu_num']):
  36. attr = {
  37. 'free': gpu_info['mem_free'][i],
  38. 'used': gpu_info['mem_used'][i],
  39. 'total': gpu_info['mem_total'][i]
  40. }
  41. gpu_mem_infos.append(attr)
  42. return {'status': 1, 'gpu_mem_infos': gpu_mem_infos}
  43. def get_machine_info(data, machine_info):
  44. path = None
  45. if "path" in data:
  46. path = data['path']
  47. if path in machine_info:
  48. return {'status': 1, 'info': machine_info}
  49. from .utils import get_machine_info
  50. info = get_machine_info(path)
  51. machine_info = info
  52. return {'status': 1, 'info': machine_info}
  53. def get_gpu_memory_size(data):
  54. """获取显存大小
  55. Args:
  56. request(comm.Request): 其中request.params为dict, key包括
  57. 'path' 显卡驱动动态链接库路径
  58. """
  59. from .utils import PyNvml
  60. p = PyNvml()
  61. p.nvml_init(data['path'])
  62. count = p.nvml_device_get_count()
  63. gpu_mem_infos = []
  64. for i in range(count):
  65. handler = p.nvml_device_get_handle_by_index(i)
  66. mem = p.nvml_device_get_memory_info(handler)
  67. attr = {'free': mem.free, 'used': mem.used, 'total': mem.total}
  68. gpu_mem_infos.append(attr)
  69. return {'status': 1, 'gpu_mem_infos': gpu_mem_infos}
  70. def exit_system(monitored_processes):
  71. while not monitored_processes.empty():
  72. pid = monitored_processes.get(timeout=0.5)
  73. print("Try to kill process {}".format(pid))
  74. pkill(pid)
  75. return {'status': 1}