Updating Feature store classes name and operations for private preview

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

@@ -24,16 +24,16 @@
     load_data,
     load_datastore,
     load_environment,
-    load_feature_set,
-    load_feature_store,
-    load_feature_store_entity,
     load_job,
     load_model,
     load_online_deployment,
     load_online_endpoint,
     load_registry,
     load_workspace,
     load_workspace_connection,
+    _load_feature_set,
+    _load_feature_store,
+    _load_feature_store_entity,
 )
 
 module_logger = logging.getLogger(__name__)
@@ -60,15 +60,15 @@
     "load_datastore",
     "load_model",
     "load_environment",
-    "load_feature_set",
-    "load_feature_store",
-    "load_feature_store_entity",
     "load_job",
     "load_online_deployment",
     "load_online_endpoint",
     "load_workspace",
     "load_registry",
     "load_workspace_connection",
+    "_load_feature_set",
+    "_load_feature_store",
+    "_load_feature_store_entity",
 ]
 
 __version__ = VERSION

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,
@@ -95,6 +95,9 @@
     WorkspaceConnectionsOperations,
     WorkspaceOperations,
     WorkspaceOutboundRuleOperations,
+    _FeatureSetOperations,
+    _FeatureStoreEntityOperations,
+    _FeatureStoreOperations,
 )
 from azure.ai.ml.operations._code_operations import CodeOperations
 from azure.ai.ml.operations._local_deployment_helper import _LocalDeploymentHelper
@@ -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,8 @@
 from ._artifacts.code import Code
 from ._artifacts.data import Data
 from ._artifacts.model import Model
-from ._artifacts.feature_set import FeatureSet
+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", "_FeatureSet", "WorkspaceAssetReference"]