Updating Feature store classes name and operations for private preview

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

@@ -362,7 +362,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/_ml_client.py

@@ -59,7 +59,7 @@
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml._utils._http_utils import HttpPipeline
 from azure.ai.ml._utils._registry_utils import get_registry_client
-from azure.ai.ml._utils.utils import _is_https_url
+from azure.ai.ml._utils.utils import _is_https_url, is_private_preview_enabled
 from azure.ai.ml.constants._common import AzureMLResourceType
 from azure.ai.ml.entities import (
     BatchDeployment,
@@ -100,6 +100,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
 
 module_logger = logging.getLogger(__name__)
 
@@ -478,6 +481,31 @@ def __init__(
 
         self._virtual_clusters = VirtualClusterOperations(self._operation_scope, self._credential, **ops_kwargs)
 
+        self._featurestores = _FeatureStoreOperations(
+            self._operation_scope,
+            self._rp_service_client,
+            self._operation_container,
+            self._credential,
+            **app_insights_handler_kwargs,
+        )
+
+        self._featuresets = _FeatureSetOperations(
+            self._operation_scope,
+            self._operation_config,
+            self._service_client_02_2023_preview,
+            self._datastores,
+            **ops_kwargs,
+        )
+
+        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)
+
     @classmethod
     def from_config(
         cls,
@@ -603,6 +631,39 @@ def registries(self) -> RegistryOperations:
         """
         return self._registries
 
+    @property
+    @experimental
+    def _feature_stores(self) -> _FeatureStoreOperations:
+        """A collection of feature-store related operations.
+        :return: Featurestore operations
+        :rtype: _FeatureStoreOperations
+        """
+        if is_private_preview_enabled():
+            return self._featurestores
+        raise Exception("feature store operations not supported")
+
+    @property
+    @experimental
+    def _feature_sets(self) -> _FeatureSetOperations:
+        """A collection of feature set related operations.
+        :return: FeatureSet operations
+        :rtype: _FeatureSetOperations
+        """
+        if is_private_preview_enabled():
+            return self._featuresets
+        raise Exception("feature set operations not supported")
+
+    @property
+    @experimental
+    def _feature_store_entities(self) -> _FeatureStoreEntityOperations:
+        """A collection of feature store entity related operations.
+        :return: FeatureStoreEntity operations
+        :rtype: _FeatureStoreEntityOperations
+        """
+        if is_private_preview_enabled():
+            return self._featurestoreentities
+        raise Exception("feature store entity operations not supported")
+
     @property
     def connections(self) -> WorkspaceConnectionsOperations:
         """A collection of workspace connection related operations.

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

@@ -8,7 +8,7 @@
 from .feature_schema import FeatureSchema
 from .feature_set_schema import FeatureSetSchema
 from .featureset_spec_schema import FeaturesetSpecSchema
-from .featureset_specification_schema import FeaturesetSpecificationSchema
+from .feature_set_specification_schema import FeatureSetSpecificationSchema
 from .materialization_settings_schema import MaterializationSettingsSchema
 from .source_metadata_schema import SourceMetadataSchema
 from .timestamp_column_metadata_schema import TimestampColumnMetadataSchema
@@ -18,7 +18,7 @@
     "FeatureSchema",
     "FeatureSetSchema",
     "FeaturesetSpecSchema",
-    "FeaturesetSpecificationSchema",
+    "FeatureSetSpecificationSchema",
     "MaterializationSettingsSchema",
     "SourceMetadataSchema",
     "TimestampColumnMetadataSchema",

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

@@ -23,6 +23,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(data_type=type, description=data.pop("description", None), **data)
+        return _Feature(data_type=type, description=data.pop("description", None), **data)

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

@@ -10,13 +10,13 @@
 from azure.ai.ml._schema.core.schema import YamlFileSchema
 
 from .materialization_settings_schema import MaterializationSettingsSchema
-from .featureset_specification_schema import FeaturesetSpecificationSchema
+from .feature_set_specification_schema import FeatureSetSpecificationSchema
 
 
 class FeatureSetSchema(YamlFileSchema):
     name = fields.Str(required=True, allow_none=False)
     version = fields.Str(required=True, allow_none=False)
-    specification = NestedField(FeaturesetSpecificationSchema, required=True, allow_none=False)
+    specification = NestedField(FeatureSetSpecificationSchema, required=True, allow_none=False)
     entities = fields.List(fields.Str, required=True, allow_none=False)
     stage = fields.Str()
     description = fields.Str()

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

@@ -9,11 +9,11 @@
 from azure.ai.ml._schema.core.schema import PatchedSchemaMeta
 
 
-class FeaturesetSpecificationSchema(metaclass=PatchedSchemaMeta):
+class FeatureSetSpecificationSchema(metaclass=PatchedSchemaMeta):
     path = fields.Str(required=True, allow_none=False)
 
     @post_load
     def make(self, data, **kwargs):
-        from azure.ai.ml.entities._feature_set.featureset_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/feature_transformation_code_schema.py

@@ -15,6 +15,6 @@ class FeatureTransformationCodeSchema(metaclass=PatchedSchemaMeta):
 
     @post_load
     def make(self, data, **kwargs):
-        from azure.ai.ml.entities._feature_set.delay_metadata import DelayMetadata
+        from azure.ai.ml.entities._feature_set.feature_transformation_code import FeatureTransformationCode
 
-        return DelayMetadata(**data)
+        return FeatureTransformationCode(**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"),
         )

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_feature_store_entity/data_column_schema.py

@@ -22,6 +22,6 @@ class DataColumnSchema(metaclass=PatchedSchemaMeta):
 
     @post_load
     def make(self, data, **kwargs):
-        from azure.ai.ml.entities._feature_store_entity.data_column import DataColumn
+        from azure.ai.ml.entities._feature_store_entity.data_column import _DataColumn
 
-        return DataColumn(**data)
+        return _DataColumn(**data)

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_notification/notification_schema.py

@@ -19,6 +19,6 @@ class NotificationSchema(metaclass=PatchedSchemaMeta):
 
     @post_load
     def make(self, data, **kwargs):
-        from azure.ai.ml.entities._notification.notification import Notification
+        from azure.ai.ml.entities._notification.notification import _Notification
 
-        return Notification(**data)
+        return _Notification(**data)

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

@@ -137,6 +137,20 @@
 from ._workspace.workspace import Workspace
 from ._workspace.workspace_keys import ContainerRegistryCredential, NotebookAccessKeys, WorkspaceKeys
 
+from ._assets._artifacts.feature_set import _FeatureSet
+from ._workspace.compute_runtime import _ComputeRuntime
+from ._workspace.feature_store_settings import _FeatureStoreSettings
+from ._feature_store_entity.feature_store_entity import _FeatureStoreEntity
+from ._feature_store_entity.data_column import _DataColumn
+from ._feature_store_entity.data_column_type import _DataColumnType
+from ._feature_set.feature_set_specification import _FeatureSetSpecification
+from ._feature_set.materialization_compute_resource import _MaterializationComputeResource
+from ._feature_set.materialization_settings import _MaterializationSettings
+from ._feature_set.materialization_type import _MaterializationType
+from ._feature_store.feature_store import _FeatureStore
+from ._feature_store.materialization_store import _MaterializationStore
+from ._notification.notification import _Notification
+
 # TODO: enable in PuP
 # from ._job.import_job import ImportJob
 # from ._component.import_component import ImportComponent
@@ -272,6 +286,19 @@
     "AutoScaleSettings",
     "AutoPauseSettings",
     "WorkspaceModelReference",
+    "_FeatureSet",
+    "_ComputeRuntime",
+    "_FeatureStoreSettings",
+    "_FeatureStoreEntity",
+    "_DataColumn",
+    "_DataColumnType",
+    "_FeatureSetSpecification",
+    "_MaterializationComputeResource",
+    "_MaterializationSettings",
+    "_MaterializationType",
+    "_FeatureStore",
+    "_MaterializationStore",
+    "_Notification",
     # builders
     "Command",
     "Parallel",

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/__init__.py

@@ -9,8 +9,7 @@
 from ._artifacts.code import Code
 from ._artifacts.data import Data
 from ._artifacts.model import Model
-from ._artifacts.feature_set import FeatureSet
 from .environment import Environment
 from .workspace_asset_reference import WorkspaceAssetReference
 
-__all__ = ["Artifact", "Model", "Code", "Data", "Environment", "FeatureSet", "WorkspaceAssetReference"]
+__all__ = ["Artifact", "Model", "Code", "Data", "Environment", "WorkspaceAssetReference"]

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/feature_set.py

@@ -25,24 +25,24 @@
     PARAMS_OVERRIDE_KEY,
 )
 from azure.ai.ml.entities._assets import Artifact
-from azure.ai.ml.entities._feature_set.featureset_specification import FeaturesetSpecification
-from azure.ai.ml.entities._feature_set.materialization_settings import MaterializationSettings
+from azure.ai.ml.entities._feature_set.feature_set_specification import _FeatureSetSpecification
+from azure.ai.ml.entities._feature_set.materialization_settings import _MaterializationSettings
 
 from .artifact import ArtifactStorageInfo
 
 
 @experimental
-class FeatureSet(Artifact):
+class _FeatureSet(Artifact):
     def __init__(
         self,
         *,
         name: str,
         version: str,
         entities: List[str],
-        specification: FeaturesetSpecification,
+        specification: _FeatureSetSpecification,
         stage: Optional[str] = None,
         description: Optional[str] = None,
-        materialization_settings: Optional[MaterializationSettings] = None,
+        materialization_settings: Optional[_MaterializationSettings] = None,
         tags: Optional[Dict] = None,
         properties: Optional[Dict[str, str]] = None,
         **kwargs,
@@ -56,15 +56,15 @@ def __init__(
         :param entities: Specifies list of entities.
         :type entities: list[str]
         :param specification: Specifies the feature spec details.
-        :type specification: ~azure.ai.ml.entities.FeaturesetSpecification
+        :type specification: ~azure.ai.ml.entities._FeatureSetSpecification
         :param description: Description of the resource.
         :type description: str
         :param tags: Tag dictionary. Tags can be added, removed, and updated.
         :type tags: dict[str, str]
         :param properties: The asset property dictionary.
         :type properties: dict[str, str]
         :param materialization_settings: Specifies the materialization settings.
-        :type materialization_settings: ~azure.ai.ml.entities.MaterializationSettings
+        :type materialization_settings: ~azure.ai.ml.entities._MaterializationSettings
         :param kwargs: A dictionary of additional configuration parameters.
         :type kwargs: dict
         """
@@ -98,38 +98,38 @@ def _to_rest_object(self) -> FeaturesetVersion:
         return FeaturesetVersion(name=self.name, properties=featureset_version_properties)
 
     @classmethod
-    def _from_rest_object(cls, featureset_rest_object: FeaturesetVersion) -> "FeatureSet":
+    def _from_rest_object(cls, featureset_rest_object: FeaturesetVersion) -> "_FeatureSet":
         if not featureset_rest_object:
             return None
         featureset_rest_object_details: FeaturesetVersionProperties = featureset_rest_object.properties
         arm_id_object = get_arm_id_object_from_id(featureset_rest_object.id)
-        featureset = FeatureSet(
+        featureset = _FeatureSet(
             id=featureset_rest_object.id,
             name=arm_id_object.asset_name,
             version=arm_id_object.asset_version,
             description=featureset_rest_object_details.description,
             tags=featureset_rest_object_details.tags,
             properties=featureset_rest_object_details.properties,
             entities=featureset_rest_object_details.entities,
-            materialization_settings=MaterializationSettings._from_rest_object(
+            materialization_settings=_MaterializationSettings._from_rest_object(
                 featureset_rest_object_details.materialization_settings
             ),
-            specification=FeaturesetSpecification._from_rest_object(featureset_rest_object_details.specification),
+            specification=_FeatureSetSpecification._from_rest_object(featureset_rest_object_details.specification),
             stage=featureset_rest_object_details.stage,
         )
         return featureset
 
     @classmethod
-    def _from_container_rest_object(cls, rest_obj: FeaturesetContainer) -> "FeatureSet":
+    def _from_container_rest_object(cls, rest_obj: FeaturesetContainer) -> "_FeatureSet":
         rest_object_details: FeaturesetContainerProperties = rest_obj.properties
         arm_id_object = get_arm_id_object_from_id(rest_obj.id)
-        featureset = FeatureSet(
+        featureset = _FeatureSet(
             name=arm_id_object.asset_name,
             description=rest_object_details.description,
             tags=rest_object_details.tags,
             properties=rest_object_details.properties,
             entities=[],
-            specification=FeaturesetSpecification(),
+            specification=_FeatureSetSpecification(),
             version="",
         )
         featureset.latest_version = rest_object_details.latest_version
@@ -142,15 +142,15 @@ def _load(
         yaml_path: Optional[Union[PathLike, str]] = None,
         params_override: Optional[list] = None,
         **kwargs,
-    ) -> "FeatureSet":
+    ) -> "_FeatureSet":
         data = data or {}
         params_override = params_override or []
         context = {
             BASE_PATH_CONTEXT_KEY: Path(yaml_path).parent if yaml_path else Path("./"),
             PARAMS_OVERRIDE_KEY: params_override,
         }
         loaded_schema = load_from_dict(FeatureSetSchema, data, context, **kwargs)
-        return FeatureSet(**loaded_schema)
+        return _FeatureSet(**loaded_schema)
 
     def _to_dict(self) -> Dict:
         # pylint: disable=no-member