Merge branch 'main' into taegyunkim/profiling-flaky-tests

ddtrace/debugging/_safety.py

@@ -1,5 +1,6 @@
 from inspect import CO_VARARGS
 from inspect import CO_VARKEYWORDS
+from itertools import chain
 from types import FrameType
 from typing import Any
 from typing import Dict
@@ -23,11 +24,11 @@ def get_args(frame: FrameType) -> Iterator[Tuple[str, Any]]:
 
 def get_locals(frame: FrameType) -> Iterator[Tuple[str, Any]]:
     code = frame.f_code
+    _locals = frame.f_locals
     nargs = code.co_argcount + bool(code.co_flags & CO_VARARGS) + bool(code.co_flags & CO_VARKEYWORDS)
-    names = code.co_varnames[nargs:]
-    values = (frame.f_locals.get(name) for name in names)
-
-    return zip(names, values)
+    return (
+        (name, _locals.get(name)) for name in chain(code.co_varnames[nargs:], code.co_freevars, code.co_cellvars)
+    )  # include freevars and cellvars
 
 
 def get_globals(frame: FrameType) -> Iterator[Tuple[str, Any]]:

ddtrace/llmobs/_llmobs.py

@@ -111,9 +111,9 @@ 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_start(self, span):
+        if self.enabled and span.span_type == SpanTypes.LLM:
+            self._do_annotations(span)
 
     def _on_span_finish(self, span):
         if self.enabled and span.span_type == SpanTypes.LLM:
@@ -272,6 +272,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", self._on_span_start)
+        core.reset_listeners("trace.span_finish", self._on_span_finish)
+
         try:
             self._evaluator_runner.stop()
             # flush remaining evaluation spans & evaluations
@@ -366,6 +370,10 @@ def enable(
         cls.enabled = True
         cls._instance.start()
 
+        # Register hooks for span events
+        core.on("trace.span_start", cls._instance._on_span_start)
+        core.on("trace.span_finish", cls._instance._on_span_finish)
+
         atexit.register(cls.disable)
         telemetry_writer.product_activated(TELEMETRY_APM_PRODUCT.LLMOBS, True)
 

ddtrace/profiling/collector/_lock.py

@@ -185,6 +185,7 @@ def _release(self, inner_func, *args, **kwargs):
             # _self_acquired_at is only set when the acquire was captured
             # if it's not set, we're not capturing the release
             start = self._self_acquired_at
+            del self._self_acquired_at
 
         try:
             return inner_func(*args, **kwargs)

releasenotes/notes/fix-er-include-nonlocals-bbeecfbbbde35496.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    exception replay: include missing nonlocal variables in snapshot log messages.

releasenotes/notes/fix-llmobs-default-writer-hooks-5e456c2f7dfd4381.yaml

@@ -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.

releasenotes/notes/profiling-lock-acquired-at-remove-c8b5b96130a46ca8.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    profiling: resolves an issue where lock release would have been captured
+    with a wrong acquire timestamp

tests/debugging/exception/test_replay.py

@@ -294,3 +294,23 @@ def c(foo=42):
             self.assert_span_count(6)
             # no new snapshots
             assert len(uploader.collector.queue) == 3
+
+    def test_debugger_exception_in_closure(self):
+        def b():
+            with self.trace("b"):
+                nonloc = 4
+
+                def a(v):
+                    if nonloc:
+                        raise ValueError("hello", v)
+
+                a(nonloc)
+
+        with exception_replay() as uploader:
+            with with_rate_limiter(RateLimiter(limit_rate=1, raise_on_exceed=False)):
+                with pytest.raises(ValueError):
+                    b()
+
+            assert all(
+                s.line_capture["locals"]["nonloc"] == {"type": "int", "value": "4"} for s in uploader.collector.queue
+            )

tests/debugging/test_safety.py

@@ -15,7 +15,10 @@ def assert_args(args):
         assert set(dict(_safety.get_args(inspect.currentframe().f_back)).keys()) == args
 
     def assert_locals(_locals):
-        assert set(dict(_safety.get_locals(inspect.currentframe().f_back)).keys()) == _locals
+        assert set(dict(_safety.get_locals(inspect.currentframe().f_back)).keys()) == _locals | {
+            "assert_args",
+            "assert_locals",
+        }
 
     def assert_globals(_globals):
         assert set(dict(_safety.get_globals(inspect.currentframe().f_back)).keys()) == _globals

tests/llmobs/test_llmobs_evaluator_runner.py

@@ -64,15 +64,7 @@ def test_evaluator_runner_on_exit(mock_writer_logs, run_python_code_in_subproces
     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": os.getenv("DD_API_KEY", "dummy-api-key"),
-            "DD_SITE": "datad0g.com",
-            "PYTHONPATH": ":".join(pypath),
-            "DD_LLMOBS_ML_APP": "unnamed-ml-app",
-            "_DD_LLMOBS_EVALUATOR_INTERVAL": "5",
-        }
-    )
+    env.update({"PYTHONPATH": ":".join(pypath), "_DD_LLMOBS_EVALUATOR_INTERVAL": "5"})
     out, err, status, pid = run_python_code_in_subprocess(
         """
 import os
@@ -87,7 +79,7 @@ def test_evaluator_runner_on_exit(mock_writer_logs, run_python_code_in_subproces
 ctx = logs_vcr.use_cassette("tests.llmobs.test_llmobs_evaluator_runner.send_score_metric.yaml")
 ctx.__enter__()
 atexit.register(lambda: ctx.__exit__())
-LLMObs.enable()
+LLMObs.enable(api_key="dummy-api-key", site="datad0g.com", ml_app="unnamed-ml-app")
 LLMObs._instance._evaluator_runner.evaluators.append(DummyEvaluator(llmobs_service=LLMObs))
 LLMObs._instance._evaluator_runner.start()
 LLMObs._instance._evaluator_runner.enqueue({"span_id": "123", "trace_id": "1234"}, None)

tests/llmobs/test_llmobs_ragas_evaluators.py

@@ -206,14 +206,11 @@ def test_llmobs_with_faithfulness_emits_traces_and_evals_on_exit(mock_writer_log
         pypath.append(env["PYTHONPATH"])
     env.update(
         {
-            "DD_API_KEY": os.getenv("DD_API_KEY", "dummy-api-key"),
-            "DD_SITE": "datad0g.com",
             "PYTHONPATH": ":".join(pypath),
             "OPENAI_API_KEY": os.getenv("OPENAI_API_KEY", "dummy-openai-api-key"),
-            "DD_LLMOBS_ML_APP": "unnamed-ml-app",
             "_DD_LLMOBS_EVALUATOR_INTERVAL": "5",
             "_DD_LLMOBS_EVALUATORS": "ragas_faithfulness",
-            "DD_LLMOBS_AGENTLESS_ENABLED": "1",
+            "DD_TRACE_ENABLED": "0",
         }
     )
     out, err, status, pid = run_python_code_in_subprocess(
@@ -232,14 +229,8 @@ def test_llmobs_with_faithfulness_emits_traces_and_evals_on_exit(mock_writer_log
 )
 ctx.__enter__()
 atexit.register(lambda: ctx.__exit__())
-with mock.patch(
-    "ddtrace.internal.writer.HTTPWriter._send_payload",
-    return_value=Response(
-        status=200,
-        body="{}",
-    ),
-):
-    LLMObs.enable()
+with mock.patch("ddtrace.internal.writer.HTTPWriter._send_payload", return_value=Response(status=200, body="{}")):
+    LLMObs.enable(api_key="dummy-api-key", site="datad0g.com", ml_app="unnamed-ml-app", agentless_enabled=True)
     LLMObs._instance._evaluator_runner.enqueue(_llm_span_with_expected_ragas_inputs_in_messages(), None)
     """,
         env=env,

tests/llmobs/test_llmobs_service.py

@@ -1310,6 +1310,34 @@ 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({"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, api_key="foobar.baz", site="datad0g.com")
+with LLMObs.agent("dummy"):
+    pass
+""",
+        env=env,
+    )
+    assert status == 0, err
+    assert out == b""
+    agentless_writer_log = 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
+    agent_proxy_log = b"failed to send, dropping 1 traces to intake at http://localhost:8126/evp_proxy/v2/api/v2/llmobs after 5 retries"  # noqa: E501
+    assert err == agentless_writer_log
+    assert agent_proxy_log not in err
+
+
 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"):