Skip to content

Commit fdb5d63

Browse files
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
1 parent e66894f commit fdb5d63

File tree

2 files changed

+128
-1
lines changed

2 files changed

+128
-1
lines changed
Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
"""Tests for sampling_rate parameter functionality."""
2+
3+
import pytest
4+
from opentelemetry.sdk.trace.sampling import TraceIdRatioBased
5+
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
6+
from traceloop.sdk import Traceloop
7+
from traceloop.sdk.tracing.tracing import TracerWrapper
8+
9+
10+
@pytest.fixture
11+
def clean_tracer_wrapper():
12+
"""Fixture to manage TracerWrapper global state."""
13+
original_instance = None
14+
if hasattr(TracerWrapper, "instance"):
15+
original_instance = TracerWrapper.instance
16+
del TracerWrapper.instance
17+
18+
yield
19+
20+
if hasattr(TracerWrapper, "instance"):
21+
del TracerWrapper.instance
22+
if original_instance is not None:
23+
TracerWrapper.instance = original_instance
24+
25+
26+
class TestSamplingRate:
27+
"""Test class for sampling_rate parameter functionality."""
28+
29+
def test_init_with_sampling_rate(self, clean_tracer_wrapper):
30+
"""Test initialization with sampling_rate parameter."""
31+
exporter = InMemorySpanExporter()
32+
33+
client = Traceloop.init(
34+
app_name="test-sampling-rate",
35+
sampling_rate=0.5,
36+
exporter=exporter,
37+
disable_batch=True
38+
)
39+
40+
assert client is None
41+
assert hasattr(TracerWrapper, "instance")
42+
assert TracerWrapper.instance is not None
43+
44+
# Verify that a tracer provider was created with a sampler
45+
tracer_provider = TracerWrapper.instance._TracerWrapper__tracer_provider
46+
assert tracer_provider is not None
47+
assert tracer_provider.sampler is not None
48+
49+
def test_sampling_rate_validation(self, clean_tracer_wrapper):
50+
"""Test that sampling_rate validates input range."""
51+
exporter = InMemorySpanExporter()
52+
53+
with pytest.raises(ValueError, match="sampling_rate must be between 0.0 and 1.0"):
54+
Traceloop.init(
55+
app_name="test-invalid-sampling-rate",
56+
sampling_rate=1.5,
57+
exporter=exporter,
58+
disable_batch=True
59+
)
60+
61+
with pytest.raises(ValueError, match="sampling_rate must be between 0.0 and 1.0"):
62+
Traceloop.init(
63+
app_name="test-invalid-sampling-rate",
64+
sampling_rate=-0.1,
65+
exporter=exporter,
66+
disable_batch=True
67+
)
68+
69+
def test_sampling_rate_and_sampler_conflict(self, clean_tracer_wrapper):
70+
"""Test that providing both sampling_rate and sampler raises an error."""
71+
exporter = InMemorySpanExporter()
72+
sampler = TraceIdRatioBased(0.5)
73+
74+
with pytest.raises(ValueError, match="Cannot specify both 'sampler' and 'sampling_rate'"):
75+
Traceloop.init(
76+
app_name="test-conflict",
77+
sampling_rate=0.5,
78+
sampler=sampler,
79+
exporter=exporter,
80+
disable_batch=True
81+
)
82+
83+
def test_sampling_rate_edge_cases(self, clean_tracer_wrapper):
84+
"""Test sampling_rate with edge case values (0.0 and 1.0)."""
85+
exporter = InMemorySpanExporter()
86+
87+
# Test 0.0 (no sampling)
88+
client = Traceloop.init(
89+
app_name="test-sampling-zero",
90+
sampling_rate=0.0,
91+
exporter=exporter,
92+
disable_batch=True
93+
)
94+
assert client is None
95+
assert hasattr(TracerWrapper, "instance")
96+
97+
tracer_provider = TracerWrapper.instance._TracerWrapper__tracer_provider
98+
assert tracer_provider.sampler is not None
99+
100+
del TracerWrapper.instance
101+
102+
# Test 1.0 (full sampling)
103+
client = Traceloop.init(
104+
app_name="test-sampling-one",
105+
sampling_rate=1.0,
106+
exporter=exporter,
107+
disable_batch=True
108+
)
109+
assert client is None
110+
assert hasattr(TracerWrapper, "instance")
111+
112+
tracer_provider = TracerWrapper.instance._TracerWrapper__tracer_provider
113+
assert tracer_provider.sampler is not None

packages/traceloop-sdk/traceloop/sdk/__init__.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from typing import Callable, List, Optional, Set, Union
66
from colorama import Fore
77
from opentelemetry.sdk.trace import SpanProcessor, ReadableSpan
8-
from opentelemetry.sdk.trace.sampling import Sampler
8+
from opentelemetry.sdk.trace.sampling import Sampler, TraceIdRatioBased
99
from opentelemetry.sdk.trace.export import SpanExporter
1010
from opentelemetry.sdk.metrics.export import MetricExporter
1111
from opentelemetry.sdk._logs.export import LogExporter
@@ -62,6 +62,7 @@ def init(
6262
processor: Optional[Union[SpanProcessor, List[SpanProcessor]]] = None,
6363
propagator: TextMapPropagator = None,
6464
sampler: Optional[Sampler] = None,
65+
sampling_rate: Optional[float] = None,
6566
traceloop_sync_enabled: bool = False,
6667
should_enrich_metrics: bool = True,
6768
resource_attributes: dict = {},
@@ -136,6 +137,19 @@ def init(
136137

137138
print(Fore.RESET)
138139

140+
if sampling_rate is not None and sampler is not None:
141+
raise ValueError(
142+
"Cannot specify both 'sampler' and 'sampling_rate'. "
143+
"Please use only one of these parameters."
144+
)
145+
146+
if sampling_rate is not None:
147+
if not 0.0 <= sampling_rate <= 1.0:
148+
raise ValueError(
149+
f"sampling_rate must be between 0.0 and 1.0, got {sampling_rate}"
150+
)
151+
sampler = TraceIdRatioBased(sampling_rate)
152+
139153
# Tracer init
140154
resource_attributes.update({SERVICE_NAME: app_name})
141155
TracerWrapper.set_static_params(

0 commit comments

Comments
 (0)