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'])