[ml] Fix missing parent ids in trace logs (#34405)

* Revert change to acs_connection_id type

* Implement working POC for trace logging

* Refactor to take ops_logger as argument

* Remove debug statements and update tests

* Add check for OpsLogger type

* Add check for OpsLogger type in monitor_with_telemetry_mixing

* Update tests

* Address mypy errors

* Address mypy errors

* Suppress mypy error

* Refactor log_record_to_telemetry and save trace_id, span_id, and activity_id

* Remove name from span

* Remove whitespace to eliminate pylint error

Author: diondrapeck

sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py

@@ -10,7 +10,7 @@
 from functools import singledispatch
 from itertools import product
 from pathlib import Path
-from typing import Any, Optional, Tuple, TypeVar, Union
+from typing import Any, Dict, Optional, Tuple, TypeVar, Union
 
 from azure.ai.ml._azure_environments import (
     CloudArgumentKeys,
@@ -269,12 +269,12 @@ def __init__(
 
         user_agent = kwargs.get("user_agent", None)
 
-        app_insights_handler = get_appinsights_log_handler(
+        app_insights_handler: Tuple = get_appinsights_log_handler(
             user_agent,
             **{"properties": properties},
             enable_telemetry=self._operation_config.enable_telemetry,
         )
-        app_insights_handler_kwargs = {"app_insights_handler": app_insights_handler}
+        app_insights_handler_kwargs: Dict[str, Tuple] = {"app_insights_handler": app_insights_handler}
 
         base_url = _get_base_url_from_metadata(cloud_name=cloud_name, is_local_mfe=True)
         self._base_url = base_url
@@ -463,7 +463,7 @@ def __init__(
             self._service_client_10_2022_preview,
             self._operation_container,
             self._credential,
-            **app_insights_handler_kwargs,
+            **app_insights_handler_kwargs,  # type: ignore[arg-type]
         )
         self._operation_container.add(AzureMLResourceType.REGISTRY, self._registries)  # type: ignore[arg-type]
 
@@ -484,15 +484,15 @@ def __init__(
             self._operation_scope,
             self._operation_config,
             self._service_client_08_2023_preview,
-            **app_insights_handler_kwargs,
+            **app_insights_handler_kwargs,  # type: ignore[arg-type]
         )
         self._operation_container.add(AzureMLResourceType.COMPUTE, self._compute)
         self._datastores = DatastoreOperations(
             operation_scope=self._operation_scope,
             operation_config=self._operation_config,
             serviceclient_2023_04_01_preview=self._service_client_04_2023_preview,
             serviceclient_2024_01_01_preview=self._service_client_01_2024_preview,
-            **ops_kwargs,
+            **ops_kwargs,  # type: ignore[arg-type]
         )
         self._operation_container.add(AzureMLResourceType.DATASTORE, self._datastores)
         self._models = ModelOperations(
@@ -510,23 +510,23 @@ def __init__(
             workspace_rg=self._ws_rg,
             workspace_sub=self._ws_sub,
             registry_reference=registry_reference,
-            **app_insights_handler_kwargs,
+            **app_insights_handler_kwargs,  # type: ignore[arg-type]
         )
         self._operation_container.add(AzureMLResourceType.MODEL, self._models)
         self._code = CodeOperations(
             self._ws_operation_scope if registry_reference else self._operation_scope,
             self._operation_config,
             self._service_client_10_2021_dataplanepreview if registry_name else self._service_client_04_2023,
             self._datastores,
-            **ops_kwargs,
+            **ops_kwargs,  # type: ignore[arg-type]
         )
         self._operation_container.add(AzureMLResourceType.CODE, self._code)
         self._environments = EnvironmentOperations(
             self._ws_operation_scope if registry_reference else self._operation_scope,
             self._operation_config,
             self._service_client_10_2021_dataplanepreview if registry_name else self._service_client_04_2023_preview,
             self._operation_container,
-            **ops_kwargs,
+            **ops_kwargs,  # type: ignore[arg-type]
         )
         self._operation_container.add(AzureMLResourceType.ENVIRONMENT, self._environments)
         self._local_endpoint_helper = _LocalEndpointHelper(requests_pipeline=self._requests_pipeline)
@@ -539,7 +539,7 @@ def __init__(
             self._local_endpoint_helper,
             self._credential,
             requests_pipeline=self._requests_pipeline,
-            **ops_kwargs,
+            **ops_kwargs,  # type: ignore[arg-type]
         )
         self._batch_endpoints = BatchEndpointOperations(
             self._operation_scope,
@@ -549,7 +549,7 @@ def __init__(
             self._credential,
             requests_pipeline=self._requests_pipeline,
             service_client_09_2020_dataplanepreview=self._service_client_09_2020_dataplanepreview,
-            **ops_kwargs,
+            **ops_kwargs,  # type: ignore[arg-type]
         )
         self._operation_container.add(AzureMLResourceType.BATCH_ENDPOINT, self._batch_endpoints)
         self._operation_container.add(AzureMLResourceType.ONLINE_ENDPOINT, self._online_endpoints)
@@ -560,7 +560,7 @@ def __init__(
             self._operation_container,
             self._local_deployment_helper,
             self._credential,
-            **ops_kwargs,
+            **ops_kwargs,  # type: ignore[arg-type]
         )
         self._batch_deployments = BatchDeploymentOperations(
             self._operation_scope,
@@ -592,7 +592,7 @@ def __init__(
             self._service_client_10_2021_dataplanepreview if registry_name else self._service_client_01_2024_preview,
             self._operation_container,
             self._preflight,
-            **ops_kwargs,
+            **ops_kwargs,  # type: ignore[arg-type]
         )
         self._operation_container.add(AzureMLResourceType.COMPONENT, self._components)
         self._jobs = JobOperations(
@@ -626,7 +626,7 @@ def __init__(
                 self._operation_scope,
                 self._credential,
                 _service_client_kwargs=kwargs,
-                **ops_kwargs,
+                **ops_kwargs,  # type: ignore[arg-type]
             )
             self._operation_container.add(
                 AzureMLResourceType.VIRTUALCLUSTER, self._virtual_clusters  # type: ignore[arg-type]
@@ -639,7 +639,7 @@ def __init__(
             self._service_client_08_2023_preview,
             self._operation_container,
             self._credential,
-            **app_insights_handler_kwargs,
+            **app_insights_handler_kwargs,  # type: ignore[arg-type]
         )
 
         self._featuresets = FeatureSetOperations(
@@ -648,22 +648,22 @@ def __init__(
             self._service_client_10_2023,
             self._service_client_08_2023_preview,
             self._datastores,
-            **ops_kwargs,
+            **ops_kwargs,  # type: ignore[arg-type]
         )
 
         self._featurestoreentities = FeatureStoreEntityOperations(
             self._operation_scope,
             self._operation_config,
             self._service_client_10_2023,
-            **ops_kwargs,
+            **ops_kwargs,  # type: ignore[arg-type]
         )
 
         self._workspace_hubs = WorkspaceHubOperations(
             self._operation_scope,
             self._service_client_08_2023_preview,
             self._operation_container,
             self._credential,
-            **app_insights_handler_kwargs,
+            **app_insights_handler_kwargs,  # type: ignore[arg-type]
         )
         self._operation_container.add(AzureMLResourceType.WORKSPACE_HUB, self._workspace_hubs)  # type: ignore[arg-type]
 

sdk/ml/azure-ai-ml/azure/ai/ml/_telemetry/activity.py

@@ -23,6 +23,7 @@
 
 from marshmallow import ValidationError
 
+from azure.ai.ml._utils._logger_utils import OpsLogger
 from azure.ai.ml._utils.utils import _is_user_error_from_exception_type, _is_user_error_from_status_code, _str_to_bool
 from azure.ai.ml.exceptions import ErrorCategory, MlException
 from azure.core.exceptions import HttpResponseError
@@ -260,8 +261,8 @@ def monitor_with_activity(
     To monitor, use the ``@monitor_with_activity`` decorator. As an alternative, you can also wrap the
     logical block of code with the ``log_activity()`` method.
 
-    :param logger: The logger adapter.
-    :type logger: logging.LoggerAdapter
+    :param logger: The operations logging class, containing loggers and tracer for the package and module
+    :type logger: ~azure.ai.ml._utils._logger_utils.OpsLogger
     :param activity_name: The name of the activity. The name should be unique per the wrapped logical code block.
     :type activity_name: str
     :param activity_type: One of PUBLICAPI, INTERNALCALL, or CLIENTPROXY which represent an incoming API call,
@@ -275,8 +276,16 @@ def monitor_with_activity(
     def monitor(f):
         @functools.wraps(f)
         def wrapper(*args, **kwargs):
-            with log_activity(logger, activity_name or f.__name__, activity_type, custom_dimensions):
-                return f(*args, **kwargs)
+            tracer = logger.package_tracer if isinstance(logger, OpsLogger) else None
+            if tracer:
+                with tracer.span():
+                    with log_activity(
+                        logger.package_logger, activity_name or f.__name__, activity_type, custom_dimensions
+                    ):
+                        return f(*args, **kwargs)
+            else:
+                with log_activity(logger.package_logger, activity_name or f.__name__, activity_type, custom_dimensions):
+                    return f(*args, **kwargs)
 
         return wrapper
 
@@ -300,7 +309,7 @@ def monitor_with_telemetry_mixin(
     will collect from return value.
     To monitor, use the ``@monitor_with_telemetry_mixin`` decorator.
 
-    :param logger: The logger adapter.
+    :param logger: The operations logging class, containing loggers and tracer for the package and module
     :type logger: logging.LoggerAdapter
     :param activity_name: The name of the activity. The name should be unique per the wrapped logical code block.
     :type activity_name: str
@@ -315,6 +324,8 @@ def monitor_with_telemetry_mixin(
     :return:
     """
 
+    logger = logger.package_logger if isinstance(logger, OpsLogger) else logger
+
     def monitor(f):
         def _collect_from_parameters(f, args, kwargs, extra_keys):
             dimensions = {}

sdk/ml/azure-ai-ml/azure/ai/ml/_telemetry/logging_handler.py

@@ -9,11 +9,15 @@
 import logging
 import platform
 import traceback
+from typing import Optional, Tuple, Union
 
 from opencensus.ext.azure.common import utils
 from opencensus.ext.azure.common.protocol import Data, Envelope, ExceptionData, Message
 from opencensus.ext.azure.log_exporter import AzureLogHandler
+from opencensus.ext.azure.trace_exporter import AzureExporter
 from opencensus.trace import config_integration
+from opencensus.trace.samplers import ProbabilitySampler
+from opencensus.trace.tracer import Tracer
 
 from azure.ai.ml._user_agent import USER_AGENT
 
@@ -87,7 +91,7 @@ def get_appinsights_log_handler(
     component_name=None,
     enable_telemetry=True,
     **kwargs,
-):
+) -> Tuple[Union["AzureMLSDKLogHandler", logging.NullHandler], Optional[Tracer]]:
     """Enable the OpenCensus logging handler for specified logger and instrumentation key to send info to AppInsights.
 
     :param user_agent: Information about the user's browser.
@@ -102,22 +106,23 @@ def get_appinsights_log_handler(
     :paramtype enable_telemetry: bool
     :keyword kwargs: Optional keyword arguments for adding additional information to messages.
     :paramtype kwargs: dict
-    :return: The logging handler.
-    :rtype: AzureMLSDKLogHandler
+    :return: The logging handler and tracer.
+    :rtype: Tuple[Union[AzureMLSDKLogHandler, logging.NullHandler], Optional[opencensus.trace.tracer.Tracer]]
     """
     try:
         if instrumentation_key is None:
             instrumentation_key = INSTRUMENTATION_KEY
 
         if not in_jupyter_notebook() or not enable_telemetry:
-            return logging.NullHandler()
+            return (logging.NullHandler(), None)
 
         if not user_agent or not user_agent.lower() == USER_AGENT.lower():
-            return logging.NullHandler()
+            return (logging.NullHandler(), None)
 
-        if "properties" in kwargs and "subscription_id" in kwargs.get("properties"):
-            if kwargs.get("properties")["subscription_id"] in test_subscriptions:
-                return logging.NullHandler()
+        if kwargs:
+            if "properties" in kwargs and "subscription_id" in kwargs.get("properties"):  # type: ignore[operator]
+                if kwargs.get("properties")["subscription_id"] in test_subscriptions:  # type: ignore[index]
+                    return (logging.NullHandler(), None)
 
         child_namespace = component_name or __name__
         current_logger = logging.getLogger(AML_INTERNAL_LOGGER_NAMESPACE).getChild(child_namespace)
@@ -135,10 +140,15 @@ def get_appinsights_log_handler(
         )
         current_logger.addHandler(handler)
 
-        return handler
+        tracer = Tracer(
+            exporter=AzureExporter(connection_string=f"InstrumentationKey={instrumentation_key}"),
+            sampler=ProbabilitySampler(1.0),
+        )
+
+        return (handler, tracer)
     except Exception:  # pylint: disable=broad-except
         # ignore any exceptions, telemetry collection errors shouldn't block an operation
-        return logging.NullHandler()
+        return (logging.NullHandler(), None)
 
 
 # cspell:ignore AzureMLSDKLogHandler
@@ -173,14 +183,20 @@ def log_record_to_envelope(self, record):
         if self._is_telemetry_collection_disabled:
             return None
 
-        envelope = create_envelope(self.options.instrumentation_key, record)
+        envelope = Envelope(
+            iKey=self.options.instrumentation_key,
+            tags=dict(utils.azure_monitor_context),
+            time=utils.timestamp_to_iso_str(record.created),
+        )
 
         properties = {
             "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"),
+            "activity_id": record.properties.get("activity_id", "00000000-0000-0000-0000-000000000000"),
+            "client-request-id": record.properties.get("client_request_id", "00000000-0000-0000-0000-000000000000"),
+            "span_id": record.spanId,
+            "trace_id": record.traceId,
         }
 
         if hasattr(record, "custom_dimensions") and isinstance(record.custom_dimensions, dict):
@@ -237,21 +253,3 @@ def log_record_to_envelope(self, record):
             )
             envelope.data = Data(baseData=data, baseType="MessageData")
         return envelope
-
-
-def create_envelope(instrumentation_key, record):
-    envelope = Envelope(
-        iKey=instrumentation_key,
-        tags=dict(utils.azure_monitor_context),
-        time=utils.timestamp_to_iso_str(record.created),
-    )
-    envelope.tags["ai.ml.operation.id"] = getattr(
-        record,
-        "traceId",
-        "00000000000000000000000000000000",
-    )
-    envelope.tags[
-        "ai.ml.operation.parentId"
-    ] = f"|{envelope.tags.get('ai.ml.operation.id')}.{getattr(record, 'spanId', '0000000000000000')}"
-
-    return envelope

sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_logger_utils.py

@@ -4,6 +4,7 @@
 
 import logging
 import sys
+from typing import Dict
 
 from azure.ai.ml._telemetry.logging_handler import AML_INTERNAL_LOGGER_NAMESPACE
 
@@ -22,9 +23,12 @@ class OpsLogger:
     def __init__(self, name: str):
         self.package_logger: logging.Logger = logging.getLogger(AML_INTERNAL_LOGGER_NAMESPACE + name)
         self.package_logger.propagate = False
+        self.package_tracer = None
         self.module_logger = logging.getLogger(name)
         self.custom_dimensions = {}
 
-    def update_info(self, data: dict) -> None:
+    def update_info(self, data: Dict) -> None:
         if "app_insights_handler" in data:
-            self.package_logger.addHandler(data.pop("app_insights_handler"))
+            logger, tracer = data.pop("app_insights_handler")
+            self.package_logger.addHandler(logger)
+            self.package_tracer = tracer