[ml] Address azure sdk review comments (Azure#30294)

* Address azure sdk review comments

* Fix

Author: bastrik

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_feature_set/featureset_spec_properties_schema.py

@@ -31,11 +31,17 @@ class FeaturePropertiesSchema(metaclass=PatchedSchemaMeta):
     tags = fields.Dict(keys=fields.Str(), values=fields.Str(), data_key="Tags")
 
 
+class ColumnPropertiesSchema(metaclass=PatchedSchemaMeta):
+    name = fields.Str(data_key="ColumnName")
+    type = fields.Str(data_key="DataType")
+
+
 class FeaturesetSpecPropertiesSchema(YamlFileSchema):
     source = fields.Nested(SourceMetadataSchema, data_key="source")
     feature_transformation_code = fields.Nested(
         FeatureTransformationCodePropertiesSchema, data_key="featureTransformationCode"
     )
     features = fields.List(NestedField(FeaturePropertiesSchema), data_key="features")
+    index_columns = fields.List(NestedField(ColumnPropertiesSchema), data_key="indexColumns")
     source_lookback = fields.Nested(DelayMetadataPropertiesSchema, data_key="sourceLookback")
     temporal_join_lookback = fields.Nested(DelayMetadataPropertiesSchema, data_key="temporalJoinLookback")

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

@@ -3,7 +3,7 @@
 # ---------------------------------------------------------
 
 # pylint: disable=protected-access
-
+from datetime import datetime
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Dict, Union
@@ -101,3 +101,7 @@ def _archive_or_restore(
         body=version_resource,
         **kwargs,
     )
+
+
+def _datetime_to_str(datetime_obj: Union[str, datetime]):
+    return datetime_obj if isinstance(datetime_obj, str) else datetime_obj.isoformat()

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

@@ -8,7 +8,7 @@
 import json
 from pathlib import Path
 from datetime import datetime
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
 
 from marshmallow.exceptions import ValidationError as SchemaValidationError
 
@@ -26,11 +26,12 @@
 from azure.ai.ml._artifacts._artifact_utilities import _check_and_upload_path
 from azure.ai.ml.operations._datastore_operations import DatastoreOperations
 
-# from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
+from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
 from azure.ai.ml._utils._feature_store_utils import (
     _archive_or_restore,
     read_feature_set_metadata_contents,
     read_remote_feature_set_spec_metadata_contents,
+    _datetime_to_str,
 )
 from azure.ai.ml._utils._logger_utils import OpsLogger
 from azure.ai.ml.entities._assets._artifacts.feature_set import FeatureSet
@@ -41,9 +42,10 @@
 from azure.ai.ml.entities._feature_set.feature import Feature
 from azure.core.polling import LROPoller
 from azure.core.paging import ItemPaged
+from azure.core.tracing.decorator import distributed_trace
 
 ops_logger = OpsLogger(__name__)
-module_logger = ops_logger.module_logger
+logger, module_logger = ops_logger.package_logger, ops_logger.module_logger
 
 
 class FeatureSetOperations(_ScopeDependentOperations):
@@ -71,7 +73,8 @@ def __init__(
         self._datastore_operation = datastore_operations
         self._init_kwargs = kwargs
 
-    # @monitor_with_activity(logger, "FeatureSet.List", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureSet.List", ActivityType.PUBLICAPI)
     def list(
         self,
         name: Optional[str] = None,
@@ -116,7 +119,8 @@ def _get(self, name: str, version: str = None, **kwargs: Dict) -> FeaturesetVers
             **kwargs,
         )
 
-    # @monitor_with_activity(logger, "FeatureSet.Get", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureSet.Get", ActivityType.PUBLICAPI)
     def get(self, name: str, version: str, **kwargs: Dict) -> FeatureSet:
         """Get the specified FeatureSet asset.
 
@@ -135,7 +139,8 @@ def get(self, name: str, version: str, **kwargs: Dict) -> FeatureSet:
         except (ValidationException, SchemaValidationError) as ex:
             log_and_raise_error(ex)
 
-    # @monitor_with_activity(logger, "FeatureSet.BeginCreateOrUpdate", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureSet.BeginCreateOrUpdate", ActivityType.PUBLICAPI)
     def begin_create_or_update(self, featureset: FeatureSet, **kwargs: Dict) -> LROPoller[FeatureSet]:
         """Create or update FeatureSet
 
@@ -166,7 +171,8 @@ def begin_create_or_update(self, featureset: FeatureSet, **kwargs: Dict) -> LROP
             cls=lambda response, deserialized, headers: FeatureSet._from_rest_object(deserialized),
         )
 
-    # @monitor_with_activity(logger, "FeatureSet.BeginBackFill", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureSet.BeginBackFill", ActivityType.PUBLICAPI)
     def begin_backfill(
         self,
         *,
@@ -224,14 +230,15 @@ def begin_backfill(
             cls=lambda response, deserialized, headers: FeatureSetBackfillMetadata._from_rest_object(deserialized),
         )
 
-    # @monitor_with_activity(logger, "FeatureSet.ListMaterializationOperation", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureSet.ListMaterializationOperation", ActivityType.PUBLICAPI)
     def list_materialization_operations(
         self,
         name: str,
         version: str,
         *,
-        feature_window_start_time: Optional[str] = None,
-        feature_window_end_time: Optional[str] = None,
+        feature_window_start_time: Optional[Union[str, datetime]] = None,
+        feature_window_end_time: Optional[Union[str, datetime]] = None,
         filters: Optional[str] = None,
         **kwargs: Dict,
     ) -> ItemPaged[FeatureSetMaterializationMetadata]:
@@ -242,15 +249,16 @@ def list_materialization_operations(
         :param version: Feature set version.
         :type version: str
         :param feature_window_start_time: Start time of the feature window to filter materialization jobs.
-        :type feature_window_start_time: str
+        :type feature_window_start_time: Union[str, datetime]
         :param feature_window_end_time: End time of the feature window to filter materialization jobs.
-        :type feature_window_end_time: str
+        :type feature_window_end_time: Union[str, datetime]
         :param filters: Comma-separated list of tag names (and optionally values). Example: tag1,tag2=value2.
         :type filters: str
         :return: An iterator like instance of ~azure.ai.ml.entities.FeatureSetMaterializationMetadata objects
         :rtype: ~azure.core.paging.ItemPaged[FeatureSetMaterializationMetadata]
         """
-
+        feature_window_start_time = _datetime_to_str(feature_window_start_time) if feature_window_start_time else None
+        feature_window_end_time = _datetime_to_str(feature_window_end_time) if feature_window_end_time else None
         materialization_jobs = self._operation.list_materialization_jobs(
             resource_group_name=self._resource_group_name,
             workspace_name=self._workspace_name,
@@ -264,7 +272,8 @@ def list_materialization_operations(
         )
         return materialization_jobs
 
-    # @monitor_with_activity(logger, "FeatureSet.ListFeatures", ActivityType.INTERNALCALL)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureSet.ListFeatures", ActivityType.PUBLICAPI)
     def list_features(
         self,
         feature_set_name: str,
@@ -303,7 +312,8 @@ def list_features(
         )
         return features
 
-    # @monitor_with_activity(logger, "FeatureSet.GetFeature", ActivityType.INTERNALCALL)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureSet.GetFeature", ActivityType.PUBLICAPI)
     def get_feature(self, feature_set_name: str, version: str, *, feature_name: str, **kwargs: Dict) -> "Feature":
         """Get Feature
 
@@ -329,7 +339,8 @@ def get_feature(self, feature_set_name: str, version: str, *, feature_name: str,
 
         return Feature._from_rest_object(feature)
 
-    # @monitor_with_activity(logger, "FeatureSet.Archive", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureSet.Archive", ActivityType.PUBLICAPI)
     def archive(
         self,
         name: str,
@@ -354,7 +365,8 @@ def archive(
             **kwargs,
         )
 
-    # @monitor_with_activity(logger, "FeatureSet.Restore", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureSet.Restore", ActivityType.PUBLICAPI)
     def restore(
         self,
         name: str,

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

@@ -15,17 +15,18 @@
 from azure.ai.ml.exceptions import ValidationException
 
 
-# from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
+from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
 from azure.ai.ml._utils._feature_store_utils import (
     _archive_or_restore,
 )
 from azure.ai.ml._utils._logger_utils import OpsLogger
 from azure.ai.ml.entities._feature_store_entity.feature_store_entity import FeatureStoreEntity
 from azure.core.polling import LROPoller
 from azure.core.paging import ItemPaged
+from azure.core.tracing.decorator import distributed_trace
 
 ops_logger = OpsLogger(__name__)
-module_logger = ops_logger.module_logger
+logger, module_logger = ops_logger.package_logger, ops_logger.module_logger
 
 
 class FeatureStoreEntityOperations(_ScopeDependentOperations):
@@ -50,7 +51,8 @@ def __init__(
         self._service_client = service_client
         self._init_kwargs = kwargs
 
-    # @monitor_with_activity(logger, "FeatureStoreEntity.List", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureStoreEntity.List", ActivityType.PUBLICAPI)
     def list(
         self,
         name: Optional[str] = None,
@@ -95,7 +97,8 @@ def _get(self, name: str, version: str = None, **kwargs: Dict) -> FeaturestoreEn
             **kwargs,
         )
 
-    # @monitor_with_activity(logger, "FeatureStoreEntity.Get", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureStoreEntity.Get", ActivityType.PUBLICAPI)
     def get(self, name: str, version: str, **kwargs: Dict) -> FeatureStoreEntity:
         """Get the specified FeatureStoreEntity asset.
 
@@ -114,7 +117,8 @@ def get(self, name: str, version: str, **kwargs: Dict) -> FeatureStoreEntity:
         except (ValidationException, SchemaValidationError) as ex:
             log_and_raise_error(ex)
 
-    # @monitor_with_activity(logger, "FeatureStoreEntity.BeginCreateOrUpdate", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureStoreEntity.BeginCreateOrUpdate", ActivityType.PUBLICAPI)
     def begin_create_or_update(
         self, feature_store_entity: FeatureStoreEntity, **kwargs: Dict
     ) -> LROPoller[FeatureStoreEntity]:
@@ -137,7 +141,8 @@ def begin_create_or_update(
             **kwargs,
         )
 
-    # @monitor_with_activity(logger, "FeatureStoreEntity.Archive", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureStoreEntity.Archive", ActivityType.PUBLICAPI)
     def archive(
         self,
         name: str,
@@ -162,7 +167,8 @@ def archive(
             **kwargs,
         )
 
-    # @monitor_with_activity(logger, "FeatureStoreEntity.Restore", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureStoreEntity.Restore", ActivityType.PUBLICAPI)
     def restore(
         self,
         name: str,

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

@@ -11,7 +11,7 @@
 from azure.ai.ml._restclient.v2023_04_01_preview import AzureMachineLearningWorkspaces as ServiceClient042023Preview
 from azure.ai.ml._scope_dependent_operations import OperationsContainer, OperationScope
 
-# from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
+from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
 from azure.core.credentials import TokenCredential
 from azure.core.polling import LROPoller
 from azure.core.tracing.decorator import distributed_trace
@@ -40,7 +40,7 @@
 from ._workspace_operations_base import WorkspaceOperationsBase
 
 ops_logger = OpsLogger(__name__)
-module_logger = ops_logger.module_logger
+logger, module_logger = ops_logger.package_logger, ops_logger.module_logger
 
 
 class FeatureStoreOperations(WorkspaceOperationsBase):
@@ -69,7 +69,9 @@ def __init__(
         )
         self._workspace_connection_operation = service_client.workspace_connections
 
-    # @monitor_with_activity(logger, "FeatureStore.List", ActivityType.PUBLICAPI)
+    @distributed_trace
+    @monitor_with_activity(logger, "FeatureStore.List", ActivityType.PUBLICAPI)
+    # pylint: disable=unused-argument
     def list(self, *, scope: str = Scope.RESOURCE_GROUP, **kwargs: Dict) -> Iterable[FeatureStore]:
         """List all feature stores that the user has access to in the current
         resource group or subscription.
@@ -82,23 +84,21 @@ def list(self, *, scope: str = Scope.RESOURCE_GROUP, **kwargs: Dict) -> Iterable
 
         if scope == Scope.SUBSCRIPTION:
             return self._operation.list_by_subscription(
-                **kwargs,
                 cls=lambda objs: [
                     FeatureStore._from_rest_object(filterObj)
                     for filterObj in filter(lambda ws: ws.kind.lower() == FEATURE_STORE_KIND, objs)
                 ],
             )
         return self._operation.list_by_resource_group(
             self._resource_group_name,
-            **kwargs,
             cls=lambda objs: [
                 FeatureStore._from_rest_object(filterObj)
                 for filterObj in filter(lambda ws: ws.kind.lower() == FEATURE_STORE_KIND, objs)
             ],
         )
 
-    # @monitor_with_activity(logger, "FeatureStore.Get", ActivityType.PUBLICAPI)
     @distributed_trace
+    @monitor_with_activity(logger, "FeatureStore.Get", ActivityType.PUBLICAPI)
     # pylint: disable=arguments-renamed
     def get(self, name: str, **kwargs: Dict) -> FeatureStore:
         """Get a feature store by name.
@@ -111,7 +111,7 @@ def get(self, name: str, **kwargs: Dict) -> FeatureStore:
 
         feature_store = None
         resource_group = kwargs.get("resource_group") or self._resource_group_name
-        rest_workspace_obj = self._operation.get(resource_group, name, **kwargs)
+        rest_workspace_obj = self._operation.get(resource_group, name)
         if rest_workspace_obj and rest_workspace_obj.kind and rest_workspace_obj.kind.lower() == FEATURE_STORE_KIND:
             feature_store = FeatureStore._from_rest_object(rest_workspace_obj)
 
@@ -171,8 +171,8 @@ def get(self, name: str, **kwargs: Dict) -> FeatureStore:
 
         return feature_store
 
-    # @monitor_with_activity(logger, "FeatureStore.BeginCreate", ActivityType.PUBLICAPI)
     @distributed_trace
+    @monitor_with_activity(logger, "FeatureStore.BeginCreate", ActivityType.PUBLICAPI)
     # pylint: disable=arguments-differ
     def begin_create(
         self,
@@ -214,8 +214,8 @@ def get_callback():
             **kwargs,
         )
 
-    # @monitor_with_activity(logger, "FeatureStore.BeginUpdate", ActivityType.PUBLICAPI)
     @distributed_trace
+    @monitor_with_activity(logger, "FeatureStore.BeginUpdate", ActivityType.PUBLICAPI)
     # pylint: disable=arguments-renamed
     def begin_update(
         self,
@@ -364,9 +364,9 @@ def deserialize_callback(rest_obj):
             **kwargs,
         )
 
-    # @monitor_with_activity(logger, "FeatureStore.BeginDelete", ActivityType.PUBLICAPI)
     @distributed_trace
-    def begin_delete(self, name: str, *, delete_dependent_resources: bool, **kwargs: Dict) -> LROPoller[None]:
+    @monitor_with_activity(logger, "FeatureStore.BeginDelete", ActivityType.PUBLICAPI)
+    def begin_delete(self, name: str, *, delete_dependent_resources: bool = False, **kwargs: Dict) -> LROPoller[None]:
         """Delete a FeatureStore.
 
         :param name: Name of the FeatureStore

sdk/ml/azure-ai-ml/tests/feature_store/unittests/test_feature_store_operations.py

@@ -66,10 +66,10 @@ def test_update(self, mock_feature_store_operation: FeatureStoreOperations) -> N
             description="description",
         )
 
-        def outgoing_get_call(rg, name, **kwargs):
+        def outgoing_get_call(rg, name):
             return Workspace(name=name, kind="featurestore")._to_rest_object()
 
-        def outgoing_call(rg, name, params, polling, cls, **kwargs):
+        def outgoing_call(rg, name, params, polling, cls):
             assert rg == "test_resource_group"
             assert name == "name"
             assert params.description == "description"
@@ -83,7 +83,7 @@ def outgoing_call(rg, name, params, polling, cls, **kwargs):
         mock_feature_store_operation._operation.begin_update.assert_called()
 
     def test_delete(self, mock_feature_store_operation: FeatureStoreOperations, mocker: MockFixture) -> None:
-        def outgoing_call(rg, name, **kwargs):
+        def outgoing_call(rg, name):
             return Workspace(name=name, kind="featurestore")._to_rest_object()
 
         mock_feature_store_operation._operation.get.side_effect = outgoing_call
@@ -94,7 +94,7 @@ def outgoing_call(rg, name, **kwargs):
     def test_delete_non_feature_store_kind(
         self, mock_feature_store_operation: FeatureStoreOperations, mocker: MockFixture
     ) -> None:
-        def outgoing_call(rg, name, **kwargs):
+        def outgoing_call(rg, name):
             return Workspace(name=name)._to_rest_object()
 
         mock_feature_store_operation._operation.get.side_effect = outgoing_call