Skip to content

Commit

Permalink
Adds OpenTelemetry integration.
Browse files Browse the repository at this point in the history
Two approaches:
1. Log Burr -> Otel -- simple adapter
2. Log any Otel traces inside a burr step to the Burr UI

This is overall fairly sound, although there is a bit of global state
that might be dealt with.

Design decisions:

1. Use a custom span processor
2. Store state of a map of span_id -> step
3. Propogate that when processing spans

It's a bit odd as the span processor should just be writing/queuing, but
the way I'm thinking of this is that the writing is converting OTel
spans to Burr spans, so we have to do that somewhere. Why not when
processing? Then we can just cascade down.

Any spans outside of a Burr step are ignored, as we are not a full OTel
provider.
  • Loading branch information
elijahbenizzy committed Aug 11, 2024
1 parent fa2d028 commit f9883e3
Show file tree
Hide file tree
Showing 4 changed files with 20 additions and 7 deletions.
11 changes: 6 additions & 5 deletions burr/cli/demo_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,16 +66,20 @@ def _modify(app_id: str) -> str:

def _run_conversation(app_id, prompts):
tracker = (
LocalTrackingClient(project=project_id, storage_dir=data_dir)
LocalTrackingClient(project=project_id + "_otel", storage_dir=data_dir)
if not s3_bucket
else S3TrackingClient(project=project_id, bucket=s3_bucket)
)
from opentelemetry.instrumentation.openai import OpenAIInstrumentor

# TODO -- get this auto-registered
OpenAIInstrumentor().instrument()
graph = (chatbot_application_with_traces if use_traces else chatbot_application).graph
app = (
ApplicationBuilder()
.with_graph(graph)
.with_identifiers(app_id=app_id)
.with_tracker(tracker)
.with_tracker(tracker, use_otel_tracing=True)
.with_entrypoint("prompt")
.build()
)
Expand Down Expand Up @@ -201,6 +205,3 @@ def generate_all(
generate_counter_data(data_dir=data_dir, s3_bucket=s3_bucket, unique_app_names=unique_app_names)
logger.info("Generating RAG data")
generate_rag_data(data_dir=data_dir, s3_bucket=s3_bucket, unique_app_names=unique_app_names)


#
10 changes: 8 additions & 2 deletions burr/core/application.py
Original file line number Diff line number Diff line change
Expand Up @@ -1812,6 +1812,7 @@ def with_tracker(
tracker: Union[Literal["local"], "TrackingClient"] = "local",
project: str = "default",
params: Dict[str, Any] = None,
use_otel_tracing: bool = False,
):
"""Adds a "tracker" to the application. The tracker specifies
a project name (used for disambiguating groups of tracers), and plugs into the
Expand All @@ -1826,6 +1827,8 @@ def with_tracker(
:param tracker: Tracker to use. ``local`` creates one, else pass one in.
:param project: Project name -- used if the tracker is string-specified (local).
:param params: Parameters to pass to the tracker if it's string-specified (local).
:param use_otel_tracing: Whether to log opentelemetry traces to the Burr UI. This is experimental but we will be adding
full support shortly. This requires burr[opentelemetry] installed. Note you can also log burr to OpenTelemetry using the OpenTelemetry adapter
:return: The application builder for future chaining.
"""
# if it's a lifecycle adapter, just add it
Expand All @@ -1839,15 +1842,18 @@ def with_tracker(
kwargs = {"project": project}
kwargs.update(params)
instantiated_tracker = LocalTrackingClient(**kwargs)
self.lifecycle_adapters.append(instantiated_tracker)
else:
raise ValueError(f"Tracker {tracker}:{project} not supported")
else:
self.lifecycle_adapters.append(instantiated_tracker)
if params is not None:
raise ValueError(
"Params are not supported for object-specified trackers, these are already initialized!"
)
if use_otel_tracing:
from burr.integrations.opentelemetry import OpenTelemetryTracker

instantiated_tracker = OpenTelemetryTracker(burr_tracker=instantiated_tracker)
self.lifecycle_adapters.append(instantiated_tracker)
self.tracker = instantiated_tracker
return self

Expand Down
2 changes: 2 additions & 0 deletions burr/tracking/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
PreRunStepHook,
PreStartSpanHook,
)
from burr.lifecycle.base import DoLogAttributeHook


class SyncTrackingClient(
Expand All @@ -15,6 +16,7 @@ class SyncTrackingClient(
PostRunStepHook,
PreStartSpanHook,
PostEndSpanHook,
DoLogAttributeHook,
abc.ABC,
):
"""Base class for synchronous tracking clients. All tracking clients must implement from this
Expand Down
4 changes: 4 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,10 @@ developer = [
"pre-commit",
]

opentelemetry = [
"opentelemetry-api",
"opentelemetry-sdk",
]
[tool.setuptools]
include-package-data = true

Expand Down

0 comments on commit f9883e3

Please sign in to comment.