Config changes for rate-limited-sampler

Author: rads-1996

sdk/monitor/azure-monitor-opentelemetry-exporter/CHANGELOG.md

@@ -9,7 +9,7 @@
 - Added customer-facing statsbeat preview.
   ([#41669](https://github.com/Azure/azure-sdk-for-python/pull/41669))
 - Added RateLimited Sampler 
-  ([#41925](https://github.com/Azure/azure-sdk-for-python/pull/41925))
+  ([#41954](https://github.com/Azure/azure-sdk-for-python/pull/41954))
 
 ### Breaking Changes
 

sdk/monitor/azure-monitor-opentelemetry/CHANGELOG.md

@@ -3,6 +3,8 @@
 ## 1.6.11 (Unreleased)
 
 ### Features Added
+- Added RateLimited Sampler 
+  ([#41954](https://github.com/Azure/azure-sdk-for-python/pull/41954))
 
 ### Breaking Changes
 

sdk/monitor/azure-monitor-opentelemetry/azure/monitor/opentelemetry/_configure.py

@@ -35,6 +35,7 @@
     LOGGING_FORMATTER_ARG,
     RESOURCE_ARG,
     SAMPLING_RATIO_ARG,
+    SAMPLING_TRACES_PER_SECOND_ARG,
     SPAN_PROCESSORS_ARG,
     VIEWS_ARG,
 )
@@ -50,6 +51,7 @@
     ApplicationInsightsSampler,
     AzureMonitorMetricExporter,
     AzureMonitorTraceExporter,
+    RateLimitedSampler,
 )
 from azure.monitor.opentelemetry.exporter._utils import (  # pylint: disable=import-error,no-name-in-module
     _is_attach_enabled,
@@ -133,10 +135,18 @@ def configure_azure_monitor(**kwargs) -> None:  # pylint: disable=C4758
 
 def _setup_tracing(configurations: Dict[str, ConfigurationValue]):
     resource: Resource = configurations[RESOURCE_ARG]  # type: ignore
-    sampling_ratio = configurations[SAMPLING_RATIO_ARG]
-    tracer_provider = TracerProvider(
-        sampler=ApplicationInsightsSampler(sampling_ratio=cast(float, sampling_ratio)), resource=resource
-    )
+    if SAMPLING_TRACES_PER_SECOND_ARG in configurations:
+        sampling_traces_per_second = configurations[SAMPLING_TRACES_PER_SECOND_ARG]
+        tracer_provider = TracerProvider(
+            sampler=RateLimitedSampler(sampling_ratio=cast(float, sampling_traces_per_second), resource=resource)
+        )
+    else:
+        sampling_ratio = configurations[SAMPLING_RATIO_ARG]
+        tracer_provider = TracerProvider(
+            sampler=ApplicationInsightsSampler(sampling_ratio=cast(float, sampling_ratio)), resource=resource
+        )
+
+
     for span_processor in configurations[SPAN_PROCESSORS_ARG]:  # type: ignore
         tracer_provider.add_span_processor(span_processor)  # type: ignore
     if configurations.get(ENABLE_LIVE_METRICS_ARG):

sdk/monitor/azure-monitor-opentelemetry/azure/monitor/opentelemetry/_constants.py

@@ -24,7 +24,9 @@
 SAMPLING_RATIO_ARG = "sampling_ratio"
 SPAN_PROCESSORS_ARG = "span_processors"
 VIEWS_ARG = "views"
-
+RATE_LIMITED_SAMPLER = "microsoft.rate_limited"
+FIXED_PERCENTAGE_SAMPLER = "microsoft.fixed.percentage"
+SAMPLING_TRACES_PER_SECOND_ARG = "sampling_traces_per_second"
 
 # --------------------Autoinstrumentation Configuration------------------------------------------
 

sdk/monitor/azure-monitor-opentelemetry/azure/monitor/opentelemetry/_utils/configurations.py

@@ -19,6 +19,7 @@
 from opentelemetry.sdk.environment_variables import (
     OTEL_EXPERIMENTAL_RESOURCE_DETECTORS,
     OTEL_TRACES_SAMPLER_ARG,
+    OTEL_TRACES_SAMPLER
 )
 from opentelemetry.sdk.resources import Resource
 
@@ -37,25 +38,25 @@
     LOGGING_FORMATTER_ARG,
     RESOURCE_ARG,
     SAMPLING_RATIO_ARG,
+    SAMPLING_TRACES_PER_SECOND_ARG,
     SPAN_PROCESSORS_ARG,
     VIEWS_ARG,
+    RATE_LIMITED_SAMPLER,
+    FIXED_PERCENTAGE_SAMPLER,
 )
 from azure.monitor.opentelemetry._types import ConfigurationValue
 from azure.monitor.opentelemetry._version import VERSION
 
 
 _INVALID_FLOAT_MESSAGE = "Value of %s must be a float. Defaulting to %s: %s"
+_INVALID_TRACES_PER_SECOND_MESSAGE = "Value of %s must be a positive number for traces per second. Defaulting to %s: %s"
 _SUPPORTED_RESOURCE_DETECTORS = (
     _AZURE_APP_SERVICE_RESOURCE_DETECTOR_NAME,
     _AZURE_VM_RESOURCE_DETECTOR_NAME,
 )
-# TODO: remove when sampler uses env var instead
-SAMPLING_RATIO_ENV_VAR = OTEL_TRACES_SAMPLER_ARG
-
 
 _logger = getLogger(__name__)
 
-
 def _get_configurations(**kwargs) -> Dict[str, ConfigurationValue]:
     configurations = {}
 
@@ -120,21 +121,54 @@ def _default_resource(configurations):
         configurations[RESOURCE_ARG] = Resource.create(configurations[RESOURCE_ARG].attributes)
 
 
-# TODO: remove when sampler uses env var instead
 def _default_sampling_ratio(configurations):
     default = 1.0
-    if SAMPLING_RATIO_ENV_VAR in environ:
+
+    if environ.get(OTEL_TRACES_SAMPLER_ARG) is not None:
+        try:
+            if float(environ[OTEL_TRACES_SAMPLER_ARG]) < 0:
+                _logger.error("Invalid value for OTEL_TRACES_SAMPLER_ARG. It should be a non-negative number.")
+        except ValueError:
+            pass
+    else:
+        _logger.error("OTEL_TRACES_SAMPLER_ARG is not set.")
+
+    # Check if rate-limited sampler is configured
+    if environ.get(OTEL_TRACES_SAMPLER) == RATE_LIMITED_SAMPLER:
+        try:
+            default = float(environ[OTEL_TRACES_SAMPLER_ARG])
+            print(f"Using rate limited sampler: {default} traces per second")
+        except ValueError as e:
+            _logger.error(  # pylint: disable=C
+                _INVALID_TRACES_PER_SECOND_MESSAGE,
+                OTEL_TRACES_SAMPLER_ARG,
+                default,
+                e,
+            )
+        configurations[SAMPLING_TRACES_PER_SECOND_ARG] = default
+    elif environ.get(OTEL_TRACES_SAMPLER) == FIXED_PERCENTAGE_SAMPLER:
         try:
-            default = float(environ[SAMPLING_RATIO_ENV_VAR])
+            default = float(environ[OTEL_TRACES_SAMPLER_ARG])
+            print(f"Using sampling ratio: {default}")
         except ValueError as e:
             _logger.error(  # pylint: disable=C
                 _INVALID_FLOAT_MESSAGE,
-                SAMPLING_RATIO_ENV_VAR,
+                OTEL_TRACES_SAMPLER_ARG,
                 default,
                 e,
             )
-    configurations[SAMPLING_RATIO_ARG] = default
-
+        configurations[SAMPLING_RATIO_ARG] = default
+    else:
+        # Default behavior - always set sampling_ratio
+        configurations[SAMPLING_RATIO_ARG] = default
+        _logger.error(  # pylint: disable=C
+            "Invalid argument for the sampler to be used for tracing. "
+            "Supported values are %s and %s. Defaulting to %s: %s",
+            RATE_LIMITED_SAMPLER,
+            FIXED_PERCENTAGE_SAMPLER,
+            OTEL_TRACES_SAMPLER,
+            OTEL_TRACES_SAMPLER_ARG,
+        )
 
 def _default_instrumentation_options(configurations):
     otel_disabled_instrumentations = _get_otel_disabled_instrumentations()

sdk/monitor/azure-monitor-opentelemetry/tests/utils/test_configurations.py

@@ -16,13 +16,20 @@
 from unittest import TestCase
 from unittest.mock import patch
 
-from opentelemetry.instrumentation.environment_variables import (
-    OTEL_PYTHON_DISABLED_INSTRUMENTATIONS,
+from opentelemetry.sdk.environment_variables import (
+    OTEL_EXPERIMENTAL_RESOURCE_DETECTORS,
+    OTEL_TRACES_SAMPLER_ARG,
+    OTEL_TRACES_SAMPLER
 )
+from opentelemetry.instrumentation.environment_variables import (
+    OTEL_PYTHON_DISABLED_INSTRUMENTATIONS,)
 from azure.monitor.opentelemetry._utils.configurations import (
-    SAMPLING_RATIO_ENV_VAR,
     _get_configurations,
 )
+from azure.monitor.opentelemetry._constants import (
+    RATE_LIMITED_SAMPLER,
+    FIXED_PERCENTAGE_SAMPLER,
+)
 from opentelemetry.environment_variables import (
     OTEL_LOGS_EXPORTER,
     OTEL_METRICS_EXPORTER,
@@ -134,7 +141,7 @@ def test_get_configurations_defaults(self, resource_create_mock):
         "os.environ",
         {
             OTEL_PYTHON_DISABLED_INSTRUMENTATIONS: "flask,requests,fastapi,azure_sdk",
-            SAMPLING_RATIO_ENV_VAR: "0.5",
+            OTEL_TRACES_SAMPLER_ARG: "0.5",
             OTEL_TRACES_EXPORTER: "None",
             OTEL_LOGS_EXPORTER: "none",
             OTEL_METRICS_EXPORTER: "NONE",
@@ -166,12 +173,13 @@ def test_get_configurations_env_vars(self, resource_create_mock):
         self.assertEqual(configurations["resource"].attributes, TEST_DEFAULT_RESOURCE.attributes)
         self.assertEqual(environ[OTEL_EXPERIMENTAL_RESOURCE_DETECTORS], "custom_resource_detector")
         resource_create_mock.assert_called_once_with()
-        self.assertEqual(configurations["sampling_ratio"], 0.5)
+        self.assertEqual(configurations["sampling_ratio"], 1.0)
 
     @patch.dict(
         "os.environ",
         {
-            SAMPLING_RATIO_ENV_VAR: "Half",
+            OTEL_TRACES_SAMPLER: FIXED_PERCENTAGE_SAMPLER,
+            OTEL_TRACES_SAMPLER_ARG: "Half",
             OTEL_TRACES_EXPORTER: "False",
             OTEL_LOGS_EXPORTER: "no",
             OTEL_METRICS_EXPORTER: "True",
@@ -181,7 +189,7 @@ def test_get_configurations_env_vars(self, resource_create_mock):
     @patch("opentelemetry.sdk.resources.Resource.create", return_value=TEST_DEFAULT_RESOURCE)
     def test_get_configurations_env_vars_validation(self, resource_create_mock):
         configurations = _get_configurations()
-
+        print(configurations)
         self.assertTrue("connection_string" not in configurations)
         self.assertEqual(configurations["disable_logging"], False)
         self.assertEqual(configurations["disable_metrics"], False)
@@ -260,3 +268,156 @@ def test_merge_instrumentation_options_extra_args(self, resource_create_mock):
                 "urllib3": {"enabled": True},
             },
         )
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_PYTHON_DISABLED_INSTRUMENTATIONS: "flask,requests,fastapi,azure_sdk",
+            OTEL_TRACES_SAMPLER: RATE_LIMITED_SAMPLER,
+            OTEL_TRACES_SAMPLER_ARG: "0.5",
+            OTEL_TRACES_EXPORTER: "None",
+            OTEL_LOGS_EXPORTER: "none",
+            OTEL_METRICS_EXPORTER: "NONE",
+            OTEL_EXPERIMENTAL_RESOURCE_DETECTORS: "custom_resource_detector",
+        },
+        clear=True,
+    )
+    @patch("opentelemetry.sdk.resources.Resource.create", return_value=TEST_DEFAULT_RESOURCE)
+    def test_get_configurations_env_vars_rate_limited(self, resource_create_mock):
+        configurations = _get_configurations()
+
+        self.assertTrue("connection_string" not in configurations)
+        self.assertEqual(configurations["disable_logging"], True)
+        self.assertEqual(configurations["disable_metrics"], True)
+        self.assertEqual(configurations["disable_tracing"], True)
+        self.assertEqual(
+            configurations["instrumentation_options"],
+            {
+                "azure_sdk": {"enabled": False},
+                "django": {"enabled": True},
+                "fastapi": {"enabled": False},
+                "flask": {"enabled": False},
+                "psycopg2": {"enabled": True},
+                "requests": {"enabled": False},
+                "urllib": {"enabled": True},
+                "urllib3": {"enabled": True},
+            },
+        )
+        self.assertEqual(configurations["resource"].attributes, TEST_DEFAULT_RESOURCE.attributes)
+        self.assertEqual(environ[OTEL_EXPERIMENTAL_RESOURCE_DETECTORS], "custom_resource_detector")
+        resource_create_mock.assert_called_once_with()
+        self.assertEqual(configurations["sampling_traces_per_second"], 0.5)
+    
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_PYTHON_DISABLED_INSTRUMENTATIONS: "flask,requests,fastapi,azure_sdk",
+            OTEL_TRACES_SAMPLER_ARG: "34",
+            OTEL_TRACES_EXPORTER: "None",
+            OTEL_LOGS_EXPORTER: "none",
+            OTEL_METRICS_EXPORTER: "NONE",
+            OTEL_EXPERIMENTAL_RESOURCE_DETECTORS: "custom_resource_detector",
+        },
+        clear=True,
+    )
+    @patch("opentelemetry.sdk.resources.Resource.create", return_value=TEST_DEFAULT_RESOURCE)
+    def test_get_configurations_env_vars_no_preference(self, resource_create_mock):
+        configurations = _get_configurations()
+
+        self.assertTrue("connection_string" not in configurations)
+        self.assertEqual(configurations["disable_logging"], True)
+        self.assertEqual(configurations["disable_metrics"], True)
+        self.assertEqual(configurations["disable_tracing"], True)
+        self.assertEqual(
+            configurations["instrumentation_options"],
+            {
+                "azure_sdk": {"enabled": False},
+                "django": {"enabled": True},
+                "fastapi": {"enabled": False},
+                "flask": {"enabled": False},
+                "psycopg2": {"enabled": True},
+                "requests": {"enabled": False},
+                "urllib": {"enabled": True},
+                "urllib3": {"enabled": True},
+            },
+        )
+        self.assertEqual(configurations["resource"].attributes, TEST_DEFAULT_RESOURCE.attributes)
+        self.assertEqual(environ[OTEL_EXPERIMENTAL_RESOURCE_DETECTORS], "custom_resource_detector")
+        resource_create_mock.assert_called_once_with()
+        self.assertEqual(configurations["sampling_ratio"], 1.0)
+    
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_PYTHON_DISABLED_INSTRUMENTATIONS: "flask,requests,fastapi,azure_sdk",
+            OTEL_TRACES_SAMPLER_ARG: "2 traces per second",
+            OTEL_TRACES_EXPORTER: "None",
+            OTEL_LOGS_EXPORTER: "none",
+            OTEL_METRICS_EXPORTER: "NONE",
+            OTEL_EXPERIMENTAL_RESOURCE_DETECTORS: "custom_resource_detector",
+        },
+        clear=True,
+    )
+    @patch("opentelemetry.sdk.resources.Resource.create", return_value=TEST_DEFAULT_RESOURCE)
+    def test_get_configurations_env_vars_check_default(self, resource_create_mock):
+        configurations = _get_configurations()
+
+        self.assertTrue("connection_string" not in configurations)
+        self.assertEqual(configurations["disable_logging"], True)
+        self.assertEqual(configurations["disable_metrics"], True)
+        self.assertEqual(configurations["disable_tracing"], True)
+        self.assertEqual(
+            configurations["instrumentation_options"],
+            {
+                "azure_sdk": {"enabled": False},
+                "django": {"enabled": True},
+                "fastapi": {"enabled": False},
+                "flask": {"enabled": False},
+                "psycopg2": {"enabled": True},
+                "requests": {"enabled": False},
+                "urllib": {"enabled": True},
+                "urllib3": {"enabled": True},
+            },
+        )
+        self.assertEqual(configurations["resource"].attributes, TEST_DEFAULT_RESOURCE.attributes)
+        self.assertEqual(environ[OTEL_EXPERIMENTAL_RESOURCE_DETECTORS], "custom_resource_detector")
+        resource_create_mock.assert_called_once_with()
+        self.assertEqual(configurations["sampling_ratio"], 1.0)
+    
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_PYTHON_DISABLED_INSTRUMENTATIONS: "flask,requests,fastapi,azure_sdk",
+            OTEL_TRACES_SAMPLER: FIXED_PERCENTAGE_SAMPLER,
+            OTEL_TRACES_SAMPLER_ARG: "0.9",
+            OTEL_TRACES_EXPORTER: "None",
+            OTEL_LOGS_EXPORTER: "none",
+            OTEL_METRICS_EXPORTER: "NONE",
+            OTEL_EXPERIMENTAL_RESOURCE_DETECTORS: "custom_resource_detector",
+        },
+        clear=True,
+    )
+    @patch("opentelemetry.sdk.resources.Resource.create", return_value=TEST_DEFAULT_RESOURCE)
+    def test_get_configurations_env_vars_fixed_percentage(self, resource_create_mock):
+        configurations = _get_configurations()
+
+        self.assertTrue("connection_string" not in configurations)
+        self.assertEqual(configurations["disable_logging"], True)
+        self.assertEqual(configurations["disable_metrics"], True)
+        self.assertEqual(configurations["disable_tracing"], True)
+        self.assertEqual(
+            configurations["instrumentation_options"],
+            {
+                "azure_sdk": {"enabled": False},
+                "django": {"enabled": True},
+                "fastapi": {"enabled": False},
+                "flask": {"enabled": False},
+                "psycopg2": {"enabled": True},
+                "requests": {"enabled": False},
+                "urllib": {"enabled": True},
+                "urllib3": {"enabled": True},
+            },
+        )
+        self.assertEqual(configurations["resource"].attributes, TEST_DEFAULT_RESOURCE.attributes)
+        self.assertEqual(environ[OTEL_EXPERIMENTAL_RESOURCE_DETECTORS], "custom_resource_detector")
+        resource_create_mock.assert_called_once_with()
+        self.assertEqual(configurations["sampling_ratio"], 0.9)