[ML] Re-enable sdk logging (#29046)

* Re-implement a custom logging handler with emit method

---------

Co-authored-by: nthande <[email protected]>

sdk/ml/azure-ai-ml/README.md

@@ -98,7 +98,7 @@ The Azure ML Python SDK includes a telemetry feature that collects usage and fai
 <u>Telemetry will **not** be collected for any use of the Python SDK outside of a Jupyter Notebook.</u>
 
 Telemetry data helps the SDK team understand how the SDK is used so it can be improved and the information about failures helps the team resolve problems and fix bugs.
-The SDK telemetry feature is enabled by default for Jupyter Notebook usage. To opt out of the telemetry feature, set the AZUREML_SDKV2_TELEMETRY_OPTOUT environment variable to '1' or 'true'.
+The SDK telemetry feature is enabled by default for Jupyter Notebook usage and cannot be enabled for non-Jupyter scenarios. To opt out of the telemetry feature in a Jupyter scenario, pass in `enable_telemetry=False` when constructing your MLClient object.
 
 ## Next steps
 

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.core.credentials import TokenCredential
 from azure.core.polling import LROPoller
@@ -58,7 +58,7 @@
     OperationScope,
 )
 
-# from azure.ai.ml._telemetry.logging_handler import get_appinsights_log_handler
+from azure.ai.ml._telemetry.logging_handler import get_appinsights_log_handler
 from azure.ai.ml._user_agent import USER_AGENT
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml._utils._http_utils import HttpPipeline
@@ -191,7 +191,8 @@ def __init__(
         self._credential = credential
 
         show_progress = kwargs.pop("show_progress", True)
-        self._operation_config = OperationConfig(show_progress=show_progress)
+        enable_telemetry = kwargs.pop("enable_telemetry", True)
+        self._operation_config = OperationConfig(show_progress=show_progress, enable_telemetry=enable_telemetry)
 
         if "cloud" in kwargs:
             cloud_name = kwargs["cloud"]
@@ -252,7 +253,12 @@ def __init__(
         if registry_name:
             properties.update({"registry_name": registry_name})
 
-        app_insights_handler_kwargs: Dict[Any, Any] = {}
+        user_agent = kwargs.get("user_agent", None)
+
+        app_insights_handler = 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}
 
         base_url = _get_base_url_from_metadata(cloud_name=cloud_name, is_local_mfe=True)
         self._base_url = base_url

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

@@ -20,8 +20,10 @@ class OperationConfig(object):
     :type object: _type_
     """
 
-    def __init__(self, show_progress: bool) -> None:
+    def __init__(self, show_progress: bool, enable_telemetry: bool) -> None:
+
         self._show_progress = show_progress
+        self._enable_telemetry = enable_telemetry
 
     @property
     def show_progress(self) -> bool:
@@ -32,6 +34,15 @@ def show_progress(self) -> bool:
         """
         return self._show_progress
 
+    @property
+    def enable_telemetry(self) -> bool:
+        """Decide whether to enable telemetry for Jupyter Notebooks - telemetry cannot be enabled for other contexts.
+
+        :return: enable_telemetry
+        :rtype: bool
+        """
+        return self._enable_telemetry
+
 
 class OperationScope(object):
     def __init__(
@@ -99,6 +110,10 @@ def _resource_group_name(self) -> str:
     def _show_progress(self) -> bool:
         return self._operation_config.show_progress
 
+    @property
+    def _enable_telemetry(self) -> bool:
+        return self._operation_config.enable_telemetry
+
 
 class OperationsContainer(object):
     def __init__(self):

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

@@ -4,14 +4,14 @@
 
 __path__ = __import__("pkgutil").extend_path(__path__, __name__)
 
-# from .activity import ActivityType, log_activity, monitor_with_activity, monitor_with_telemetry_mixin
-# from .logging_handler import AML_INTERNAL_LOGGER_NAMESPACE, get_appinsights_log_handler
+from .activity import ActivityType, log_activity, monitor_with_activity, monitor_with_telemetry_mixin
+from .logging_handler import AML_INTERNAL_LOGGER_NAMESPACE, get_appinsights_log_handler
 
 __all__ = [
-    #     "monitor_with_activity",
-    #     "monitor_with_telemetry_mixin",
-    #     "log_activity",
-    #     "ActivityType",
-    #     "get_appinsights_log_handler",
-    #     "AML_INTERNAL_LOGGER_NAMESPACE",
+    "monitor_with_activity",
+    "monitor_with_telemetry_mixin",
+    "log_activity",
+    "ActivityType",
+    "get_appinsights_log_handler",
+    "AML_INTERNAL_LOGGER_NAMESPACE",
 ]

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

@@ -7,33 +7,25 @@
 """Contains functionality for sending telemetry to Application Insights via OpenCensus Azure Monitor Exporter."""
 
 import logging
+import platform
+import traceback
 
-# import platform
-from os import getenv
-
-# from opencensus.ext.azure.log_exporter import AzureLogHandler
-
-# from azure.ai.ml._user_agent import USER_AGENT
+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
 
 
 AML_INTERNAL_LOGGER_NAMESPACE = "azure.ai.ml._telemetry"
 
 # vienna-sdk-unitedstates
 INSTRUMENTATION_KEY = "71b954a8-6b7d-43f5-986c-3d3a6605d803"
 
-AZUREML_SDKV2_TELEMETRY_OPTOUT_ENV_VAR = "AZUREML_SDKV2_TELEMETRY_OPTOUT"
-
-# application insight logger name
-LOGGER_NAME = "ApplicationInsightLogger"
-
-SUCCESS = True
-FAILURE = False
-
-TRACEBACK_LOOKUP_STR = "Traceback (most recent call last)"
-
-# extract traceback path from message
-reformat_traceback = True
-
 test_subscriptions = [
     "b17253fa-f327-42d6-9686-f3e553e24763",
     "test_subscription",
@@ -79,65 +71,177 @@ def in_jupyter_notebook() -> bool:
     return True
 
 
-def is_telemetry_collection_disabled():
-    telemetry_disabled = getenv(AZUREML_SDKV2_TELEMETRY_OPTOUT_ENV_VAR)
-    if telemetry_disabled and (telemetry_disabled.lower() == "true" or telemetry_disabled == "1"):
-        return True
-    if not in_jupyter_notebook:
-        return True
-    return False
-
-
-# def get_appinsights_log_handler(
-#      user_agent,
-#      *args, # pylint: disable=unused-argument
-#      instrumentation_key=None,
-#      component_name=None,
-#      **kwargs
-#  ):
-#     """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.
-#     :type user_agent: Dict[str, str]
-#     :param instrumentation_key: The Application Insights instrumentation key.
-#     :type instrumentation_key: str
-#     :param component_name: The component name.
-#     :type component_name: str
-#     :param args: Optional arguments for formatting messages.
-#     :type args: list
-#     :param kwargs: Optional keyword arguments for adding additional information to messages.
-#     :type kwargs: dict
-#     :return: The logging handler.
-#     :rtype: opencensus.ext.azure.log_exporter.AzureLogHandler
-#     """
-#     try:
-#         if instrumentation_key is None:
-#             instrumentation_key = INSTRUMENTATION_KEY
-
-#     if is_telemetry_collection_disabled():
-#         return logging.NullHandler()
-
-#     if not user_agent or not user_agent.lower() == USER_AGENT.lower():
-#         return logging.NullHandler()
-
-#     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)
-#     current_logger.propagate = False
-#     current_logger.setLevel(logging.CRITICAL)
-
-#     custom_properties = {"PythonVersion": platform.python_version()}
-#     custom_properties.update({"user_agent": user_agent})
-#     if "properties" in kwargs:
-#         custom_properties.update(kwargs.pop("properties"))
-#     handler = AzureLogHandler(connection_string=f'InstrumentationKey={instrumentation_key}')
-#     current_logger.addHandler(handler)
-#     handler.addFilter(CustomDimensionsFilter(custom_properties))
-
-#     return handler
-# except Exception: # pylint: disable=broad-except
-#     # ignore exceptions, telemetry should not block
-#     return logging.NullHandler()
+# cspell:ignore overriden
+def get_appinsights_log_handler(
+    user_agent,
+    *args,  # pylint: disable=unused-argument
+    instrumentation_key=None,
+    component_name=None,
+    enable_telemetry=True,
+    **kwargs,
+):
+    """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.
+    :type user_agent: Dict[str, str]
+    :param instrumentation_key: The Application Insights instrumentation key.
+    :type instrumentation_key: str
+    :param component_name: The component name.
+    :type component_name: str
+    :param enable_telemetry: Whether to enable telemetry. Will be overriden to False if not in a Jupyter Notebook.
+    :type enable_telemetry: bool
+    :param args: Optional arguments for formatting messages.
+    :type args: list
+    :param kwargs: Optional keyword arguments for adding additional information to messages.
+    :type kwargs: dict
+    :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()
+
+        if not user_agent or not user_agent.lower() == USER_AGENT.lower():
+            return logging.NullHandler()
+
+        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)
+        current_logger.propagate = False
+        current_logger.setLevel(logging.CRITICAL)
+
+        custom_properties = {"PythonVersion": platform.python_version()}
+        custom_properties.update({"user_agent": user_agent})
+        if "properties" in kwargs:
+            custom_properties.update(kwargs.pop("properties"))
+        handler = AzureMLSDKLogHandler(
+            connection_string=f"InstrumentationKey={instrumentation_key}",
+            custom_properties=custom_properties,
+            enable_telemetry=enable_telemetry,
+        )
+        current_logger.addHandler(handler)
+
+        return handler
+    except Exception:  # pylint: disable=broad-except
+        # ignore any exceptions, telemetry collection errors shouldn't block an operation
+        return logging.NullHandler()
+
+
+# cspell:ignore AzureMLSDKLogHandler
+class AzureMLSDKLogHandler(AzureLogHandler):
+    """Customized AzureLogHandler for AzureML SDK"""
+
+    def __init__(self, custom_properties, enable_telemetry, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._is_telemetry_collection_disabled = not enable_telemetry
+        self._custom_properties = custom_properties
+        self.addFilter(CustomDimensionsFilter(self._custom_properties))
+
+    def emit(self, record):
+        if self._is_telemetry_collection_disabled:
+            return
+
+        try:
+            self._queue.put(record, block=False)
+
+            # log the record immediately if it is an error
+            if record.exc_info and not all(item is None for item in record.exc_info):
+                self._queue.flush()
+        except Exception:  # pylint: disable=broad-except
+            # ignore any exceptions, telemetry collection errors shouldn't block an operation
+            return
+
+    # The code below is vendored from opencensus-ext-azure's AzureLogHandler base class, but the telemetry disabling
+    # logic has been added to the beginning. Without this, the base class would still send telemetry even if
+    # enable_telemetry had been set to true.
+    def log_record_to_envelope(self, record):
+        if self._is_telemetry_collection_disabled:
+            return
+
+        envelope = create_envelope(self.options.instrumentation_key, record)
+
+        properties = {
+            "process": record.processName,
+            "module": record.module,
+            "level": record.levelname,
+        }
+        if hasattr(record, "custom_dimensions") and isinstance(record.custom_dimensions, dict):
+            properties.update(record.custom_dimensions)
+
+        if record.exc_info:
+            exctype, _value, tb = record.exc_info
+            callstack = []
+            level = 0
+            has_full_stack = False
+            exc_type = "N/A"
+            message = self.format(record)
+            if tb is not None:
+                has_full_stack = True
+                for _, line, method, _text in traceback.extract_tb(tb):
+                    callstack.append(
+                        {
+                            "level": level,
+                            "method": method,
+                            "line": line,
+                        }
+                    )
+                    level += 1
+                callstack.reverse()
+            elif record.message:
+                message = record.message
+
+            if exctype is not None:
+                exc_type = exctype.__name__
+
+            envelope.name = "Microsoft.ApplicationInsights.Exception"
+
+            data = ExceptionData(
+                exceptions=[
+                    {
+                        "id": 1,
+                        "outerId": 0,
+                        "typeName": exc_type,
+                        "message": message,
+                        "hasFullStack": has_full_stack,
+                        "parsedStack": callstack,
+                    }
+                ],
+                severityLevel=max(0, record.levelno - 1) // 10,
+                properties=properties,
+            )
+            envelope.data = Data(baseData=data, baseType="ExceptionData")
+        else:
+            envelope.name = "Microsoft.ApplicationInsights.Message"
+            data = Message(
+                message=self.format(record),
+                severityLevel=max(0, record.levelno - 1) // 10,
+                properties=properties,
+            )
+            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.operation.id"] = getattr(
+        record,
+        "traceId",
+        "00000000000000000000000000000000",
+    )
+    envelope.tags["ai.operation.parentId"] = "|{}.{}.".format(
+        envelope.tags["ai.operation.id"],
+        getattr(record, "spanId", "0000000000000000"),
+    )
+
+    return envelope