import os class GPUManager(): queries = ('index', 'gpu_name', 'memory.free', 'memory.used', 'memory.total', 'power.draw', 'power.limit') def __init__(self): all_gpus = self.query_gpu(False) def get_info(self, ctype): cmd = 'nvidia-smi --query-gpu={} --format=csv,noheader'.format(ctype) lines = os.popen(cmd).readlines() lines = [line.strip('\n') for line in lines] return lines def query_gpu(self, show=True): num_gpus = len( self.get_info('index') ) all_gpus = [ {} for i in range(num_gpus) ] for query in self.queries: infos = self.get_info(query) for idx, info in enumerate(infos): all_gpus[idx][query] = info if 'CUDA_VISIBLE_DEVICES' in os.environ: CUDA_VISIBLE_DEVICES = os.environ['CUDA_VISIBLE_DEVICES'].split(',') selected_gpus = [] for idx, CUDA_VISIBLE_DEVICE in enumerate(CUDA_VISIBLE_DEVICES): find = False for gpu in all_gpus: if gpu['index'] == CUDA_VISIBLE_DEVICE: assert not find, 'Duplicate cuda device index : {}'.format(CUDA_VISIBLE_DEVICE) find = True selected_gpus.append( gpu.copy() ) selected_gpus[-1]['index'] = '{}'.format(idx) assert find, 'Does not find the device : {}'.format(CUDA_VISIBLE_DEVICE) all_gpus = selected_gpus if show: allstrings = '' for gpu in all_gpus: string = '| ' for query in self.queries: if query.find('memory') == 0: xinfo = '{:>9}'.format(gpu[query]) else: xinfo = gpu[query] string = string + query + ' : ' + xinfo + ' | ' allstrings = allstrings + string + '\n' return allstrings else: return all_gpus def select_by_memory(self, numbers=1): all_gpus = self.query_gpu(False) assert numbers <= len(all_gpus), 'Require {} gpus more than you have'.format(numbers) alls = [] for idx, gpu in enumerate(all_gpus): free_memory = gpu['memory.free'] free_memory = free_memory.split(' ')[0] free_memory = int(free_memory) index = gpu['index'] alls.append((free_memory, index)) alls.sort(reverse = True) alls = [ int(alls[i][1]) for i in range(numbers) ] return sorted(alls) """ if __name__ == '__main__': manager = GPUManager() manager.query_gpu(True) indexes = manager.select_by_memory(3) print (indexes) """