Skip to content

Commit

Permalink
Move listener hooks to enable instead of on init
Browse files Browse the repository at this point in the history
  • Loading branch information
Yun-Kim committed Jan 9, 2025
1 parent 04ee68f commit e1b6ee1
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 4 deletions.
12 changes: 8 additions & 4 deletions ddtrace/llmobs/_llmobs.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,10 +109,6 @@ def __init__(self, tracer=None):
self._annotations = []
self._annotation_context_lock = forksafe.RLock()

# Register hooks for span events
core.on("trace.span_start", self._do_annotations)
core.on("trace.span_finish", self._on_span_finish)

def _on_span_finish(self, span):
if self.enabled and span.span_type == SpanTypes.LLM:
self._submit_llmobs_span(span)
Expand Down Expand Up @@ -270,6 +266,10 @@ def _start_service(self) -> None:
log.debug("Error starting evaluator runner")

def _stop_service(self) -> None:
# Remove listener hooks for span events
core.reset_listeners("trace.span_start")
core.reset_listeners("trace.span_finish")

try:
self._evaluator_runner.stop()
# flush remaining evaluation spans & evaluations
Expand Down Expand Up @@ -364,6 +364,10 @@ def enable(
cls.enabled = True
cls._instance.start()

# Register hooks for span events
core.on("trace.span_start", cls._instance._do_annotations)
core.on("trace.span_finish", cls._instance._on_span_finish)

atexit.register(cls.disable)
telemetry_writer.product_activated(TELEMETRY_APM_PRODUCT.LLMOBS, True)

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
---
fixes:
- |
LLM Observability: Resolves an issue where enabling LLM Observability in agentless mode would result in traces also being sent to the agent proxy endpoint.
29 changes: 29 additions & 0 deletions tests/llmobs/test_llmobs_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -1319,6 +1319,35 @@ def test_activate_distributed_headers_activates_context(llmobs, mock_llmobs_logs
mock_activate.assert_called_once_with(dummy_context)


def test_listener_hooks_enqueue_correct_writer(run_python_code_in_subprocess):
"""
Regression test that ensures that listener hooks enqueue span events to the correct writer,
not the default writer created at startup.
"""
env = os.environ.copy()
pypath = [os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))]
if "PYTHONPATH" in env:
pypath.append(env["PYTHONPATH"])
env.update(
{"DD_API_KEY": "foobar.baz", "DD_SITE": "datad0g.com", "PYTHONPATH": ":".join(pypath), "DD_TRACE_ENABLED": "0"}
)

out, err, status, pid = run_python_code_in_subprocess(
"""
from ddtrace.llmobs import LLMObs
LLMObs.enable(ml_app="repro-issue", agentless_enabled=True)
with LLMObs.agent("dummy"):
pass
""",
env=env,
)
assert status == 0, err
assert out == b""
assert err == b"failed to send traces to intake at https://llmobs-intake.datad0g.com/api/v2/llmobs: HTTP error status 403, reason Forbidden\n" # noqa: E501
assert b"failed to send, dropping 1 traces to intake at http://localhost:8126/evp_proxy/v2/api/v2/llmobs after 5 retries" not in err # noqa: E501


def test_llmobs_fork_recreates_and_restarts_span_writer():
"""Test that forking a process correctly recreates and restarts the LLMObsSpanWriter."""
with mock.patch("ddtrace.internal.writer.HTTPWriter._send_payload"):
Expand Down

0 comments on commit e1b6ee1

Please sign in to comment.