wehub-resource-sync
117 lines
3.9 KiB
Python
117 lines
3.9 KiB
Python
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
|