Update featurestore classes for PuP (Azure#29884)

* Update featurestore classes for PuP

Author: bastrik

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

@@ -6,6 +6,7 @@
 - Added experimental scatter gather node to DSL package. This node has a unique mldesigner dependency.
 - Added support to make JobService and ServiceInstance objects serializable when printed
 - Support Singularity compute in pipeline job
+- Added Feature Store, its dedicated classes and updated the docstrings, now available in public interface. The classes added are `FeatureStoreOperations, FeatureSetOperations, FeatureStoreEntityOperations` with properties classes specific to the new features.
 
 ### Bugs Fixed
 

sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_artifact_utilities.py

@@ -361,7 +361,7 @@ def _update_gen2_metadata(name, version, indicator_file, storage_client) -> None
 
 def _check_and_upload_path(
     artifact: T,
-    asset_operations: Union["DataOperations", "ModelOperations", "CodeOperations", "_FeatureSetOperations"],
+    asset_operations: Union["DataOperations", "ModelOperations", "CodeOperations", "FeatureSetOperations"],
     artifact_type: str,
     datastore_name: Optional[str] = None,
     sas_uri: Optional[str] = None,

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/_input_outputs.py

@@ -114,6 +114,7 @@ def _from_base(cls, _input: Union[Input, Dict]) -> Optional["InternalInput"]:
 
 def _map_v1_io_type(output_type: str) -> str:
     """Map v1 IO type to v2."""
+
     # TODO: put it in a common place
     def _map_primitive_type(_type):
         """Convert double and float to number type."""

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/runsettings/ai_super_computer_configuration.py

@@ -8,7 +8,6 @@
 
 
 class PascalCaseProperty(BaseProperty):
-
     _KEY_MAPPING = {}
 
     def items(self):
@@ -21,7 +20,6 @@ def items(self):
 
 
 class AISuperComputerStorageReferenceConfiguration(PascalCaseProperty):
-
     _KEY_MAPPING = {
         "container_name": "ContainerName",
         "relative_path": "RelativePath",
@@ -45,7 +43,6 @@ def __init__(
 
 
 class AISuperComputerScalePolicy(PascalCaseProperty):
-
     _KEY_MAPPING = {
         "auto_scale_instance_type_count_set": "AutoScaleInstanceTypeCountSet",
         "auto_scale_interval_in_sec": "AutoScaleIntervalInSec",

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

@@ -103,9 +103,9 @@
 from azure.ai.ml.operations._local_deployment_helper import _LocalDeploymentHelper
 from azure.ai.ml.operations._local_endpoint_helper import _LocalEndpointHelper
 from azure.ai.ml.operations._schedule_operations import ScheduleOperations
-from azure.ai.ml.operations._feature_set_operations import _FeatureSetOperations
-from azure.ai.ml.operations._feature_store_operations import _FeatureStoreOperations
-from azure.ai.ml.operations._feature_store_entity_operations import _FeatureStoreEntityOperations
+from azure.ai.ml.operations._feature_set_operations import FeatureSetOperations
+from azure.ai.ml.operations._feature_store_operations import FeatureStoreOperations
+from azure.ai.ml.operations._feature_store_entity_operations import FeatureStoreEntityOperations
 
 module_logger = logging.getLogger(__name__)
 
@@ -510,30 +510,29 @@ def __init__(
         except Exception as ex:  # pylint: disable=broad-except
             module_logger.debug("Virtual Cluster operations could not be initialized due to %s ", ex)
 
-        self._featurestores = _FeatureStoreOperations(
+        self._featurestores = FeatureStoreOperations(
             self._operation_scope,
             self._rp_service_client,
             self._operation_container,
             self._credential,
             **app_insights_handler_kwargs,
         )
 
-        self._featuresets = _FeatureSetOperations(
+        self._featuresets = FeatureSetOperations(
             self._operation_scope,
             self._operation_config,
-            self._service_client_02_2023_preview,
+            self._service_client_04_2023_preview,
             self._datastores,
             **ops_kwargs,
         )
 
-        self._featurestoreentities = _FeatureStoreEntityOperations(
+        self._featurestoreentities = FeatureStoreEntityOperations(
             self._operation_scope, self._operation_config, self._service_client_02_2023_preview, **ops_kwargs
         )
 
-        if is_private_preview_enabled():
-            self._operation_container.add(AzureMLResourceType.FEATURE_STORE, self._featurestores)
-            self._operation_container.add(AzureMLResourceType.FEATURE_SET, self._featuresets)
-            self._operation_container.add(AzureMLResourceType.FEATURE_STORE_ENTITY, self._featurestoreentities)
+        self._operation_container.add(AzureMLResourceType.FEATURE_STORE, self._featurestores)
+        self._operation_container.add(AzureMLResourceType.FEATURE_SET, self._featuresets)
+        self._operation_container.add(AzureMLResourceType.FEATURE_STORE_ENTITY, self._featurestoreentities)
 
     @classmethod
     def from_config(
@@ -662,36 +661,30 @@ def registries(self) -> RegistryOperations:
 
     @property
     @experimental
-    def _feature_stores(self) -> _FeatureStoreOperations:
+    def feature_stores(self) -> FeatureStoreOperations:
         """A collection of feature-store related operations.
         :return: Featurestore operations
-        :rtype: _FeatureStoreOperations
+        :rtype: FeatureStoreOperations
         """
-        if is_private_preview_enabled():
-            return self._featurestores
-        raise Exception("feature store operations not supported")
+        return self._featurestores
 
     @property
     @experimental
-    def _feature_sets(self) -> _FeatureSetOperations:
+    def feature_sets(self) -> FeatureSetOperations:
         """A collection of feature set related operations.
         :return: FeatureSet operations
-        :rtype: _FeatureSetOperations
+        :rtype: FeatureSetOperations
         """
-        if is_private_preview_enabled():
-            return self._featuresets
-        raise Exception("feature set operations not supported")
+        return self._featuresets
 
     @property
     @experimental
-    def _feature_store_entities(self) -> _FeatureStoreEntityOperations:
+    def feature_store_entities(self) -> FeatureStoreEntityOperations:
         """A collection of feature store entity related operations.
         :return: FeatureStoreEntity operations
-        :rtype: _FeatureStoreEntityOperations
+        :rtype: FeatureStoreEntityOperations
         """
-        if is_private_preview_enabled():
-            return self._featurestoreentities
-        raise Exception("feature store entity operations not supported")
+        return self._featurestoreentities
 
     @property
     def connections(self) -> WorkspaceConnectionsOperations:

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

@@ -25,6 +25,6 @@ class FeatureSchema(metaclass=PatchedSchemaMeta):
 
     @post_load
     def make(self, data, **kwargs):
-        from azure.ai.ml.entities._feature_set.feature import _Feature
+        from azure.ai.ml.entities._feature_set.feature import Feature
 
-        return _Feature(description=data.pop("description", None), **data)
+        return Feature(description=data.pop("description", None), **data)

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

@@ -14,6 +14,6 @@ class FeatureSetSpecificationSchema(metaclass=PatchedSchemaMeta):
 
     @post_load
     def make(self, data, **kwargs):
-        from azure.ai.ml.entities._feature_set.feature_set_specification import _FeatureSetSpecification
+        from azure.ai.ml.entities._feature_set.feature_set_specification import FeatureSetSpecification
 
-        return _FeatureSetSpecification(**data)
+        return FeatureSetSpecification(**data)

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

@@ -17,9 +17,9 @@ class MaterializationComputeResourceSchema(metaclass=PatchedSchemaMeta):
 
     @post_load
     def make(self, data, **kwargs):
-        from azure.ai.ml.entities._feature_set.materialization_compute_resource import _MaterializationComputeResource
+        from azure.ai.ml.entities._feature_set.materialization_compute_resource import MaterializationComputeResource
 
-        return _MaterializationComputeResource(instance_type=data.pop("instance_type"), **data)
+        return MaterializationComputeResource(instance_type=data.pop("instance_type"), **data)
 
 
 class MaterializationSettingsSchema(metaclass=PatchedSchemaMeta):
@@ -32,6 +32,6 @@ class MaterializationSettingsSchema(metaclass=PatchedSchemaMeta):
 
     @post_load
     def make(self, data, **kwargs):
-        from azure.ai.ml.entities._feature_set.materialization_settings import _MaterializationSettings
+        from azure.ai.ml.entities._feature_set.materialization_settings import MaterializationSettings
 
-        return _MaterializationSettings(**data)
+        return MaterializationSettings(**data)

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_feature_store/compute_runtime_schema.py

@@ -14,6 +14,6 @@ class ComputeRuntimeSchema(metaclass=PatchedSchemaMeta):
 
     @post_load
     def make(self, data, **kwargs):
-        from azure.ai.ml.entities._workspace.compute_runtime import _ComputeRuntime
+        from azure.ai.ml.entities._workspace.compute_runtime import ComputeRuntime
 
-        return _ComputeRuntime(spark_runtime_version=data.pop("spark_runtime_version"))
+        return ComputeRuntime(spark_runtime_version=data.pop("spark_runtime_version"))

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_feature_store/materialization_store_schema.py

@@ -15,9 +15,9 @@ class MaterializationStoreSchema(metaclass=PatchedSchemaMeta):
 
     @post_load
     def make(self, data, **kwargs):
-        from azure.ai.ml.entities._feature_store.materialization_store import _MaterializationStore
+        from azure.ai.ml.entities._feature_store.materialization_store import MaterializationStore
 
-        return _MaterializationStore(
+        return MaterializationStore(
             type=data.pop("type"),
             target=data.pop("target"),
         )