From 7651f8201bf181ed2b73dee467cff30ade96ba59 Mon Sep 17 00:00:00 2001 From: Martin Liska Date: Thu, 14 Mar 2024 13:59:52 +0100 Subject: [PATCH] usage-wrapper: add --skip-gpu argument --- usage-wrapper.py | 55 +++++++++++++++++++++++++----------------------- 1 file changed, 29 insertions(+), 26 deletions(-) diff --git a/usage-wrapper.py b/usage-wrapper.py index d766037..4bb6631 100755 --- a/usage-wrapper.py +++ b/usage-wrapper.py @@ -144,6 +144,7 @@ def difference_in_gb(self): help='Scale up CPU data to used CPUs ' 'instead of available CPUs', ) parser.add_argument('--y-scale', type=int, help='Minimal y-scale (in GiB)') +parser.add_argument('--skip-gpu', action='store_false', help='Skip collecting statistics for GPU') args = parser.parse_args() @@ -182,32 +183,34 @@ def difference_in_gb(self): print('WARNING: disk IO counters not supported by the system') pass -try: - import GPUtil - - def collect_gpu(): - try: - return 100 * GPUtil.getGPUs()[0].load - except IndexError: - # sometimes GPUtil.getGPUs() returns [] if we are terminating - return 0 - - def collect_gpu_memory(): - try: - return GPUtil.getGPUs()[0].memoryUsed / 1024 - except IndexError: - # sometimes GPUtil.getGPUs() returns [] if we are terminating - return 0 - - gpu_stats = DataStatistic(collect_gpu) - # the memory consumption is reported in MiBs - gpu_mem_stats = DataStatistic(collect_gpu_memory) - collectors.append(gpu_stats) - collectors.append(gpu_mem_stats) -except ImportError: - gpu_stats = None - gpu_mem_stats = None - print('WARNING: missing GPUtil package (pip install GPUtil)') +gpu_stats = None +gpu_mem_stats = None + +if args.skip_gpu: + try: + import GPUtil + + def collect_gpu(): + try: + return 100 * GPUtil.getGPUs()[0].load + except IndexError: + # sometimes GPUtil.getGPUs() returns [] if we are terminating + return 0 + + def collect_gpu_memory(): + try: + return GPUtil.getGPUs()[0].memoryUsed / 1024 + except IndexError: + # sometimes GPUtil.getGPUs() returns [] if we are terminating + return 0 + + gpu_stats = DataStatistic(collect_gpu) + # the memory consumption is reported in MiBs + gpu_mem_stats = DataStatistic(collect_gpu_memory) + collectors.append(gpu_stats) + collectors.append(gpu_mem_stats) + except ImportError: + print('WARNING: missing GPUtil package (pip install GPUtil)') def get_process_name(proc):