chore(weave): fix some metrics (#1032)

* fix some metrics * fix set_metadata * remove stuff * uncomment logging * add set meta stubs * lint
wandb · Feb 1, 2024 · 83ac2fb · 83ac2fb
1 parent 4140d02
commit 83ac2fb
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 11 deletions.
diff --git a/weave/engine_trace.py b/weave/engine_trace.py
@@ -43,6 +43,9 @@ def set_tag(self, key, unredacted_val, redacted_val=None):
     def set_meta(self, *args, **kwargs):
         pass
 
+    def set_metric(self, *args, **kwargs):
+        pass
+
     def finish(self, *args, **kwargs):
         pass
 
@@ -152,18 +155,26 @@ def attributes(self):
     def set_tag(self, key, unredacted_val, redacted_val=None):
         if "tags" not in self.attributes:
             self.attributes["tags"] = {}
-        if os.getenv("DISABLE_WEAVE_PII"):
-            self.attributes["tags"][key] = redacted_val
-        else:
+        if not os.getenv("DISABLE_WEAVE_PII"):
             self.attributes["tags"][key] = unredacted_val
+        elif redacted_val is not None:
+            self.attributes["tags"][key] = redacted_val
 
     def set_meta(self, key, unredacted_val, redacted_val=None):
         if "metadata" not in self.attributes:
             self.attributes["metadata"] = {}
-        if os.getenv("DISABLE_WEAVE_PII"):
-            self.attributes["metadata"][key] = redacted_val
-        else:
+        if not os.getenv("DISABLE_WEAVE_PII"):
             self.attributes["metadata"][key] = unredacted_val
+        elif redacted_val is not None:
+            self.attributes["metadata"][key] = redacted_val
+
+    def set_metric(self, key, unredacted_val, redacted_val=None):
+        if "metrics" not in self.attributes:
+            self.attributes["metrics"] = {}
+        if not os.getenv("DISABLE_WEAVE_PII"):
+            self.attributes["metrics"][key] = unredacted_val
+        elif redacted_val is not None:
+            self.attributes["metrics"][key] = redacted_val
 
     def finish(self, *args, **kwargs):
         pass

diff --git a/weave/execute.py b/weave/execute.py
@@ -386,10 +386,10 @@ def do_one(
                             span.finish()
 
                     if span is not None:
-                        span.set_tag(
+                        span.set_metric(
                             "bytes_read_to_arrow",
                             report["bytes_read_to_arrow"],
-                            report["bytes_read_to_arrow"],
+                            True,
                         )
 
                     stats.add_node(

diff --git a/weave/weave_server.py b/weave/weave_server.py
@@ -353,13 +353,12 @@ def execute():
                     "http://localhost:8080/snakeviz/"
                     + urllib.parse.quote(profile_filename),
                 )
-    if root_span is not None:
-        root_span.set_tag("request_size", len(req_bytes), len(req_bytes))
-    fixed_response = response.results.safe_map(weavejs_fixes.fixup_data)
 
+    fixed_response = response.results.safe_map(weavejs_fixes.fixup_data)
     response_payload = _value_or_errors_to_response(fixed_response)
 
     if root_span is not None:
+        root_span.set_metric("request_size", len(req_bytes), True)
         root_span.set_metric("node_count", len(response_payload["data"]), True)
         root_span.set_metric(
             "error_count",