def getGPUNodes(pyslurmNodes): #TODO: need to change max_gpu_cnt if no-GPU node add other gres gpu_nodes = [ n_name for n_name, node in pyslurmNodes.items() if 'gpu' in node['features'] ] max_gpu_cnt = max([ MyTool.getNodeGresGPUCount(pyslurmNodes[n]['gres']) for n in gpu_nodes ]) return gpu_nodes, max_gpu_cnt
def getJobGPUNodes(jobs, pyslurmNodes): gpu_nodes = reduce(lambda rlt, curr: rlt.union(curr), [ set(job['gpus_allocated'].keys()) for job in jobs.values() if 'gpus_allocated' in job and job['gpus_allocated'] ], set()) if gpu_nodes: max_gpu_cnt = max([ MyTool.getNodeGresGPUCount(pyslurmNodes[n]['gres']) for n in gpu_nodes ]) else: max_gpu_cnt = 0 return gpu_nodes, max_gpu_cnt
def getNodeAllocGPU(node_name, node_dict=None): if not node_dict: node_dict = pyslurm.node().get() if node_name not in node_dict or 'gres' not in node_dict[node_name]: return None return MyTool.getNodeGresGPUCount(node_dict[node_name]['gres'])