Iperf chart generator (Salt Runner)

pyperf_diagram.py

### PyPerf Diagram Generator
## Salt Runner, 
## gets bandwidth results from each minion and generates a diagram

import ast
import pygraphviz as pgv
import logging
import time
from dictor import dictor
import salt.client
from salt.exceptions import AuthorizationError

log = logging.getLogger(__name__)
client = salt.client.LocalClient()

timestamp = time.strftime("%Y%m%d_%H%M", time.localtime())
label = f"{timestamp}\n\n\n"

# you can give each host an alias for the datacenter its located in,
minion_list = [
    "nyweb1,ec2us",
    "ukweb1,ec2uk",
    "sydweb1,ec2aus",
    "db1,dc123"
]

def _get_color(bw):
    """ create colored arrows between hosts, red to green depending on bw speed """
    r1 = range(0, 10, 1)
    r2 = range(10, 20, 1)
    r3 = range(20, 40, 1)
    r4 = range(40, 80, 1)
    r5 = range(80, 150, 1)
    r6 = range(150, 5000, 1)

    color = "#B0B0B0"

    if bw in r1:
        color = "#D51000"
    if bw in r2:
        color = "#D95200"
    if bw in r3:
        color = "#E1DC00"
    if bw in r4:
        color = "#00D2A3"
    if bw in r5:
        color = "#A4E500"
    if bw in r6:
        color = "#00D926"

    return color


def _get_snapshot():
    """connect to each minion, run a pyperf snapshot, get results as JSON"""

    ret = {}

    for row in sorted(minion_list):
        minion = row.split(",")[0]
        log.warning(f"perfing {minion}")

        # generate snapshot
        __salt__["salt.execute"](minion, "cmd.run", arg=[f"/opt/pyperf/snapshot.py {timestamp}"])

        file_exists = client.cmd(minion, "file.file_exists", ["/opt/pyperf/output/snapshot.json"])

        if file_exists[minion]:
            contents = client.cmd(minion, "file.read", ["/opt/pyperf/output/snapshot.json"])

            for key, val in ast.literal_eval(contents[minion]).items():
                if key == timestamp:
                    ret[minion] = val
    return ret


def _get_region(minion):
    """ provide a minion hostname, get its geo region or datacenter name """
    for row in minion_list:
        if minion == row.split(",")[0]:
            return row.split(",")[1]


def _generate_diagram(snapshot, label):
    """uses graphviz to generate bwidth diagram"""
    weight = 1
    A = pgv.AGraph(directed=True, strict=False)
    A.graph_attr.update(landscape="false", ranksep=4)

    for k, v in snapshot.items():

        # get region of minion
        region = _get_region(k)

        A.add_node(region, color="gray", shape="box")
        for target, bw in v.items():
            color = _get_color(int(bw))
            target_region = _get_region(target)
            A.add_edge(
                region,
                target_region,
                weight=weight,
                label=bw,
                labelloc="c",
                labelfloat=True,
                color=color,
                fontcolor=color,
                fontsize=8,
            )
            label = label + f"{region} > {target_region} = {bw}\n"

    A.layout(prog="circo")
    A.graph_attr.update(label=label, labelloc="b", fontsize=9, fontname="Arial")
    A.draw(f"/tmp/bw_{timestamp}.png")
    return f"diagram generated: /tmp/bw_{timestamp}.png"


def diagram():
    log.warning("running pyperf diagram")
    snapshot = _get_snapshot()
    log.warning(snapshot)
    return _generate_diagram(snapshot, label)