observatory: toggle default lenses with config

Author: quic-boyuc

devtools/observatory/lenses/accuracy.py

@@ -479,14 +479,19 @@ def _clear_state(cls) -> None:
 
     @classmethod
     def observe(cls, artifact: Any, context: ObservationContext) -> Any:
+
+        acc_config = context.config.get("accuracy", {})
+        if not acc_config.get("enabled", True):
+            return None
+
         record_name = context.shared_state.get("record_name", "")
 
         # Lazily configure evaluator on first "Exported Float" record
         if record_name == "Exported Float" and cls._evaluator is None:
             cls._configure_from_float_model(artifact)
 
-        acc_config = context.config.get("accuracy", {})
         evaluator = acc_config.get("evaluator") or cls._evaluator
+
         if not evaluator:
             return None
 

devtools/observatory/lenses/graph.py

@@ -47,6 +47,11 @@ def _to_graph_module(cls, artifact: Any) -> Optional[torch.fx.GraphModule]:
 
     @classmethod
     def observe(cls, artifact: Any, context: ObservationContext) -> Any:
+
+        graph_config = context.config.get("graph", {})
+        if not graph_config.get("enabled", True):
+            return None
+
         graph_module = cls._to_graph_module(artifact)
         if graph_module is None:
             return None

devtools/observatory/lenses/metadata.py

@@ -36,6 +36,10 @@ def get_name(cls) -> str:
 
     @classmethod
     def observe(cls, artifact: Any, context: ObservationContext) -> Any:
+        metadata_config = context.config.get("metadata", {})
+        if not metadata_config.get("enabled", True):
+            return None
+
         artifact_type = str(type(artifact).__name__)
         node_count: Any = "N/A"
 

devtools/observatory/lenses/per_layer_accuracy.py

@@ -430,6 +430,10 @@ def _values(k: str) -> List[float]:
 
     @classmethod
     def observe(cls, artifact: Any, context: ObservationContext) -> Any:
+        acc_config = context.config.get("accuracy", {})
+        if not acc_config.get("enabled", True):
+            return None
+
         graph_module = cls._to_graph_module(artifact)
         if graph_module is None:
             return None

devtools/observatory/lenses/stack_trace.py

@@ -63,6 +63,10 @@ def _get_stack_trace(cls) -> List[Dict[str, Any]]:
 
     @classmethod
     def observe(cls, artifact: Any, context: ObservationContext) -> Any:
+        stacktrace_config = context.config.get("stacktrace", {})
+        if not stacktrace_config.get("enabled", True):
+            return None
+
         try:
             return cls._get_stack_trace()
         except Exception as exc: