Revert "[ml] Fix parent ID logging and add activity ID logging for operations…" (#34602)

This reverts commit c409ced.

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, Dict, Optional, Tuple, TypeVar, Union
+from typing import Any, 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: Tuple = get_appinsights_log_handler(
+        app_insights_handler = get_appinsights_log_handler(
             user_agent,
             **{"properties": properties},
             enable_telemetry=self._operation_config.enable_telemetry,
         )
-        app_insights_handler_kwargs: Dict[str, Tuple] = {"app_insights_handler": app_insights_handler}
+        app_insights_handler_kwargs = {"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
@@ -442,7 +442,7 @@ def __init__(
             self._service_client_10_2022_preview,
             self._operation_container,
             self._credential,
-            **app_insights_handler_kwargs,  # type: ignore[arg-type]
+            **app_insights_handler_kwargs,
         )
         self._operation_container.add(AzureMLResourceType.REGISTRY, self._registries)  # type: ignore[arg-type]
 
@@ -463,15 +463,15 @@ def __init__(
             self._operation_scope,
             self._operation_config,
             self._service_client_08_2023_preview,
-            **app_insights_handler_kwargs,  # type: ignore[arg-type]
+            **app_insights_handler_kwargs,
         )
         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,  # type: ignore[arg-type]
+            **ops_kwargs,
         )
         self._operation_container.add(AzureMLResourceType.DATASTORE, self._datastores)
         self._models = ModelOperations(
@@ -487,23 +487,23 @@ def __init__(
             workspace_rg=self._ws_rg,
             workspace_sub=self._ws_sub,
             registry_reference=registry_reference,
-            **app_insights_handler_kwargs,  # type: ignore[arg-type]
+            **app_insights_handler_kwargs,
         )
         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,  # type: ignore[arg-type]
+            **ops_kwargs,
         )
         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,  # type: ignore[arg-type]
+            **ops_kwargs,
         )
         self._operation_container.add(AzureMLResourceType.ENVIRONMENT, self._environments)
         self._local_endpoint_helper = _LocalEndpointHelper(requests_pipeline=self._requests_pipeline)
@@ -516,7 +516,7 @@ def __init__(
             self._local_endpoint_helper,
             self._credential,
             requests_pipeline=self._requests_pipeline,
-            **ops_kwargs,  # type: ignore[arg-type]
+            **ops_kwargs,
         )
         self._batch_endpoints = BatchEndpointOperations(
             self._operation_scope,
@@ -526,7 +526,7 @@ def __init__(
             self._credential,
             requests_pipeline=self._requests_pipeline,
             service_client_09_2020_dataplanepreview=self._service_client_09_2020_dataplanepreview,
-            **ops_kwargs,  # type: ignore[arg-type]
+            **ops_kwargs,
         )
         self._operation_container.add(AzureMLResourceType.BATCH_ENDPOINT, self._batch_endpoints)
         self._operation_container.add(AzureMLResourceType.ONLINE_ENDPOINT, self._online_endpoints)
@@ -537,7 +537,7 @@ def __init__(
             self._operation_container,
             self._local_deployment_helper,
             self._credential,
-            **ops_kwargs,  # type: ignore[arg-type]
+            **ops_kwargs,
         )
         self._batch_deployments = BatchDeploymentOperations(
             self._operation_scope,
@@ -569,7 +569,7 @@ def __init__(
             self._service_client_10_2021_dataplanepreview if registry_name else self._service_client_10_2022,
             self._operation_container,
             self._preflight,
-            **ops_kwargs,  # type: ignore[arg-type]
+            **ops_kwargs,
         )
         self._operation_container.add(AzureMLResourceType.COMPONENT, self._components)
         self._jobs = JobOperations(
@@ -602,7 +602,7 @@ def __init__(
                 self._operation_scope,
                 self._credential,
                 _service_client_kwargs=kwargs,
-                **ops_kwargs,  # type: ignore[arg-type]
+                **ops_kwargs,
             )
             self._operation_container.add(
                 AzureMLResourceType.VIRTUALCLUSTER, self._virtual_clusters  # type: ignore[arg-type]
@@ -615,7 +615,7 @@ def __init__(
             self._service_client_08_2023_preview,
             self._operation_container,
             self._credential,
-            **app_insights_handler_kwargs,  # type: ignore[arg-type]
+            **app_insights_handler_kwargs,
         )
 
         self._featuresets = FeatureSetOperations(
@@ -624,22 +624,22 @@ def __init__(
             self._service_client_10_2023,
             self._service_client_08_2023_preview,
             self._datastores,
-            **ops_kwargs,  # type: ignore[arg-type]
+            **ops_kwargs,
         )
 
         self._featurestoreentities = FeatureStoreEntityOperations(
             self._operation_scope,
             self._operation_config,
             self._service_client_10_2023,
-            **ops_kwargs,  # type: ignore[arg-type]
+            **ops_kwargs,
         )
 
         self._workspace_hubs = WorkspaceHubOperations(
             self._operation_scope,
             self._service_client_08_2023_preview,
             self._operation_container,
             self._credential,
-            **app_insights_handler_kwargs,  # type: ignore[arg-type]
+            **app_insights_handler_kwargs,
         )
         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,7 +23,6 @@
 
 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
@@ -261,8 +260,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 operations logging class, containing loggers and tracer for the package and module
-    :type logger: ~azure.ai.ml._utils._logger_utils.OpsLogger
+    :param logger: The logger adapter.
+    :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
     :param activity_type: One of PUBLICAPI, INTERNALCALL, or CLIENTPROXY which represent an incoming API call,
@@ -276,16 +275,8 @@ def monitor_with_activity(
     def monitor(f):
         @functools.wraps(f)
         def wrapper(*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)
+            with log_activity(logger, activity_name or f.__name__, activity_type, custom_dimensions):
+                return f(*args, **kwargs)
 
         return wrapper
 
@@ -309,7 +300,7 @@ def monitor_with_telemetry_mixin(
     will collect from return value.
     To monitor, use the ``@monitor_with_telemetry_mixin`` decorator.
 
-    :param logger: The operations logging class, containing loggers and tracer for the package and module
+    :param logger: The logger adapter.
     :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
@@ -324,8 +315,6 @@ 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,15 +9,11 @@
 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
 
@@ -91,7 +87,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.
@@ -106,23 +102,22 @@ 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 and tracer.
-    :rtype: Tuple[Union[AzureMLSDKLogHandler, logging.NullHandler], Optional[opencensus.trace.tracer.Tracer]]
+    :return: The logging handler.
+    :rtype: AzureMLSDKLogHandler
     """
     try:
         if instrumentation_key is None:
             instrumentation_key = INSTRUMENTATION_KEY
 
         if not in_jupyter_notebook() or not enable_telemetry:
-            return (logging.NullHandler(), None)
+            return logging.NullHandler()
 
         if not user_agent or not user_agent.lower() == USER_AGENT.lower():
-            return (logging.NullHandler(), None)
+            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)
+        if "properties" in kwargs and "subscription_id" in kwargs.get("properties"):
+            if kwargs.get("properties")["subscription_id"] in test_subscriptions:
+                return logging.NullHandler()
 
         child_namespace = component_name or __name__
         current_logger = logging.getLogger(AML_INTERNAL_LOGGER_NAMESPACE).getChild(child_namespace)
@@ -140,15 +135,10 @@ def get_appinsights_log_handler(
         )
         current_logger.addHandler(handler)
 
-        tracer = Tracer(
-            exporter=AzureExporter(connection_string=f"InstrumentationKey={instrumentation_key}"),
-            sampler=ProbabilitySampler(1.0),
-        )
-
-        return (handler, tracer)
+        return handler
     except Exception:  # pylint: disable=broad-except
         # ignore any exceptions, telemetry collection errors shouldn't block an operation
-        return (logging.NullHandler(), None)
+        return logging.NullHandler()
 
 
 # cspell:ignore AzureMLSDKLogHandler
@@ -183,20 +173,14 @@ def log_record_to_envelope(self, record):
         if self._is_telemetry_collection_disabled:
             return None
 
-        envelope = Envelope(
-            iKey=self.options.instrumentation_key,
-            tags=dict(utils.azure_monitor_context),
-            time=utils.timestamp_to_iso_str(record.created),
-        )
+        envelope = create_envelope(self.options.instrumentation_key, record)
 
         properties = {
             "process": record.processName,
             "module": record.module,
             "level": record.levelname,
-            "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,
+            "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):
@@ -253,3 +237,21 @@ 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,7 +4,6 @@
 
 import logging
 import sys
-from typing import Dict
 
 from azure.ai.ml._telemetry.logging_handler import AML_INTERNAL_LOGGER_NAMESPACE
 
@@ -23,12 +22,9 @@ 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:
-            logger, tracer = data.pop("app_insights_handler")
-            self.package_logger.addHandler(logger)
-            self.package_tracer = tracer
+            self.package_logger.addHandler(data.pop("app_insights_handler"))