feat(weave): support resource attributes for wandb vars

Author: zbirenbaum

tests/trace_server/isolated_client_executor/cross_process_trace_server.py

@@ -269,7 +269,7 @@ def ensure_project_exists(
         raise NotImplementedError("ensure_project_exists is not implemented")
 
     # Regular method implementations (reduced duplication)
-    def otel_export(self, req: tsi.OtelExportReq) -> tsi.OtelExportRes:
+    def otel_export(self, req: tsi.OTelExportReq) -> tsi.OTelExportRes:
         """Export OTEL traces."""
         return self._send_request("otel_export", req)
 

tests/trace_server/test_opentelemetry.py

@@ -6,9 +6,6 @@
 from typing import Any
 
 import pytest
-from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
-    ExportTraceServiceRequest,
-)
 from opentelemetry.proto.common.v1.common_pb2 import (
     AnyValue,
     InstrumentationScope,
@@ -105,8 +102,8 @@ def create_test_span():
     return span
 
 
-def create_test_export_request(project_id="test_project"):
-    """Create a test ExportTraceServiceRequest with one span."""
+def create_test_export_request(project_id="test_project") -> tsi.OTelExportReq:
+    """Create a test OTelExportReq with one span."""
     span = create_test_span()
 
     # Create instrumentation scope
@@ -131,15 +128,17 @@ def create_test_export_request(project_id="test_project"):
     resource_spans.resource.CopyFrom(resource)
     resource_spans.scope_spans.append(scope_spans)
 
-    # Create traces data
-    traces_data = TracesData()
-    traces_data.resource_spans.append(resource_spans)
-
-    # Create export request
-    request = ExportTraceServiceRequest()
-    request.resource_spans.append(resource_spans)
+    # Create processed resource spans
+    processed_span = tsi.ProcessedResourceSpans(
+        entity="test-entity",
+        project="test-project",
+        run_id=None,
+        resource_spans=resource_spans,
+    )
 
-    return tsi.OtelExportReq(project_id=project_id, traces=request, wb_user_id=None)
+    return tsi.OTelExportReq(
+        project_id=project_id, processed_spans=[processed_span], wb_user_id=None
+    )
 
 
 def test_otel_export_clickhouse(client: weave_client.WeaveClient):
@@ -152,7 +151,7 @@ def test_otel_export_clickhouse(client: weave_client.WeaveClient):
     # Export the otel traces
     response = client.server.otel_export(export_req)
     # Verify the response is of the correct type
-    assert isinstance(response, tsi.OtelExportRes)
+    assert isinstance(response, tsi.OTelExportRes)
 
     # Query the calls
     res = client.server.calls_query(
@@ -164,7 +163,7 @@ def test_otel_export_clickhouse(client: weave_client.WeaveClient):
     assert len(res.calls) == 1
 
     call = res.calls[0]
-    export_span = export_req.traces.resource_spans[0].scope_spans[0].spans[0]
+    export_span = export_req.processed_spans[0].resource_spans.scope_spans[0].spans[0]
     decoded_trace = hexlify(export_span.trace_id).decode("ascii")
     decoded_span = hexlify(export_span.span_id).decode("ascii")
 
@@ -213,8 +212,10 @@ def test_otel_export_with_turn_and_thread(client: weave_client.WeaveClient):
     export_req.wb_user_id = "abcd123"
 
     # Add turn and thread attributes to the span
+    # Materialize processed_spans to avoid iterator exhaustion
     test_thread_id = str(uuid.uuid4())
-    span = export_req.traces.resource_spans[0].scope_spans[0].spans[0]
+    processed_spans_list = export_req.processed_spans
+    span = processed_spans_list[0].resource_spans.scope_spans[0].spans[0]
 
     kv_is_turn = KeyValue()
     kv_is_turn.key = "wandb.is_turn"
@@ -226,9 +227,11 @@ def test_otel_export_with_turn_and_thread(client: weave_client.WeaveClient):
     kv_thread.value.string_value = test_thread_id
     span.attributes.append(kv_thread)
 
+    export_req.processed_spans = processed_spans_list
+
     # Export the otel traces
     response = client.server.otel_export(export_req)
-    assert isinstance(response, tsi.OtelExportRes)
+    assert isinstance(response, tsi.OTelExportRes)
 
     # Query the calls
     res = client.server.calls_query(
@@ -263,16 +266,20 @@ def test_otel_export_with_turn_no_thread(client: weave_client.WeaveClient):
     export_req.wb_user_id = "abcd123"
 
     # Add only is_turn attribute (no thread_id)
-    span = export_req.traces.resource_spans[0].scope_spans[0].spans[0]
+    # Materialize processed_spans to avoid iterator exhaustion
+    processed_spans_list = export_req.processed_spans
+    span = processed_spans_list[0].resource_spans.scope_spans[0].spans[0]
 
     kv_is_turn = KeyValue()
     kv_is_turn.key = "wandb.is_turn"
     kv_is_turn.value.bool_value = True
     span.attributes.append(kv_is_turn)
 
+    export_req.processed_spans = processed_spans_list
+
     # Export the otel traces
     response = client.server.otel_export(export_req)
-    assert isinstance(response, tsi.OtelExportRes)
+    assert isinstance(response, tsi.OTelExportRes)
 
     # Query the calls
     res = client.server.calls_query(
@@ -462,8 +469,9 @@ def test_span_to_call_with_turn_and_thread(self):
     def test_traces_data_from_proto(self):
         """Test converting protobuf TracesData to Python TracesData."""
         export_req = create_test_export_request()
+        resource_spans_list = [ps.resource_spans for ps in export_req.processed_spans]
         traces_data = PyTracesData.from_proto(
-            TracesData(resource_spans=export_req.traces.resource_spans)
+            TracesData(resource_spans=resource_spans_list)
         )
 
         assert len(traces_data.resource_spans) == 1
@@ -1160,7 +1168,12 @@ def test_opentelemetry_cost_calculation(self, client: weave_client.WeaveClient):
 
         # Create export request
         export_req = create_test_export_request(project_id=project_id)
-        export_req.traces.resource_spans[0].scope_spans[0].spans[0].CopyFrom(span_gpt4)
+        # Materialize processed_spans to avoid iterator exhaustion
+        processed_spans_list = export_req.processed_spans
+        processed_spans_list[0].resource_spans.scope_spans[0].spans[0].CopyFrom(
+            span_gpt4
+        )
+        export_req.processed_spans = processed_spans_list
         export_req.wb_user_id = "test_user"
 
         # Export the trace
@@ -1444,7 +1457,11 @@ def test_otel_export_partial_success_on_attribute_conflict(
     export_req.wb_user_id = "abcd123"
 
     # Good span (already present at index 0)
-    good_span = export_req.traces.resource_spans[0].scope_spans[0].spans[0]
+    # Materialize processed_spans to avoid iterator exhaustion
+    processed_spans_list = export_req.processed_spans
+
+    # Good span (already present at index 0)
+    good_span = processed_spans_list[0].resource_spans.scope_spans[0].spans[0]
     good_span_id = hexlify(good_span.span_id).decode("ascii")
 
     # Add a conflicting span to the same batch
@@ -1461,12 +1478,13 @@ def test_otel_export_partial_success_on_attribute_conflict(
     kv_child.value.string_value = "Hello"
     bad_span.attributes.append(kv_child)
 
-    export_req.traces.resource_spans[0].scope_spans[0].spans.append(bad_span)
+    processed_spans_list[0].resource_spans.scope_spans[0].spans.append(bad_span)
+    export_req.processed_spans = processed_spans_list
     bad_span_id = hexlify(bad_span.span_id).decode("ascii")
 
     # Export
     res = client.server.otel_export(export_req)
-    assert isinstance(res, tsi.OtelExportRes)
+    assert isinstance(res, tsi.OTelExportRes)
     assert res.partial_success is not None
     assert res.partial_success.rejected_spans == 1
     # Error message should mention the conflicting key and guidance
@@ -1522,7 +1540,9 @@ def test_otel_span_wandb_attributes_and_data_routing(
     export_req.wb_user_id = "abcd123"
 
     # Get the span to add custom attributes
-    span = export_req.traces.resource_spans[0].scope_spans[0].spans[0]
+    # Materialize processed_spans to avoid iterator exhaustion
+    processed_spans_list = export_req.processed_spans
+    span = processed_spans_list[0].resource_spans.scope_spans[0].spans[0]
 
     # Clear default test attributes
     del span.attributes[:]
@@ -1591,9 +1611,11 @@ def test_otel_span_wandb_attributes_and_data_routing(
     kv_provider.value.string_value = "openai"
     span.attributes.append(kv_provider)
 
+    export_req.processed_spans = processed_spans_list
+
     # Export the OTEL traces
     response = client.server.otel_export(export_req)
-    assert isinstance(response, tsi.OtelExportRes)
+    assert isinstance(response, tsi.OTelExportRes)
 
     # Query the calls
     res = client.server.calls_query(

weave/trace_server/clickhouse_trace_server_batched.py

@@ -21,9 +21,7 @@
 from clickhouse_connect.driver.httputil import get_pool_manager
 from clickhouse_connect.driver.query import QueryResult
 from clickhouse_connect.driver.summary import QuerySummary
-from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
-    ExportTraceServiceRequest,
-)
+from opentelemetry.proto.trace.v1.trace_pb2 import ResourceSpans
 from tenacity import (
     retry,
     retry_if_exception_type,
@@ -382,20 +380,23 @@ def _create_or_get_placeholder_ops_digest(
         return self.file_create(source_file_req).digest
 
     @ddtrace.tracer.wrap(name="clickhouse_trace_server_batched.otel_export")
-    def otel_export(self, req: tsi.OtelExportReq) -> tsi.OtelExportRes:
+    def otel_export(self, req: tsi.OTelExportReq) -> tsi.OTelExportRes:
         assert_non_null_wb_user_id(req)
-
-        if not isinstance(req.traces, ExportTraceServiceRequest):
-            raise TypeError(
-                "Expected traces as ExportTraceServiceRequest, got {type(req.traces)}"
-            )
         calls: list[
             tuple[tsi.StartedCallSchemaForInsert, tsi.EndedCallSchemaForInsert]
         ] = []
         rejected_spans = 0
         error_messages: list[str] = []
+        for processed_span in req.processed_spans:
+            # Extract wb_run_id from the processed span
+            wb_run_id = processed_span.run_id
+
+            if not isinstance(processed_span.resource_spans, ResourceSpans):
+                raise TypeError(
+                    f"Expected resource_spans as ResourceSpans, got {type(processed_span.resource_spans)}"
+                )
 
-        for proto_resource_spans in req.traces.resource_spans:
+            proto_resource_spans = processed_span.resource_spans
             resource = Resource.from_proto(proto_resource_spans.resource)
             for proto_scope_spans in proto_resource_spans.scope_spans:
                 for proto_span in proto_scope_spans.spans:
@@ -423,69 +424,69 @@ def otel_export(self, req: tsi.OtelExportReq) -> tsi.OtelExportRes:
                         span.to_call(
                             req.project_id,
                             wb_user_id=req.wb_user_id,
-                            wb_run_id=req.wb_run_id,
+                            wb_run_id=wb_run_id,
                         )
                     )
 
-        obj_id_idx_map = defaultdict(list)
-        for idx, (start_call, _) in enumerate(calls):
-            op_name = object_creation_utils.make_safe_name(start_call.op_name)
-            obj_id_idx_map[op_name].append(idx)
+            obj_id_idx_map = defaultdict(list)
+            for idx, (start_call, _) in enumerate(calls):
+                op_name = object_creation_utils.make_safe_name(start_call.op_name)
+                obj_id_idx_map[op_name].append(idx)
 
-        existing_objects = self._get_existing_ops_from_spans(
-            seen_ids=set(obj_id_idx_map.keys()),
-            project_id=req.project_id,
-            limit=len(calls),
-        )
-        # We know that OTel will always use the placeholder source.
-        # We can instead just reuse the existing file if we know it is present
-        # and create it just once if we are not sure.
-        if len(existing_objects) == 0:
-            digest = self._create_or_get_placeholder_ops_digest(
-                project_id=req.project_id, create=True
-            )
-        else:
-            digest = self._create_or_get_placeholder_ops_digest(
-                project_id=req.project_id, create=False
+            existing_objects = self._get_existing_ops_from_spans(
+                seen_ids=set(obj_id_idx_map.keys()),
+                project_id=req.project_id,
+                limit=len(calls),
             )
+            # We know that OTel will always use the placeholder source.
+            # We can instead just reuse the existing file if we know it is present
+            # and create it just once if we are not sure.
+            if len(existing_objects) == 0:
+                digest = self._create_or_get_placeholder_ops_digest(
+                    project_id=req.project_id, create=True
+                )
+            else:
+                digest = self._create_or_get_placeholder_ops_digest(
+                    project_id=req.project_id, create=False
+                )
 
-        for obj in existing_objects:
-            op_ref_uri = ri.InternalOpRef(
-                project_id=req.project_id,
-                name=obj.object_id,
-                version=obj.digest,
-            ).uri()
-
-            # Modify each of the matched start calls in place
-            for idx in obj_id_idx_map[obj.object_id]:
-                calls[idx][0].op_name = op_ref_uri
-            # Remove this ID from the mapping so that once the for loop is done we are left with only new objects
-            obj_id_idx_map.pop(obj.object_id)
-
-        obj_creation_batch = []
-        for op_obj_id in obj_id_idx_map.keys():
-            op_val = object_creation_utils.build_op_val(digest)
-            obj_creation_batch.append(
-                tsi.ObjSchemaForInsert(
+            for obj in existing_objects:
+                op_ref_uri = ri.InternalOpRef(
                     project_id=req.project_id,
-                    object_id=op_obj_id,
-                    val=op_val,
-                    wb_user_id=req.wb_user_id,
+                    name=obj.object_id,
+                    version=obj.digest,
+                ).uri()
+
+                # Modify each of the matched start calls in place
+                for idx in obj_id_idx_map[obj.object_id]:
+                    calls[idx][0].op_name = op_ref_uri
+                # Remove this ID from the mapping so that once the for loop is done we are left with only new objects
+                obj_id_idx_map.pop(obj.object_id)
+
+            obj_creation_batch = []
+            for op_obj_id in obj_id_idx_map.keys():
+                op_val = object_creation_utils.build_op_val(digest)
+                obj_creation_batch.append(
+                    tsi.ObjSchemaForInsert(
+                        project_id=req.project_id,
+                        object_id=op_obj_id,
+                        val=op_val,
+                        wb_user_id=req.wb_user_id,
+                    )
                 )
-            )
-        res = self.obj_create_batch(obj_creation_batch)
+            res = self.obj_create_batch(obj_creation_batch)
 
-        for result in res:
-            if result.object_id is None:
-                raise RuntimeError("Otel Export - Expected object_id but got None")
+            for result in res:
+                if result.object_id is None:
+                    raise RuntimeError("Otel Export - Expected object_id but got None")
 
-            op_ref_uri = ri.InternalOpRef(
-                project_id=req.project_id,
-                name=result.object_id,
-                version=result.digest,
-            ).uri()
-            for idx in obj_id_idx_map[result.object_id]:
-                calls[idx][0].op_name = op_ref_uri
+                op_ref_uri = ri.InternalOpRef(
+                    project_id=req.project_id,
+                    name=result.object_id,
+                    version=result.digest,
+                ).uri()
+                for idx in obj_id_idx_map[result.object_id]:
+                    calls[idx][0].op_name = op_ref_uri
 
         write_target = self.table_routing_resolver.resolve_v2_write_target(
             req.project_id,
@@ -531,13 +532,13 @@ def otel_export(self, req: tsi.OtelExportReq) -> tsi.OtelExportRes:
             joined_errors = "; ".join(error_messages[:20]) + (
                 "; ..." if len(error_messages) > 20 else ""
             )
-            return tsi.OtelExportRes(
+            return tsi.OTelExportRes(
                 partial_success=tsi.ExportTracePartialSuccess(
                     rejected_spans=rejected_spans,
                     error_message=joined_errors,
                 )
             )
-        return tsi.OtelExportRes()
+        return tsi.OTelExportRes()
 
     @ddtrace.tracer.wrap(name="clickhouse_trace_server_batched.kafka_producer.flush")
     def _flush_kafka_producer(self) -> None: