chore: import upstream snapshot with attribution

dev/benchmarks/tracing/_data.py

@@ -0,0 +1,116 @@
+import json
+import random
+import time
+import uuid
+
+from opentelemetry import trace as trace_api
+from opentelemetry.sdk.resources import Resource as _OTelResource
+from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
+from opentelemetry.trace import SpanContext
+
+from mlflow.entities.span import Span, SpanType, create_mlflow_span
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_location import TraceLocation
+from mlflow.entities.trace_state import TraceState
+from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
+from mlflow.tracing.constant import SpanAttributeKey, TraceTagKey
+from mlflow.tracing.utils import TraceJSONEncoder
+
+ENV_CHOICES = ["prod", "staging", "dev"]
+NAME_PREFIXES = ["agent_run", "qa_chain", "rag_pipeline", "summarizer"]
+WEEK_MS = 7 * 24 * 60 * 60 * 1000
+
+SEED_TRACES = 1000
+SEED_SPANS_PER_TRACE = 10
+
+
+def generate_trace_data(
+    experiment_id: str,
+    num_spans: int,
+    rng: random.Random,
+) -> tuple[TraceInfo, list[Span]]:
+    trace_id = f"tr-{uuid.uuid4().hex}"
+    request_time = int(time.time() * 1000) - rng.randint(0, WEEK_MS)
+    name_prefix = rng.choice(NAME_PREFIXES)
+    trace_info = TraceInfo(
+        trace_id=trace_id,
+        trace_location=TraceLocation.from_experiment_id(experiment_id),
+        request_time=request_time,
+        state=rng.choice([TraceState.OK, TraceState.OK, TraceState.OK, TraceState.ERROR]),
+        execution_duration=rng.randint(100, 5000),
+        tags={
+            TraceTagKey.TRACE_NAME: f"{name_prefix}_{trace_id[-4:]}",
+            "env": rng.choice(ENV_CHOICES),
+        },
+    )
+
+    span_types = [SpanType.LLM, SpanType.RETRIEVER, SpanType.TOOL, SpanType.CHAIN]
+    base_ns = 1_000_000_000_000
+    spans: list[Span] = []
+
+    for i in range(num_spans):
+        is_root = i == 0
+        span_type = SpanType.AGENT if is_root else rng.choice(span_types)
+        parent_id = None if is_root else rng.choice(range(max(0, i - 3), i))
+
+        trace_num = rng.randint(1, 2**63)
+        ctx = SpanContext(
+            trace_id=trace_num,
+            span_id=i + 1,
+            is_remote=False,
+            trace_flags=trace_api.TraceFlags(1),
+            trace_state=trace_api.TraceState(),
+        )
+
+        parent_ctx = None
+        if parent_id is not None:
+            parent_ctx = SpanContext(
+                trace_id=trace_num,
+                span_id=parent_id + 1,
+                is_remote=False,
+                trace_flags=trace_api.TraceFlags(1),
+                trace_state=trace_api.TraceState(),
+            )
+
+        attrs: dict[str, object] = {}
+        if is_root:
+            attrs[SpanAttributeKey.INPUTS] = json.dumps(
+                {"query": "What is ML?"}, cls=TraceJSONEncoder
+            )
+            attrs[SpanAttributeKey.OUTPUTS] = json.dumps(
+                {"response": "ML is..."}, cls=TraceJSONEncoder
+            )
+
+        otel_span = OTelReadableSpan(
+            name=f"{span_type.lower()}_{i}" if not is_root else "agent_run",
+            context=ctx,
+            parent=parent_ctx,
+            attributes={
+                "mlflow.traceRequestId": json.dumps(trace_id),
+                "mlflow.spanType": json.dumps(span_type, cls=TraceJSONEncoder),
+                **attrs,
+            },
+            start_time=base_ns + i * 10_000_000,
+            end_time=base_ns + i * 10_000_000 + rng.randint(5_000_000, 50_000_000),
+            status=trace_api.Status(trace_api.StatusCode.OK),
+            resource=_OTelResource.get_empty(),
+        )
+        spans.append(create_mlflow_span(otel_span, trace_id, span_type))
+
+    return trace_info, spans
+
+
+def seed_traces(
+    store: SqlAlchemyStore,
+    experiment_id: str,
+    count: int,
+    spans_per_trace: int,
+) -> list[str]:
+    rng = random.Random(123)
+    trace_ids: list[str] = []
+    for _ in range(count):
+        ti, sp = generate_trace_data(experiment_id, spans_per_trace, rng)
+        store.start_trace(ti)
+        store.log_spans(experiment_id, sp)
+        trace_ids.append(ti.trace_id)
+    return trace_ids