Distinguish Azure AI SDKs in statsbeats

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

@@ -194,6 +194,7 @@ class _RP_Names(Enum):
 # Special constant for azure-sdk opentelemetry instrumentation
 _AZURE_SDK_OPENTELEMETRY_NAME = "azure-sdk-opentelemetry"
 _AZURE_SDK_NAMESPACE_NAME = "az.namespace"
+_AZURE_AI_SDK_NAME = "azure-ai-opentelemetry"
 
 _BASE = 2
 
@@ -253,6 +254,7 @@ class _RP_Names(Enum):
     "openai_v2",
     "vertexai",
     # Instrumentations below this line have not been added to statsbeat report yet
+    _AZURE_AI_SDK_NAME
 ]
 
 _INSTRUMENTATIONS_BIT_MAP = {_INSTRUMENTATIONS_LIST[i]: _BASE**i for i in range(len(_INSTRUMENTATIONS_LIST))}

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/_quickpulse/_processor.py

@@ -9,12 +9,22 @@
 
 # pylint: disable=protected-access
 class _QuickpulseLogRecordProcessor(LogRecordProcessor):
+    def __init__(self):
+        super().__init__()
+        self.call_on_emit = hasattr(super(), 'on_emit')
 
-    def emit(self, log_data: LogData) -> None:  # type: ignore
+    def on_emit(self, log_data: LogData) -> None:  # type: ignore
         qpm = _QuickpulseManager._instance
         if qpm:
             qpm._record_log_record(log_data)
-        super().emit(log_data)  # type: ignore[safe-super]
+        if self.call_on_emit:
+            super().on_emit(log_data)  # type: ignore[safe-super]
+        else:
+            # this method was removed in opentelemetry-sdk and replaced with on_emit
+            super().emit(log_data)  # type: ignore[safe-super,misc] # pylint: disable=no-member
+
+    def emit(self, log_data: LogData) -> None:
+        self.on_emit(log_data)
 
     def shutdown(self):
         pass

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

@@ -124,6 +124,7 @@ def _getlocale():
             # by continuing to use getdefaultlocale() even though it has been deprecated.
             # we ignore the deprecation warnings to reduce noise
             warnings.simplefilter("ignore", category=DeprecationWarning)
+            # pylint: disable=deprecated-method
             return locale.getdefaultlocale()[0]
     except AttributeError:
         # locale.getlocal() has issues on Windows: https://github.com/python/cpython/issues/82986

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/export/trace/_exporter.py

@@ -23,6 +23,7 @@
     _APPLICATIONINSIGHTS_OPENTELEMETRY_RESOURCE_METRIC_DISABLED,
     _AZURE_SDK_NAMESPACE_NAME,
     _AZURE_SDK_OPENTELEMETRY_NAME,
+    _AZURE_AI_SDK_NAME,
     _INSTRUMENTATION_SUPPORTING_METRICS_LIST,
     _SAMPLE_RATE_KEY,
     _METRIC_ENVELOPE_NAME,
@@ -525,13 +526,18 @@ def _convert_span_events_to_envelopes(span: ReadableSpan) -> Sequence[TelemetryI
 
 
 def _check_instrumentation_span(span: ReadableSpan) -> None:
-    # Special use-case for spans generated from azure-sdk services
-    # Identified by having az.namespace as a span attribute
-    if span.attributes and _AZURE_SDK_NAMESPACE_NAME in span.attributes:
-        _utils.add_instrumentation(_AZURE_SDK_OPENTELEMETRY_NAME)
-        return
     if span.instrumentation_scope is None:
         return
+
+    # Special use-case for spans generated from azure-sdk services
+    # `azure-` or `azure.` is a prefix
+    if span.instrumentation_scope.name.startswith("azure"):
+        # spec-case for Azure AI SDKs - identified by `az.namespace` attribute
+        if span.attributes and span.attributes.get(_AZURE_SDK_NAMESPACE_NAME) == "Microsoft.CognitiveServices":
+            _utils.add_instrumentation(_AZURE_AI_SDK_NAME)
+        else:
+            _utils.add_instrumentation(_AZURE_SDK_OPENTELEMETRY_NAME)
+        return
     # All instrumentation scope names from OpenTelemetry instrumentations have
     # `opentelemetry.instrumentation.` as a prefix
     if span.instrumentation_scope.name.startswith("opentelemetry.instrumentation."):

sdk/monitor/azure-monitor-opentelemetry-exporter/dev_requirements.txt

@@ -1,4 +1,5 @@
 -e ../../../tools/azure-sdk-tools
 ../../core/azure-core
+../../core/azure-core-tracing-opentelemetry
 -e ../../identity/azure-identity
 aiohttp>=3.0; python_version >= '3.7'

sdk/monitor/azure-monitor-opentelemetry-exporter/tests/quickpulse/test_processor.py

@@ -23,7 +23,7 @@ def tearDownClass(cls) -> None:
     def test_emit(self):
         processor = _QuickpulseLogRecordProcessor()
         log_data = mock.Mock()
-        processor.emit(log_data)
+        processor.on_emit(log_data)
         self.qpm._record_log_record.assert_called_once_with(log_data)
 
 

sdk/monitor/azure-monitor-opentelemetry-exporter/tests/trace/test_trace.py

@@ -11,7 +11,11 @@
 # pylint: disable=import-error
 from opentelemetry.trace import get_tracer_provider, set_tracer_provider
 from opentelemetry.sdk import trace, resources
+from opentelemetry.sdk.trace import TracerProvider
+
 from opentelemetry.sdk.trace.export import SpanExportResult
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.util.instrumentation import InstrumentationScope
 from opentelemetry.semconv.attributes.exception_attributes import (
     EXCEPTION_ESCAPED,
@@ -22,6 +26,14 @@
 from opentelemetry.trace import Link, SpanContext, SpanKind
 from opentelemetry.trace.status import Status, StatusCode
 
+from azure.core.settings import settings
+from azure.core.tracing.ext.opentelemetry_span import OpenTelemetrySpan
+
+try:
+    from azure.core.instrumentation import get_tracer as get_azure_sdk_tracer
+except ImportError:
+    # azure.core.instrumentation is not available in older versions of azure-core
+    get_azure_sdk_tracer = None
 from azure.monitor.opentelemetry.exporter.export._base import ExportResult
 from azure.monitor.opentelemetry.exporter.export.trace._exporter import (
     AzureMonitorTraceExporter,
@@ -31,6 +43,7 @@
 from azure.monitor.opentelemetry.exporter._constants import (
     _AZURE_SDK_NAMESPACE_NAME,
     _AZURE_SDK_OPENTELEMETRY_NAME,
+    _AZURE_AI_SDK_NAME,
 )
 from azure.monitor.opentelemetry.exporter._generated.models import ContextTagKeys
 from azure.monitor.opentelemetry.exporter._utils import azure_monitor_context
@@ -1128,7 +1141,7 @@ def test_span_envelope_server_http(self):
             "url.path": "/path",
             "url.query": "query",
             "server.address": "www.example.org",
-            "server.port": "80"
+            "server.port": "80",
         }
         envelope = exporter._span_to_envelope(span)
         self.assertEqual(envelope.data.base_data.url, "https://www.example.org:80/path?query")
@@ -1687,8 +1700,89 @@ def test_check_instrumentation_span_not_instrumentation(self):
 
     def test_check_instrumentation_span_azure_sdk(self):
         span = mock.Mock()
-        span.attributes = {_AZURE_SDK_NAMESPACE_NAME: "Microsoft.EventHub"}
-        span.instrumentation_scope.name = "__main__"
+        span.attributes = {}
+        span.instrumentation_scope.name = "azure.foo.bar.__init__"
+
         with mock.patch("azure.monitor.opentelemetry.exporter._utils.add_instrumentation") as add:
             _check_instrumentation_span(span)
             add.assert_called_once_with(_AZURE_SDK_OPENTELEMETRY_NAME)
+
+    @mock.patch("opentelemetry.trace.get_tracer_provider")
+    def test_check_instrumentation_span_azure_sdk_otel_span(self, mock_get_tracer_provider):
+        mock_get_tracer_provider.return_value = self.get_tracer_provider()
+
+        with OpenTelemetrySpan() as azure_sdk_span:
+            with mock.patch("azure.monitor.opentelemetry.exporter._utils.add_instrumentation") as add:
+                _check_instrumentation_span(azure_sdk_span.span_instance)
+                add.assert_called_once_with(_AZURE_SDK_OPENTELEMETRY_NAME)
+
+            with mock.patch("azure.monitor.opentelemetry.exporter._utils.add_instrumentation") as add:
+                azure_sdk_span.add_attribute(_AZURE_SDK_NAMESPACE_NAME, "Microsoft.ServiceBus")
+                _check_instrumentation_span(azure_sdk_span.span_instance)
+                add.assert_called_once_with(_AZURE_SDK_OPENTELEMETRY_NAME)
+
+            with mock.patch("azure.monitor.opentelemetry.exporter._utils.add_instrumentation") as add:
+                azure_sdk_span.add_attribute(_AZURE_SDK_NAMESPACE_NAME, "Microsoft.CognitiveServices")
+                _check_instrumentation_span(azure_sdk_span.span_instance)
+                add.assert_called_once_with(_AZURE_AI_SDK_NAME)
+
+    @mock.patch("opentelemetry.trace.get_tracer_provider")
+    def test_check_instrumentation_span_azure_sdk_span_impl(self, mock_get_tracer_provider):
+        mock_get_tracer_provider.return_value = self.get_tracer_provider()
+
+        settings.tracing_implementation = "opentelemetry"
+        try:
+            azure_sdk_span = settings.tracing_implementation()(name="test")
+
+            with mock.patch("azure.monitor.opentelemetry.exporter._utils.add_instrumentation") as add:
+                _check_instrumentation_span(azure_sdk_span.span_instance)
+                add.assert_called_once_with(_AZURE_SDK_OPENTELEMETRY_NAME)
+
+            with mock.patch("azure.monitor.opentelemetry.exporter._utils.add_instrumentation") as add:
+                azure_sdk_span.add_attribute(_AZURE_SDK_NAMESPACE_NAME, "Microsoft.ServiceBus")
+                _check_instrumentation_span(azure_sdk_span.span_instance)
+                add.assert_called_once_with(_AZURE_SDK_OPENTELEMETRY_NAME)
+
+            with mock.patch("azure.monitor.opentelemetry.exporter._utils.add_instrumentation") as add:
+                azure_sdk_span.add_attribute(_AZURE_SDK_NAMESPACE_NAME, "Microsoft.CognitiveServices")
+                _check_instrumentation_span(azure_sdk_span.span_instance)
+                add.assert_called_once_with(_AZURE_AI_SDK_NAME)
+
+        finally:
+            settings.tracing_implementation = None
+
+    @mock.patch("opentelemetry.trace.get_tracer_provider")
+    def test_check_instrumentation_span_azure_sdk_get_tracer(self, mock_get_tracer_provider):
+        mock_get_tracer_provider.return_value = self.get_tracer_provider()
+
+        if not get_azure_sdk_tracer:
+            self.skipTest("azure.core.instrumentation is not available")
+
+        azure_sdk_tracer = get_azure_sdk_tracer(library_name="azure-foo-bar")
+        azure_sdk_span = azure_sdk_tracer.start_span(name="test")
+
+        with mock.patch("azure.monitor.opentelemetry.exporter._utils.add_instrumentation") as add:
+            _check_instrumentation_span(azure_sdk_span)
+            add.assert_called_once_with(_AZURE_SDK_OPENTELEMETRY_NAME)
+
+        with mock.patch("azure.monitor.opentelemetry.exporter._utils.add_instrumentation") as add:
+            azure_sdk_span.set_attribute(_AZURE_SDK_NAMESPACE_NAME, "Microsoft.ServiceBus")
+            _check_instrumentation_span(azure_sdk_span)
+            add.assert_called_once_with(_AZURE_SDK_OPENTELEMETRY_NAME)
+
+        with mock.patch("azure.monitor.opentelemetry.exporter._utils.add_instrumentation") as add:
+            azure_sdk_span.set_attribute(_AZURE_SDK_NAMESPACE_NAME, "Microsoft.CognitiveServices")
+            _check_instrumentation_span(azure_sdk_span)
+            add.assert_called_once_with(_AZURE_AI_SDK_NAME)
+
+        not_azure_sdk_span = get_azure_sdk_tracer(library_name="not-azure-foo-bar").start_span(name="test")
+        with mock.patch("azure.monitor.opentelemetry.exporter._utils.add_instrumentation") as add:
+            _check_instrumentation_span(not_azure_sdk_span)
+            add.assert_not_called()
+
+    def get_tracer_provider(self):
+        tracer_provider = TracerProvider()
+        span_exporter = InMemorySpanExporter()
+        processor = SimpleSpanProcessor(span_exporter)
+        tracer_provider.add_span_processor(processor)
+        return tracer_provider