Display a table of GPU availability on a Slurm server.

gpu_avail.py

#!/usr/bin/env python3

"""
Display a table of GPU availability on a Slurm server.
Usage:
  $ scontrol -M gpu show node | gpu_avail.py
"""

import sys


class Node:
    pass


try:
    lines = [line.rstrip('\r\n') for line in sys.stdin]
except KeyboardInterrupt:
    sys.exit(1)
lines.reverse()
nodes = []
while lines:
    line = lines.pop().strip()
    if not line:
        continue
    node = Node()
    node.arch = '-'
    nodes.append(node)
    assert line.startswith('NodeName=')
    for pair in line.split():
        key, val = pair.split('=')
        if key == 'NodeName':
            node.name = val
        elif key == 'Arch':
            node.arch = val
    assert hasattr(node, 'name')
    while lines and lines[-1].startswith('   '):
        line = lines.pop().strip()
        if line.startswith('ActiveFeatures='):
            node.features = line.split('=', maxsplit=1)[1]
        elif line.startswith('Partitions='):
            node.partitions = line.split('=', maxsplit=1)[1].rstrip('-')
        elif line.startswith('CfgTRES') or line.startswith('AllocTRES'):
            attr = 'gpu_total' if line.startswith('CfgTRES') else 'gpu_alloc'
            line = line.split('=', maxsplit=1)[1]
            count = 0
            if line:
                for pair in line.split(','):
                    key, val = pair.split('=')
                    if key == 'gres/gpu':
                        count = int(val)
                        break
            setattr(node, attr, count)
    for attr in ('features', 'partitions', 'gpu_total', 'gpu_alloc'):
        assert hasattr(node, attr)
    assert None not in (node.features, node.partitions, node.gpu_total, node.gpu_alloc)

header = ['name', 'arch', 'features', 'partitions', 'gpu_total', 'gpu_avail']
widths = [len(x) for x in header]
for node in nodes:
    for idx, key in enumerate(header):
        val = node.gpu_total - node.gpu_alloc if key == 'gpu_avail' else getattr(node, key)
        widths[idx] = max(widths[idx], len(str(val)))
fmt = ' '.join([f'{{: <{w}}}' for w in widths])
print(fmt.format(*header))
for node in nodes:
    gpu_avail = node.gpu_total - node.gpu_alloc
    print(fmt.format(node.name, node.arch, node.features, node.partitions, node.gpu_total, gpu_avail))