fix(vis): avoid graphviz.Digraph memory bomb when ELK fails on large graphs

When ELK layout fails (OOM/timeout) on 1M+ node graphs, the fallback path
previously built a graphviz.Digraph in Python — nested subgraph body-list
copies exploded memory and hung indefinitely. Now render_elk_direct handles
the failure internally: reuses already-collected Phase 1 data to generate
DOT text without positions and renders directly with sfdp, bypassing
graphviz.Digraph entirely.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: johnmarktaylor91

torchlens/visualization/elk_layout.py

@@ -1113,30 +1113,45 @@ def _patch_sizes(elk_node):
         _seed_stress_positions(elk_graph, all_edges)
 
     elk_timeout = max(_ELK_TIMEOUT, int(num_elk_nodes * 0.015))
-    positioned = run_elk_layout(elk_graph, timeout=elk_timeout)
+
+    # If ELK layout fails (OOM, timeout, etc.), fall back to generating DOT
+    # without positions and rendering with sfdp.  This avoids the catastrophic
+    # graphviz.Digraph construction path in rendering.py which explodes on
+    # very large graphs (1M+ nodes) due to nested subgraph body-list copies.
+    elk_failed = False
+    try:
+        positioned = run_elk_layout(elk_graph, timeout=elk_timeout)
+    except RuntimeError as e:
+        warnings.warn(
+            f"ELK layout failed ({e}), generating DOT without positions and rendering with sfdp."
+        )
+        elk_failed = True
 
     # Collect leaf node centers and compound node bounding boxes from ELK output.
     positions = {}  # leaf_id -> (center_x, center_y) in ELK coords
     compound_bboxes = {}  # "group_<mod>" -> (x, y, w, h) in ELK coords (absolute)
+    max_y = 0
+
+    if not elk_failed:
+
+        def _collect_pos(elk_node, ox=0, oy=0):
+            for ch in elk_node.get("children", []):
+                ax = ox + ch.get("x", 0)
+                ay = oy + ch.get("y", 0)
+                if ch["id"].startswith("group_"):
+                    w = ch.get("width", 0)
+                    h = ch.get("height", 0)
+                    compound_bboxes[ch["id"]] = (ax, ay, w, h)
+                    _collect_pos(ch, ax, ay)
+                else:
+                    w = ch.get("width", _DEFAULT_NODE_WIDTH)
+                    h = ch.get("height", _DEFAULT_NODE_HEIGHT)
+                    positions[ch["id"]] = (ax + w / 2, ay + h / 2)
 
-    def _collect_pos(elk_node, ox=0, oy=0):
-        for ch in elk_node.get("children", []):
-            ax = ox + ch.get("x", 0)
-            ay = oy + ch.get("y", 0)
-            if ch["id"].startswith("group_"):
-                w = ch.get("width", 0)
-                h = ch.get("height", 0)
-                compound_bboxes[ch["id"]] = (ax, ay, w, h)
-                _collect_pos(ch, ax, ay)
-            else:
-                w = ch.get("width", _DEFAULT_NODE_WIDTH)
-                h = ch.get("height", _DEFAULT_NODE_HEIGHT)
-                positions[ch["id"]] = (ax + w / 2, ay + h / 2)
-
-    _collect_pos(positioned)
-    # Use the root node's full height as the y-flip reference.
-    root_h = positioned.get("height", 0)
-    max_y = max(root_h, max((y for _, y in positions.values()), default=0))
+        _collect_pos(positioned)
+        # Use the root node's full height as the y-flip reference.
+        root_h = positioned.get("height", 0)
+        max_y = max(root_h, max((y for _, y in positions.values()), default=0))
 
     # ── Phase 3: Generate DOT with clusters and positions ──
 
@@ -1265,7 +1280,9 @@ def _write_cluster(mod_key, depth, indent):
     lines.append("}")
     dot_source = "\n".join(lines)
 
-    # ── Phase 4: Render with neato -n ──
+    # ── Phase 4: Render ──
+    # When ELK succeeded, render with neato -n (pre-positioned layout).
+    # When ELK failed, render with sfdp (spring-embedder, computes own layout).
 
     if num_elk_nodes > 25000 and vis_fileformat != "svg":
         warnings.warn(
@@ -1279,24 +1296,40 @@ def _write_cluster(mod_key, depth, indent):
         f.write(dot_source)
 
     rendered_path = f"{vis_outpath}.{vis_fileformat}"
-    # Spline routing is O(n^2) — use straight lines for large graphs.
     num_nodes = len(node_data)
-    spline_mode = "true" if num_nodes < 1000 else "line"
-    cmd = [
-        "neato",
-        "-n",
-        f"-Gsplines={spline_mode}",
-        f"-T{vis_fileformat}",
-        "-o",
-        rendered_path,
-        source_path,
-    ]
+
+    if elk_failed:
+        # sfdp computes its own layout — no positions needed in DOT.
+        # Use overlap removal for readability.
+        cmd = [
+            "sfdp",
+            "-Goverlap=prism",
+            f"-T{vis_fileformat}",
+            "-o",
+            rendered_path,
+            source_path,
+        ]
+    else:
+        # neato -n uses ELK-computed positions.
+        # Spline routing is O(n^2) — use straight lines for large graphs.
+        spline_mode = "true" if num_nodes < 1000 else "line"
+        cmd = [
+            "neato",
+            "-n",
+            f"-Gsplines={spline_mode}",
+            f"-T{vis_fileformat}",
+            "-o",
+            rendered_path,
+            source_path,
+        ]
+
     render_timeout = max(_SFDP_TIMEOUT, int(num_nodes * 0.01))
     try:
         result = subprocess.run(cmd, timeout=render_timeout, capture_output=True, text=True)
         if result.returncode != 0:
+            engine_name = "sfdp" if elk_failed else "neato"
             raise RuntimeError(
-                f"neato rendering failed (exit {result.returncode}):\n{result.stderr}"
+                f"{engine_name} rendering failed (exit {result.returncode}):\n{result.stderr}"
             )
         if not vis_save_only:
             import graphviz

torchlens/visualization/rendering.py

@@ -204,28 +204,26 @@ def render_graph(
 
     # ELK fast path: skip graphviz.Digraph construction entirely.
     # Generates DOT directly with ELK positions and cluster subgraphs (module boxes).
+    # If ELK layout fails (OOM, timeout), render_elk_direct falls back internally
+    # to sfdp — still using the fast DOT-text path, never graphviz.Digraph.
     if engine == "elk":
-        from .elk_layout import render_elk_direct, render_with_sfdp
+        from .elk_layout import render_elk_direct
 
-        try:
-            result = render_elk_direct(
-                self,
-                entries_to_plot,
-                vis_mode,
-                vis_nesting_depth,
-                show_buffer_layers,
-                overrides,
-                vis_outpath,
-                vis_fileformat,
-                vis_save_only,
-                graph_caption,
-                rankdir,
-            )
-            _vprint(self, f"Graph saved to {vis_outpath}.{vis_fileformat}")
-            return result
-        except RuntimeError as e:
-            warnings.warn(f"ELK layout failed ({e}), falling back to sfdp.")
-            engine = "sfdp"  # fall through to build graphviz.Digraph for sfdp
+        result = render_elk_direct(
+            self,
+            entries_to_plot,
+            vis_mode,
+            vis_nesting_depth,
+            show_buffer_layers,
+            overrides,
+            vis_outpath,
+            vis_fileformat,
+            vis_save_only,
+            graph_caption,
+            rankdir,
+        )
+        _vprint(self, f"Graph saved to {vis_outpath}.{vis_fileformat}")
+        return result
 
     dot = graphviz.Digraph(
         name=self.model_name,