From ab91407cb706057f717aeee5d29c6f38e8e7af42 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 1 Mar 2025 10:53:36 -0800 Subject: [PATCH] (bug fix) - dd tracer, only send traces when user opts into sending dd-trace (#8928) * fix dd tracing null tracer bug * fix dd tracing * fix base aws llm * test_should_use_dd_tracer --- litellm/litellm_core_utils/dd_tracing.py | 97 ++++++++++--------- .../litellm_core_utils/test_dd_tracing.py | 45 ++++++--- 2 files changed, 81 insertions(+), 61 deletions(-) diff --git a/litellm/litellm_core_utils/dd_tracing.py b/litellm/litellm_core_utils/dd_tracing.py index 4b33b2c4239..fe2c96000c2 100644 --- a/litellm/litellm_core_utils/dd_tracing.py +++ b/litellm/litellm_core_utils/dd_tracing.py @@ -9,57 +9,58 @@ from contextlib import contextmanager from litellm.secret_managers.main import get_secret_bool +class NullSpan: + """A no-op span implementation.""" + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def finish(self): + pass + + +@contextmanager +def null_tracer(name, **kwargs): + """Context manager that yields a no-op span.""" + yield NullSpan() + + +class NullTracer: + """A no-op tracer implementation.""" + + def trace(self, name, **kwargs): + return NullSpan() + + def wrap(self, name=None, **kwargs): + # If called with no arguments (as @tracer.wrap()) + if callable(name): + return name + + # If called with arguments (as @tracer.wrap(name="something")) + def decorator(f): + return f + + return decorator + + def _should_use_dd_tracer(): - """ - Returns True if `USE_DDTRACE` is set to True in .env - """ + """Returns True if `USE_DDTRACE` is set to True in .env""" return get_secret_bool("USE_DDTRACE", False) is True -has_ddtrace = False -try: - from ddtrace import tracer as dd_tracer +# Initialize tracer +should_use_dd_tracer = _should_use_dd_tracer() +tracer = None - if _should_use_dd_tracer(): - has_ddtrace = True -except ImportError: - has_ddtrace = False +if should_use_dd_tracer: + try: + from ddtrace import tracer as dd_tracer - @contextmanager - def null_tracer(name, **kwargs): - class NullSpan: - def __enter__(self): - return self - - def __exit__(self, *args): - pass - - def finish(self): - pass - - yield NullSpan() - - class NullTracer: - def trace(self, name, **kwargs): - class NullSpan: - def __enter__(self): - return self - - def __exit__(self, *args): - pass - - def finish(self): - pass - - return NullSpan() - - def wrap(self, name=None, **kwargs): - def decorator(f): - return f - - return decorator - - dd_tracer = NullTracer() - -# Export the tracer instance -tracer = dd_tracer + tracer = dd_tracer + except ImportError: + tracer = NullTracer() +else: + tracer = NullTracer() diff --git a/tests/litellm/litellm_core_utils/test_dd_tracing.py b/tests/litellm/litellm_core_utils/test_dd_tracing.py index 75318c83ab7..9efb1fe7faa 100644 --- a/tests/litellm/litellm_core_utils/test_dd_tracing.py +++ b/tests/litellm/litellm_core_utils/test_dd_tracing.py @@ -9,26 +9,39 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.litellm_core_utils.dd_tracing import _should_use_dd_tracer, dd_tracer +from litellm.litellm_core_utils.dd_tracing import _should_use_dd_tracer +from litellm.litellm_core_utils.dd_tracing import tracer as dd_tracer def test_dd_tracer_when_package_exists(): - with patch("litellm.litellm_core_utils.dd_tracing.has_ddtrace", True): - # Test the trace context manager - with dd_tracer.trace("test_operation") as span: - assert span is not None + with patch("litellm.litellm_core_utils.dd_tracing.should_use_dd_tracer", True): + """ + Operations works as expected when ddtrace is installed and USE_DDTRACE is True + """ + # Create a mock tracer with the necessary methods + mock_tracer = MagicMock() + mock_span = MagicMock() + mock_tracer.trace.return_value.__enter__.return_value = mock_span - # Test the wrapper decorator - @dd_tracer.wrap(name="test_function") - def sample_function(): - return "test" + with patch("litellm.litellm_core_utils.dd_tracing.tracer", mock_tracer): + # Test the trace context manager + with dd_tracer.trace("test_operation") as span: + assert span is not None - result = sample_function() - assert result == "test" + # Test the wrapper decorator + @dd_tracer.wrap(name="test_function") + def sample_function(): + return "test" + + result = sample_function() + assert result == "test" def test_dd_tracer_when_package_not_exists(): - with patch("litellm.litellm_core_utils.dd_tracing.has_ddtrace", False): + """ + Operations works as expected when ddtrace is not installed and USE_DDTRACE is False + """ + with patch("litellm.litellm_core_utils.dd_tracing.should_use_dd_tracer", False): # Test the trace context manager with null tracer with dd_tracer.trace("test_operation") as span: assert span is not None @@ -45,7 +58,10 @@ def test_dd_tracer_when_package_not_exists(): def test_null_tracer_context_manager(): - with patch("litellm.litellm_core_utils.dd_tracing.has_ddtrace", False): + """ + Test that the context manager works without raising exceptions when should_use_dd_tracer is False + """ + with patch("litellm.litellm_core_utils.dd_tracing.should_use_dd_tracer", False): # Test that the context manager works without raising exceptions with dd_tracer.trace("test_operation") as span: # Test that we can call methods on the null span @@ -54,6 +70,9 @@ def test_null_tracer_context_manager(): def test_should_use_dd_tracer(): + """ + Test that the should_use_dd_tracer function works as expected + """ with patch( "litellm.litellm_core_utils.dd_tracing.get_secret_bool" ) as mock_get_secret: