feat: graph visualization

In order to ease flow_graph inspection, we added a visualization tool. It accepts a job config as an input and it generates a visual graph. In order to save it as PNG, user must specify folder path and the file name will be {job_id}.png

Author: raresgaia123

pyproject.toml

@@ -16,7 +16,9 @@ dependencies = [
     "datasets >= 2.14.6",
     "fastapi >= 0.109.0",
     "grpcio >= 1.60.0",
+    "matplotlib>=3.7",
     "multiworld >= 0.2.2",
+    "networkx>=3.2",
     "protobuf >= 4.25.1",
     "psutil >= 5.9.5",
     "pydantic >= 2.4.2",
@@ -42,6 +44,9 @@ dev = [
     "python-lsp-server",
 ]
 
+[tool.setuptools.packages.find]
+exclude = ["tests*", "examples*", "integration-tests*", "tools*"]
+
 # [tool.setuptools]
 # packages = ["infscale"]
 

tools/visualise_flow_graph.py

@@ -0,0 +1,194 @@
+import argparse
+import os
+
+import matplotlib.pyplot as plt
+import networkx as nx
+import yaml
+
+from infscale.configs.job import JobConfig
+
+
+def load_job_config(path: str) -> JobConfig:
+    with open(path) as f:
+        data = yaml.safe_load(f)
+
+    return JobConfig(**data)
+
+
+def build_graph(job: JobConfig) -> tuple[nx.DiGraph, dict[str, int]]:
+    """Build graph using job configuration."""
+    graph = nx.DiGraph()
+
+    # Add all worker nodes
+    worker_stage = {w.id: w.stage["start"] for w in job.workers}
+
+    for wid in job.flow_graph:
+        graph.add_node(wid, stage=worker_stage.get(wid, 0))
+
+    # Add edges (peer -> wid) with labels from worlds
+    for wid, worlds in job.flow_graph.items():
+        for world in worlds:
+            for peer in world.peers:
+                label = f"{world.name}"
+                if world.addr:
+                    label += f"\n{world.addr}"
+                if world.backend:
+                    label += f"\n{world.backend}"
+                graph.add_edge(peer, wid, label=label)
+
+    return graph, worker_stage
+
+
+def draw_graph(
+    graph: nx.DiGraph,
+    worker_stage: dict[str, int],
+    job_id: str,
+    output_path: str = "",
+) -> None:
+    """Draw graph where worker_stage maps node -> stage (start)."""
+    # build positions (horizontal by stage, vertical stacked)
+    stage_to_nodes = {}
+    for node, stage in worker_stage.items():
+        stage_to_nodes.setdefault(stage, []).append(node)
+
+    sorted_stages = sorted(stage_to_nodes.keys())
+    pos = {}
+    x_spacing = 4
+    y_spacing = 2
+
+    for x, stage in enumerate(sorted_stages):
+        nodes = sorted(stage_to_nodes[stage])  # sort for stable layout
+        for y, node in enumerate(nodes):
+            pos[node] = (x * x_spacing, -y * y_spacing)
+
+    plt.figure(figsize=(12, 7))
+    ax = plt.gca()
+
+    nx.draw_networkx_nodes(
+        graph, pos, node_size=2000, node_color="#5bf4a7", edgecolors="black", ax=ax
+    )
+    nx.draw_networkx_labels(graph, pos, font_size=10, font_weight="bold", ax=ax)
+
+    # draw edges and capture artists. Use explicit edgelist to ensure order.
+    edgelist = list(graph.edges())
+    edge_artists = list(
+        nx.draw_networkx_edges(
+            graph,
+            pos,
+            edgelist=edgelist,
+            arrows=True,
+            arrowstyle="-|>",
+            arrowsize=20,
+            width=1.5,
+            connectionstyle=f"arc3,rad={0.15}",
+            min_source_margin=25,
+            min_target_margin=25,
+            ax=ax,
+        )
+    )
+
+    # get edge labels from graph attributes (we created them earlier as "label")
+    edge_labels = nx.get_edge_attributes(graph, "label")
+
+    # For each edge artist, get its path and compute the point at t along the curve.
+    for (src, dst), artist in zip(edgelist, edge_artists):
+        label = edge_labels.get((src, dst), "")
+        if not label:
+            continue
+
+        # try to extract the actual path used by the artist.
+        # for FancyArrowPatch, get_path() returns a Path containing
+        # MOVETO and CURVE3 (quadratic bezier) segments for arc3.
+        try:
+            path = artist.get_path()
+            verts = path.vertices
+        except Exception:
+            verts = None
+
+        label_x, label_y = None, None
+
+        if verts is not None and len(verts) >= 3:
+            # Common pattern for quad bezier: [P0, P1, P2] or path with MOVETO + CURVE3 + CURVE3...
+            # We'll attempt to find a quadratic bezier triplet P0, P1, P2.
+            # Find first MOVETO index
+            try:
+                # Find indices for MOVETO and first CURVE3 sequence
+                # Fallback: take the first three vertices
+                P0 = verts[0]
+                P1 = verts[1]
+                P2 = verts[2]
+                # If the path contains more points, using the first three typically represents the curve.
+                # Evaluate quadratic Bezier at t
+                t = float(0.2)
+                one_minus_t = 1 - t
+                bx = (
+                    (one_minus_t**2) * P0[0]
+                    + 2 * one_minus_t * t * P1[0]
+                    + (t**2) * P2[0]
+                )
+                by = (
+                    (one_minus_t**2) * P0[1]
+                    + 2 * one_minus_t * t * P1[1]
+                    + (t**2) * P2[1]
+                )
+                label_x, label_y = bx, by
+            except Exception:
+                label_x, label_y = None, None
+
+        if label_x is None:
+            # Fallback: straight-line t% between source and target positions
+            x1, y1 = pos[src]
+            x2, y2 = pos[dst]
+            label_x = x1 + (x2 - x1) * 0.2
+            label_y = y1 + (y2 - y1) * 0.2
+
+        # Draw the label at computed position
+        ax.text(
+            label_x,
+            label_y,
+            label,
+            fontsize=8,
+            ha="center",
+            va="center",
+            rotation=0,
+            bbox=dict(boxstyle="round,pad=0.2", fc="white", alpha=0.9),
+            zorder=10,
+        )
+
+    ax.set_axis_off()
+    plt.tight_layout()
+    if output_path:
+        os.makedirs(output_path, exist_ok=True)
+        output_file = os.path.join(output_path, f"{job_id}.png")
+        plt.savefig(output_file, dpi=300, bbox_inches="tight")
+        print(f"Graph saved at: {output_file}")
+    else:
+        print("Graph opened in a new window.")
+        plt.show()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Visualize JobConfig flow graph")
+    parser.add_argument("config_path", help="Path to job YAML config")
+    parser.add_argument(
+        "-o", "--output", help="Directory to save output image (optional)", default=None
+    )
+    args = parser.parse_args()
+
+    try:
+        config = load_job_config(args.config_path)
+    except FileNotFoundError as e:
+        print(f"Error while loading file: {e}")
+        return
+
+    graph, worker_stage = build_graph(config)
+    try:
+        draw_graph(graph, worker_stage, config.job_id, args.output)
+    except nx.exception.NetworkXError as e:
+        print(f"Error while drawing graph: {e}")
+    except KeyboardInterrupt:
+        pass
+
+
+if __name__ == "__main__":
+    main()