observatory: add observe_pass decorator and collect() name deduplication

Add observe_pass, a decorator that wraps any PassBase subclass or callable
pass to automatically collect input/output graphs via Observatory. Names
are derived from the class/function name and deduplicated by collect()
itself (#2, pytorch#3, ...) so repeated calls never silently overwrite records.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: quic-boyuc

devtools/observatory/README.md

@@ -71,19 +71,38 @@ This produces:
 ### Python API
 
 ```python
-from executorch.devtools.observatory import Observatory
+from executorch.devtools.observatory import Observatory, observe_pass
+
+# Wrap passes for automatic graph collection
+pass_a = observe_pass(SomePass())
 
 Observatory.clear()
 with Observatory.enable_context():
     # Auto: Lenses can auto-insert collection points by monkey patching when entering context
     # Manual: Insert the collection point anywhere
     Observatory.collect("step_0", graph_module)
 
+    # observe_pass: auto-collects input and output graphs
+    result = pass_a(graph_module)
+    # collects "SomePass/input" and "SomePass/output"
+
 Observatory.export_html_report("/tmp/report.html")
 ```
 
 ## Core concepts
 
+### `observe_pass` decorator
+
+The `observe_pass` decorator wraps any pass (PassBase subclass or callable) to automatically collect graphs via Observatory. By default it captures both input and output graphs, making pass debugging a one-line change:
+
+```python
+observed = observe_pass(SomePass())  # wrap once
+result = observed(graph_module)       # auto-collects input + output
+```
+
+Record names are derived from the class name and auto-deduplicated on repeat calls.
+See [USAGE.md](USAGE.md) for the full decorator reference.
+
 ### Lenses
 
 A **Lens** is a modular extension that adds domain-specific debugging logic. Each lens can participate in capture, analysis, and visualization. This is what makes Observatory a framework rather than a fixed tool.

devtools/observatory/USAGE.md

@@ -159,20 +159,63 @@ Observatory.export_json("pass_debug.json")
 
 ### Pass transform debugging
 
-To compare graphs before and after a specific pass:
+Use `observe_pass` to automatically collect graphs before and after a pass.
+Wrap any `PassBase` subclass instance, callable, or use it as a class decorator:
 
 ```python
+from executorch.devtools.observatory import Observatory, observe_pass
+from executorch.exir.passes.remove_graph_asserts_pass import RemoveGraphAssertsPass
 
-from executorch.devtools.observatory import Observatory
+# Wrap pass instances — default collects both input and output graphs
+pass_a = observe_pass(RemoveGraphAssertsPass())
+pass_b = observe_pass(MyCustomPass())
+
+Observatory.clear()
 with Observatory.enable_context():
-    for name, module in model.named_modules():
-        before = torch.fx.symbolic_trace(module)
-        Observatory.collect(f"{name}/before", before)
+    result_a = pass_a(graph_module)
+    # collects "RemoveGraphAssertsPass/input" and "RemoveGraphAssertsPass/output"
+
+    result_b = pass_b(result_a.graph_module)
+    # collects "MyCustomPass/input" and "MyCustomPass/output"
+
+    # Call again — names auto-deduplicate
+    result_c = pass_a(result_b.graph_module)
+    # collects "RemoveGraphAssertsPass/input #2" and "RemoveGraphAssertsPass/output #2"
 
-        after = my_transform(before)
-        Observatory.collect(f"{name}/after", after)
+Observatory.export_html_report("pass_debug.html")
 ```
 
+Control what is collected with boolean flags:
+
+```python
+# Collect only the output graph
+observed = observe_pass(SomePass(), collect_input=False)
+
+# Collect only the input graph
+observed = observe_pass(SomePass(), collect_output=False)
+
+# Override the record name
+observed = observe_pass(SomePass(), name="step_1")
+```
+
+Use as a class decorator to make all instances observable:
+
+```python
+@observe_pass
+class MyPass(PassBase):
+    def call(self, gm):
+        # ... transform logic ...
+        return PassResult(gm, True)
+
+# Or with parameters:
+@observe_pass(name="Quantize", collect_input=False)
+class QuantizePass(PassBase):
+    def call(self, gm):
+        ...
+```
+
+`observe_pass` is a no-op when no Observatory context is active.
+
 ### Inside the CLI-wrapped script (zero-code-change)
 
 When running via the CLI, `Observatory.enable_context()` is already active.

devtools/observatory/__init__.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from .observe_pass import observe_pass
 from .observatory import Observatory
 
-__all__ = ["Observatory"]
+__all__ = ["Observatory", "observe_pass"]

devtools/observatory/observatory.py

@@ -165,6 +165,7 @@ def collect(cls, name: str, artifact: Any) -> None:
         Notes:
         1. No-op when context is disabled.
         2. Record name is exposed via `context.shared_state['record_name']`.
+        3. Duplicate names are auto-suffixed with #2, #3, … to prevent overwrites.
         """
 
         if any(ignored in name for ignored in cls._ignored_graphs):
@@ -173,6 +174,13 @@ def collect(cls, name: str, artifact: Any) -> None:
         if not cls._config_stack:
             return
 
+        # Deduplicate: if name exists, suffix with #2, #3, ...
+        if name in cls._records:
+            n = 2
+            while f"{name} #{n}" in cls._records:
+                n += 1
+            name = f"{name} #{n}"
+
         active_config = cls._config_stack[-1]
         ctx = ObservationContext(config=active_config)
         ctx.shared_state["record_name"] = name

devtools/observatory/observe_pass.py

@@ -0,0 +1,111 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""observe_pass — decorator to auto-collect pass input/output via Observatory."""
+
+from __future__ import annotations
+
+import functools
+import logging
+from typing import Any, Callable, Optional, Type, Union
+
+from torch.fx.passes.infra.pass_base import PassResult
+
+
+def observe_pass(
+    target: Union[Callable, Type, None] = None,
+    *,
+    name: Optional[str] = None,
+    collect_input: bool = True,
+    collect_output: bool = True,
+) -> Any:
+    """Wrap a pass to auto-collect its graph via Observatory.
+
+    Works as a class decorator, instance wrapper, or function wrapper.
+
+    Args:
+        target: A PassBase subclass (class decorator), a pass instance, or a
+            callable pass. When ``None``, returns a parameterized decorator.
+        name: Override the record name (default: derived from class/function name).
+        collect_input: Collect the input graph before the pass runs (default True).
+        collect_output: Collect the output graph after the pass runs (default True).
+    """
+    collect_both = collect_input and collect_output
+
+    def _base_name(obj: Any) -> str:
+        if name:
+            return name
+        if isinstance(obj, type):
+            return obj.__name__
+        cls_name = type(obj).__name__
+        if cls_name == "function":
+            return getattr(obj, "__name__", "pass")
+        return cls_name
+
+    def _collect_artifact(record_name: str, artifact: Any) -> None:
+        from .observatory import Observatory
+
+        try:
+            Observatory.collect(record_name, artifact)
+        except Exception as exc:
+            logging.debug("[observe_pass] collection failed for %s: %s", record_name, exc)
+
+    def _output_graph(gm: Any, result: Any) -> Any:
+        if isinstance(result, tuple) and hasattr(result, "graph_module"):
+            return result.graph_module
+        if result is None:
+            return gm
+        return result
+
+    def _wrap_callable(fn: Callable) -> Callable:
+        @functools.wraps(fn)
+        def wrapper(gm: Any, *args: Any, **kwargs: Any) -> Any:
+            base = _base_name(fn)
+            if collect_input:
+                _collect_artifact(f"{base}/input" if collect_both else base, gm)
+            result = fn(gm, *args, **kwargs)
+            if collect_output:
+                _collect_artifact(
+                    f"{base}/output" if collect_both else base,
+                    _output_graph(gm, result),
+                )
+            return result
+
+        return wrapper
+
+    def _wrap_class(cls: Type) -> Type:
+        original_call = cls.__call__
+
+        @functools.wraps(original_call)
+        def patched_call(self: Any, gm: Any, *args: Any, **kwargs: Any) -> Any:
+            base = _base_name(self)
+            if collect_input:
+                _collect_artifact(f"{base}/input" if collect_both else base, gm)
+            result = original_call(self, gm, *args, **kwargs)
+            if collect_output:
+                _collect_artifact(
+                    f"{base}/output" if collect_both else base,
+                    _output_graph(gm, result),
+                )
+            return result
+
+        cls.__call__ = patched_call
+        return cls
+
+    # Dispatch based on how observe_pass was called.
+    if target is None:
+        # Parameterized: @observe_pass(name="X", collect_output=False)
+        def decorator(t: Any) -> Any:
+            if isinstance(t, type):
+                return _wrap_class(t)
+            return _wrap_callable(t)
+
+        return decorator
+
+    if isinstance(target, type):
+        return _wrap_class(target)
+
+    return _wrap_callable(target)