[ml] Fix missing parent ids in trace logs

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
@@ -449,7 +449,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]
 
@@ -470,15 +470,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(
@@ -494,23 +494,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)
@@ -523,7 +523,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,
@@ -533,7 +533,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)
@@ -544,7 +544,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,
@@ -576,7 +576,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(
@@ -610,7 +610,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]
@@ -623,7 +623,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(
@@ -632,22 +632,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(name=f.__name__):
+                    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

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

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_batch_deployment_operations.py

@@ -34,7 +34,7 @@
 from ._operation_orchestrator import OperationOrchestrator
 
 ops_logger = OpsLogger(__name__)
-logger, module_logger = ops_logger.package_logger, ops_logger.module_logger
+module_logger = ops_logger.module_logger
 DeploymentType = TypeVar(
     "DeploymentType", bound=Union[BatchDeployment, PipelineComponentBatchDeployment, ModelBatchDeployment]
 )
@@ -84,7 +84,7 @@ def __init__(
         self._requests_pipeline: HttpPipeline = kwargs.pop("requests_pipeline")
 
     @distributed_trace
-    @monitor_with_activity(logger, "BatchDeployment.BeginCreateOrUpdate", ActivityType.PUBLICAPI)
+    @monitor_with_activity(ops_logger, "BatchDeployment.BeginCreateOrUpdate", ActivityType.PUBLICAPI)
     def begin_create_or_update(
         self,
         deployment: DeploymentType,
@@ -171,7 +171,7 @@ def begin_create_or_update(
             raise ex
 
     @distributed_trace
-    @monitor_with_activity(logger, "BatchDeployment.Get", ActivityType.PUBLICAPI)
+    @monitor_with_activity(ops_logger, "BatchDeployment.Get", ActivityType.PUBLICAPI)
     def get(self, name: str, endpoint_name: str) -> BatchDeployment:
         """Get a deployment resource.
 
@@ -205,7 +205,7 @@ def get(self, name: str, endpoint_name: str) -> BatchDeployment:
         return deployment
 
     @distributed_trace
-    @monitor_with_activity(logger, "BatchDeployment.BeginDelete", ActivityType.PUBLICAPI)
+    @monitor_with_activity(ops_logger, "BatchDeployment.BeginDelete", ActivityType.PUBLICAPI)
     def begin_delete(self, name: str, endpoint_name: str) -> LROPoller[None]:
         """Delete a batch deployment.
 
@@ -247,7 +247,7 @@ def begin_delete(self, name: str, endpoint_name: str) -> LROPoller[None]:
         return delete_poller
 
     @distributed_trace
-    @monitor_with_activity(logger, "BatchDeployment.List", ActivityType.PUBLICAPI)
+    @monitor_with_activity(ops_logger, "BatchDeployment.List", ActivityType.PUBLICAPI)
     def list(self, endpoint_name: str) -> ItemPaged[BatchDeployment]:
         """List a deployment resource.
 
@@ -274,7 +274,7 @@ def list(self, endpoint_name: str) -> ItemPaged[BatchDeployment]:
         )
 
     @distributed_trace
-    @monitor_with_activity(logger, "BatchDeployment.ListJobs", ActivityType.PUBLICAPI)
+    @monitor_with_activity(ops_logger, "BatchDeployment.ListJobs", ActivityType.PUBLICAPI)
     def list_jobs(self, endpoint_name: str, *, name: Optional[str] = None) -> ItemPaged[BatchJob]:
         """List jobs under the provided batch endpoint deployment. This is only valid for batch endpoint.