Merge pull request #123 from johnmarktaylor91/feat/verbose-mode

Author: johnmarktaylor91

scripts/CLAUDE.md

@@ -5,8 +5,7 @@
 | File | Purpose |
 |------|---------|
 | `check_flops_coverage.py` | Reports FLOPs module coverage against all decorated torch functions |
-| `run_250k.py` | Render a 250k-node graph with ELK layout |
-| `run_1M.py` | Render a 1M-node graph with ELK (phased: construct → log → render) |
+| `render_large_graph.py` | Render a large random graph with ELK layout (any node count) |
 
 ## check_flops_coverage.py
 
@@ -20,3 +19,12 @@ Run with:
 ```bash
 python scripts/check_flops_coverage.py
 ```
+
+## render_large_graph.py
+
+Renders a large random graph using the ELK layout engine. Accepts any target node count.
+
+```bash
+python scripts/render_large_graph.py 250000
+python scripts/render_large_graph.py 1000000 --format png --seed 123
+```

scripts/render_large_graph.py

@@ -0,0 +1,75 @@
+"""Render a large random graph with ELK layout.
+
+Usage:
+    python scripts/render_large_graph.py NUM_NODES [OPTIONS]
+
+Examples:
+    python scripts/render_large_graph.py 250000
+    python scripts/render_large_graph.py 1000000 --format png --seed 123
+    python scripts/render_large_graph.py 50000 --outdir /tmp/graphs
+"""
+
+import argparse
+import os
+import sys
+import time
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "tests"))
+
+import torch
+from example_models import RandomGraphModel
+from torchlens import log_forward_pass
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Render a large random graph with ELK layout.")
+    parser.add_argument("num_nodes", type=int, help="Target number of nodes in the graph")
+    parser.add_argument("--format", default="svg", help="Output format (default: svg)")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed (default: 42)")
+    parser.add_argument(
+        "--outdir",
+        default=os.path.join("tests", "test_outputs", "visualizations", "large"),
+        help="Output directory",
+    )
+    args = parser.parse_args()
+
+    os.makedirs(args.outdir, exist_ok=True)
+    label = f"elk_{args.num_nodes // 1000}k"
+
+    # Phase 1: Model construction
+    t0 = time.time()
+    model = RandomGraphModel(target_nodes=args.num_nodes, seed=args.seed)
+    x = torch.randn(2, 64)
+    t1 = time.time()
+    print(f"Phase 1 — Model construction: {t1 - t0:.1f}s", flush=True)
+
+    # Phase 2: log_forward_pass (logging + postprocessing)
+    ml = log_forward_pass(model, x, layers_to_save=None, detect_loops=False, verbose=True)
+    t2 = time.time()
+    print(f"Phase 2 — log_forward_pass:    {t2 - t1:.1f}s ({len(ml)} layers)", flush=True)
+
+    # Phase 3: Render
+    ml.render_graph(
+        vis_opt="unrolled",
+        vis_nesting_depth=1000,
+        vis_outpath=os.path.join(args.outdir, label),
+        save_only=True,
+        vis_fileformat=args.format,
+        vis_node_placement="elk",
+    )
+    t3 = time.time()
+    print(f"Phase 3 — ELK render:          {t3 - t2:.1f}s", flush=True)
+
+    total = t3 - t0
+    print(f"Total: {total:.1f}s ({total / 60:.1f} min)", flush=True)
+
+    out_path = os.path.join(args.outdir, f"{label}.{args.format}")
+    if os.path.exists(out_path):
+        size_mb = os.path.getsize(out_path) / 1024 / 1024
+        print(f"Output: {out_path} ({size_mb:.1f} MB)")
+
+    ml.cleanup()
+
+
+if __name__ == "__main__":
+    main()

scripts/run_1M.py

@@ -49,6 +49,7 @@
 from ..utils.arg_handling import safe_copy_args, safe_copy_kwargs, normalize_input_args
 from .source_tensors import log_source_tensor
 from ..data_classes.interface import _give_user_feedback_about_lookup_key
+from ..utils.display import _vprint, _vtimed
 
 
 def save_new_activations(
@@ -133,6 +134,7 @@ def save_new_activations(
         self._raw_layer_dict.pop(label, None)
 
     # Now run and log the new inputs.
+    _vprint(self, "Running fast pass (saving requested activations)")
     self._run_and_log_inputs_through_model(
         model, input_args, input_kwargs, layers_to_save, random_seed
     )
@@ -431,12 +433,23 @@ def run_and_log_inputs_through_model(
         # Per-session model preparation
         _prepare_model_session(self, model, self._optimizer)
         self.elapsed_time_setup = time.time() - self.pass_start_time
+        _vprint(self, f"Model prepared ({self.elapsed_time_setup:.2f}s)")
+
+        # Print input summary
+        if getattr(self, "verbose", False):
+            devices = set()
+            for t in input_tensors:
+                if hasattr(t, "device"):
+                    devices.add(str(t.device))
+            device_str = ", ".join(sorted(devices)) if devices else "unknown"
+            _vprint(self, f"Inputs: {len(input_tensors)} tensor(s) on {device_str}")
 
         # Turn on the logging toggle and run the forward pass.
         # Inside this context, every decorated torch function will log its
         # inputs/outputs.  Source tensors (model inputs) are logged explicitly
         # before invoking the model; all subsequent operations are captured
         # automatically by the decorated wrappers.
+        _vprint(self, f"Running {self.logging_mode} forward pass...")
         with _state.active_logging(self):
             for i, t in enumerate(input_tensors):
                 log_source_tensor(self, t, "input", input_tensor_addresses[i])
@@ -446,10 +459,16 @@ def run_and_log_inputs_through_model(
         self.elapsed_time_forward_pass = (
             time.time() - self.pass_start_time - self.elapsed_time_setup
         )
+        _vprint(
+            self,
+            f"Forward pass complete ({self.elapsed_time_forward_pass:.2f}s, "
+            f"{len(self._raw_layer_dict)} raw operations)",
+        )
 
         output_tensors, output_tensor_addresses = _extract_and_mark_outputs(self, outputs)
 
         _cleanup_model_session(model, input_tensors)
+        _vprint(self, f"Postprocessing {len(self._raw_layer_dict)} operations...")
         self._postprocess(output_tensors, output_tensor_addresses)
 
     except Exception as e:

scripts/run_250k.py

@@ -112,6 +112,7 @@ def __init__(
         save_source_context: bool = False,
         save_rng_states: bool = False,
         detect_loops: bool = True,
+        verbose: bool = False,
     ):
         """Initialise a fresh ModelLog for a new logging session.
 
@@ -130,6 +131,7 @@ def __init__(
                 around each function call (used by FuncCallLocation).
             optimizer: Optional torch optimizer, used to annotate which params
                 have optimizers attached.
+            verbose: If True, print timed progress messages at each major pipeline stage.
         """
         # Callables are effectively immutable — deepcopy is unnecessary.
 
@@ -159,6 +161,7 @@ def __init__(
         self.save_source_context = save_source_context
         self.save_rng_states = save_rng_states
         self.detect_loops = detect_loops
+        self.verbose = verbose
         self.has_saved_gradients = False
         self.mark_input_output_distances = mark_input_output_distances