From fdb5d632750c35a01e2d18bac9a6f6b150452597 Mon Sep 17 00:00:00 2001 From: Nikhil172913832 Date: Sat, 11 Oct 2025 18:37:55 +0530 Subject: [PATCH] feat(traceloop-sdk): add sampling_rate parameter to control trace volume Add simplified sampling_rate parameter to Traceloop.init() for easy trace sampling configuration without manually creating sampler objects. Closes #2109 --- .../traceloop-sdk/tests/test_sampling_rate.py | 113 ++++++++++++++++++ .../traceloop-sdk/traceloop/sdk/__init__.py | 16 ++- 2 files changed, 128 insertions(+), 1 deletion(-) create mode 100644 packages/traceloop-sdk/tests/test_sampling_rate.py diff --git a/packages/traceloop-sdk/tests/test_sampling_rate.py b/packages/traceloop-sdk/tests/test_sampling_rate.py new file mode 100644 index 0000000000..045d48f6bf --- /dev/null +++ b/packages/traceloop-sdk/tests/test_sampling_rate.py @@ -0,0 +1,113 @@ +"""Tests for sampling_rate parameter functionality.""" + +import pytest +from opentelemetry.sdk.trace.sampling import TraceIdRatioBased +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from traceloop.sdk import Traceloop +from traceloop.sdk.tracing.tracing import TracerWrapper + + +@pytest.fixture +def clean_tracer_wrapper(): + """Fixture to manage TracerWrapper global state.""" + original_instance = None + if hasattr(TracerWrapper, "instance"): + original_instance = TracerWrapper.instance + del TracerWrapper.instance + + yield + + if hasattr(TracerWrapper, "instance"): + del TracerWrapper.instance + if original_instance is not None: + TracerWrapper.instance = original_instance + + +class TestSamplingRate: + """Test class for sampling_rate parameter functionality.""" + + def test_init_with_sampling_rate(self, clean_tracer_wrapper): + """Test initialization with sampling_rate parameter.""" + exporter = InMemorySpanExporter() + + client = Traceloop.init( + app_name="test-sampling-rate", + sampling_rate=0.5, + exporter=exporter, + disable_batch=True + ) + + assert client is None + assert hasattr(TracerWrapper, "instance") + assert TracerWrapper.instance is not None + + # Verify that a tracer provider was created with a sampler + tracer_provider = TracerWrapper.instance._TracerWrapper__tracer_provider + assert tracer_provider is not None + assert tracer_provider.sampler is not None + + def test_sampling_rate_validation(self, clean_tracer_wrapper): + """Test that sampling_rate validates input range.""" + exporter = InMemorySpanExporter() + + with pytest.raises(ValueError, match="sampling_rate must be between 0.0 and 1.0"): + Traceloop.init( + app_name="test-invalid-sampling-rate", + sampling_rate=1.5, + exporter=exporter, + disable_batch=True + ) + + with pytest.raises(ValueError, match="sampling_rate must be between 0.0 and 1.0"): + Traceloop.init( + app_name="test-invalid-sampling-rate", + sampling_rate=-0.1, + exporter=exporter, + disable_batch=True + ) + + def test_sampling_rate_and_sampler_conflict(self, clean_tracer_wrapper): + """Test that providing both sampling_rate and sampler raises an error.""" + exporter = InMemorySpanExporter() + sampler = TraceIdRatioBased(0.5) + + with pytest.raises(ValueError, match="Cannot specify both 'sampler' and 'sampling_rate'"): + Traceloop.init( + app_name="test-conflict", + sampling_rate=0.5, + sampler=sampler, + exporter=exporter, + disable_batch=True + ) + + def test_sampling_rate_edge_cases(self, clean_tracer_wrapper): + """Test sampling_rate with edge case values (0.0 and 1.0).""" + exporter = InMemorySpanExporter() + + # Test 0.0 (no sampling) + client = Traceloop.init( + app_name="test-sampling-zero", + sampling_rate=0.0, + exporter=exporter, + disable_batch=True + ) + assert client is None + assert hasattr(TracerWrapper, "instance") + + tracer_provider = TracerWrapper.instance._TracerWrapper__tracer_provider + assert tracer_provider.sampler is not None + + del TracerWrapper.instance + + # Test 1.0 (full sampling) + client = Traceloop.init( + app_name="test-sampling-one", + sampling_rate=1.0, + exporter=exporter, + disable_batch=True + ) + assert client is None + assert hasattr(TracerWrapper, "instance") + + tracer_provider = TracerWrapper.instance._TracerWrapper__tracer_provider + assert tracer_provider.sampler is not None diff --git a/packages/traceloop-sdk/traceloop/sdk/__init__.py b/packages/traceloop-sdk/traceloop/sdk/__init__.py index 1b91cb9ff6..2db40242bf 100644 --- a/packages/traceloop-sdk/traceloop/sdk/__init__.py +++ b/packages/traceloop-sdk/traceloop/sdk/__init__.py @@ -5,7 +5,7 @@ from typing import Callable, List, Optional, Set, Union from colorama import Fore from opentelemetry.sdk.trace import SpanProcessor, ReadableSpan -from opentelemetry.sdk.trace.sampling import Sampler +from opentelemetry.sdk.trace.sampling import Sampler, TraceIdRatioBased from opentelemetry.sdk.trace.export import SpanExporter from opentelemetry.sdk.metrics.export import MetricExporter from opentelemetry.sdk._logs.export import LogExporter @@ -62,6 +62,7 @@ def init( processor: Optional[Union[SpanProcessor, List[SpanProcessor]]] = None, propagator: TextMapPropagator = None, sampler: Optional[Sampler] = None, + sampling_rate: Optional[float] = None, traceloop_sync_enabled: bool = False, should_enrich_metrics: bool = True, resource_attributes: dict = {}, @@ -136,6 +137,19 @@ def init( print(Fore.RESET) + if sampling_rate is not None and sampler is not None: + raise ValueError( + "Cannot specify both 'sampler' and 'sampling_rate'. " + "Please use only one of these parameters." + ) + + if sampling_rate is not None: + if not 0.0 <= sampling_rate <= 1.0: + raise ValueError( + f"sampling_rate must be between 0.0 and 1.0, got {sampling_rate}" + ) + sampler = TraceIdRatioBased(sampling_rate) + # Tracer init resource_attributes.update({SERVICE_NAME: app_name}) TracerWrapper.set_static_params(