Enable custom sampler configuration via env vars

Author: jeremydvoss

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.13.0...HEAD)
 
+- Enabled custom samplers via entry points
+  ([#2972](https://github.com/open-telemetry/opentelemetry-python/pull/2972))
 - Update explicit histogram bucket boundaries
   ([#2947](https://github.com/open-telemetry/opentelemetry-python/pull/2947))
 

opentelemetry-sdk/src/opentelemetry/sdk/_configuration/__init__.py

@@ -23,7 +23,6 @@
 from os import environ
 from typing import Dict, Optional, Sequence, Tuple, Type
 
-from pkg_resources import iter_entry_points
 from typing_extensions import Literal
 
 from opentelemetry.environment_variables import (
@@ -55,6 +54,7 @@
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
 from opentelemetry.sdk.trace.id_generator import IdGenerator
+from opentelemetry.sdk.util import _import_config_components
 from opentelemetry.semconv.resource import ResourceAttributes
 from opentelemetry.trace import set_tracer_provider
 
@@ -228,26 +228,6 @@ def _init_logging(
     logging.getLogger().addHandler(handler)
 
 
-def _import_config_components(
-    selected_components, entry_point_name
-) -> Sequence[Tuple[str, object]]:
-    component_entry_points = {
-        ep.name: ep for ep in iter_entry_points(entry_point_name)
-    }
-    component_impls = []
-    for selected_component in selected_components:
-        entry_point = component_entry_points.get(selected_component, None)
-        if not entry_point:
-            raise RuntimeError(
-                f"Requested component '{selected_component}' not found in entry points for '{entry_point_name}'"
-            )
-
-        component_impl = entry_point.load()
-        component_impls.append((selected_component, component_impl))
-
-    return component_impls
-
-
 def _import_exporters(
     trace_exporter_names: Sequence[str],
     metric_exporter_names: Sequence[str],

opentelemetry-sdk/src/opentelemetry/sdk/trace/sampling.py

@@ -64,7 +64,7 @@
         ...
 
 The tracer sampler can also be configured via environment variables ``OTEL_TRACES_SAMPLER`` and ``OTEL_TRACES_SAMPLER_ARG`` (only if applicable).
-The list of known values for ``OTEL_TRACES_SAMPLER`` are:
+The list of built-in values for ``OTEL_TRACES_SAMPLER`` are:
 
     * always_on - Sampler that always samples spans, regardless of the parent span's sampling decision.
     * always_off - Sampler that never samples spans, regardless of the parent span's sampling decision.
@@ -73,7 +73,24 @@
     * parentbased_always_off - Sampler that respects its parent span's sampling decision, but otherwise never samples.
     * parentbased_traceidratio - Sampler that respects its parent span's sampling decision, but otherwise samples probabalistically based on rate.
 
-Sampling probability can be set with ``OTEL_TRACES_SAMPLER_ARG`` if the sampler is traceidratio or parentbased_traceidratio, when not provided rate will be set to 1.0 (maximum rate possible).
+In order to configure a custom sampler via environment variables, create an entry point for the custom sampler class under the entry point group, ``opentelemtry_traces_sampler``. Then, set the ``OTEL_TRACES_SAMPLER`` environment variable to the key name of the entry point. For example, set ``OTEL_TRACES_SAMPLER=custom_sampler_name`` and ``OTEL_TRACES_SAMPLER_ARG=0.5`` after creating the following entry point:
+
+.. code:: python
+
+    setup(
+        ...
+        entry_points={
+            ...
+            "opentelemtry_traces_sampler": [
+                "custom_sampler_name = path.to.sampler.module:CustomSampler"
+            ]
+        }
+    )
+    ...
+    class CustomSampler(TraceIdRatioBased):
+        ...
+
+Sampling probability can be set with ``OTEL_TRACES_SAMPLER_ARG`` if the sampler is a ``TraceIdRatioBased`` Sampler, such as ``traceidratio`` and ``parentbased_traceidratio``. When not provided rate will be set to 1.0 (maximum rate possible).
 
 
 Prev example but with environment variables. Please make sure to set the env ``OTEL_TRACES_SAMPLER=traceidratio`` and ``OTEL_TRACES_SAMPLER_ARG=0.001``.
@@ -111,10 +128,13 @@
     OTEL_TRACES_SAMPLER,
     OTEL_TRACES_SAMPLER_ARG,
 )
+from opentelemetry.sdk.util import _import_config_components
 from opentelemetry.trace import Link, SpanKind, get_current_span
 from opentelemetry.trace.span import TraceState
 from opentelemetry.util.types import Attributes
 
+# from opentelemetry.sdk._configuration import _import_config_components
+
 _logger = getLogger(__name__)
 
 
@@ -161,6 +181,9 @@ def __init__(
         self.trace_state = trace_state
 
 
+_OTEL_SAMPLER_ENTRY_POINT_GROUP = "opentelemtry_traces_sampler"
+
+
 class Sampler(abc.ABC):
     @abc.abstractmethod
     def should_sample(
@@ -350,7 +373,7 @@ def get_description(self):
 """Sampler that respects its parent span's sampling decision, but otherwise always samples."""
 
 
-class ParentBasedTraceIdRatio(ParentBased):
+class ParentBasedTraceIdRatio(ParentBased, TraceIdRatioBased):
     """
     Sampler that respects its parent span's sampling decision, but otherwise
     samples probabalistically based on `rate`.
@@ -361,37 +384,64 @@ def __init__(self, rate: float):
         super().__init__(root=root)
 
 
-_KNOWN_SAMPLERS = {
+_KNOWN_INITIALIZED_SAMPLERS = {
     "always_on": ALWAYS_ON,
     "always_off": ALWAYS_OFF,
     "parentbased_always_on": DEFAULT_ON,
     "parentbased_always_off": DEFAULT_OFF,
+}
+
+_KNOWN_SAMPLER_CLASSES = {
     "traceidratio": TraceIdRatioBased,
     "parentbased_traceidratio": ParentBasedTraceIdRatio,
 }
 
 
 def _get_from_env_or_default() -> Sampler:
-    trace_sampler = os.getenv(
+    trace_sampler_name = os.getenv(
         OTEL_TRACES_SAMPLER, "parentbased_always_on"
     ).lower()
-    if trace_sampler not in _KNOWN_SAMPLERS:
-        _logger.warning("Couldn't recognize sampler %s.", trace_sampler)
-        trace_sampler = "parentbased_always_on"
 
-    if trace_sampler in ("traceidratio", "parentbased_traceidratio"):
+    if trace_sampler_name in _KNOWN_INITIALIZED_SAMPLERS:
+        return _KNOWN_INITIALIZED_SAMPLERS[trace_sampler_name]
+
+    trace_sampler_impl = None
+    if trace_sampler_name in _KNOWN_SAMPLER_CLASSES:
+        trace_sampler_impl = _KNOWN_SAMPLER_CLASSES[trace_sampler_name]
+    else:
+        try:
+            trace_sampler_impl = _import_sampler(trace_sampler_name)
+        except RuntimeError as err:
+            _logger.warning(
+                "Unable to recognize sampler %s: %s", trace_sampler_name, err
+            )
+            return _KNOWN_INITIALIZED_SAMPLERS["parentbased_always_on"]
+
+    if issubclass(trace_sampler_impl, TraceIdRatioBased):
         try:
             rate = float(os.getenv(OTEL_TRACES_SAMPLER_ARG))
         except ValueError:
             _logger.warning("Could not convert TRACES_SAMPLER_ARG to float.")
             rate = 1.0
-        return _KNOWN_SAMPLERS[trace_sampler](rate)
+        return trace_sampler_impl(rate)
 
-    return _KNOWN_SAMPLERS[trace_sampler]
+    return trace_sampler_impl()
 
 
 def _get_parent_trace_state(parent_context) -> Optional["TraceState"]:
     parent_span_context = get_current_span(parent_context).get_span_context()
     if parent_span_context is None or not parent_span_context.is_valid:
         return None
     return parent_span_context.trace_state
+
+
+def _import_sampler(sampler_name: str) -> Sampler:
+    # pylint: disable=unbalanced-tuple-unpacking
+    [(sampler_name, sampler_impl)] = _import_config_components(
+        [sampler_name.strip()], _OTEL_SAMPLER_ENTRY_POINT_GROUP
+    )
+
+    if issubclass(sampler_impl, Sampler):
+        return sampler_impl
+
+    raise RuntimeError(f"{sampler_name} is not an Sampler")

opentelemetry-sdk/src/opentelemetry/sdk/util/__init__.py

@@ -15,10 +15,11 @@
 import datetime
 import threading
 from collections import OrderedDict, deque
-from collections.abc import MutableMapping, Sequence
-from typing import Optional
+from collections import abc
+from typing import Optional, Sequence, Tuple
 
 from deprecated import deprecated
+from pkg_resources import iter_entry_points
 
 
 def ns_to_iso_str(nanoseconds):
@@ -41,7 +42,27 @@ def get_dict_as_key(labels):
     )
 
 
-class BoundedList(Sequence):
+def _import_config_components(
+    selected_components, entry_point_name
+) -> Sequence[Tuple[str, object]]:
+    component_entry_points = {
+        ep.name: ep for ep in iter_entry_points(entry_point_name)
+    }
+    component_impls = []
+    for selected_component in selected_components:
+        entry_point = component_entry_points.get(selected_component, None)
+        if not entry_point:
+            raise RuntimeError(
+                f"Requested component '{selected_component}' not found in entry points for '{entry_point_name}'"
+            )
+
+        component_impl = entry_point.load()
+        component_impls.append((selected_component, component_impl))
+
+    return component_impls
+
+
+class BoundedList(abc.Sequence):
     """An append only list with a fixed max size.
 
     Calls to `append` and `extend` will drop the oldest elements if there is
@@ -92,7 +113,7 @@ def from_seq(cls, maxlen, seq):
 
 
 @deprecated(version="1.4.0")  # type: ignore
-class BoundedDict(MutableMapping):
+class BoundedDict(abc.MutableMapping):
     """An ordered dict with a fixed max capacity.
 
     Oldest elements are dropped when the dict is full and a new element is

opentelemetry-sdk/tests/test_configurator.py

@@ -257,7 +257,7 @@ def test_trace_init_otlp(self):
 
     @patch.dict(environ, {OTEL_PYTHON_ID_GENERATOR: "custom_id_generator"})
     @patch("opentelemetry.sdk._configuration.IdGenerator", new=IdGenerator)
-    @patch("opentelemetry.sdk._configuration.iter_entry_points")
+    @patch("opentelemetry.sdk.util.iter_entry_points")
     def test_trace_init_custom_id_generator(self, mock_iter_entry_points):
         mock_iter_entry_points.configure_mock(
             return_value=[

opentelemetry-sdk/tests/trace/test_trace.py

@@ -139,6 +139,63 @@ def test_tracer_provider_accepts_concurrent_multi_span_processor(self):
         )
 
 
+class CustomSampler(sampling.Sampler):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def get_description(self) -> str:
+        return super().get_description()
+
+    def should_sample(
+        self,
+        parent_context,
+        trace_id,
+        name,
+        kind,
+        attributes,
+        links,
+        trace_state,
+    ):
+        return super().should_sample(
+            parent_context,
+            trace_id,
+            name,
+            kind,
+            attributes,
+            links,
+            trace_state,
+        )
+
+
+class CustomRatioSampler(sampling.TraceIdRatioBased):
+    def __init__(self, ratio):
+        self.ratio = ratio
+        super().__init__(ratio)
+
+    def get_description(self) -> str:
+        return super().get_description()
+
+    def should_sample(
+        self,
+        parent_context,
+        trace_id,
+        name,
+        kind,
+        attributes,
+        links,
+        trace_state,
+    ):
+        return super().should_sample(
+            parent_context,
+            trace_id,
+            name,
+            kind,
+            attributes,
+            links,
+            trace_state,
+        )
+
+
 class TestTracerSampling(unittest.TestCase):
     def tearDown(self):
         reload(trace)
@@ -219,6 +276,50 @@ def test_ratio_sampler_with_env(self):
         self.assertIsInstance(tracer_provider.sampler, sampling.ParentBased)
         self.assertEqual(tracer_provider.sampler._root.rate, 0.25)
 
+    @mock.patch.dict(
+        "os.environ", {OTEL_TRACES_SAMPLER: "non_existent_entry_point"}
+    )
+    def test_sampler_with_env_non_existent_entry_point(self):
+        # pylint: disable=protected-access
+        reload(trace)
+        tracer_provider = trace.TracerProvider()
+        self.assertIsInstance(tracer_provider.sampler, sampling.ParentBased)
+        # pylint: disable=protected-access
+        self.assertEqual(tracer_provider.sampler._root, sampling.ALWAYS_ON)
+
+    @mock.patch("opentelemetry.sdk.trace.sampling._import_config_components")
+    @mock.patch.dict("os.environ", {OTEL_TRACES_SAMPLER: "custom_sampler"})
+    def test_custom_sampler_with_env(
+        self, mock_sampling_import_config_components
+    ):
+        mock_sampling_import_config_components.return_value = [
+            ("custom_sampler", CustomSampler)
+        ]
+        # pylint: disable=protected-access
+        reload(trace)
+        tracer_provider = trace.TracerProvider()
+        self.assertIsInstance(tracer_provider.sampler, CustomSampler)
+
+    @mock.patch("opentelemetry.sdk.trace.sampling._import_config_components")
+    @mock.patch.dict(
+        "os.environ",
+        {
+            OTEL_TRACES_SAMPLER: "custom_ratio_sampler",
+            OTEL_TRACES_SAMPLER_ARG: "0.5",
+        },
+    )
+    def test_custom_ratio_sampler_with_env(
+        self, mock_sampling_import_config_components
+    ):
+        mock_sampling_import_config_components.return_value = [
+            ("custom_ratio_sampler", CustomRatioSampler)
+        ]
+        # pylint: disable=protected-access
+        reload(trace)
+        tracer_provider = trace.TracerProvider()
+        self.assertIsInstance(tracer_provider.sampler, CustomRatioSampler)
+        self.assertEqual(tracer_provider.sampler.ratio, 0.5)
+
 
 class TestSpanCreation(unittest.TestCase):
     def test_start_span_invalid_spancontext(self):