import math
from dataclasses import dataclass
from typing import *
import networkx as nx
import numpy as np
import tensorkit as tk
from tensorkit import tensor as T
from tracegnn.data import *
from tracegnn.utils import *
from .constants import *
from .tensor_utils import *
import dgl
import torch
__all__ = [
'flat_to_nx_graphs',
'p_net_to_trace_graphs',
'GraphNodeMatch', 'GraphNodeDiff',
'diff_graph',
]
# util to reshape an array
def reshape_to(x, ndims):
shape = T.shape(x)
return T.reshape(x, [-1] + shape[len(shape) - ndims + 1:])
def to_scalar(x):
return T.to_numpy(x).tolist()
def flat_to_nx_graphs(p: tk.BayesianNet,
id_manager: TraceGraphIDManager,
latency_range: TraceGraphLatencyRangeFile,
min_edge_weight: float = 0.2,
) -> List[nx.Graph]:
"""Convert `p` net sampled from a flat TraceVAE to nx.Graph."""
# extract features
adjs = reshape_to(p['adj'].distribution.probs, 2)
node_counts = T.to_numpy(reshape_to(p['node_count'].tensor, 1))
node_types = T.to_numpy(reshape_to(p['node_type'].tensor, 2))
# span_counts = reshape_to(p['span_count'].tensor, 2)
if 'latency' in p:
latency_src = T.to_numpy(reshape_to(p['latency'].distribution.base_distribution.mean, 3))
latencies = np.zeros(latency_src.shape, dtype=np.float32)
for i in range(node_types.shape[0]):
for j in range(node_types.shape[1]):
try:
node_type = int(node_types[i, j])
mu, std = latency_range[node_type]
latencies[i, j] = latency_src[i, j] * std + mu
except KeyError:
latencies[i, j] = -1. # todo: is this okay?
else:
latencies = None
# build the graph
ret = []
for i, node_count in enumerate(node_counts):
g = nx.Graph()
# add nodes
for j in range(node_count):
g.add_node(j)
# add edges
adj = triu_to_dense(adjs[i: i+1], MAX_NODE_COUNT)
for u in range(node_count):
for v in range(u + 1, node_count):
w = float(to_scalar(adj[u, v]))
if w >= min_edge_weight:
g.add_edge(u, v, weight=w)
# add node attributes
for j in range(node_count):
node_type = int(node_types[i, j])
g.nodes[j]['node_type'] = node_type
g.nodes[j]['operation'] = id_manager.operation_id.reverse_map(node_type)
if latencies is not None:
for k, pfx in enumerate(('avg_', 'max_', 'min_')):
if k < LATENCY_DIM:
g.nodes[j][f'{pfx}latency'] = latencies[i, j, k]
# g.nodes[j]['span_count'] = to_scalar(span_counts[i, j])
# for pfx in ('avg_', 'max_', 'min_'):
# g.nodes[j][f'{pfx}latency'] = latencies[f'{pfx}latency'][i, j]
ret.append(g)
# return the graphs
return ret
def p_net_to_trace_graphs(p: tk.BayesianNet,
id_manager: TraceGraphIDManager,
latency_range: TraceGraphLatencyRangeFile,
discard_node_with_type_0: bool = True,
discard_node_with_unknown_latency_range: bool = True,
discard_graph_with_error_node_count: bool = False,
keep_front_shape: bool = False,
) -> Union[List[Optional[TraceGraph]], np.ndarray]:
"""Convert `p` net sampled from a flat TraceVAE to TraceGraph."""
if USE_MULTI_DIM_LATENCY_CODEC:
raise RuntimeError(f'`USE_MULTI_DIM_LATENCY_CODEC` is not supported.')
# find the base distribution (Normal, Categorical, OneHotCategorical)
def find_base(t: tk.StochasticTensor):
d = t.distribution
while not isinstance(d, (tk.Normal,
tk.Bernoulli,
tk.Categorical,
tk.OneHotCategorical)):
d = d.base_distribution
return d
# extract features
def get_adj(t, pad_value=0):
t = reshape_to(t, 2)
return np.stack(
[
T.to_numpy(triu_to_dense(
t[i: i + 1],
MAX_NODE_COUNT,
pad_value=pad_value
))
for i in range(len(t))
],
axis=0
)
def bernoulli_log_prob(l):
# log(1 / (1 + exp(-l)) = log(exp(l) / (1 + exp(l)))
return T.where(
l >= 0,
-T.log1p(T.exp(-l)),
l - T.log1p(T.exp(l)),
)
def softmax_log_prob(l):
# log(exp(l) / sum(exp(l))
return l - T.log_sum_exp(l, axis=[-1], keepdims=True)
front_shape = T.shape(p['adj'].tensor)[:-1]
adjs = get_adj(p['adj'].tensor)
adj_probs = get_adj(find_base(p['adj']).probs)
adj_logits = get_adj(bernoulli_log_prob(find_base(p['adj']).logits), pad_value=-100000)
node_counts = T.to_numpy(reshape_to(p['node_count'].tensor, 1))
node_types = T.to_numpy(reshape_to(p['node_type'].tensor, 2))
node_count_logits = T.to_numpy(reshape_to(softmax_log_prob(find_base(p['node_count']).logits), 2))
node_type_logits = T.to_numpy(reshape_to(softmax_log_prob(find_base(p['node_type']).logits), 3))
if 'latency' in p:
latencies = T.to_numpy(reshape_to(p['latency'].tensor, 3))
avg_latencies = latencies[..., 0]
latency_means = T.to_numpy(reshape_to(find_base(p['latency']).mean, 3))
latency_stds = T.to_numpy(reshape_to(find_base(p['latency']).std, 3))
# build the graph
ret = []
for i, node_count in enumerate(node_counts):
# extract the arrays
adj = adjs[i][:node_count][:, :node_count]
adj_prob = adj_probs[i][:node_count][:, :node_count]
adj_logit = adj_logits[i] # [:node_count][:, :node_count]
node_type = node_types[i] # [:node_count]
node_mask = np.full([node_count], True, dtype=np.bool)
node_count_logit = node_count_logits[i]
node_type_logit = node_type_logits[i]
if 'latency' in p:
avg_latency = avg_latencies[i]
latency_mean = latency_means[i]
latency_std = latency_stds[i]
# if `discard_node_with_type_0`, set all adjs that from / to `node_type == 0` as 0
node_count_new = node_count
for j in range(node_count):
n_type = int(node_type[j])
if (discard_node_with_type_0 and n_type == 0) or \
(discard_node_with_unknown_latency_range and n_type not in latency_range):
node_mask[j] = False
node_count_new -= 1
adj[:, j] = 0
adj[j, :] = 0
adj_prob[:, j] = 0
adj_prob[j, :] = 0
# for each column in `adj`, if there are more than 2 candidate in-edges,
# or no in-edge, then choose an edge sampled w.r.t. to adj_prob
for j in range(node_count):
if node_mask[j] and np.sum(adj[:, j]) != 1:
prob_vec = adj_prob[:, j]
prob_sum = np.sum(prob_vec)
if prob_sum > 1e-7:
pvals = prob_vec / np.sum(prob_vec)
pvals_mask = pvals > 1e-7
indices = np.arange(len(pvals))[pvals_mask]
k = indices[np.argmax(np.random.multinomial(1, pvals[pvals_mask]))]
adj[:, j] = 0
adj[k, j] = 1
# select the edges
edges = list(zip(*np.where(adj)))
if len(edges) < node_count_new - 1:
# pick out the root sub-graph
union_set = {j: -1 for j in range(node_count) if node_mask[j]}
def find_root(s):
t = union_set[s]
if t == -1:
return s
r = find_root(t)
if r != t:
union_set[s] = r
return r
def link_edge(s, t):
union_set[t] = s
edges_new = []
for s, t in edges:
link_edge(s, t)
for s, t in edges:
if s == 0 or find_root(s) == 0:
edges_new.append((s, t))
edges = edges_new
node_count_new = len(edges_new) + 1
if discard_graph_with_error_node_count and (node_count_new != node_count):
ret.append(None)
continue
# build the trace graph
def get_node(s):
if s not in nodes:
n_type = node_type[s]
if 'latency' in p:
latency = avg_latency[s]
if n_type in latency_range:
mu, std = latency_range[n_type]
latency = latency * std + mu
features = TraceGraphNodeFeatures(
span_count=1,
avg_latency=latency,
max_latency=latency,
min_latency=latency,
)
avg_latency_nstd = float(
abs(avg_latency[s] - latency_mean[s, 0]) /
latency_std[s, 0]
)
else:
features = TraceGraphNodeFeatures(
span_count=1,
avg_latency=math.nan,
max_latency=math.nan,
min_latency=math.nan,
)
avg_latency_nstd = 0
nodes[s] = TraceGraphNode.new_sampled(
node_id=s,
operation_id=node_type[s],
features=features,
scores=TraceGraphNodeReconsScores(
edge_logit=0,
operation_logit=node_type_logit[s, n_type],
avg_latency_nstd=avg_latency_nstd,
)
)
return nodes[s]
nodes = {}
edges.sort()
for u, v in edges:
if node_mask[u] and node_mask[v]:
v_node = get_node(v)
get_node(u).children.append(v_node)
v_node.scores.edge_logit = adj_logit[u, v]
if 0 in nodes:
g = TraceGraph.new_sampled(nodes[0], len(nodes), -1)
g.merge_spans_and_assign_id()
ret.append(g)
else:
ret.append(None)
# return the graphs
if keep_front_shape:
ret = np.array(ret).reshape(front_shape)
return ret
@dataclass(init=False)
class GraphNodeMatch(object):
__slots__ = [
'g1_to_g2',
'g2_to_g1',
]
g1_to_g2: Dict[TraceGraphNode, TraceGraphNode]
g2_to_g1: Dict[TraceGraphNode, TraceGraphNode]
def __init__(self):
self.g1_to_g2 = {}
self.g2_to_g1 = {}
def add_match(self, node1, node2):
self.g1_to_g2[node1] = node2
self.g2_to_g1[node2] = node1
@dataclass(init=False)
class GraphNodeDiff(object):
__slots__ = [
'parent', 'depth', 'node', 'offset', 'node_count',
]
parent: Optional[TraceGraphNode]
depth: int
node: TraceGraphNode
offset: int # -1: present in g but absent in g2; 1: present in g2 but absent in g1
node_count: int # count of nodes in this branch
def __init__(self, parent, depth, node, offset):
self.parent = parent
self.depth = depth
self.node = node
self.offset = offset
self.node_count = node.count_nodes()
def __repr__(self):
return f'GraphNodeDiff(depth={self.depth}, offset={self.offset})'
def diff_graph(g1: TraceGraph,
g2: TraceGraph
) -> Tuple[GraphNodeMatch, List[GraphNodeDiff]]:
m = GraphNodeMatch()
ret = []
def match_node(depth: int,
parent1: Optional[TraceGraphNode],
parent2: Optional[TraceGraphNode],
node1: Optional[TraceGraphNode],
node2: Optional[TraceGraphNode]):
if node1 is None:
if node2 is None:
pass
else:
ret.append(GraphNodeDiff(parent=parent2, depth=depth, node=node2, offset=1))
else:
if node2 is None:
ret.append(GraphNodeDiff(parent=parent1, depth=depth, node=node1, offset=-1))
elif node1.operation_id != node2.operation_id:
ret.append(GraphNodeDiff(parent=parent1, depth=depth, node=node1, offset=-1))
ret.append(GraphNodeDiff(parent=parent2, depth=depth, node=node2, offset=1))
else:
m.add_match(node1, node2)
c_depth = depth + 1
i, j = 0, 0
while i < len(node1.children) and j < len(node2.children):
c1 = node1.children[i]
c2 = node2.children[j]
if c1.operation_id < c2.operation_id:
match_node(c_depth, node1, None, c1, None)
i += 1
elif c2.operation_id < c1.operation_id:
match_node(c_depth, None, node2, None, c2)
j += 1
else:
match_node(c_depth, node1, node2, c1, c2)
i += 1
j += 1
while i < len(node1.children):
c1 = node1.children[i]
match_node(c_depth, node1, None, c1, None)
i += 1
while j < len(node2.children):
c2 = node2.children[j]
match_node(c_depth, None, node2, None, c2)
j += 1
match_node(0, None, None, g1.root, g2.root)
return m, ret
def dgl_graph_key(graph: dgl.DGLGraph) -> str:
return edges_to_key(graph.ndata['operation_id'], *graph.edges())
@torch.jit.script
def edges_to_key(operation_id: torch.Tensor, u_list: torch.Tensor, v_list: torch.Tensor) -> str:
mask = u_list != v_list
u_id: List[int] = operation_id[u_list][mask].tolist()
v_id: List[int] = operation_id[v_list][mask].tolist()
graph_key = f'0,{operation_id[0].item()};' + ';'.join(sorted([f'{u},{v}' for (u, v) in zip(u_id, v_id)]))
return graph_key
def trace_graph_key(graph: TraceGraph) -> str:
def dfs(nd: TraceGraphNode, pa_id: int, cnt: int=1):
cur_cnt = cnt * len(nd.spans)
spans = [f'{pa_id},{nd.operation_id}'] * cur_cnt
for child in nd.children:
spans += dfs(child, nd.operation_id, cur_cnt)
return spans
spans = dfs(graph.root, 0, 1)
return ';'.join(sorted(spans))