Skip to content

Commit b02ff47

Browse files
jeremydvosssrikanthccvlzchenocelotl
authored
Custom sampler fix (#3026)
* Fixed circular dependency that can happen when injecting custom samplers * lint * Deleted duplicate tests * lint * lint * lint * lint * lint * Update opentelemetry-sdk/src/opentelemetry/sdk/_configuration/__init__.py Co-authored-by: Leighton Chen <[email protected]> * Update opentelemetry-sdk/src/opentelemetry/sdk/_configuration/__init__.py Co-authored-by: Leighton Chen <[email protected]> * typing * Update opentelemetry-sdk/src/opentelemetry/sdk/_configuration/__init__.py Co-authored-by: Leighton Chen <[email protected]> * Update opentelemetry-sdk/src/opentelemetry/sdk/_configuration/__init__.py Co-authored-by: Srikanth Chekuri <[email protected]> * Retry tests * Fixed circular dependency that can happen when injecting custom samplers * lint * Deleted duplicate tests * lint * lint * lint * lint * lint * Update opentelemetry-sdk/src/opentelemetry/sdk/_configuration/__init__.py Co-authored-by: Leighton Chen <[email protected]> * Update opentelemetry-sdk/src/opentelemetry/sdk/_configuration/__init__.py Co-authored-by: Leighton Chen <[email protected]> * typing * Retry tests * Update opentelemetry-sdk/src/opentelemetry/sdk/_configuration/__init__.py Co-authored-by: Leighton Chen <[email protected]> * Update opentelemetry-sdk/src/opentelemetry/sdk/_configuration/__init__.py Co-authored-by: Srikanth Chekuri <[email protected]> * Updated contrib sha Co-authored-by: Srikanth Chekuri <[email protected]> Co-authored-by: Leighton Chen <[email protected]> Co-authored-by: Diego Hurtado <[email protected]>
1 parent 74ced85 commit b02ff47

File tree

8 files changed

+412
-329
lines changed

8 files changed

+412
-329
lines changed

Diff for: .github/workflows/test.yml

+1-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ env:
1010
# Otherwise, set variable to the commit of your branch on
1111
# opentelemetry-python-contrib which is compatible with these Core repo
1212
# changes.
13-
CONTRIB_REPO_SHA: 66edf69811e142c397d8500cafe6eddeb5565d6e
13+
CONTRIB_REPO_SHA: c6134843900e2eeb1b8b3383a897b38cc0905c38
1414
# This is needed because we do not clone the core repo in contrib builds anymore.
1515
# When running contrib builds as part of core builds, we use actions/checkout@v2 which
1616
# does not set an environment variable (simply just runs tox), which is different when

Diff for: CHANGELOG.md

+2
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
77

88
## Unreleased
99

10+
- Fixed circular dependency issue with custom samplers
11+
([#3026](https://github.com/open-telemetry/opentelemetry-python/pull/3026))
1012
- Add missing entry points for OTLP/HTTP exporter
1113
([#3027](https://github.com/open-telemetry/opentelemetry-python/pull/3027))
1214

Diff for: opentelemetry-sdk/src/opentelemetry/sdk/_configuration/__init__.py

+72-6
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,9 @@
2121
import os
2222
from abc import ABC, abstractmethod
2323
from os import environ
24-
from typing import Dict, Optional, Sequence, Tuple, Type
24+
from typing import Callable, Dict, List, Optional, Sequence, Tuple, Type
2525

26+
from pkg_resources import iter_entry_points
2627
from typing_extensions import Literal
2728

2829
from opentelemetry.environment_variables import (
@@ -44,6 +45,8 @@
4445
OTEL_EXPORTER_OTLP_METRICS_PROTOCOL,
4546
OTEL_EXPORTER_OTLP_PROTOCOL,
4647
OTEL_EXPORTER_OTLP_TRACES_PROTOCOL,
48+
OTEL_TRACES_SAMPLER,
49+
OTEL_TRACES_SAMPLER_ARG,
4750
)
4851
from opentelemetry.sdk.metrics import MeterProvider
4952
from opentelemetry.sdk.metrics.export import (
@@ -54,7 +57,7 @@
5457
from opentelemetry.sdk.trace import TracerProvider
5558
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
5659
from opentelemetry.sdk.trace.id_generator import IdGenerator
57-
from opentelemetry.sdk.util import _import_config_components
60+
from opentelemetry.sdk.trace.sampling import Sampler
5861
from opentelemetry.semconv.resource import ResourceAttributes
5962
from opentelemetry.trace import set_tracer_provider
6063

@@ -82,9 +85,35 @@
8285
_RANDOM_ID_GENERATOR = "random"
8386
_DEFAULT_ID_GENERATOR = _RANDOM_ID_GENERATOR
8487

88+
_OTEL_SAMPLER_ENTRY_POINT_GROUP = "opentelemetry_traces_sampler"
89+
8590
_logger = logging.getLogger(__name__)
8691

8792

93+
def _import_config_components(
94+
selected_components: List[str], entry_point_name: str
95+
) -> Sequence[Tuple[str, object]]:
96+
component_entry_points = {
97+
ep.name: ep for ep in iter_entry_points(entry_point_name)
98+
}
99+
component_impls = []
100+
for selected_component in selected_components:
101+
entry_point = component_entry_points.get(selected_component, None)
102+
if not entry_point:
103+
raise RuntimeError(
104+
f"Requested component '{selected_component}' not found in entry points for '{entry_point_name}'"
105+
)
106+
107+
component_impl = entry_point.load()
108+
component_impls.append((selected_component, component_impl))
109+
110+
return component_impls
111+
112+
113+
def _get_sampler() -> Optional[str]:
114+
return environ.get(OTEL_TRACES_SAMPLER, None)
115+
116+
88117
def _get_id_generator() -> str:
89118
return environ.get(OTEL_PYTHON_ID_GENERATOR, _DEFAULT_ID_GENERATOR)
90119

@@ -149,7 +178,8 @@ def _get_exporter_names(
149178

150179
def _init_tracing(
151180
exporters: Dict[str, Type[SpanExporter]],
152-
id_generator: IdGenerator,
181+
id_generator: IdGenerator = None,
182+
sampler: Sampler = None,
153183
auto_instrumentation_version: Optional[str] = None,
154184
):
155185
# if env var OTEL_RESOURCE_ATTRIBUTES is given, it will read the service_name
@@ -161,7 +191,8 @@ def _init_tracing(
161191
ResourceAttributes.TELEMETRY_AUTO_VERSION
162192
] = auto_instrumentation_version
163193
provider = TracerProvider(
164-
id_generator=id_generator(),
194+
id_generator=id_generator,
195+
sampler=sampler,
165196
resource=Resource.create(auto_resource),
166197
)
167198
set_tracer_provider(provider)
@@ -266,13 +297,41 @@ def _import_exporters(
266297
return trace_exporters, metric_exporters, log_exporters
267298

268299

300+
def _import_sampler_factory(sampler_name: str) -> Callable[[str], Sampler]:
301+
_, sampler_impl = _import_config_components(
302+
[sampler_name.strip()], _OTEL_SAMPLER_ENTRY_POINT_GROUP
303+
)[0]
304+
return sampler_impl
305+
306+
307+
def _import_sampler(sampler_name: str) -> Optional[Sampler]:
308+
if not sampler_name:
309+
return None
310+
try:
311+
sampler_factory = _import_sampler_factory(sampler_name)
312+
sampler_arg = os.getenv(OTEL_TRACES_SAMPLER_ARG, "")
313+
sampler = sampler_factory(sampler_arg)
314+
if not isinstance(sampler, Sampler):
315+
message = f"Sampler factory, {sampler_factory}, produced output, {sampler}, which is not a Sampler."
316+
_logger.warning(message)
317+
raise ValueError(message)
318+
return sampler
319+
except Exception as exc: # pylint: disable=broad-except
320+
_logger.warning(
321+
"Using default sampler. Failed to initialize custom sampler, %s: %s",
322+
sampler_name,
323+
exc,
324+
)
325+
return None
326+
327+
269328
def _import_id_generator(id_generator_name: str) -> IdGenerator:
270329
id_generator_name, id_generator_impl = _import_config_components(
271330
[id_generator_name.strip()], "opentelemetry_id_generator"
272331
)[0]
273332

274333
if issubclass(id_generator_impl, IdGenerator):
275-
return id_generator_impl
334+
return id_generator_impl()
276335

277336
raise RuntimeError(f"{id_generator_name} is not an IdGenerator")
278337

@@ -283,9 +342,16 @@ def _initialize_components(auto_instrumentation_version):
283342
_get_exporter_names("metrics"),
284343
_get_exporter_names("logs"),
285344
)
345+
sampler_name = _get_sampler()
346+
sampler = _import_sampler(sampler_name)
286347
id_generator_name = _get_id_generator()
287348
id_generator = _import_id_generator(id_generator_name)
288-
_init_tracing(trace_exporters, id_generator, auto_instrumentation_version)
349+
_init_tracing(
350+
exporters=trace_exporters,
351+
id_generator=id_generator,
352+
sampler=sampler,
353+
auto_instrumentation_version=auto_instrumentation_version,
354+
)
289355
_init_metrics(metric_exporters, auto_instrumentation_version)
290356
logging_enabled = os.getenv(
291357
_OTEL_PYTHON_LOGGING_AUTO_INSTRUMENTATION_ENABLED, "false"

Diff for: opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py

+25-12
Original file line numberDiff line numberDiff line change
@@ -78,9 +78,6 @@
7878

7979
_ENV_VALUE_UNSET = ""
8080

81-
# pylint: disable=protected-access
82-
_TRACE_SAMPLER = sampling._get_from_env_or_default()
83-
8481

8582
class SpanProcessor:
8683
"""Interface which allows hooks for SDK's `Span` start and end method
@@ -334,7 +331,7 @@ def _check_span_ended(func):
334331
def wrapper(self, *args, **kwargs):
335332
already_ended = False
336333
with self._lock: # pylint: disable=protected-access
337-
if self._end_time is None:
334+
if self._end_time is None: # pylint: disable=protected-access
338335
func(self, *args, **kwargs)
339336
else:
340337
already_ended = True
@@ -519,7 +516,11 @@ def _format_events(events):
519516
f_event = OrderedDict()
520517
f_event["name"] = event.name
521518
f_event["timestamp"] = util.ns_to_iso_str(event.timestamp)
522-
f_event["attributes"] = Span._format_attributes(event.attributes)
519+
f_event[
520+
"attributes"
521+
] = Span._format_attributes( # pylint: disable=protected-access
522+
event.attributes
523+
)
523524
f_events.append(f_event)
524525
return f_events
525526

@@ -528,8 +529,16 @@ def _format_links(links):
528529
f_links = []
529530
for link in links:
530531
f_link = OrderedDict()
531-
f_link["context"] = Span._format_context(link.context)
532-
f_link["attributes"] = Span._format_attributes(link.attributes)
532+
f_link[
533+
"context"
534+
] = Span._format_context( # pylint: disable=protected-access
535+
link.context
536+
)
537+
f_link[
538+
"attributes"
539+
] = Span._format_attributes( # pylint: disable=protected-access
540+
link.attributes
541+
)
533542
f_links.append(f_link)
534543
return f_links
535544

@@ -691,10 +700,12 @@ def _from_env_if_absent(
691700
)
692701

693702
# not removed for backward compat. please use SpanLimits instead.
694-
SPAN_ATTRIBUTE_COUNT_LIMIT = SpanLimits._from_env_if_absent(
695-
None,
696-
OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT,
697-
_DEFAULT_OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT,
703+
SPAN_ATTRIBUTE_COUNT_LIMIT = (
704+
SpanLimits._from_env_if_absent( # pylint: disable=protected-access
705+
None,
706+
OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT,
707+
_DEFAULT_OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT,
708+
)
698709
)
699710

700711

@@ -1115,7 +1126,7 @@ class TracerProvider(trace_api.TracerProvider):
11151126

11161127
def __init__(
11171128
self,
1118-
sampler: sampling.Sampler = _TRACE_SAMPLER,
1129+
sampler: sampling.Sampler = None,
11191130
resource: Resource = Resource.create({}),
11201131
shutdown_on_exit: bool = True,
11211132
active_span_processor: Union[
@@ -1132,6 +1143,8 @@ def __init__(
11321143
else:
11331144
self.id_generator = id_generator
11341145
self._resource = resource
1146+
if not sampler:
1147+
sampler = sampling._get_from_env_or_default()
11351148
self.sampler = sampler
11361149
self._span_limits = span_limits or SpanLimits()
11371150
self._atexit_handler = None

Diff for: opentelemetry-sdk/src/opentelemetry/sdk/trace/sampling.py

+20-44
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,8 @@
7373
* parentbased_always_off - Sampler that respects its parent span's sampling decision, but otherwise never samples.
7474
* parentbased_traceidratio - Sampler that respects its parent span's sampling decision, but otherwise samples probabalistically based on rate.
7575
76-
Sampling probability can be set with ``OTEL_TRACES_SAMPLER_ARG`` if the sampler is traceidratio or parentbased_traceidratio. Rate must be in the range [0.0,1.0]. When not provided rate will be set to 1.0 (maximum rate possible).
76+
Sampling probability can be set with ``OTEL_TRACES_SAMPLER_ARG`` if the sampler is traceidratio or parentbased_traceidratio. Rate must be in the range [0.0,1.0]. When not provided rate will be set to
77+
1.0 (maximum rate possible).
7778
7879
Prev example but with environment variables. Please make sure to set the env ``OTEL_TRACES_SAMPLER=traceidratio`` and ``OTEL_TRACES_SAMPLER_ARG=0.001``.
7980
@@ -97,9 +98,10 @@
9798
with trace.get_tracer(__name__).start_as_current_span("Test Span"):
9899
...
99100
100-
In order to create a configurable custom sampler, create an entry point for the custom sampler factory method under the entry point group, ``opentelemetry_traces_sampler``. The custom sampler factory
101-
method must be of type ``Callable[[str], Sampler]``, taking a single string argument and returning a Sampler object. The single input will come from the string value of the
102-
``OTEL_TRACES_SAMPLER_ARG`` environment variable. If ``OTEL_TRACES_SAMPLER_ARG`` is not configured, the input will be an empty string. For example:
101+
When utilizing a configurator, you can configure a custom sampler. In order to create a configurable custom sampler, create an entry point for the custom sampler
102+
factory method or function under the entry point group, ``opentelemetry_traces_sampler``. The custom sampler factory method must be of type ``Callable[[str], Sampler]``, taking a single string argument and
103+
returning a Sampler object. The single input will come from the string value of the ``OTEL_TRACES_SAMPLER_ARG`` environment variable. If ``OTEL_TRACES_SAMPLER_ARG`` is not configured, the input will
104+
be an empty string. For example:
103105
104106
.. code:: python
105107
@@ -134,15 +136,14 @@ class CustomSamplerFactory:
134136
import os
135137
from logging import getLogger
136138
from types import MappingProxyType
137-
from typing import Callable, Optional, Sequence
139+
from typing import Optional, Sequence
138140

139141
# pylint: disable=unused-import
140142
from opentelemetry.context import Context
141143
from opentelemetry.sdk.environment_variables import (
142144
OTEL_TRACES_SAMPLER,
143145
OTEL_TRACES_SAMPLER_ARG,
144146
)
145-
from opentelemetry.sdk.util import _import_config_components
146147
from opentelemetry.trace import Link, SpanKind, get_current_span
147148
from opentelemetry.trace.span import TraceState
148149
from opentelemetry.util.types import Attributes
@@ -193,9 +194,6 @@ def __init__(
193194
self.trace_state = trace_state
194195

195196

196-
_OTEL_SAMPLER_ENTRY_POINT_GROUP = "opentelemetry_traces_sampler"
197-
198-
199197
class Sampler(abc.ABC):
200198
@abc.abstractmethod
201199
def should_sample(
@@ -407,48 +405,26 @@ def __init__(self, rate: float):
407405

408406

409407
def _get_from_env_or_default() -> Sampler:
410-
traces_sampler_name = os.getenv(
408+
trace_sampler = os.getenv(
411409
OTEL_TRACES_SAMPLER, "parentbased_always_on"
412410
).lower()
411+
if trace_sampler not in _KNOWN_SAMPLERS:
412+
_logger.warning("Couldn't recognize sampler %s.", trace_sampler)
413+
trace_sampler = "parentbased_always_on"
413414

414-
if traces_sampler_name in _KNOWN_SAMPLERS:
415-
if traces_sampler_name in ("traceidratio", "parentbased_traceidratio"):
416-
try:
417-
rate = float(os.getenv(OTEL_TRACES_SAMPLER_ARG))
418-
except ValueError:
419-
_logger.warning(
420-
"Could not convert TRACES_SAMPLER_ARG to float."
421-
)
422-
rate = 1.0
423-
return _KNOWN_SAMPLERS[traces_sampler_name](rate)
424-
return _KNOWN_SAMPLERS[traces_sampler_name]
425-
try:
426-
traces_sampler_factory = _import_sampler_factory(traces_sampler_name)
427-
sampler_arg = os.getenv(OTEL_TRACES_SAMPLER_ARG, "")
428-
traces_sampler = traces_sampler_factory(sampler_arg)
429-
if not isinstance(traces_sampler, Sampler):
430-
message = f"Traces sampler factory, {traces_sampler_factory}, produced output, {traces_sampler}, which is not a Sampler object."
431-
_logger.warning(message)
432-
raise ValueError(message)
433-
return traces_sampler
434-
except Exception as exc: # pylint: disable=broad-except
435-
_logger.warning(
436-
"Using default sampler. Failed to initialize custom sampler, %s: %s",
437-
traces_sampler_name,
438-
exc,
439-
)
440-
return _KNOWN_SAMPLERS["parentbased_always_on"]
415+
if trace_sampler in ("traceidratio", "parentbased_traceidratio"):
416+
try:
417+
rate = float(os.getenv(OTEL_TRACES_SAMPLER_ARG))
418+
except ValueError:
419+
_logger.warning("Could not convert TRACES_SAMPLER_ARG to float.")
420+
rate = 1.0
421+
return _KNOWN_SAMPLERS[trace_sampler](rate)
422+
423+
return _KNOWN_SAMPLERS[trace_sampler]
441424

442425

443426
def _get_parent_trace_state(parent_context) -> Optional["TraceState"]:
444427
parent_span_context = get_current_span(parent_context).get_span_context()
445428
if parent_span_context is None or not parent_span_context.is_valid:
446429
return None
447430
return parent_span_context.trace_state
448-
449-
450-
def _import_sampler_factory(sampler_name: str) -> Callable[[str], Sampler]:
451-
_, sampler_impl = _import_config_components(
452-
[sampler_name.strip()], _OTEL_SAMPLER_ENTRY_POINT_GROUP
453-
)[0]
454-
return sampler_impl

0 commit comments

Comments
 (0)