[ml] Add operation id tracking to client-side logging (#32923)

* Add operation id tracking

* import instrumentation key

* Remove logging format

* Set upper bound for opencensus-ext-logging

* Remove azureml-telemetry dependency

* Set upperbound for opencensus and use .get for safe envelope tag indexing

* Update to use f string instead of format

* Reduce duplicate declaration of instrumentation key

Author: diondrapeck

Diff for: sdk/ai/azure-ai-generative/azure/ai/generative/_telemetry/logging_handler.py

@@ -21,9 +21,7 @@
 )
 from opencensus.trace import config_integration
 
-from azureml.telemetry import INSTRUMENTATION_KEY
-
-from azure.ai.ml._telemetry.logging_handler import in_jupyter_notebook, CustomDimensionsFilter
+from azure.ai.ml._telemetry.logging_handler import in_jupyter_notebook, CustomDimensionsFilter, INSTRUMENTATION_KEY
 
 from azure.ai.generative._user_agent import USER_AGENT
 
@@ -152,8 +150,8 @@ def log_record_to_envelope(self, record):
             "process": record.processName,
             "module": record.module,
             "level": record.levelname,
-            "operation_id": envelope.tags["gen.ai.operation.id"],
-            "operation_parent_id": envelope.tags["gen.ai.operation.parentId"],
+            "operation_id": envelope.tags.get("ai.generative.operation.id"),
+            "operation_parent_id": envelope.tags.get("ai.generative.operation.parentId"),
         }
         if hasattr(record, "custom_dimensions") and isinstance(record.custom_dimensions, dict):
             properties.update(record.custom_dimensions)
@@ -217,14 +215,11 @@ def create_envelope(instrumentation_key, record):
         tags=dict(utils.azure_monitor_context),
         time=utils.timestamp_to_iso_str(record.created),
     )
-    envelope.tags["gen.ai.operation.id"] = getattr(
+    envelope.tags["ai.generative.operation.id"] = getattr(
         record,
         "traceId",
         "00000000000000000000000000000000",
     )
-    envelope.tags["gen.ai.operation.parentId"] = "|{}.{}.".format(
-        envelope.tags["gen.ai.operation.id"],
-        getattr(record, "spanId", "0000000000000000"),
-    )
+    envelope.tags["ai.generative.operation.parentId"] = f"|{envelope.tags.get('ai.generative.operation.id')}.{getattr(record, 'spanId', '0000000000000000')}"
 
     return envelope

Diff for: sdk/ai/azure-ai-generative/setup.py

@@ -74,10 +74,9 @@
     install_requires=[
         # NOTE: To avoid breaking changes in a major version bump, all dependencies should pin an upper bound if possible.
         "azure-ai-resources<2.0.0,>=1.0.0b1",
-        "azureml-telemetry~=1.0,>=1.51.0",
         "mlflow-skinny<3",
         "opencensus-ext-azure~=1.0",
-        "opencensus-ext-logging",
+        "opencensus-ext-logging<=0.1.1",
     ],
     extras_require={
         "evaluate": [

Diff for: sdk/ai/azure-ai-resources/azure/ai/resources/_telemetry/logging_handler.py

@@ -21,9 +21,7 @@
 )
 from opencensus.trace import config_integration
 
-from azureml.telemetry import INSTRUMENTATION_KEY
-
-from azure.ai.ml._telemetry.logging_handler import in_jupyter_notebook, CustomDimensionsFilter
+from azure.ai.ml._telemetry.logging_handler import in_jupyter_notebook, CustomDimensionsFilter, INSTRUMENTATION_KEY
 from azure.ai.resources._version import VERSION
 
 
@@ -157,8 +155,8 @@ def log_record_to_envelope(self, record):
             "process": record.processName,
             "module": record.module,
             "level": record.levelname,
-            "operation_id": envelope.tags["gen.ai.operation.id"],
-            "operation_parent_id": envelope.tags["gen.ai.operation.parentId"],
+            "operation_id": envelope.tags.get("ai.resources.operation.id"),
+            "operation_parent_id": envelope.tags.get("ai.resources.operation.parentId"),
         }
         if hasattr(record, "custom_dimensions") and isinstance(record.custom_dimensions, dict):
             properties.update(record.custom_dimensions)
@@ -222,14 +220,11 @@ def create_envelope(instrumentation_key, record):
         tags=dict(utils.azure_monitor_context),
         time=utils.timestamp_to_iso_str(record.created),
     )
-    envelope.tags["gen.ai.operation.id"] = getattr(
+    envelope.tags["ai.resources.operation.id"] = getattr(
         record,
         "traceId",
         "00000000000000000000000000000000",
     )
-    envelope.tags["gen.ai.operation.parentId"] = "|{}.{}.".format(
-        envelope.tags["gen.ai.operation.id"],
-        getattr(record, "spanId", "0000000000000000"),
-    )
+    envelope.tags["ai.resources.operation.parentId"] = f"|{envelope.tags.get('ai.resources.operation.id')}.{getattr(record, 'spanId', '0000000000000000')}"
 
     return envelope

Diff for: sdk/ai/azure-ai-resources/setup.py

@@ -68,7 +68,7 @@
         # NOTE: To avoid breaking changes in a major version bump, all dependencies should pin an upper bound if possible.
         "azure-ai-ml>=1.12.0",
         "mlflow-skinny<3",
-        "azureml-telemetry~=1.0,>=1.51.0",
+        "opencensus-ext-logging<=0.1.1",
         "azure-mgmt-resource<23.0.0,>=22.0.0",
     ],
     project_urls={

Diff for: sdk/ml/azure-ai-ml/azure/ai/ml/_telemetry/logging_handler.py

@@ -10,20 +10,14 @@
 import platform
 import traceback
 
-from opencensus.ext.azure.log_exporter import AzureLogHandler
 from opencensus.ext.azure.common import utils
-from opencensus.ext.azure.common.protocol import (
-    Data,
-    ExceptionData,
-    Message,
-    Envelope,
-)
-from azure.ai.ml._user_agent import USER_AGENT
+from opencensus.ext.azure.common.protocol import Data, Envelope, ExceptionData, Message
+from opencensus.ext.azure.log_exporter import AzureLogHandler
+from opencensus.trace import config_integration
 
+from azure.ai.ml._user_agent import USER_AGENT
 
 AML_INTERNAL_LOGGER_NAMESPACE = "azure.ai.ml._telemetry"
-
-# vienna-sdk-unitedstates
 INSTRUMENTATION_KEY = "71b954a8-6b7d-43f5-986c-3d3a6605d803"
 
 test_subscriptions = [
@@ -37,6 +31,10 @@
 ]
 
 
+# activate operation id tracking
+config_integration.trace_integrations(["logging"])
+
+
 class CustomDimensionsFilter(logging.Filter):
     """Add application-wide properties to AzureLogHandler records"""
 
@@ -181,7 +179,10 @@ def log_record_to_envelope(self, record):
             "process": record.processName,
             "module": record.module,
             "level": record.levelname,
+            "operation_id": envelope.tags.get("ai.ml.operation.id"),
+            "operation_parent_id": envelope.tags.get("ai.ml.operation.parentId"),
         }
+
         if hasattr(record, "custom_dimensions") and isinstance(record.custom_dimensions, dict):
             properties.update(record.custom_dimensions)
 
@@ -244,14 +245,13 @@ def create_envelope(instrumentation_key, record):
         tags=dict(utils.azure_monitor_context),
         time=utils.timestamp_to_iso_str(record.created),
     )
-    envelope.tags["ai.operation.id"] = getattr(
+    envelope.tags["ai.ml.operation.id"] = getattr(
         record,
         "traceId",
         "00000000000000000000000000000000",
     )
-    envelope.tags["ai.operation.parentId"] = "|{}.{}.".format(
-        envelope.tags["ai.operation.id"],
-        getattr(record, "spanId", "0000000000000000"),
-    )
+    envelope.tags[
+        "ai.ml.operation.parentId"
+    ] = f"|{envelope.tags.get('ai.ml.operation.id')}.{getattr(record, 'spanId', '0000000000000000')}"
 
     return envelope

Diff for: sdk/ml/azure-ai-ml/setup.py

@@ -86,6 +86,7 @@
         "azure-common<2.0.0,>=1.1",
         "typing-extensions<5.0.0",
         "opencensus-ext-azure<2.0.0",
+        "opencensus-ext-logging<=0.1.1",
     ],
     extras_require={
         # user can run `pip install azure-ai-ml[designer]` to install mldesigner alone with this package