Show statistics on PBS computing nodes.

pbstat.py

#!/bin/sh
""":" .

exec python "$0" "$@"
"""

# This file is distributed under the MIT License:
#
# MIT License
#
# Copyright (c) 2018-2020 Takahiro Ueda
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

__doc__ = """Show statistics on PBS computing nodes."""

import itertools
import re
import subprocess
import xml.etree.ElementTree

if False:
    from typing import Any, Dict, List, Sequence  # noqa: F401


def is_consecutive(a, b):
    # type: (Any, Any) -> bool
    """Return True if the two IDs are considered as consecutive."""
    try:
        if int(a) + 1 == int(b):
            return True
    except ValueError:
        pass
    return False


def sorted_numerically(a):
    # type: (Sequence[str]) -> Sequence[str]
    """Perform numerical sorting."""
    try:
        b = [int(x) for x in a]
        b.sort()
        return tuple(str(x) for x in b)
    except ValueError:
        return a


def group_by_id(a):
    # type: (Sequence[str]) -> List[str]
    """Group job IDs.

    >>> group_by_id(['98', '98', '98', '99', '100', '101', '102'])
    ['98*3', '99..102']

    >>> group_by_id(['5', '5', '6', '6', '8', '8'])
    ['5*2..6*2', '8*2']

    >>> group_by_id(['9', '9', '9', '10', '10', '10', '11', '11', '11'])
    ['9*3..11*3']

    >>> group_by_id(['100[99]', '100[100]', '100[101]', '100[102]'])
    ['100[99..102]']

    """
    b = [[key, key, len(list(group))]
         for key, group in itertools.groupby(sorted_numerically(a))]

    # consecutive job ids
    loop = True
    while loop:
        loop = False
        for i in range(0, len(b) - 1):
            b1 = b[i]
            b2 = b[i + 1]
            if b1[2] == b2[2] and is_consecutive(b1[1], b2[0]):
                b1[1] = str(int(str(b1[1])) + 1)
                del b[i + 1]
                loop = True
                break

    # array jobs
    loop = True
    while loop:
        loop = False
        for i in range(0, len(b) - 1):
            b1 = b[i]
            b2 = b[i + 1]
            if b1[0] == b1[1] and b2[0] == b2[1] and b1[2] == b2[2]:
                a1 = re.split(r'[\[\].]+', str(b1[0]))
                a2 = re.split(r'[\[\].]+', str(b2[0]))
                if len(a1) >= 3 and len(a2) >= 3 and not a1[-1] and not a2[-1]:
                    if ((len(a1) == 3 and is_consecutive(a1[1], a2[1]))
                        or (len(a1) == 4 and is_consecutive(a1[2], a2[1]))):
                            if len(a2) == 3:
                                b1[0] = b1[1] = '{0}[{1}..{2}]'.format(
                                    a1[0], a1[1], a2[1]
                                )
                            else:
                                b1[0] = b1[1] = '{0}[{1}..{2}]'.format(
                                    a1[0], a1[1], a2[2]
                                )
                            del b[i + 1]
                            loop = True
                            break

    return [str(key1) if key1 == key2 and count == 1 else
            '{0}*{1}'.format(key1, count) if key1 == key2 else
            '{0}..{1}'.format(key1, key2) if count == 1 else
            '{0}*{2}..{1}*{2}'.format(key1, key2, count)
            for key1, key2, count in b]


output = subprocess.Popen(['pbsnodes', '-x'],
                          stdout=subprocess.PIPE).communicate()[0]

root = xml.etree.ElementTree.fromstring(output)

nodes = []  # type: List[Dict[str, Any]]
for child in root:
    assert child.tag == 'Node'
    node = {}  # type: Dict[str, Any]
    for item in child:
        text = item.text if item.text is not None else ''
        if item.tag == 'status':
            status = {}
            subitems = text.split(',')  # type: List[str]
            for subitem in subitems:
                a = subitem.split('=', 1)
                status[a[0]] = a[1]
            node[item.tag] = status
        elif item.tag == 'jobs':
            subitems = text.split(',')
            subitems = [b.split('.', 1)[0] for b in subitems]
            subitems = [b.split('/', 1)[1] for b in subitems]
            subitems = group_by_id(subitems)
            node[item.tag] = ','.join(subitems)
        else:
            node[item.tag] = text
    nodes.append(node)

rows = [('name', 'state', 'np', 'mem', 'prop', 'jobs')] + [(
    n['name'],
    n['state'],
    n['np'] if 'np' in n else '',
    n['status']['totmem'] if 'status' in n and 'totmem' in n['status'] else '',
    n['properties'] if 'properties' in n else '',
    n['jobs'] if 'jobs' in n else '',
) for n in nodes]

fmt = ('{{0:{0}}}  {{1:>{1}}}  {{2:{2}}}  {{3:{3}}}  {{4:{4}}}  '
       '{{5:{5}}}').format(
    max([len(x[0]) for x in rows]),
    max([len(x[1]) for x in rows]),
    max([len(x[2]) for x in rows]),
    max([len(x[3]) for x in rows]),
    max([len(x[4]) for x in rows]),
    max([len(x[5]) for x in rows]),
)

for row in rows:
    print(fmt.format(*row))