Visualizer.py

import asyncio
import logging
import pickle
import time

import matplotlib
import matplotlib.pyplot as plt
import networkx as nx
import seaborn as sns
from pyvis.network import Network

from CapybaraNetty import CapybaraNetty


class Visualizer(CapybaraNetty):
    def __init__(self):
        super().__init__()
        self.pinger_data = []
        self.external_targets_for_optimizations = []
        self.optimization_threshold_ms = 0

    async def run(self):
        G_simple = await self.get_host_latency_overview()
        if len(G_simple) > 0:
            await self.draw_graphs(G_simple, "simple")
            G_optimizations = await self.get_host_latency_optimization_overview(G_simple.copy())
            await self.draw_graphs(self.get_graph_only_improved_overview(G_optimizations), "optimizations_overview")
            for target in self.external_targets_for_optimizations:
                target_filtered_graph = self.get_graph_only_improved_specific_target(G_optimizations, target)
                if target_filtered_graph.number_of_nodes() > 0:
                    await self.draw_graphs(
                        target_filtered_graph,
                        f"optimizations_target_{target}",
                    )
                else:
                    self.__logger.debug(f"Not drawing graph for {target=} since the filtered graph is empty")
        else:
            self.__logger.warning("Doing nothing since G is empty")

    async def get_host_latency_overview(self):
        G = nx.Graph()
        for result in self.pinger_data:
            host = result["from_name"]
            dest = result["to_name"]
            time = int(result["time"])
            result_kind = result["result_kind"]
            if result_kind == "external":
                attrs = {
                    "weight": time,
                }
            else:
                attrs = {f"weight_{result_kind}": time}
            G.add_edge(host, dest, **attrs, color="black")
            G.nodes[host]["type"] = result["from_type"]
            G.nodes[dest]["type"] = result["to_type"]
        return G

    async def get_host_latency_optimization_overview(self, G: nx.Graph):
        def filter_node(n1):
            return G.nodes[n1].get("type") == "router" or (G.nodes[n1].get("type") == "target" and n1 == target)

        for target in self.external_targets_for_optimizations:
            # G has all routers and targets
            # We need G routers + single target, filter out all other targets
            # otherwise the graph finds shortest path with including external targets
            G_subgraph = nx.subgraph_view(G, filter_node=filter_node)
            lengths, paths = nx.single_source_dijkstra(G_subgraph, source=target, weight="weight")

            palette = sns.color_palette("tab10")
            palette_index = 0
            for key, path in paths.items():
                if len(path) > 2 and not (
                    path[0] in self.external_targets_for_optimizations
                    and path[-1] in self.external_targets_for_optimizations
                ):
                    optimized_ms = G[key][path[0]]["weight"] - lengths[key]
                    if optimized_ms < self.optimization_threshold_ms:
                        self.__logger.debug(
                            f"VISUALIZER IGNORED OPTIMIZATION: {optimized_ms:<3d}ms {'->'.join(path[::-1])}"
                        )
                    else:
                        color = palette[palette_index]
                        palette_index += 1
                        for previous, current in zip(path, path[1:]):
                            G[previous][current]["color"] = color
                            G[previous][current]["improved_path"] = True
                            G[previous][current].setdefault("improved_path_targets", []).append(target)

                        G[path[0]][path[-1]]["original_of_improved_path"] = True
                        G[path[0]][path[-1]].setdefault("improved_path_targets", []).append(target)
                        for node in path:
                            G.nodes[node]["improved_path"] = True
                            G.nodes[node].setdefault("improved_path_targets", []).append(target)

                        self.__logger.info(
                            f"VISUALIZER VIABLE OPTIMIZATION: {optimized_ms:<3d}ms {'->'.join(path[::-1])}"
                        )
        return G

    def get_graph_only_improved_overview(self, G: nx.Graph):
        def filter_edge(n1, n2):
            return G[n1][n2].get("improved_path", False) or G[n1][n2].get("original_of_improved_path", False)

        def filter_node(n1):
            return G.nodes[n1].get("improved_path", False)

        return nx.subgraph_view(G, filter_node=filter_node, filter_edge=filter_edge)

    def get_graph_only_improved_specific_target(self, G: nx.Graph, target):
        def filter_edge(n1, n2):
            return target in G[n1][n2].get("improved_path_targets", [])

        def filter_node(n1):
            return target in G.nodes[n1].get("improved_path_targets", [])

        G = self.get_graph_only_improved_overview(G)
        return nx.subgraph_view(G, filter_node=filter_node, filter_edge=filter_edge)

    async def draw_graphs(self, G: nx.Graph, file_base_name):
        output_base_name = f'out/network_latency_{file_base_name}_{time.strftime("%Y%m%d-%H_%M_%S")}'
        await self.draw_graph_networkx(G, output_base_name)
        await self.draw_graph_pyvis(G, output_base_name)

    async def draw_graph_pyvis(self, G, output_base):
        # PyVis
        net = Network("1280px", "1280px")
        net.show_buttons(filter_=["physics"])
        net.from_nx(G)
        net.barnes_hut(spring_length=250)
        for edge in net.edges:
            edge["label"] = edge["weight"]

        output_name = f"{output_base}_pyvis.html"
        self.__logger.info(f"Writing {output_name}")
        net.save_graph(output_name)

    async def draw_graph_networkx(self, G, output_base):
        # Matplotlib + NetworkX
        pos = nx.fruchterman_reingold_layout(G, 5)
        fig = plt.figure(1, figsize=(60, 20), dpi=120)

        labels = nx.get_edge_attributes(G, "weight")
        edge_colors = nx.get_edge_attributes(G, "color").values()
        node_colors = ["red" if node in self.external_targets_for_optimizations else "green" for node in G]
        nx.draw_networkx_edge_labels(G, pos, edge_labels=labels, label_pos=0.5)
        nx.draw(G, pos, with_labels=True, edge_color=edge_colors, node_size=100, node_color=node_colors)
        output_name = f"{output_base}_networkx.png"
        self.__logger.info(f"Writing {output_name}")
        plt.savefig(output_name, bbox_inches="tight")
        plt.close()

    async def run_daemon(self, pinger_data_fn):
        while True:
            self.optimization_threshold_ms = self.config["optimization_threshold_ms"]
            self.pinger_data = pinger_data_fn()
            self.external_targets_for_optimizations = self.config["external_targets_for_optimizations"]
            await self.run()
            await asyncio.sleep(self.config["visualizer_interval"])