monitor-cluster-resources

#!/usr/bin/python3 -u

# Copyright 2023 Robert Krawitz/Red Hat
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import re
from datetime import datetime
import time
import sys
import subprocess
import selectors
import argparse
import fnmatch


node_cpu = {}
node_memory = {}
node_cpu_capacity = {}
node_memory_capacity = {}
node_pods_capacity = {}
pod_status = {}
pod_memory = {}
pod_cpu = {}
node_pods = {}
pod_nodes = {}
all_nodes = []
starting_time = None
resources_to_list = ['cpu', 'memory', 'pods']
interval = 5
first_time = True
last_row = None
last_row_base = None
last_changed = 0
tab = '\t'
nl = '\n'
n_timestamps = 0
nodes_schedulable = {}
node_pattern = None


def format_pods(num: float):
    if args.absolute:
        return str(int(num))
    else:
        return pformat(num)


def format_cpu(num: float):
    return pformat(num)


def format_memory(num: float):
    if args.absolute:
        if num < 1024 or args.exact:
            return str(int(num))
        elif num < 1024 * 1024:
            return f'{pformat(num / 1024)}Ki'
        elif num < 1024 * 1024 * 1024:
            return f'{pformat(num / (1024 * 1024))}Mi'
        elif num < 1024 * 1024 * 1024 * 1024:
            return f'{pformat(num / (1024 * 1024 * 1024))}Gi'
        else:
            return f'{pformat(num / (1024 * 1024 * 1024 * 1024))}Ti'
    else:
        return pformat(num)


print_funcs = {
    'cpu': format_cpu,
    'memory': format_memory,
    'pods': format_pods
    }


parser = argparse.ArgumentParser(description='Monitor cluster resources')
parser.add_argument('-N', '--node-pattern', help='Pattern(s) of nodes to list',
                    metavar='pattern', action='append')
parser.add_argument('-P', '--pod-pattern', help='Pattern(s) of pods to list',
                    metavar='pattern', action='append')
parser.add_argument('-R', '--resource',
                    help='Resource(s) to list (memory, cpu, pods).  Repeat for multiple.',
                    metavar='resource', action='append')
parser.add_argument('-i', '--interval', help=f'Interval between scans (default {interval})',
                    metavar='time interval', type=float)
parser.add_argument('-q', '--quiet', help='Do not print rows for individual pods',
                    action='store_true')
parser.add_argument('-Q', '--veryquiet', help='Do not print individual pods that have changed',
                    action='store_true')
parser.add_argument('-c', '--changed-rows-only', help='Only print changed rows', action='store_true')
parser.add_argument('-t', '--relative-timestamps', help='Print relative timestamps', action='store_true')
parser.add_argument('-T', '--absolute-timestamps', help='Print absolute timestamps', action='store_true')
parser.add_argument('-O', '--once', help="Run only once", action='store_true')
parser.add_argument('-A', '--absolute', help="Report absolute units", action='store_true')
parser.add_argument('-E', '--exact', help="Report exact quantities", action='store_true')
parser.add_argument('--only-matches', help="Count only matching pods", action='store_true')
parser.add_argument('--no-summary', help="Don't print summary", action='store_true')
parser.add_argument('--no-headers', help="Don't print headers", action='store_true')
parser.add_argument('rest', nargs=argparse.REMAINDER)
args = parser.parse_args()

if not args.once:
    if args.relative_timestamps:
        n_timestamps += 1
    if args.absolute_timestamps is None or args.absolute_timestamps is True:
        n_timestamps += 1

if args.resource:
    resources_to_list = [resource for resource in resources_to_list if resource in args.resource]
if args.interval:
    interval = args.interval

if args.node_pattern:
    node_pattern = args.node_pattern
elif args.rest:
    node_pattern = args.rest


def _ts(delimiter: str = ' '):
    global starting_time
    answers = []
    now_time = time.time()
    if args.once:
        return ''
    if args.absolute_timestamps is None or args.absolute_timestamps is True:
        answers.append(datetime.utcfromtimestamp(now_time).strftime('%Y-%m-%dT%T.%f'))
    if args.relative_timestamps:
        if starting_time is None:
            starting_time = now_time
        delta = now_time - starting_time
        answers.append(f'{delta:.3f}')
    return delimiter.join(answers)


def get_timestamp(string, delimiter: str = tab):
    """
    Return a string with a timestamp prepended to the first line
    and any other lines indented
    :param string: String to be timestamped
    :return: Timestamped string
    """
    string = re.sub(r'\n(.*\S.*)', r'\n            \1', string)
    if n_timestamps > 0:
        return delimiter.join([_ts(delimiter), string])
    else:
        return string


def timestamp(string, delimiter: str = ' '):
    """
    Timestamp a string and print it to stderr
    :param string: String to be printed to stderr with timestamp attached
    """
    print(get_timestamp(str(string), delimiter=delimiter))


def process_memory(token):
    if not token or token == '':
        return 0
    tokens = token.split()
    quantity = 0
    for token in tokens:
        m = re.match(r'([0-9]+)([kmg])?', token.lower())
        if m:
            mem = int(m.group(1))
            if m.group(2) == 'k':
                quantity += mem * 1024
            elif m.group(2) == 'm':
                quantity += mem * 1024 * 1024
            elif m.group(2) == 'g':
                quantity += mem * 1024 * 1024 * 1024
            else:
                quantity += mem
    return quantity


def process_cpu(token):
    if not token or token == '':
        return 0
    tokens = token.split()
    quantity = 0
    for token in tokens:
        m = re.match(r'([0-9]+)(m)?', token)
        cpu = int(m.group(1))
        if m.group(2) == 'm':
            quantity += cpu * .001
        else:
            quantity += cpu
    return quantity


def pformat(num: float, precision: int = 3):
    """
    Return a rounded representation of a number.
    :param num:
    :param precision:
    """
    if not args.absolute:
        num = num * 100
        precision = precision - 2
    try:
        if precision >= 1:
            return f'{num:.{precision}f}'
        else:
            return str(round(num))
    except (KeyboardInterrupt, BrokenPipeError):
        sys.exit(1)
    except TypeError:
        return str(num)
    except Exception:
        return str(num)


def is_match(elt: str, patterns=None):
    if patterns:
        for pattern in patterns:
            if fnmatch.fnmatch(elt, pattern):
                return True
        return False
    else:
        return True


def define_node(line: str):
    node, cpu, memory, pods = line.split('|')
    global all_nodes
    if node not in node_cpu_capacity:
        node_cpu_capacity[node] = process_cpu(cpu)
        node_memory_capacity[node] = process_memory(memory)
        node_pods_capacity[node] = int(pods)
        node_cpu[node] = 0
        node_memory[node] = 0
        node_pods[node] = 0
        all_nodes = sorted([node for node in list(node_cpu_capacity.keys()) if is_match(node, node_pattern)])


def fetch_node(node: str):
    run_command(['kubectl', 'get', 'node', node,
                 '-ojsonpath={%s}|{%s}|{%s}|{%s}{"\\n"}' % ('.metadata.name', '.status.allocatable.cpu',
                                                            '.status.allocatable.memory', '.status.allocatable.pods')],
                define_node)


def mk_header():
    line1_base = ''
    line2_base = ''
    line3_base = ''
    if args.once:
        return tab.join(['Node', tab.join(resources_to_list)])
    if n_timestamps > 0:
        line1_base = ''
        line2_base = f'Node{tab * n_timestamps}{(tab * len(resources_to_list)).join(all_nodes)}'
        if not args.no_summary:
            line2_base += f'{tab}%s%sPods changed' % ("%Capacity" if not args.absolute else "Total",
                                                      {(tab * len(resources_to_list))})
        line3_base = f'{"Timestamp"}{(tab * (n_timestamps - 1)) + ((tab + tab.join(resources_to_list)) * (len(all_nodes) + 1))}'
    else:
        line1a_base = ''
        if not args.no_summary:
            line1a_base = (tab * (len(resources_to_list) * len(all_nodes))) + "Summary"
        line1_base = f'Node{line1a_base}{nl}'
        line2_base = f'{(tab * len(resources_to_list)).join(all_nodes)}'
        if not args.no_summary:
            line2_base += '%s%s%sPods changed' % (tab * len(resources_to_list),
                                                  "%Capacity" if not args.absolute else "Total",
                                                  tab * len(resources_to_list))
        line3_base = f'{((tab.join(resources_to_list) + tab) * (len(all_nodes) + (0 if args.no_summary else 1)))}'
    summary_base = ''
    if not args.no_summary:
        summary_base = f'{tab}add{tab}remove{tab}change{tab}total'
    if args.veryquiet:
        return f'''{line1_base}{line2_base}
{line3_base}{summary_base}'''
    else:
        return f'''{line1_base}{line2_base}
{line3_base}{summary_base}'''


def get_node_data():
    answer = []
    totals = {'cpu': 0, 'memory': 0, 'pods': 0}
    for node in all_nodes:
        if 'cpu' in resources_to_list:
            num = node_cpu[node] / (node_cpu_capacity[node] if not args.absolute else 1)
            answer.append(format_cpu(num))
            if node in nodes_schedulable:
                totals['cpu'] += num
        if 'memory' in resources_to_list:
            num = node_memory[node] / (node_memory_capacity[node] if not args.absolute else 1)
            answer.append(format_memory(num))
            if node in nodes_schedulable:
                totals['memory'] += num
        if 'pods' in resources_to_list:
            num = node_pods[node] / (node_pods_capacity[node] if not args.absolute else 1)
            answer.append(format_pods(num))
            if node in nodes_schedulable:
                totals['pods'] += num
    if not args.no_summary:
        for resource in ['cpu', 'memory', 'pods']:
            if resource in resources_to_list:
                if args.exact:
                    answer.append(str(print_funcs[resource](totals[resource])))
                else:
                    answer.append(str(print_funcs[resource](totals[resource] / (len(nodes_schedulable) if not args.absolute else 1)
                                                            if len(nodes_schedulable) > 0 else 0)))
    return answer


def process_lines(lines: list):
    def node_status(lines: list):
        global nodes_schedulable
        nodes_schedulable = {}
        for line in lines:
            node, unschedulable = line.split('|')
            if node in all_nodes and not unschedulable.lower().endswith('true'):
                nodes_schedulable[node] = True

    global first_time, last_row, last_row_base
    new_pod_cpu = {}
    new_pod_memory = {}
    new_pod_nodes = {}
    pods_changed = []
    npods_added = 0
    npods_removed = 0
    npods_changed = 0
    npods_total = 0
    node_status(run_command(['kubectl', 'get', 'node',
                             '-ojsonpath={range .items[*]}{.metadata.name}|u={.spec.unschedulable}{"\\n"}{end}'],
                            process_stdout=True))
    for line in lines:
        pod, status, node, cpu, memory = line.split('|')
        if status != 'Running':
            continue
        new_pod_cpu[pod] = process_cpu(cpu)
        new_pod_memory[pod] = process_memory(memory)
        new_pod_nodes[pod] = node
    for pod in [spod for spod in pod_nodes if spod not in {pod: None for pod in new_pod_nodes}]:
        node = pod_nodes[pod]
        node_memory[node] -= pod_memory[pod]
        del pod_memory[pod]
        node_cpu[node] -= pod_cpu[pod]
        del pod_cpu[pod]
        node_pods[node] -= 1
        del pod_nodes[pod]
        if is_match(pod, args.pod_pattern):
            pods_changed.append(f'-{node}.{pod}')
            npods_removed += 1
            if not args.quiet and not args.veryquiet:
                timestamp(f'{tab.join(get_node_data())}{tab}-{node}.{pod}')

    for pod, node in new_pod_nodes.items():
        cpu = new_pod_cpu[pod]
        memory = new_pod_memory[pod]
        prefix = '+'
        pod_node = None
        if is_match(pod, args.pod_pattern):
            npods_total += 1
        elif args.only_matches:
            continue
        if pod in pod_nodes:
            if node == pod_nodes[pod] and cpu == pod_cpu[pod] and memory == pod_memory[pod]:
                continue
            else:
                pod_node = pod_nodes[pod]
                node_memory[pod_node] -= pod_memory[pod]
                node_cpu[pod_node] -= pod_cpu[pod]
                node_pods[pod_node] -= 1
                prefix = '!'
        pod_memory[pod] = memory
        pod_cpu[pod] = cpu
        pod_nodes[pod] = node
        node_memory[node] += memory
        node_cpu[node] += cpu
        node_pods[node] += 1
        if is_match(pod, args.pod_pattern):
            if prefix == '+':
                npods_added += 1
            else:
                npods_changed += 1
            pods_changed.append(f'{prefix}{node}.{pod}')
            if ((not args.once and not args.quiet and not args.veryquiet and
                 (not pod_node or pod_node in all_nodes) and node in all_nodes)):
                timestamp(f'{tab.join(get_node_data())}{tab}{prefix}{node}.{pod}')
    if args.once:
        data = get_node_data()
        element_counter = 0
        nodes = all_nodes
        if not args.no_summary:
            nodes.append('Summary')
        for node in nodes:
            print(tab.join([node, tab.join(data[element_counter:element_counter + len(resources_to_list)])]))
            element_counter += len(resources_to_list)
        if args.quiet and not args.veryquiet:
            print(" ".join(sorted(pods_changed)))
    elif args.quiet or args.veryquiet:
        row_base = tab.join(get_node_data())
        if not args.no_summary:
            row_base += f'{tab}{npods_added}{tab}{npods_removed}{tab}{npods_changed}{tab}{npods_total}'
        if args.quiet and not args.veryquiet:
            row_base += f'{tab}{" ".join(sorted(pods_changed))}'
        row = get_timestamp(row_base)
        if first_time or row_base != last_row_base or not args.changed_rows_only:
            if last_row:
                print(last_row)
                last_row = None
            print(row)
            first_time = False
        else:
            last_row = row
        last_row_base = row_base


def run_command(cmd, process_stdout=None, process_stderr=None):
    """ Run specified command, capturing stdout and stderr as array of timestamped lines.
        Optionally fail if return status is non-zero.  Also optionally report
        stdout and/or stderr to the appropriate file descriptors
    """
    with subprocess.Popen(cmd, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE) as command:
        sel = selectors.DefaultSelector()
        sel.register(command.stdout, selectors.EVENT_READ)
        sel.register(command.stderr, selectors.EVENT_READ)
        foundSomething = True
        stdout_lines = []
        while foundSomething:
            # Keep reading until we reach EOF on both channels.
            # command.poll() is not a good criterion because the process
            # might complete before everything has been read.
            foundSomething = False
            for key, _ in sel.select():
                data = key.fileobj.readline()
                if len(data) > 0:
                    foundSomething = True
                    data = data.decode().rstrip()
                    if key.fileobj is command.stdout:
                        if process_stdout is True:
                            stdout_lines.append(data)
                        elif process_stdout:
                            process_stdout(data)
                        else:
                            timestamp(data)
                    elif key.fileobj is command.stderr:
                        if process_stderr:
                            process_stderr(data)
                        else:
                            timestamp(data)
        if process_stdout is True:
            return stdout_lines


try:
    run_command(['kubectl', 'get', 'node',
                 ('-ojsonpath={range .items[*]}{%s}{%s}{%s}{%s}{"\\n"}{end}' %
                  ('.metadata.name', '.status.allocatable.cpu', '.status.allocatable.memory', '.status.allocatable.pods'))],
                define_node)

    if not args.no_headers:
        print(mk_header())
    exit_next = False
    while True:
        # Unfortunately, watching doesn't work here because we don't normally
        # receive any notification that the pod has gone away.
        curtime = time.time()
        process_lines(run_command(['kubectl', 'get', 'pod', '-A', '-o',
                                   ('jsonpath='
                                    '{range .items[*]}'
                                    '{.metadata.name}.{.metadata.namespace}|'
                                    '{.status.phase}|'
                                    '{.spec.nodeName}|'
                                    '{.spec.containers[*].resources.requests.cpu}|'
                                    '{.spec.containers[*].resources.requests.memory}{"\\n"}{end}')],
                                  process_stdout=True))
        if exit_next or args.once:
            sys.exit()
        try:
            now = time.time()
            if now - curtime < interval:
                time.sleep(interval - (now - curtime))
        except KeyboardInterrupt:
            first_time = True
            last_row = None
            exit_next = True
except (KeyboardInterrupt, BrokenPipeError):
    sys.exit()