# ------------------------------------------------------------- # # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you 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 sys import re import networkx as nx import matplotlib.pyplot as plt try: import pygraphviz from networkx.drawing.nx_agraph import graphviz_layout HAS_PYGRAPHVIZ = True except ImportError: HAS_PYGRAPHVIZ = False print("[WARNING] pygraphviz not found. Please install via 'pip install pygraphviz'.\n" "If not installed, we will use an alternative layout (spring_layout).") def parse_line(line: str): """ Parse a single line from the trace file to extract: - Node ID - Operation (hop name) - Kind (e.g., FOUT, LOUT, NREF) - Total cost - Weight - Refs (list of IDs that this node depends on) """ # 1) Match a node ID in the form of "(R)" or "()" match_id = re.match(r'^\((R|\d+)\)', line) if not match_id: return None node_id = match_id.group(1) # 2) The remaining string after the node ID after_id = line[match_id.end():].strip() # Extract operation (hop name) before the first "[" match_label = re.search(r'^(.*?)\s*\[', after_id) if match_label: operation = match_label.group(1).strip() else: operation = after_id.strip() # 3) Extract the kind (content inside the first pair of brackets "[]") match_bracket = re.search(r'\[([^\]]+)\]', after_id) if match_bracket: kind = match_bracket.group(1).strip() else: kind = "" # 4) Extract total and weight from the content inside curly braces "{}" total = "" weight = "" match_curly = re.search(r'\{([^}]+)\}', line) if match_curly: curly_content = match_curly.group(1) m_total = re.search(r'Total:\s*([\d\.]+)', curly_content) m_weight = re.search(r'Weight:\s*([\d\.]+)', curly_content) if m_total: total = m_total.group(1) if m_weight: weight = m_weight.group(1) # 5) Extract reference nodes: look for the first parenthesis containing numbers after the hop name match_refs = re.search(r'\(\s*(\d+(?:,\d+)*)\s*\)', after_id) if match_refs: ref_str = match_refs.group(1) refs = [r.strip() for r in ref_str.split(',') if r.strip().isdigit()] else: refs = [] return { 'node_id': node_id, 'operation': operation, 'kind': kind, 'total': total, 'weight': weight, 'refs': refs } def build_dag_from_file(filename: str): """ Read a trace file line by line and build a directed acyclic graph (DAG) using NetworkX. """ G = nx.DiGraph() with open(filename, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line: continue info = parse_line(line) if not info: continue node_id = info['node_id'] operation = info['operation'] kind = info['kind'] total = info['total'] weight = info['weight'] refs = info['refs'] # Add node with attributes G.add_node(node_id, label=operation, kind=kind, total=total, weight=weight) # Add edges from references to this node for r in refs: if r not in G: G.add_node(r, label=r, kind="", total="", weight="") G.add_edge(r, node_id) return G def main(): """ Main function that: - Reads a filename from command-line arguments - Builds a DAG from the file - Draws and displays the DAG using matplotlib """ # Get filename from command-line argument if len(sys.argv) < 2: print("[ERROR] No filename provided.\nUsage: python plot_federated_dag.py ") sys.exit(1) filename = sys.argv[1] print(f"[INFO] Running with filename '{filename}'") # Build the DAG G = build_dag_from_file(filename) # Print debug info: nodes and edges print("Nodes:", G.nodes(data=True)) print("Edges:", list(G.edges())) # Decide on layout if HAS_PYGRAPHVIZ: # graphviz_layout with rankdir=BT (bottom to top), etc. pos = graphviz_layout(G, prog='dot', args='-Grankdir=BT -Gnodesep=0.5 -Granksep=0.8') else: # Fallback layout if pygraphviz is not installed pos = nx.spring_layout(G, seed=42) # Dynamically adjust figure size based on number of nodes node_count = len(G.nodes()) fig_width = 10 + node_count / 10.0 fig_height = 6 + node_count / 10.0 plt.figure(figsize=(fig_width, fig_height), facecolor='white', dpi=300) ax = plt.gca() ax.set_facecolor('white') # Generate labels for each node in the format: # node_id: operation_name # C (W) labels = { n: f"{n}: {G.nodes[n].get('label', n)}\n C{G.nodes[n].get('total', '')} (W{G.nodes[n].get('weight', '')})" for n in G.nodes() } # Function to determine color based on 'kind' def get_color(n): k = G.nodes[n].get('kind', '').lower() if k == 'fout': return 'tomato' elif k == 'lout': return 'dodgerblue' elif k == 'nref': return 'mediumpurple' else: return 'mediumseagreen' # Determine node shapes based on operation name: # - '^' (triangle) if the label contains "twrite" # - 's' (square) if the label contains "tread" # - 'o' (circle) otherwise triangle_nodes = [n for n in G.nodes() if 'twrite' in G.nodes[n].get('label', '').lower()] square_nodes = [n for n in G.nodes() if 'tread' in G.nodes[n].get('label', '').lower()] other_nodes = [ n for n in G.nodes() if 'twrite' not in G.nodes[n].get('label', '').lower() and 'tread' not in G.nodes[n].get('label', '').lower() ] # Colors for each group triangle_colors = [get_color(n) for n in triangle_nodes] square_colors = [get_color(n) for n in square_nodes] other_colors = [get_color(n) for n in other_nodes] # Draw nodes group-wise node_collection_triangle = nx.draw_networkx_nodes( G, pos, nodelist=triangle_nodes, node_size=800, node_color=triangle_colors, node_shape='^', ax=ax ) node_collection_square = nx.draw_networkx_nodes( G, pos, nodelist=square_nodes, node_size=800, node_color=square_colors, node_shape='s', ax=ax ) node_collection_other = nx.draw_networkx_nodes( G, pos, nodelist=other_nodes, node_size=800, node_color=other_colors, node_shape='o', ax=ax ) # Set z-order for nodes, edges, and labels node_collection_triangle.set_zorder(1) node_collection_square.set_zorder(1) node_collection_other.set_zorder(1) edge_collection = nx.draw_networkx_edges(G, pos, arrows=True, arrowstyle='->', ax=ax) if isinstance(edge_collection, list): for ec in edge_collection: ec.set_zorder(2) else: edge_collection.set_zorder(2) label_dict = nx.draw_networkx_labels(G, pos, labels=labels, font_size=9, ax=ax) for text in label_dict.values(): text.set_zorder(3) # Set the title plt.title("Program Level Federated Plan", fontsize=14, fontweight="bold") # Provide a small legend on the top-right or top-left plt.text(1, 1, "[LABEL]\n hopID: hopName\n C(Total) (W(Weight))", fontsize=12, ha='right', va='top', transform=ax.transAxes) # Example mini-legend for different 'kind' values plt.scatter(0.05, 0.95, color='dodgerblue', s=200, transform=ax.transAxes) plt.scatter(0.18, 0.95, color='tomato', s=200, transform=ax.transAxes) plt.scatter(0.31, 0.95, color='mediumpurple', s=200, transform=ax.transAxes) plt.text(0.08, 0.95, "LOUT", fontsize=12, va='center', transform=ax.transAxes) plt.text(0.21, 0.95, "FOUT", fontsize=12, va='center', transform=ax.transAxes) plt.text(0.34, 0.95, "NREF", fontsize=12, va='center', transform=ax.transAxes) plt.axis("off") # Save the plot to a file with the same name as the input file, but with a .png extension output_filename = f"{filename.rsplit('.', 1)[0]}.png" plt.savefig(output_filename, format='png', dpi=300, bbox_inches='tight') plt.show() if __name__ == '__main__': main()