Visualizer.py 6.99 KB
Newer Older
Alberts S's avatar
Alberts S committed
1
import asyncio
2
import logging
Alberts S's avatar
Alberts S committed
3
4
5
6
7
8
9
import pickle
import time

import matplotlib
import matplotlib.pyplot as plt
import networkx as nx
import seaborn as sns
Alberts S's avatar
Alberts S committed
10
from pyvis.network import Network
Alberts S's avatar
Alberts S committed
11

12
from CapybaraNetty import CapybaraNetty
Alberts S's avatar
Alberts S committed
13

14
15

class Visualizer(CapybaraNetty):
Alberts S's avatar
Alberts S committed
16
    def __init__(self):
17
        super().__init__()
Alberts S's avatar
Alberts S committed
18
19
        self.pinger_data = []
        self.external_targets_for_optimizations = []
20
        self.optimization_threshold_ms = 0
Alberts S's avatar
Alberts S committed
21
22

    async def run(self):
23
24
        G_simple = await self.get_host_latency_overview()
        if len(G_simple) > 0:
Alberts S's avatar
Alberts S committed
25
            await self.draw_graphs(G_simple, "simple")
26
            G_optimizations = await self.get_host_latency_optimization_overview(G_simple.copy())
Alberts S's avatar
Alberts S committed
27
            await self.draw_graphs(self.get_graph_only_improved_overview(G_optimizations), "optimizations_overview")
28
            for target in self.external_targets_for_optimizations:
Alberts S's avatar
Alberts S committed
29
30
31
32
33
34
35
36
                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")
37
38
        else:
            self.__logger.warning("Doing nothing since G is empty")
Alberts S's avatar
Alberts S committed
39
40
41
42

    async def get_host_latency_overview(self):
        G = nx.Graph()
        for result in self.pinger_data:
43
44
            host = result["from_name"]
            dest = result["to_name"]
Alberts S's avatar
Alberts S committed
45
            time = int(result["time"])
Alberts S's avatar
Alberts S committed
46
47
48
49
50
51
52
53
            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")
54
55
            G.nodes[host]["type"] = result["from_type"]
            G.nodes[dest]["type"] = result["to_type"]
Alberts S's avatar
Alberts S committed
56
57
        return G

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

62
        for target in self.external_targets_for_optimizations:
63
64
65
66
67
            # 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")
68
69
70
71
72
73

            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
Alberts S's avatar
Alberts S committed
74
                    and path[-1] in self.external_targets_for_optimizations
75
76
77
                ):
                    optimized_ms = G[key][path[0]]["weight"] - lengths[key]
                    if optimized_ms < self.optimization_threshold_ms:
Alberts S's avatar
Alberts S committed
78
79
80
                        self.__logger.debug(
                            f"VISUALIZER IGNORED OPTIMIZATION: {optimized_ms:<3d}ms {'->'.join(path[::-1])}"
                        )
81
82
83
84
85
86
87
88
89
90
91
92
93
94
                    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)

Alberts S's avatar
Alberts S committed
95
96
97
                        self.__logger.info(
                            f"VISUALIZER VIABLE OPTIMIZATION: {optimized_ms:<3d}ms {'->'.join(path[::-1])}"
                        )
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
        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)

Alberts S's avatar
Alberts S committed
119
120
121
122
123
124
125
126
    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")
127
        net.show_buttons(filter_=["physics"])
Alberts S's avatar
Alberts S committed
128
        net.from_nx(G)
129
        net.barnes_hut(spring_length=250)
Alberts S's avatar
Alberts S committed
130
131
132
133
134
135
136
137
138
        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
Alberts S's avatar
Alberts S committed
139
140
141
142
143
144
145
146
        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)
Alberts S's avatar
Alberts S committed
147
        output_name = f"{output_base}_networkx.png"
148
149
        self.__logger.info(f"Writing {output_name}")
        plt.savefig(output_name, bbox_inches="tight")
Alberts S's avatar
Alberts S committed
150
151
        plt.close()

Alberts S's avatar
Alberts S committed
152
    async def run_daemon(self, pinger_data_fn):
Alberts S's avatar
Alberts S committed
153
        while True:
Alberts S's avatar
Alberts S committed
154
            self.optimization_threshold_ms = self.config["optimization_threshold_ms"]
Alberts S's avatar
Alberts S committed
155
            self.pinger_data = pinger_data_fn()
Alberts S's avatar
Alberts S committed
156
            self.external_targets_for_optimizations = self.config["external_targets_for_optimizations"]
Alberts S's avatar
Alberts S committed
157
            await self.run()
Alberts S's avatar
Alberts S committed
158
            await asyncio.sleep(self.config["visualizer_interval"])