Add support for OTEL_TRACES_SAMPLER (Azure#44535)

* Add support for OTEL_TRACES_SAMPLER

* Added CHANGELOG

* Added tests and modified README

* Remove extra space

* Fix formatting

* Added tests and modified README

* Fix README

* Fix black

* Add sample for different sampling configurations

* Fix cspell

Author: rads-1996

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

@@ -3,6 +3,8 @@
 ## 1.8.4 (Unreleased)
 
 ### Features Added
+- Added support for OTEL_TRACES_SAMPLER
+  ([#44535](https://github.com/Azure/azure-sdk-for-python/pull/44535))
 - Added ability to add additional Log Record Processors and Metric Readers via configure_azure_monitor
   ([#44367](https://github.com/Azure/azure-sdk-for-python/pull/44367))
 

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

@@ -82,8 +82,8 @@ You can configure further with [OpenTelemetry environment variables][ot_env_vars
 | `OTEL_TRACES_EXPORTER` | If set to `None`, disables collection and export of distributed tracing telemetry. |
 | `OTEL_BLRP_SCHEDULE_DELAY` | Specifies the logging export interval in milliseconds. Defaults to 5000. |
 | `OTEL_BSP_SCHEDULE_DELAY` | Specifies the distributed tracing export interval in milliseconds. Defaults to 5000. |
-| `OTEL_TRACES_SAMPLER` | Specifies the sampler to be used for traces. Supports both [application_insights_sampling] and [rate_limited_sampling]. Use `microsoft.fixed.percentage` for the Application Insights sampler or `microsoft.rate_limited` for the Rate Limited sampler. |
-| `OTEL_TRACES_SAMPLER_ARG` | Specifies the sampling parameter for the configured sampler. For the Application Insights sampler, this sets the ratio of distributed tracing telemetry to be [sampled][application_insights_sampling] with accepted values in the range [0,1]. Defaults to 1.0 (no sampling). For the Rate Limited sampler, this sets the maximum traces per second to be [sampled][rate_limited_sampler]. For example, 0.5 means one trace every two seconds, while 5.0 means five traces per second. |
+| `OTEL_TRACES_SAMPLER` | Specifies the sampler to be used for traces. Supports `always_on`, `always_off`, `trace_id_ratio`, `parentbased_always_on`, `parentbased_always_off`, `parentbased_trace_id_ratio`, [application_insights_sampling] and [rate_limited_sampling]. Use `microsoft.fixed.percentage` for the Application Insights sampler or `microsoft.rate_limited` for the Rate Limited sampler. |
+| `OTEL_TRACES_SAMPLER_ARG` | Specifies the sampling parameter for the configured sampler. For the standard OpenTelemetry samplers `trace_id_ratio` and `parentbased_trace_id_ratio`, this is the sampling ratio in the range [0.0, 1.0]. Not needed to be specified for `always_on`, `always_off`, `parentbased_always_on`, or `parentbased_always_off` samplers. For the Application Insights sampler, this sets the ratio of distributed tracing telemetry to be [sampled][application_insights_sampling] with accepted values in the range [0,1]. Defaults to 1.0 (no sampling). For the Rate Limited sampler, this sets the maximum traces per second to be [sampled][rate_limited_sampler]. For example, 0.5 means one trace every two seconds, while 5.0 means five traces per second. |
 | `OTEL_PYTHON_DISABLED_INSTRUMENTATIONS` | Specifies which of the supported instrumentations to disable. Disabled instrumentations will not be instrumented as part of `configure_azure_monitor`. However, they can still be manually instrumented with `instrument()` directly. Accepts a comma-separated list of lowercase [Library Names](#officially-supported-instrumentations). For example, set to `"psycopg2,fastapi"` to disable the Psycopg2 and FastAPI instrumentations. Defaults to an empty list, enabling all supported instrumentations. |
 | `OTEL_EXPERIMENTAL_RESOURCE_DETECTORS` | An experimental OpenTelemetry environment variable used to specify Resource Detectors to be used to generate Resource Attributes. This is an experimental feature and the name of this variable and its behavior can change in a non-backwards compatible way. Defaults to "azure_app_service,azure_vm" to enable the [Azure Resource Detectors][ot_resource_detector_azure] for Azure App Service and Azure VM. To add or remove specific resource detectors, set the environment variable accordingly. See the [OpenTelemetry Python Resource Detector Documentation][ot_python_resource_detectors] for more. |
 

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

@@ -42,6 +42,8 @@
     METRIC_READERS_ARG,
     VIEWS_ARG,
     ENABLE_TRACE_BASED_SAMPLING_ARG,
+    SAMPLING_ARG,
+    SAMPLER_TYPE,
 )
 from azure.monitor.opentelemetry._types import ConfigurationValue
 from azure.monitor.opentelemetry.exporter._quickpulse import (  # pylint: disable=import-error,no-name-in-module
@@ -75,6 +77,7 @@
 from azure.monitor.opentelemetry._utils.configurations import (
     _get_configurations,
     _is_instrumentation_enabled,
+    _get_sampler_from_name,
 )
 from azure.monitor.opentelemetry._utils.instrumentation import (
     get_dist_dependency_conflicts,
@@ -155,7 +158,12 @@ def configure_azure_monitor(**kwargs) -> None:  # pylint: disable=C4758
 def _setup_tracing(configurations: Dict[str, ConfigurationValue]):
     resource: Resource = configurations[RESOURCE_ARG]  # type: ignore
     enable_performance_counters_config = configurations[ENABLE_PERFORMANCE_COUNTERS_ARG]
-    if SAMPLING_TRACES_PER_SECOND_ARG in configurations:
+    if SAMPLING_ARG in configurations:
+        sampler_arg = configurations[SAMPLING_ARG]
+        sampler_type = configurations[SAMPLER_TYPE]
+        sampler = _get_sampler_from_name(sampler_type, sampler_arg)
+        tracer_provider = TracerProvider(sampler=sampler, resource=resource)
+    elif SAMPLING_TRACES_PER_SECOND_ARG in configurations:
         traces_per_second = configurations[SAMPLING_TRACES_PER_SECOND_ARG]
         tracer_provider = TracerProvider(
             sampler=RateLimitedSampler(target_spans_per_second_limit=cast(float, traces_per_second)), resource=resource

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

@@ -31,6 +31,24 @@
 FIXED_PERCENTAGE_SAMPLER = "microsoft.fixed.percentage"
 SAMPLING_TRACES_PER_SECOND_ARG = "traces_per_second"
 ENABLE_TRACE_BASED_SAMPLING_ARG = "enable_trace_based_sampling_for_logs"
+SAMPLER_TYPE = "sampler_type"
+SAMPLING_ARG = "sampling_arg"
+ALWAYS_ON_SAMPLER = "always_on"
+ALWAYS_OFF_SAMPLER = "always_off"
+TRACE_ID_RATIO_SAMPLER = "trace_id_ratio"
+PARENT_BASED_ALWAYS_ON_SAMPLER = "parentbased_always_on"
+PARENT_BASED_ALWAYS_OFF_SAMPLER = "parentbased_always_off"
+PARENT_BASED_TRACE_ID_RATIO_SAMPLER = "parentbased_trace_id_ratio"
+SUPPORTED_OTEL_SAMPLERS = (
+    RATE_LIMITED_SAMPLER,
+    FIXED_PERCENTAGE_SAMPLER,
+    ALWAYS_ON_SAMPLER,
+    ALWAYS_OFF_SAMPLER,
+    TRACE_ID_RATIO_SAMPLER,
+    PARENT_BASED_ALWAYS_ON_SAMPLER,
+    PARENT_BASED_ALWAYS_OFF_SAMPLER,
+    PARENT_BASED_TRACE_ID_RATIO_SAMPLER,
+)
 
 # --------------------Autoinstrumentation Configuration------------------------------------------
 

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

@@ -23,6 +23,13 @@
 )
 from opentelemetry.sdk.resources import Resource
 
+from opentelemetry.sdk.trace.sampling import (
+    TraceIdRatioBased,
+    ALWAYS_OFF,
+    ALWAYS_ON,
+    ParentBased,
+)
+
 from azure.monitor.opentelemetry._constants import (
     _AZURE_APP_SERVICE_RESOURCE_DETECTOR_NAME,
     _AZURE_VM_RESOURCE_DETECTOR_NAME,
@@ -49,6 +56,15 @@
     RATE_LIMITED_SAMPLER,
     FIXED_PERCENTAGE_SAMPLER,
     ENABLE_TRACE_BASED_SAMPLING_ARG,
+    SUPPORTED_OTEL_SAMPLERS,
+    ALWAYS_OFF_SAMPLER,
+    ALWAYS_ON_SAMPLER,
+    TRACE_ID_RATIO_SAMPLER,
+    PARENT_BASED_ALWAYS_ON_SAMPLER,
+    PARENT_BASED_ALWAYS_OFF_SAMPLER,
+    PARENT_BASED_TRACE_ID_RATIO_SAMPLER,
+    SAMPLING_ARG,
+    SAMPLER_TYPE,
 )
 from azure.monitor.opentelemetry._types import ConfigurationValue
 from azure.monitor.opentelemetry._version import VERSION
@@ -149,6 +165,7 @@ def _default_resource(configurations):
         configurations[RESOURCE_ARG] = Resource.create(configurations[RESOURCE_ARG].attributes)
 
 
+# pylint: disable=too-many-statements,too-many-branches
 def _default_sampling_ratio(configurations):
     default_value = 1.0
     sampler_type = environ.get(OTEL_TRACES_SAMPLER)
@@ -158,7 +175,7 @@ def _default_sampling_ratio(configurations):
     if sampler_type == RATE_LIMITED_SAMPLER:
         try:
             sampler_value = float(sampler_arg)
-            if sampler_value < 0:
+            if sampler_value < 0.0:
                 _logger.error("Invalid value for OTEL_TRACES_SAMPLER_ARG. It should be a non-negative number.")
                 sampler_value = default_value
             else:
@@ -177,7 +194,7 @@ def _default_sampling_ratio(configurations):
     elif sampler_type == FIXED_PERCENTAGE_SAMPLER:
         try:
             sampler_value = float(sampler_arg)
-            if sampler_value < 0:
+            if sampler_value < 0.0:
                 _logger.error("Invalid value for OTEL_TRACES_SAMPLER_ARG. It should be a non-negative number.")
                 sampler_value = default_value
             else:
@@ -192,16 +209,75 @@ def _default_sampling_ratio(configurations):
             )
             configurations[SAMPLING_RATIO_ARG] = default_value
 
+    # Handle always_on sampler
+    elif sampler_type == ALWAYS_ON_SAMPLER:
+        configurations[SAMPLING_ARG] = 1.0
+        configurations[SAMPLER_TYPE] = ALWAYS_ON_SAMPLER
+
+    # Handle always_off sampler
+    elif sampler_type == ALWAYS_OFF_SAMPLER:
+        configurations[SAMPLING_ARG] = 0.0
+        configurations[SAMPLER_TYPE] = ALWAYS_OFF_SAMPLER
+
+    # Handle trace_id_ratio sampler
+    elif sampler_type == TRACE_ID_RATIO_SAMPLER:
+        try:
+            sampler_value = float(sampler_arg) if sampler_arg is not None else default_value
+            if sampler_value < 0.0 or sampler_value > 1.0:
+                _logger.error("Invalid value for OTEL_TRACES_SAMPLER_ARG. It should be a value between 0 and 1.")
+                sampler_value = default_value
+            else:
+                _logger.info("Using sampling value: %s", sampler_value)
+            configurations[SAMPLING_ARG] = sampler_value
+        except ValueError as e:
+            _logger.error(  # pylint: disable=C
+                _INVALID_FLOAT_MESSAGE,
+                OTEL_TRACES_SAMPLER_ARG,
+                default_value,
+                e,
+            )
+            configurations[SAMPLING_ARG] = default_value
+        configurations[SAMPLER_TYPE] = TRACE_ID_RATIO_SAMPLER
+
+    # Handle parentbased_always_on sampler
+    elif sampler_type == PARENT_BASED_ALWAYS_ON_SAMPLER:
+        configurations[SAMPLING_ARG] = 1.0
+        configurations[SAMPLER_TYPE] = PARENT_BASED_ALWAYS_ON_SAMPLER
+
+    # Handle parentbased_always_off sampler
+    elif sampler_type == PARENT_BASED_ALWAYS_OFF_SAMPLER:
+        configurations[SAMPLING_ARG] = 0.0
+        configurations[SAMPLER_TYPE] = PARENT_BASED_ALWAYS_OFF_SAMPLER
+
+    # Handle parentbased_trace_id_ratio sampler
+    elif sampler_type == PARENT_BASED_TRACE_ID_RATIO_SAMPLER:
+        try:
+            sampler_value = float(sampler_arg) if sampler_arg is not None else default_value
+            if sampler_value < 0.0 or sampler_value > 1.0:
+                _logger.error("Invalid value for OTEL_TRACES_SAMPLER_ARG. It should be a value between 0 and 1.")
+                sampler_value = default_value
+            else:
+                _logger.info("Using sampling value: %s", sampler_value)
+            configurations[SAMPLING_ARG] = sampler_value
+        except ValueError as e:
+            _logger.error(  # pylint: disable=C
+                _INVALID_FLOAT_MESSAGE,
+                OTEL_TRACES_SAMPLER_ARG,
+                default_value,
+                e,
+            )
+            configurations[SAMPLING_ARG] = default_value
+        configurations[SAMPLER_TYPE] = PARENT_BASED_TRACE_ID_RATIO_SAMPLER
+
     # Handle all other cases (no sampler type specified or unsupported sampler type)
     else:
         if configurations.get(SAMPLING_RATIO_ARG) is None:
             configurations[SAMPLING_RATIO_ARG] = default_value
         if sampler_type is not None:
             _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,
+                "Supported values are %s. Defaulting to %s: %s",
+                SUPPORTED_OTEL_SAMPLERS,
                 FIXED_PERCENTAGE_SAMPLER,
                 configurations[SAMPLING_RATIO_ARG],
             )
@@ -271,3 +347,19 @@ def _is_instrumentation_enabled(configurations, lib_name):
 
 def _default_enable_trace_based_sampling(configurations):
     configurations.setdefault(ENABLE_TRACE_BASED_SAMPLING_ARG, False)
+
+
+def _get_sampler_from_name(sampler_type, sampler_arg):
+    if sampler_type == ALWAYS_ON_SAMPLER:
+        return ALWAYS_ON
+    if sampler_type == ALWAYS_OFF_SAMPLER:
+        return ALWAYS_OFF
+    if sampler_type == TRACE_ID_RATIO_SAMPLER:
+        ratio = float(sampler_arg) if sampler_arg is not None else 1.0
+        return TraceIdRatioBased(ratio)
+    if sampler_type == PARENT_BASED_ALWAYS_OFF_SAMPLER:
+        return ParentBased(ALWAYS_OFF)
+    if sampler_type == PARENT_BASED_TRACE_ID_RATIO_SAMPLER:
+        ratio = float(sampler_arg) if sampler_arg is not None else 1.0
+        return ParentBased(TraceIdRatioBased(ratio))
+    return ParentBased(ALWAYS_ON)

sdk/monitor/azure-monitor-opentelemetry/samples/tracing/sampling_configurations.py

@@ -0,0 +1,58 @@
+from azure.monitor.opentelemetry import configure_azure_monitor
+from opentelemetry import trace
+
+# Using always_on sampler
+# Set the OTEL_TRACES_SAMPLER environment variable to "always_on"
+# The sampling rate is 1.0, so 100% of the traces are sampled.
+
+# Using always_off sampler
+# Set the OTEL_TRACES_SAMPLER environment variable to "always_off"
+# The sampling rate is 0.0, so None of the traces are sampled.
+
+# Using trace_id_ratio sampler
+# Set the OTEL_TRACES_SAMPLER environment variable to "trace_id_ratio"
+# Set the OTEL_TRACES_SAMPLER_ARG environment variable to 0.1, it has to be a number between 0 and 1, else it will throw an error and default to 1.0
+# The sampling rate is 0.1 means approximately 10% of your traces are sent
+
+# Using parentbased_always_on sampler
+# Set the OTEL_TRACES_SAMPLER environment variable to "parentbased_always_on"
+# The sampling rate is 1.0, so 100% of the traces are sampled.
+
+# Using parentbased_always_off sampler
+# Set the OTEL_TRACES_SAMPLER environment variable to "parentbased_always_off"
+# The sampling rate is 0.0, so None of the traces are sampled.
+
+# Using parentbased_trace_id_ratio sampler
+# Set the OTEL_TRACES_SAMPLER environment variable to "parentbased_trace_id_ratio"
+# Set the OTEL_TRACES_SAMPLER_ARG environment variable to 0.45, it has to be a number between 0 and 1, else it will throw an error and default to 1.0
+# The sampling rate is 0.45 means approximately 45% of your traces are sent
+
+# Using rate limited sampler
+# Set the OTEL_TRACES_SAMPLER environment variable to "microsoft.rate_limited"
+# Set the OTEL_TRACES_SAMPLER_ARG environment variable to the desired rate limit (e.g., 0.5 means one trace every two seconds, while 5.0 means five traces per second)
+
+# Using fixed percentage sampler
+# Set the OTEL_TRACES_SAMPLER environment variable to "microsoft.fixed.percentage"
+# Set the OTEL_TRACES_SAMPLER_ARG environment variable to 0.2, it has to be a number between 0 and 1, else it will throw an error and default to 1.0
+
+# Using trace_based_sampling configuration # cspell: ignore unsampled
+# Determines whether the logger should drop log records associated with unsampled traces.
+# Passing the enable_trace_based_sampling_for_logs=True argument to configure_azure_monitor ensure that log records associated with unsampled traces are dropped by the `Logger`.
+# A log record is considered associated with an unsampled trace if it has a valid `SpanId` and its `TraceFlags` indicate that the trace is unsampled.
+# The value of this config is False by default
+
+"""
+    configure_azure_monitor (
+        "enable_trace_based_sampling_for_logs": True,
+    ) 
+"""
+
+configure_azure_monitor()
+
+tracer = trace.get_tracer(__name__)
+
+for i in range(100):
+    with tracer.start_as_current_span("hello"):
+        print("Hello, World!")
+
+input()