OpenTelemetry SDK로 Unity Catalog에 전송하기

pip install --upgrade "mlflow[databricks]>=3.14.0"
pip install opentelemetry-sdk
pip install opentelemetry-exporter-otlp-proto-http

import os
import mlflow
from mlflow.entities.trace_location import UnityCatalog

os.environ["MLFLOW_TRACING_SQL_WAREHOUSE_ID"] = "<SQL_WAREHOUSE_ID>"
mlflow.set_tracking_uri("databricks")

experiment = mlflow.set_experiment(
    experiment_name="/Shared/agent-traces",
    trace_location=UnityCatalog(
        catalog_name="<UC_CATALOG>",
        schema_name="<UC_SCHEMA>",
        table_prefix="agent",
    ),
)
print(experiment.experiment_id)
print(experiment.trace_location.full_otel_spans_table_name)

GRANT USE CATALOG ON CATALOG <UC_CATALOG> TO `<SERVICE_PRINCIPAL_OR_USER>`;
GRANT USE SCHEMA ON SCHEMA <UC_CATALOG>.<UC_SCHEMA> TO `<SERVICE_PRINCIPAL_OR_USER>`;
GRANT MODIFY, SELECT ON TABLE <UC_CATALOG>.<UC_SCHEMA>.agent_otel_spans TO `<SERVICE_PRINCIPAL_OR_USER>`;
GRANT MODIFY, SELECT ON TABLE <UC_CATALOG>.<UC_SCHEMA>.agent_otel_logs TO `<SERVICE_PRINCIPAL_OR_USER>`;
GRANT MODIFY, SELECT ON TABLE <UC_CATALOG>.<UC_SCHEMA>.agent_otel_metrics TO `<SERVICE_PRINCIPAL_OR_USER>`;
GRANT MODIFY, SELECT ON TABLE <UC_CATALOG>.<UC_SCHEMA>.agent_otel_annotations TO `<SERVICE_PRINCIPAL_OR_USER>`;

import json
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter

WORKSPACE_URL = "<WORKSPACE_URL>"
UC_SPANS_TABLE = "<UC_CATALOG>.<UC_SCHEMA>.agent_otel_spans"
DATABRICKS_TOKEN = "<DATABRICKS_TOKEN>"

exporter = OTLPSpanExporter(
    endpoint=f"https://{WORKSPACE_URL}/api/2.0/otel/v1/traces",
    headers={
        "content-type": "application/x-protobuf",
        "X-Databricks-UC-Table-Name": UC_SPANS_TABLE,
        "Authorization": f"Bearer {DATABRICKS_TOKEN}",
    },
)
provider = TracerProvider()
provider.add_span_processor(BatchSpanProcessor(exporter))
trace.set_tracer_provider(provider)
tracer = trace.get_tracer("my-agent")

def run_agent(query: str) -> str:
    with tracer.start_as_current_span("agent-run") as root:
        root.set_attribute("gen_ai.operation.name", "invoke_agent")
        root.set_attribute("session.id", "conversation-123")
        root.set_attribute("user.id", "user-456")
        root.set_attribute("gen_ai.input.messages", json.dumps([{"role": "user", "content": query}]))
        with tracer.start_as_current_span("llm-call") as llm_span:
            llm_span.set_attribute("gen_ai.operation.name", "chat")
            llm_span.set_attribute("gen_ai.request.model", "my-model")
            response = call_llm(query)
            llm_span.set_attribute("gen_ai.output.messages", json.dumps([{"role": "assistant", "content": response}]))
        root.set_attribute("gen_ai.output.messages", json.dumps([{"role": "assistant", "content": response}]))
        root.set_attribute("gen_ai.usage.input_tokens", 150)
        root.set_attribute("gen_ai.usage.output_tokens", 42)
        return response

provider.shutdown()