diff --git a/google/cloud/aiplatform/__init__.py b/google/cloud/aiplatform/__init__.py index 76d9b6efe9..3e206a5538 100644 --- a/google/cloud/aiplatform/__init__.py +++ b/google/cloud/aiplatform/__init__.py @@ -33,6 +33,11 @@ from google.cloud.aiplatform import explain from google.cloud.aiplatform import gapic from google.cloud.aiplatform import hyperparameter_tuning +from google.cloud.aiplatform.featurestore import ( + EntityType, + Feature, + Featurestore, +) from google.cloud.aiplatform.metadata import metadata from google.cloud.aiplatform.models import Endpoint from google.cloud.aiplatform.models import Model @@ -92,6 +97,9 @@ "CustomContainerTrainingJob", "CustomPythonPackageTrainingJob", "Endpoint", + "EntityType", + "Feature", + "Featurestore", "ImageDataset", "HyperparameterTuningJob", "Model", diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index c4eb2e4853..2ee4bf4635 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -898,6 +898,7 @@ def _list( project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, + parent: Optional[str] = None, ) -> List[VertexAiResourceNoun]: """Private method to list all instances of this Vertex AI Resource, takes a `cls_filter` arg to filter to a particular SDK resource @@ -925,6 +926,8 @@ def _list( credentials (auth_credentials.Credentials): Optional. Custom credentials to use to retrieve list. Overrides credentials set in aiplatform.init. + parent (str): + Optional. The parent resource name if any to retrieve resource list from. Returns: List[VertexAiResourceNoun] - A list of SDK resource objects @@ -934,12 +937,13 @@ def _list( ) # Fetch credentials once and re-use for all `_empty_constructor()` calls - creds = initializer.global_config.credentials + creds = resource.credentials resource_list_method = getattr(resource.api_client, resource._list_method) list_request = { - "parent": initializer.global_config.common_location_path( + "parent": parent + or initializer.global_config.common_location_path( project=project, location=location ), "filter": filter, @@ -1029,6 +1033,7 @@ def list( project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, + parent: Optional[str] = None, ) -> List[VertexAiResourceNoun]: """List all instances of this Vertex AI Resource. @@ -1057,6 +1062,8 @@ def list( credentials (auth_credentials.Credentials): Optional. Custom credentials to use to retrieve list. Overrides credentials set in aiplatform.init. + parent (str): + Optional. The parent resource name if any to retrieve list from. Returns: List[VertexAiResourceNoun] - A list of SDK resource objects @@ -1068,6 +1075,7 @@ def list( project=project, location=location, credentials=credentials, + parent=parent, ) @optional_sync() diff --git a/google/cloud/aiplatform/compat/__init__.py b/google/cloud/aiplatform/compat/__init__.py index 057ba9344f..d435cf9578 100644 --- a/google/cloud/aiplatform/compat/__init__.py +++ b/google/cloud/aiplatform/compat/__init__.py @@ -27,6 +27,10 @@ services.dataset_service_client = services.dataset_service_client_v1beta1 services.endpoint_service_client = services.endpoint_service_client_v1beta1 + services.featurestore_online_serving_service_client = ( + services.featurestore_online_serving_service_client_v1beta1 + ) + services.featurestore_service_client = services.featurestore_service_client_v1beta1 services.job_service_client = services.job_service_client_v1beta1 services.model_service_client = services.model_service_client_v1beta1 services.pipeline_service_client = services.pipeline_service_client_v1beta1 @@ -53,11 +57,19 @@ types.encryption_spec = types.encryption_spec_v1beta1 types.endpoint = types.endpoint_v1beta1 types.endpoint_service = types.endpoint_service_v1beta1 + types.entity_type = types.entity_type_v1beta1 types.env_var = types.env_var_v1beta1 types.event = types.event_v1beta1 types.execution = types.execution_v1beta1 types.explanation = types.explanation_v1beta1 types.explanation_metadata = types.explanation_metadata_v1beta1 + types.feature = types.feature_v1beta1 + types.feature_monitoring_stats = types.feature_monitoring_stats_v1beta1 + types.feature_selector = types.feature_selector_v1beta1 + types.featurestore = types.featurestore_v1beta1 + types.featurestore_monitoring = types.featurestore_monitoring_v1beta1 + types.featurestore_online_service = types.featurestore_online_service_v1beta1 + types.featurestore_service = types.featurestore_service_v1beta1 types.hyperparameter_tuning_job = types.hyperparameter_tuning_job_v1beta1 types.io = types.io_v1beta1 types.job_service = types.job_service_v1beta1 @@ -90,6 +102,10 @@ services.dataset_service_client = services.dataset_service_client_v1 services.endpoint_service_client = services.endpoint_service_client_v1 + services.featurestore_online_serving_service_client = ( + services.featurestore_online_serving_service_client_v1 + ) + services.featurestore_service_client = services.featurestore_service_client_v1 services.job_service_client = services.job_service_client_v1 services.model_service_client = services.model_service_client_v1 services.pipeline_service_client = services.pipeline_service_client_v1 @@ -113,11 +129,18 @@ types.encryption_spec = types.encryption_spec_v1 types.endpoint = types.endpoint_v1 types.endpoint_service = types.endpoint_service_v1 + types.entity_type = types.entity_type_v1 types.env_var = types.env_var_v1 types.event = types.event_v1 types.execution = types.execution_v1 types.explanation = types.explanation_v1 types.explanation_metadata = types.explanation_metadata_v1 + types.feature = types.feature_v1 + types.feature_monitoring_stats = types.feature_monitoring_stats_v1 + types.feature_selector = types.feature_selector_v1 + types.featurestore = types.featurestore_v1 + types.featurestore_online_service = types.featurestore_online_service_v1 + types.featurestore_service = types.featurestore_service_v1 types.hyperparameter_tuning_job = types.hyperparameter_tuning_job_v1 types.io = types.io_v1 types.job_service = types.job_service_v1 diff --git a/google/cloud/aiplatform/compat/services/__init__.py b/google/cloud/aiplatform/compat/services/__init__.py index b0feaf6f0a..f8545a688c 100644 --- a/google/cloud/aiplatform/compat/services/__init__.py +++ b/google/cloud/aiplatform/compat/services/__init__.py @@ -21,6 +21,12 @@ from google.cloud.aiplatform_v1beta1.services.endpoint_service import ( client as endpoint_service_client_v1beta1, ) +from google.cloud.aiplatform_v1beta1.services.featurestore_online_serving_service import ( + client as featurestore_online_serving_service_client_v1beta1, +) +from google.cloud.aiplatform_v1beta1.services.featurestore_service import ( + client as featurestore_service_client_v1beta1, +) from google.cloud.aiplatform_v1beta1.services.job_service import ( client as job_service_client_v1beta1, ) @@ -49,6 +55,12 @@ from google.cloud.aiplatform_v1.services.endpoint_service import ( client as endpoint_service_client_v1, ) +from google.cloud.aiplatform_v1.services.featurestore_online_serving_service import ( + client as featurestore_online_serving_service_client_v1, +) +from google.cloud.aiplatform_v1.services.featurestore_service import ( + client as featurestore_service_client_v1, +) from google.cloud.aiplatform_v1.services.job_service import ( client as job_service_client_v1, ) @@ -75,6 +87,8 @@ # v1 dataset_service_client_v1, endpoint_service_client_v1, + featurestore_online_serving_service_client_v1beta1, + featurestore_service_client_v1beta1, job_service_client_v1, metadata_service_client_v1, model_service_client_v1, @@ -85,6 +99,8 @@ # v1beta1 dataset_service_client_v1beta1, endpoint_service_client_v1beta1, + featurestore_online_serving_service_client_v1, + featurestore_service_client_v1, job_service_client_v1beta1, model_service_client_v1beta1, pipeline_service_client_v1beta1, diff --git a/google/cloud/aiplatform/compat/types/__init__.py b/google/cloud/aiplatform/compat/types/__init__.py index 4d64c50236..a2a82fdca9 100644 --- a/google/cloud/aiplatform/compat/types/__init__.py +++ b/google/cloud/aiplatform/compat/types/__init__.py @@ -32,11 +32,19 @@ encryption_spec as encryption_spec_v1beta1, endpoint as endpoint_v1beta1, endpoint_service as endpoint_service_v1beta1, + entity_type as entity_type_v1beta1, env_var as env_var_v1beta1, event as event_v1beta1, execution as execution_v1beta1, explanation as explanation_v1beta1, explanation_metadata as explanation_metadata_v1beta1, + feature as feature_v1beta1, + feature_monitoring_stats as feature_monitoring_stats_v1beta1, + feature_selector as feature_selector_v1beta1, + featurestore as featurestore_v1beta1, + featurestore_monitoring as featurestore_monitoring_v1beta1, + featurestore_online_service as featurestore_online_service_v1beta1, + featurestore_service as featurestore_service_v1beta1, hyperparameter_tuning_job as hyperparameter_tuning_job_v1beta1, io as io_v1beta1, job_service as job_service_v1beta1, @@ -82,11 +90,18 @@ encryption_spec as encryption_spec_v1, endpoint as endpoint_v1, endpoint_service as endpoint_service_v1, + entity_type as entity_type_v1, env_var as env_var_v1, event as event_v1, execution as execution_v1, explanation as explanation_v1, explanation_metadata as explanation_metadata_v1, + feature as feature_v1, + feature_monitoring_stats as feature_monitoring_stats_v1, + feature_selector as feature_selector_v1, + featurestore as featurestore_v1, + featurestore_online_service as featurestore_online_service_v1, + featurestore_service as featurestore_service_v1, hyperparameter_tuning_job as hyperparameter_tuning_job_v1, io as io_v1, job_service as job_service_v1, @@ -134,11 +149,18 @@ encryption_spec_v1, endpoint_v1, endpoint_service_v1, + entity_type_v1, env_var_v1, event_v1, execution_v1, explanation_v1, explanation_metadata_v1, + feature_v1, + feature_monitoring_stats_v1, + feature_selector_v1, + featurestore_v1, + featurestore_online_service_v1, + featurestore_service_v1, hyperparameter_tuning_job_v1, io_v1, job_service_v1, @@ -182,11 +204,19 @@ encryption_spec_v1beta1, endpoint_v1beta1, endpoint_service_v1beta1, + entity_type_v1beta1, env_var_v1beta1, event_v1beta1, execution_v1beta1, explanation_v1beta1, explanation_metadata_v1beta1, + feature_v1beta1, + feature_monitoring_stats_v1beta1, + feature_selector_v1beta1, + featurestore_v1beta1, + featurestore_monitoring_v1beta1, + featurestore_online_service_v1beta1, + featurestore_service_v1beta1, hyperparameter_tuning_job_v1beta1, io_v1beta1, job_service_v1beta1, diff --git a/google/cloud/aiplatform/featurestore/__init__.py b/google/cloud/aiplatform/featurestore/__init__.py new file mode 100644 index 0000000000..883f72dd26 --- /dev/null +++ b/google/cloud/aiplatform/featurestore/__init__.py @@ -0,0 +1,26 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from google.cloud.aiplatform.featurestore.entity_type import EntityType +from google.cloud.aiplatform.featurestore.feature import Feature +from google.cloud.aiplatform.featurestore.featurestore import Featurestore + +__all__ = ( + "EntityType", + "Feature", + "Featurestore", +) diff --git a/google/cloud/aiplatform/featurestore/entity_type.py b/google/cloud/aiplatform/featurestore/entity_type.py new file mode 100644 index 0000000000..327bf1931d --- /dev/null +++ b/google/cloud/aiplatform/featurestore/entity_type.py @@ -0,0 +1,383 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import Dict, List, Optional, Sequence, Tuple + +from google.auth import credentials as auth_credentials +from google.protobuf import field_mask_pb2 + +from google.cloud.aiplatform import base +from google.cloud.aiplatform.compat.types import entity_type as gca_entity_type +from google.cloud.aiplatform import featurestore +from google.cloud.aiplatform import utils +from google.cloud.aiplatform.utils import featurestore_utils + +_LOGGER = base.Logger(__name__) +_ALL_FEATURE_IDS = "*" + + +class EntityType(base.VertexAiResourceNounWithFutureManager): + """Managed entityType resource for Vertex AI.""" + + client_class = utils.FeaturestoreClientWithOverride + + _is_client_prediction_client = False + _resource_noun = None + _getter_method = "get_entity_type" + _list_method = "list_entity_types" + _delete_method = "delete_entity_type" + + def __init__( + self, + entity_type_name: str, + featurestore_id: Optional[str] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """Retrieves an existing managed entityType given an entityType resource name or an entity_type ID. + + Example Usage: + + my_entity_type = aiplatform.EntityType( + entity_type_name='projects/123/locations/us-central1/featurestores/my_featurestore_id/\ + entityTypes/my_entity_type_id' + ) + or + my_entity_type = aiplatform.EntityType( + entity_type_name='my_entity_type_id', + featurestore_id='my_featurestore_id', + ) + + Args: + entity_type_name (str): + Required. A fully-qualified entityType resource name or an entity_type ID. + Example: "projects/123/locations/us-central1/featurestores/my_featurestore_id/entityTypes/my_entity_type_id" + or "my_entity_type_id" when project and location are initialized or passed, with featurestore_id passed. + featurestore_id (str): + Optional. Featurestore ID to retrieve entityType from, when entity_type_name is passed as entity_type ID. + project (str): + Optional. Project to retrieve entityType from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve entityType from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this EntityType. Overrides + credentials set in aiplatform.init. + """ + + ( + featurestore_id, + _, + ) = featurestore_utils.validate_and_get_entity_type_resource_ids( + entity_type_name=entity_type_name, featurestore_id=featurestore_id + ) + + # TODO(b/208269923): Temporary workaround, update when base class supports nested resource + self._resource_noun = f"featurestores/{featurestore_id}/entityTypes" + + super().__init__( + project=project, + location=location, + credentials=credentials, + resource_name=entity_type_name, + ) + self._gca_resource = self._get_gca_resource(resource_name=entity_type_name) + + @property + def featurestore_name(self) -> str: + """Full qualified resource name of the managed featurestore in which this EntityType is.""" + entity_type_name_components = featurestore_utils.CompatFeaturestoreServiceClient.parse_entity_type_path( + path=self.resource_name + ) + return featurestore_utils.CompatFeaturestoreServiceClient.featurestore_path( + project=entity_type_name_components["project"], + location=entity_type_name_components["location"], + featurestore=entity_type_name_components["featurestore"], + ) + + def get_featurestore(self) -> "featurestore.Featurestore": + """Retrieves the managed featurestore in which this EntityType is. + + Returns: + featurestore.Featurestore - The managed featurestore in which this EntityType is. + """ + return featurestore.Featurestore(self.featurestore_name) + + def get_feature(self, feature_id: str) -> "featurestore.Feature": + """Retrieves an existing managed feature in this EntityType. + + Args: + feature_id (str): + Required. The managed feature resource ID in this EntityType. + Returns: + featurestore.Feature - The managed feature resource object. + """ + entity_type_name_components = featurestore_utils.CompatFeaturestoreServiceClient.parse_entity_type_path( + path=self.resource_name + ) + + return featurestore.Feature( + feature_name=featurestore_utils.CompatFeaturestoreServiceClient.feature_path( + project=entity_type_name_components["project"], + location=entity_type_name_components["location"], + featurestore=entity_type_name_components["featurestore"], + entity_type=entity_type_name_components["entity_type"], + feature=feature_id, + ) + ) + + def update( + self, + description: Optional[str] = None, + labels: Optional[Dict[str, str]] = None, + request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + ) -> "EntityType": + """Updates an existing managed entityType resource. + + Example Usage: + + my_entity_type = aiplatform.EntityType( + entity_type_name='my_entity_type_id', + featurestore_id='my_featurestore_id', + ) + my_entity_type.update( + description='update my description', + ) + + Args: + description (str): + Optional. Description of the EntityType. + labels (Dict[str, str]): + Optional. The labels with user-defined + metadata to organize your EntityTypes. + Label keys and values can be no longer than 64 + characters (Unicode codepoints), can only + contain lowercase letters, numeric characters, + underscores and dashes. International characters + are allowed. + See https://goo.gl/xmQnxf for more information + on and examples of labels. No more than 64 user + labels can be associated with one Feature + (System labels are excluded)." + System reserved label keys are prefixed with + "aiplatform.googleapis.com/" and are immutable. + request_metadata (Sequence[Tuple[str, str]]): + Optional. Strings which should be sent along with the request as metadata. + Returns: + EntityType - The updated entityType resource object. + """ + update_mask = list() + + if description: + update_mask.append("description") + + if labels: + utils.validate_labels(labels) + update_mask.append("labels") + + update_mask = field_mask_pb2.FieldMask(paths=update_mask) + + gapic_entity_type = gca_entity_type.EntityType( + name=self.resource_name, description=description, labels=labels, + ) + + _LOGGER.log_action_start_against_resource( + "Updating", "entityType", self, + ) + + update_entity_type_lro = self.api_client.update_entity_type( + entity_type=gapic_entity_type, + update_mask=update_mask, + metadata=request_metadata, + ) + + _LOGGER.log_action_started_against_resource_with_lro( + "Update", "entityType", self.__class__, update_entity_type_lro + ) + + update_entity_type_lro.result() + + _LOGGER.log_action_completed_against_resource("entityType", "updated", self) + + return self + + @classmethod + def list( + cls, + featurestore_name: str, + filter: Optional[str] = None, + order_by: Optional[str] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> List["EntityType"]: + """Lists existing managed entityType resources in a featurestore, given a featurestore resource name or a featurestore ID. + + Example Usage: + + my_entityTypes = aiplatform.EntityType.list( + featurestore_name='projects/123/locations/us-central1/featurestores/my_featurestore_id' + ) + or + my_entityTypes = aiplatform.EntityType.list( + featurestore_name='my_featurestore_id' + ) + + Args: + featurestore_name (str): + Required. A fully-qualified featurestore resource name or a featurestore ID to list entityTypes in + Example: "projects/123/locations/us-central1/featurestores/my_featurestore_id" + or "my_featurestore_id" when project and location are initialized or passed. + filter (str): + Optional. Lists the EntityTypes that match the filter expression. The + following filters are supported: + + - ``create_time``: Supports ``=``, ``!=``, ``<``, ``>``, + ``>=``, and ``<=`` comparisons. Values must be in RFC + 3339 format. + - ``update_time``: Supports ``=``, ``!=``, ``<``, ``>``, + ``>=``, and ``<=`` comparisons. Values must be in RFC + 3339 format. + - ``labels``: Supports key-value equality as well as key + presence. + + Examples: + + - ``create_time > \"2020-01-31T15:30:00.000000Z\" OR update_time > \"2020-01-31T15:30:00.000000Z\"`` + --> EntityTypes created or updated after + 2020-01-31T15:30:00.000000Z. + - ``labels.active = yes AND labels.env = prod`` --> + EntityTypes having both (active: yes) and (env: prod) + labels. + - ``labels.env: *`` --> Any EntityType which has a label + with 'env' as the key. + order_by (str): + Optional. A comma-separated list of fields to order by, sorted in + ascending order. Use "desc" after a field name for + descending. + + Supported fields: + + - ``entity_type_id`` + - ``create_time`` + - ``update_time`` + project (str): + Optional. Project to list entityTypes in. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to list entityTypes in. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to list entityTypes. Overrides + credentials set in aiplatform.init. + + Returns: + List[EntityType] - A list of managed entityType resource objects + """ + + return cls._list( + filter=filter, + order_by=order_by, + project=project, + location=location, + credentials=credentials, + parent=utils.full_resource_name( + resource_name=featurestore_name, + resource_noun="featurestores", + project=project, + location=location, + ), + ) + + def list_features( + self, filter: Optional[str] = None, order_by: Optional[str] = None, + ) -> List["featurestore.Feature"]: + """Lists existing managed feature resources in this EntityType. + + Example Usage: + + my_entity_type = aiplatform.EntityType( + entity_type_name='my_entity_type_id', + featurestore_id='my_featurestore_id', + ) + my_entityType.list_features() + + Args: + filter (str): + Optional. Lists the Features that match the filter expression. The + following filters are supported: + + - ``value_type``: Supports = and != comparisons. + - ``create_time``: Supports =, !=, <, >, >=, and <= + comparisons. Values must be in RFC 3339 format. + - ``update_time``: Supports =, !=, <, >, >=, and <= + comparisons. Values must be in RFC 3339 format. + - ``labels``: Supports key-value equality as well as key + presence. + + Examples: + + - ``value_type = DOUBLE`` --> Features whose type is + DOUBLE. + - ``create_time > \"2020-01-31T15:30:00.000000Z\" OR update_time > \"2020-01-31T15:30:00.000000Z\"`` + --> EntityTypes created or updated after + 2020-01-31T15:30:00.000000Z. + - ``labels.active = yes AND labels.env = prod`` --> + Features having both (active: yes) and (env: prod) + labels. + - ``labels.env: *`` --> Any Feature which has a label with + 'env' as the key. + order_by (str): + Optional. A comma-separated list of fields to order by, sorted in + ascending order. Use "desc" after a field name for + descending. Supported fields: + + - ``feature_id`` + - ``value_type`` + - ``create_time`` + - ``update_time`` + + Returns: + List[featurestore.Feature] - A list of managed feature resource objects. + """ + return featurestore.Feature.list( + entity_type_name=self.resource_name, filter=filter, order_by=order_by, + ) + + @base.optional_sync() + def delete_features(self, feature_ids: List[str], sync: bool = True,) -> None: + """Deletes feature resources in this EntityType given their feature IDs. + WARNING: This deletion is permanent. + + Args: + feature_ids (List[str]): + Required. The list of feature IDs to be deleted. + sync (bool): + Optional. Whether to execute this deletion synchronously. If False, this method + will be executed in concurrent Future and any downstream object will + be immediately returned and synced when the Future has completed. + """ + features = [] + for feature_id in feature_ids: + feature = self.get_feature(feature_id=feature_id) + feature.delete(sync=False) + features.append(feature) + + for feature in features: + feature.wait() diff --git a/google/cloud/aiplatform/featurestore/feature.py b/google/cloud/aiplatform/featurestore/feature.py new file mode 100644 index 0000000000..ab199d0c57 --- /dev/null +++ b/google/cloud/aiplatform/featurestore/feature.py @@ -0,0 +1,463 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import Dict, List, Optional, Sequence, Tuple + +from google.auth import credentials as auth_credentials +from google.protobuf import field_mask_pb2 + +from google.cloud.aiplatform import base +from google.cloud.aiplatform.compat.types import feature as gca_feature +from google.cloud.aiplatform import featurestore +from google.cloud.aiplatform import initializer +from google.cloud.aiplatform import utils +from google.cloud.aiplatform.utils import featurestore_utils + +_LOGGER = base.Logger(__name__) + + +class Feature(base.VertexAiResourceNounWithFutureManager): + """Managed feature resource for Vertex AI.""" + + client_class = utils.FeaturestoreClientWithOverride + + _is_client_prediction_client = False + _resource_noun = None + _getter_method = "get_feature" + _list_method = "list_features" + _delete_method = "delete_feature" + + def __init__( + self, + feature_name: str, + featurestore_id: Optional[str] = None, + entity_type_id: Optional[str] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """Retrieves an existing managed feature given a feature resource name or a feature ID. + + Example Usage: + + my_feature = aiplatform.Feature( + feature_name='projects/123/locations/us-central1/featurestores/my_featurestore_id/\ + entityTypes/my_entity_type_id/features/my_feature_id' + ) + or + my_feature = aiplatform.Feature( + feature_name='my_feature_id', + featurestore_id='my_featurestore_id', + entity_type_id='my_entity_type_id', + ) + + Args: + feature_name (str): + Required. A fully-qualified feature resource name or a feature ID. + Example: "projects/123/locations/us-central1/featurestores/my_featurestore_id/entityTypes/my_entity_type_id/features/my_feature_id" + or "my_feature_id" when project and location are initialized or passed, with featurestore_id and entity_type_id passed. + featurestore_id (str): + Optional. Featurestore ID to retrieve feature from, when feature_name is passed as Feature ID. + entity_type_id (str): + Optional. EntityType ID to retrieve feature from, when feature_name is passed as Feature ID. + project (str): + Optional. Project to retrieve feature from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve feature from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this Feature. Overrides + credentials set in aiplatform.init. + """ + ( + featurestore_id, + entity_type_id, + _, + ) = featurestore_utils.validate_and_get_feature_resource_ids( + feature_name=feature_name, + entity_type_id=entity_type_id, + featurestore_id=featurestore_id, + ) + + # TODO(b/208269923): Temporary workaround, update when base class supports nested resource + self._resource_noun = ( + f"featurestores/{featurestore_id}/entityTypes/{entity_type_id}/features" + ) + + super().__init__( + project=project, + location=location, + credentials=credentials, + resource_name=feature_name, + ) + self._gca_resource = self._get_gca_resource(resource_name=feature_name) + + @property + def featurestore_name(self) -> str: + """Full qualified resource name of the managed featurestore in which this Feature is.""" + feature_path_components = featurestore_utils.CompatFeaturestoreServiceClient.parse_feature_path( + path=self.resource_name + ) + + return featurestore_utils.CompatFeaturestoreServiceClient.featurestore_path( + project=feature_path_components["project"], + location=feature_path_components["location"], + featurestore=feature_path_components["featurestore"], + ) + + def get_featurestore(self) -> "featurestore.Featurestore": + """Retrieves the managed featurestore in which this Feature is. + + Returns: + featurestore.Featurestore - The managed featurestore in which this Feature is. + """ + return featurestore.Featurestore(featurestore_name=self.featurestore_name) + + @property + def entity_type_name(self) -> str: + """Full qualified resource name of the managed entityType in which this Feature is.""" + feature_path_components = featurestore_utils.CompatFeaturestoreServiceClient.parse_feature_path( + path=self.resource_name + ) + + return featurestore_utils.CompatFeaturestoreServiceClient.entity_type_path( + project=feature_path_components["project"], + location=feature_path_components["location"], + featurestore=feature_path_components["featurestore"], + entity_type=feature_path_components["entity_type"], + ) + + def get_entity_type(self) -> "featurestore.EntityType": + """Retrieves the managed entityType in which this Feature is. + + Returns: + featurestore.EntityType - The managed entityType in which this Feature is. + """ + return featurestore.EntityType(entity_type_name=self.entity_type_name) + + def update( + self, + description: Optional[str] = None, + labels: Optional[Dict[str, str]] = None, + request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + ) -> "Feature": + """Updates an existing managed feature resource. + + Example Usage: + + my_feature = aiplatform.Feature( + feature_name='my_feature_id', + featurestore_id='my_featurestore_id', + entity_type_id='my_entity_type_id', + ) + my_feature.update( + description='update my description', + ) + + Args: + description (str): + Optional. Description of the Feature. + labels (Dict[str, str]): + Optional. The labels with user-defined + metadata to organize your Features. + Label keys and values can be no longer than 64 + characters (Unicode codepoints), can only + contain lowercase letters, numeric characters, + underscores and dashes. International characters + are allowed. + See https://goo.gl/xmQnxf for more information + on and examples of labels. No more than 64 user + labels can be associated with one Feature + (System labels are excluded)." + System reserved label keys are prefixed with + "aiplatform.googleapis.com/" and are immutable. + request_metadata (Sequence[Tuple[str, str]]): + Optional. Strings which should be sent along with the request as metadata. + + Returns: + Feature - The updated feature resource object. + """ + update_mask = list() + + if description: + update_mask.append("description") + + if labels: + utils.validate_labels(labels) + update_mask.append("labels") + + update_mask = field_mask_pb2.FieldMask(paths=update_mask) + + gapic_feature = gca_feature.Feature( + name=self.resource_name, description=description, labels=labels, + ) + + _LOGGER.log_action_start_against_resource( + "Updating", "feature", self, + ) + + update_feature_lro = self.api_client.update_feature( + feature=gapic_feature, update_mask=update_mask, metadata=request_metadata, + ) + + _LOGGER.log_action_started_against_resource_with_lro( + "Update", "feature", self.__class__, update_feature_lro + ) + + update_feature_lro.result() + + _LOGGER.log_action_completed_against_resource("feature", "updated", self) + + return self + + @classmethod + def list( + cls, + entity_type_name: str, + featurestore_id: Optional[str] = None, + filter: Optional[str] = None, + order_by: Optional[str] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> List["Feature"]: + """Lists existing managed feature resources in an entityType, given an entityType resource name or an entity_type ID. + + Example Usage: + + my_features = aiplatform.Feature.list( + entity_type_name='projects/123/locations/us-central1/featurestores/my_featurestore_id/\ + entityTypes/my_entity_type_id' + ) + or + my_features = aiplatform.Feature.list( + entity_type_name='my_entity_type_id', + featurestore_id='my_featurestore_id', + ) + + Args: + entity_type_name (str): + Required. A fully-qualified entityType resource name or an entity_type ID to list features in + Example: "projects/123/locations/us-central1/featurestores/my_featurestore_id/entityTypes/my_entity_type_id" + or "my_entity_type_id" when project and location are initialized or passed, with featurestore_id passed. + featurestore_id (str): + Optional. Featurestore ID to list features in, when entity_type_name is passed as entity_type ID. + filter (str): + Optional. Lists the Features that match the filter expression. The + following filters are supported: + + - ``value_type``: Supports = and != comparisons. + - ``create_time``: Supports =, !=, <, >, >=, and <= + comparisons. Values must be in RFC 3339 format. + - ``update_time``: Supports =, !=, <, >, >=, and <= + comparisons. Values must be in RFC 3339 format. + - ``labels``: Supports key-value equality as well as key + presence. + + Examples: + + - ``value_type = DOUBLE`` --> Features whose type is + DOUBLE. + - ``create_time > \"2020-01-31T15:30:00.000000Z\" OR update_time > \"2020-01-31T15:30:00.000000Z\"`` + --> EntityTypes created or updated after + 2020-01-31T15:30:00.000000Z. + - ``labels.active = yes AND labels.env = prod`` --> + Features having both (active: yes) and (env: prod) + labels. + - ``labels.env: *`` --> Any Feature which has a label with + 'env' as the key. + order_by (str): + Optional. A comma-separated list of fields to order by, sorted in + ascending order. Use "desc" after a field name for + descending. Supported fields: + + - ``feature_id`` + - ``value_type`` + - ``create_time`` + - ``update_time`` + project (str): + Optional. Project to list features in. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to list features in. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to list features. Overrides + credentials set in aiplatform.init. + + Returns: + List[Feature] - A list of managed feature resource objects + """ + ( + featurestore_id, + entity_type_id, + ) = featurestore_utils.validate_and_get_entity_type_resource_ids( + entity_type_name=entity_type_name, featurestore_id=featurestore_id, + ) + + return cls._list( + filter=filter, + order_by=order_by, + project=project, + location=location, + credentials=credentials, + parent=utils.full_resource_name( + resource_name=entity_type_name, + resource_noun=f"featurestores/{featurestore_id}/entityTypes", + project=project, + location=location, + ), + ) + + @classmethod + def search( + cls, + query: Optional[str] = None, + page_size: Optional[int] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> List["Feature"]: + """Searches existing managed Feature resources. + + Example Usage: + + my_features = aiplatform.Feature.search() + + Args: + query (str): + Optional. Query string that is a conjunction of field-restricted + queries and/or field-restricted filters. + Field-restricted queries and filters can be combined + using ``AND`` to form a conjunction. + + A field query is in the form FIELD:QUERY. This + implicitly checks if QUERY exists as a substring within + Feature's FIELD. The QUERY and the FIELD are converted + to a sequence of words (i.e. tokens) for comparison. + This is done by: + + - Removing leading/trailing whitespace and tokenizing + the search value. Characters that are not one of + alphanumeric ``[a-zA-Z0-9]``, underscore ``_``, or + asterisk ``*`` are treated as delimiters for tokens. + ``*`` is treated as a wildcard that matches + characters within a token. + - Ignoring case. + - Prepending an asterisk to the first and appending an + asterisk to the last token in QUERY. + + A QUERY must be either a singular token or a phrase. A + phrase is one or multiple words enclosed in double + quotation marks ("). With phrases, the order of the + words is important. Words in the phrase must be matching + in order and consecutively. + + Supported FIELDs for field-restricted queries: + + - ``feature_id`` + - ``description`` + - ``entity_type_id`` + + Examples: + + - ``feature_id: foo`` --> Matches a Feature with ID + containing the substring ``foo`` (eg. ``foo``, + ``foofeature``, ``barfoo``). + - ``feature_id: foo*feature`` --> Matches a Feature + with ID containing the substring ``foo*feature`` (eg. + ``foobarfeature``). + - ``feature_id: foo AND description: bar`` --> Matches + a Feature with ID containing the substring ``foo`` + and description containing the substring ``bar``. + + Besides field queries, the following exact-match filters + are supported. The exact-match filters do not support + wildcards. Unlike field-restricted queries, exact-match + filters are case-sensitive. + + - ``feature_id``: Supports = comparisons. + - ``description``: Supports = comparisons. Multi-token + filters should be enclosed in quotes. + - ``entity_type_id``: Supports = comparisons. + - ``value_type``: Supports = and != comparisons. + - ``labels``: Supports key-value equality as well as + key presence. + - ``featurestore_id``: Supports = comparisons. + + Examples: + + - ``description = "foo bar"`` --> Any Feature with + description exactly equal to ``foo bar`` + - ``value_type = DOUBLE`` --> Features whose type is + DOUBLE. + - ``labels.active = yes AND labels.env = prod`` --> + Features having both (active: yes) and (env: prod) + labels. + - ``labels.env: *`` --> Any Feature which has a label + with ``env`` as the key. + + This corresponds to the ``query`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + page_size (int): + Optional. The maximum number of Features to return. The + service may return fewer than this value. If + unspecified, at most 100 Features will be + returned. The maximum value is 100; any value + greater than 100 will be coerced to 100. + project (str): + Optional. Project to list features in. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to list features in. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to list features. Overrides + credentials set in aiplatform.init. + + Returns: + List[Feature] - A list of managed feature resource objects + """ + resource = cls._empty_constructor( + project=project, location=location, credentials=credentials + ) + + # Fetch credentials once and re-use for all `_empty_constructor()` calls + creds = resource.credentials + + search_features_request = { + "location": initializer.global_config.common_location_path( + project=project, location=location + ), + "query": query, + } + + if page_size: + search_features_request["page_size"] = page_size + + resource_list = ( + resource.api_client.search_features(request=search_features_request) or [] + ) + + return [ + cls._construct_sdk_resource_from_gapic( + gapic_resource, project=project, location=location, credentials=creds + ) + for gapic_resource in resource_list + ] diff --git a/google/cloud/aiplatform/featurestore/featurestore.py b/google/cloud/aiplatform/featurestore/featurestore.py new file mode 100644 index 0000000000..d3bb0a0c11 --- /dev/null +++ b/google/cloud/aiplatform/featurestore/featurestore.py @@ -0,0 +1,324 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import Dict, List, Optional, Sequence, Tuple + +from google.auth import credentials as auth_credentials +from google.protobuf import field_mask_pb2 + +from google.cloud.aiplatform import base +from google.cloud.aiplatform.compat.types import featurestore as gca_featurestore +from google.cloud.aiplatform import featurestore +from google.cloud.aiplatform import utils +from google.cloud.aiplatform.utils import featurestore_utils + +_LOGGER = base.Logger(__name__) + + +class Featurestore(base.VertexAiResourceNounWithFutureManager): + """Managed featurestore resource for Vertex AI.""" + + client_class = utils.FeaturestoreClientWithOverride + + _is_client_prediction_client = False + _resource_noun = "featurestores" + _getter_method = "get_featurestore" + _list_method = "list_featurestores" + _delete_method = "delete_featurestore" + + def __init__( + self, + featurestore_name: str, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """Retrieves an existing managed featurestore given a featurestore resource name or a featurestore ID. + + Example Usage: + + my_featurestore = aiplatform.Featurestore( + featurestore_name='projects/123/locations/us-central1/featurestores/my_featurestore_id' + ) + or + my_featurestore = aiplatform.Featurestore( + featurestore_name='my_featurestore_id' + ) + + Args: + featurestore_name (str): + Required. A fully-qualified featurestore resource name or a featurestore ID. + Example: "projects/123/locations/us-central1/featurestores/my_featurestore_id" + or "my_featurestore_id" when project and location are initialized or passed. + project (str): + Optional. Project to retrieve featurestore from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve featurestore from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this Featurestore. Overrides + credentials set in aiplatform.init. + """ + + super().__init__( + project=project, + location=location, + credentials=credentials, + resource_name=featurestore_name, + ) + self._gca_resource = self._get_gca_resource(resource_name=featurestore_name) + + def get_entity_type(self, entity_type_id: str) -> "featurestore.EntityType": + """Retrieves an existing managed entityType in this Featurestore. + + Args: + entity_type_id (str): + Required. The managed entityType resource ID in this Featurestore. + Returns: + featurestore.EntityType - The managed entityType resource object. + """ + featurestore_name_components = featurestore_utils.CompatFeaturestoreServiceClient.parse_featurestore_path( + path=self.resource_name + ) + + return featurestore.EntityType( + entity_type_name=featurestore_utils.CompatFeaturestoreServiceClient.entity_type_path( + project=featurestore_name_components["project"], + location=featurestore_name_components["location"], + featurestore=featurestore_name_components["featurestore"], + entity_type=entity_type_id, + ) + ) + + def update( + self, + labels: Optional[Dict[str, str]] = None, + request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + ) -> "Featurestore": + """Updates an existing managed featurestore resource. + + Example Usage: + + my_featurestore = aiplatform.Featurestore( + featurestore_name='my_featurestore_id', + ) + my_featurestore.update( + labels={'update my key': 'update my value'}, + ) + + Args: + labels (Dict[str, str]): + Optional. The labels with user-defined + metadata to organize your Featurestores. + Label keys and values can be no longer than 64 + characters (Unicode codepoints), can only + contain lowercase letters, numeric characters, + underscores and dashes. International characters + are allowed. + See https://goo.gl/xmQnxf for more information + on and examples of labels. No more than 64 user + labels can be associated with one Feature + (System labels are excluded)." + System reserved label keys are prefixed with + "aiplatform.googleapis.com/" and are immutable. + request_metadata (Sequence[Tuple[str, str]]): + Optional. Strings which should be sent along with the request as metadata. + + Returns: + Featurestore - The updated featurestore resource object. + """ + + return self._update(labels=labels, request_metadata=request_metadata) + + # TODO(b/206818784): Add enable_online_store and disable_online_store methods + def update_online_store( + self, + fixed_node_count: int, + request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + ) -> "Featurestore": + """Updates the online store of an existing managed featurestore resource. + + Example Usage: + + my_featurestore = aiplatform.Featurestore( + featurestore_name='my_featurestore_id', + ) + my_featurestore.update_online_store( + fixed_node_count=2, + ) + + Args: + fixed_node_count (int): + Required. Config for online serving resources, can only update the node count to >= 1. + request_metadata (Sequence[Tuple[str, str]]): + Optional. Strings which should be sent along with the request as metadata. + + Returns: + Featurestore - The updated featurestore resource object. + """ + return self._update( + fixed_node_count=fixed_node_count, request_metadata=request_metadata + ) + + def _update( + self, + labels: Optional[Dict[str, str]] = None, + fixed_node_count: Optional[int] = None, + request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + ) -> "Featurestore": + """Updates an existing managed featurestore resource. + + Args: + labels (Dict[str, str]): + Optional. The labels with user-defined + metadata to organize your Featurestores. + Label keys and values can be no longer than 64 + characters (Unicode codepoints), can only + contain lowercase letters, numeric characters, + underscores and dashes. International characters + are allowed. + See https://goo.gl/xmQnxf for more information + on and examples of labels. No more than 64 user + labels can be associated with one Feature + (System labels are excluded)." + System reserved label keys are prefixed with + "aiplatform.googleapis.com/" and are immutable. + fixed_node_count (int): + Optional. Config for online serving resources, can only update the node count to >= 1. + request_metadata (Sequence[Tuple[str, str]]): + Optional. Strings which should be sent along with the request as metadata. + + Returns: + Featurestore - The updated featurestore resource object. + """ + update_mask = list() + + if labels: + utils.validate_labels(labels) + update_mask.append("labels") + + if fixed_node_count is not None: + update_mask.append("online_serving_config.fixed_node_count") + + update_mask = field_mask_pb2.FieldMask(paths=update_mask) + + gapic_featurestore = gca_featurestore.Featurestore( + name=self.resource_name, + labels=labels, + online_serving_config=gca_featurestore.Featurestore.OnlineServingConfig( + fixed_node_count=fixed_node_count + ), + ) + + _LOGGER.log_action_start_against_resource( + "Updating", "featurestore", self, + ) + + update_featurestore_lro = self.api_client.update_featurestore( + featurestore=gapic_featurestore, + update_mask=update_mask, + metadata=request_metadata, + ) + + _LOGGER.log_action_started_against_resource_with_lro( + "Update", "featurestore", self.__class__, update_featurestore_lro + ) + + update_featurestore_lro.result() + + _LOGGER.log_action_completed_against_resource("featurestore", "updated", self) + + return self + + def list_entity_types( + self, filter: Optional[str] = None, order_by: Optional[str] = None, + ) -> List["featurestore.EntityType"]: + """Lists existing managed entityType resources in this Featurestore. + + Example Usage: + + my_featurestore = aiplatform.Featurestore( + featurestore_name='my_featurestore_id', + ) + my_featurestore.list_entity_types() + + Args: + filter (str): + Optional. Lists the EntityTypes that match the filter expression. The + following filters are supported: + + - ``create_time``: Supports ``=``, ``!=``, ``<``, ``>``, + ``>=``, and ``<=`` comparisons. Values must be in RFC + 3339 format. + - ``update_time``: Supports ``=``, ``!=``, ``<``, ``>``, + ``>=``, and ``<=`` comparisons. Values must be in RFC + 3339 format. + - ``labels``: Supports key-value equality as well as key + presence. + + Examples: + + - ``create_time > \"2020-01-31T15:30:00.000000Z\" OR update_time > \"2020-01-31T15:30:00.000000Z\"`` + --> EntityTypes created or updated after + 2020-01-31T15:30:00.000000Z. + - ``labels.active = yes AND labels.env = prod`` --> + EntityTypes having both (active: yes) and (env: prod) + labels. + - ``labels.env: *`` --> Any EntityType which has a label + with 'env' as the key. + order_by (str): + Optional. A comma-separated list of fields to order by, sorted in + ascending order. Use "desc" after a field name for + descending. + + Supported fields: + + - ``entity_type_id`` + - ``create_time`` + - ``update_time`` + + Returns: + List[featurestore.EntityType] - A list of managed entityType resource objects. + """ + return featurestore.EntityType.list( + featurestore_name=self.resource_name, filter=filter, order_by=order_by, + ) + + @base.optional_sync() + def delete_entity_types( + self, entity_type_ids: List[str], sync: bool = True, + ) -> None: + """Deletes entity_type resources in this Featurestore given their entity_type IDs. + WARNING: This deletion is permanent. + + Args: + entity_type_ids (List[str]): + Required. The list of entity_type IDs to be deleted. + sync (bool): + Optional. Whether to execute this deletion synchronously. If False, this method + will be executed in concurrent Future and any downstream object will + be immediately returned and synced when the Future has completed. + """ + entity_types = [] + for entity_type_id in entity_type_ids: + entity_type = self.get_entity_type(entity_type_id=entity_type_id) + entity_type.delete(sync=False) + entity_types.append(entity_type) + + for entity_type in entity_types: + entity_type.wait() diff --git a/google/cloud/aiplatform/utils/__init__.py b/google/cloud/aiplatform/utils/__init__.py index 7d49d57c1e..cac1248ee7 100644 --- a/google/cloud/aiplatform/utils/__init__.py +++ b/google/cloud/aiplatform/utils/__init__.py @@ -36,6 +36,8 @@ from google.cloud.aiplatform.compat.services import ( dataset_service_client_v1beta1, endpoint_service_client_v1beta1, + featurestore_online_serving_service_client_v1beta1, + featurestore_service_client_v1beta1, job_service_client_v1beta1, metadata_service_client_v1beta1, model_service_client_v1beta1, @@ -46,6 +48,8 @@ from google.cloud.aiplatform.compat.services import ( dataset_service_client_v1, endpoint_service_client_v1, + featurestore_online_serving_service_client_v1, + featurestore_service_client_v1, job_service_client_v1, metadata_service_client_v1, model_service_client_v1, @@ -63,6 +67,8 @@ # v1beta1 dataset_service_client_v1beta1.DatasetServiceClient, endpoint_service_client_v1beta1.EndpointServiceClient, + featurestore_online_serving_service_client_v1beta1.FeaturestoreOnlineServingServiceClient, + featurestore_service_client_v1beta1.FeaturestoreServiceClient, model_service_client_v1beta1.ModelServiceClient, prediction_service_client_v1beta1.PredictionServiceClient, pipeline_service_client_v1beta1.PipelineServiceClient, @@ -72,6 +78,8 @@ # v1 dataset_service_client_v1.DatasetServiceClient, endpoint_service_client_v1.EndpointServiceClient, + featurestore_online_serving_service_client_v1.FeaturestoreOnlineServingServiceClient, + featurestore_service_client_v1.FeaturestoreServiceClient, metadata_service_client_v1.MetadataServiceClient, model_service_client_v1.ModelServiceClient, prediction_service_client_v1.PredictionServiceClient, @@ -453,6 +461,30 @@ class EndpointClientWithOverride(ClientWithOverride): ) +class FeaturestoreClientWithOverride(ClientWithOverride): + _is_temporary = True + _default_version = compat.DEFAULT_VERSION + _version_map = ( + (compat.V1, featurestore_service_client_v1.FeaturestoreServiceClient), + (compat.V1BETA1, featurestore_service_client_v1beta1.FeaturestoreServiceClient), + ) + + +class FeaturestoreOnlineServingClientWithOverride(ClientWithOverride): + _is_temporary = False + _default_version = compat.DEFAULT_VERSION + _version_map = ( + ( + compat.V1, + featurestore_online_serving_service_client_v1.FeaturestoreOnlineServingServiceClient, + ), + ( + compat.V1BETA1, + featurestore_online_serving_service_client_v1beta1.FeaturestoreOnlineServingServiceClient, + ), + ) + + class JobClientWithOverride(ClientWithOverride): _is_temporary = True _default_version = compat.DEFAULT_VERSION @@ -520,6 +552,7 @@ class TensorboardClientWithOverride(ClientWithOverride): "VertexAiServiceClientWithOverride", DatasetClientWithOverride, EndpointClientWithOverride, + FeaturestoreClientWithOverride, JobClientWithOverride, ModelClientWithOverride, PipelineClientWithOverride, diff --git a/google/cloud/aiplatform/utils/featurestore_utils.py b/google/cloud/aiplatform/utils/featurestore_utils.py new file mode 100644 index 0000000000..c78a96d185 --- /dev/null +++ b/google/cloud/aiplatform/utils/featurestore_utils.py @@ -0,0 +1,117 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import re +from typing import Optional, Tuple + +from google.cloud.aiplatform.compat.services import featurestore_service_client + +CompatFeaturestoreServiceClient = featurestore_service_client.FeaturestoreServiceClient + +RESOURCE_ID_PATTERN_REGEX = r"[a-z_][a-z0-9_]{0,59}" + + +def validate_id(resource_id: str) -> bool: + """Validates feature store resource ID pattern.""" + return bool(re.compile(r"^" + RESOURCE_ID_PATTERN_REGEX + r"$").match(resource_id)) + + +def validate_and_get_entity_type_resource_ids( + entity_type_name: str, featurestore_id: Optional[str] = None, +) -> Tuple[str, str]: + """Validates and gets featurestore ID and entity_type ID of the entity_type resource. + + Args: + entity_type_name (str): + Required. A fully-qualified entityType resource name or an entity_type ID + Example: "projects/123/locations/us-central1/featurestores/my_featurestore_id/entityTypes/my_entity_type_id" + or "my_entity_type_id", with featurestore_id passed. + featurestore_id (str): + Optional. Featurestore ID of the entity_type resource. + + Returns: + Tuple[str, str] - featurestore ID and entity_type ID + + Raises: + ValueError if the provided entity_type_name is not in form of a fully-qualified + entityType resource name nor an entity_type ID with featurestore_id passed. + """ + match = CompatFeaturestoreServiceClient.parse_entity_type_path( + path=entity_type_name + ) + + if match: + featurestore_id = match["featurestore"] + entity_type_id = match["entity_type"] + elif ( + validate_id(entity_type_name) + and featurestore_id + and validate_id(featurestore_id) + ): + entity_type_id = entity_type_name + else: + raise ValueError( + f"{entity_type_name} is not in form of a fully-qualified entityType resource name " + f"nor an entity_type ID with featurestore_id passed." + ) + return (featurestore_id, entity_type_id) + + +def validate_and_get_feature_resource_ids( + feature_name: str, + featurestore_id: Optional[str] = None, + entity_type_id: Optional[str] = None, +) -> Tuple[str, str, str]: + """Validates and gets featurestore ID, entity_type ID, and feature ID for the feature resource. + Args: + feature_name (str): + Required. A fully-qualified feature resource name or a feature ID. + Example: "projects/123/locations/us-central1/featurestores/my_featurestore_id/entityTypes/my_entity_type_id/features/my_feature_id" + or "my_feature_id" when project and location are initialized or passed, with featurestore_id and entity_type_id passed. + featurestore_id (str): + Optional. Featurestore ID of the feature resource. + entity_type_id (str): + Optional. EntityType ID of the feature resource. + + Returns: + Tuple[str, str, str] - featurestore ID, entity_type ID, and feature ID + + Raises: + ValueError if the provided feature_name is not in form of a fully-qualified + feature resource name nor a feature ID with featurestore_id and entity_type_id passed. + """ + + match = CompatFeaturestoreServiceClient.parse_feature_path(path=feature_name) + + if match: + featurestore_id = match["featurestore"] + entity_type_id = match["entity_type"] + feature_id = match["feature"] + elif ( + validate_id(feature_name) + and featurestore_id + and entity_type_id + and validate_id(featurestore_id) + and validate_id(entity_type_id) + ): + feature_id = feature_name + else: + raise ValueError( + f"{feature_name} is not in form of a fully-qualified feature resource name " + f"nor a feature ID with featurestore_id and entity_type_id passed." + ) + return (featurestore_id, entity_type_id, feature_id) diff --git a/tests/system/aiplatform/test_featurestore.py b/tests/system/aiplatform/test_featurestore.py new file mode 100644 index 0000000000..6107f826ec --- /dev/null +++ b/tests/system/aiplatform/test_featurestore.py @@ -0,0 +1,38 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from google.cloud import aiplatform +from tests.system.aiplatform import e2e_base + + +class TestFeaturestore(e2e_base.TestEndToEnd): + + _temp_prefix = "temp-vertex-sdk-e2e-feature-store-test" + + def test_create_and_get_featurestore(self, shared_state): + + aiplatform.init( + project=e2e_base._PROJECT, location=e2e_base._LOCATION, + ) + + shared_state["resources"] = [] + + list_featurestores = aiplatform.Featurestore.list() + assert len(list_featurestores) >= 0 + + list_searched_features = aiplatform.Feature.search() + assert len(list_searched_features) >= 0 diff --git a/tests/unit/aiplatform/test_featurestores.py b/tests/unit/aiplatform/test_featurestores.py new file mode 100644 index 0000000000..4cede4ba09 --- /dev/null +++ b/tests/unit/aiplatform/test_featurestores.py @@ -0,0 +1,719 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import pytest + +from unittest import mock +from importlib import reload +from unittest.mock import patch + +from google.api_core import operation +from google.protobuf import field_mask_pb2 + +from google.cloud import aiplatform +from google.cloud.aiplatform import base +from google.cloud.aiplatform import initializer + +from google.cloud.aiplatform.utils import featurestore_utils + +from google.cloud.aiplatform_v1.services.featurestore_service import ( + client as featurestore_service_client, +) + +from google.cloud.aiplatform_v1.types import ( + featurestore as gca_featurestore, + entity_type as gca_entity_type, + feature as gca_feature, + encryption_spec as gca_encryption_spec, +) + +# project +_TEST_PROJECT = "test-project" +_TEST_LOCATION = "us-central1" +_TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" + +# featurestore +_TEST_FEATURESTORE_ID = "featurestore_id" +_TEST_FEATURESTORE_NAME = f"{_TEST_PARENT}/featurestores/{_TEST_FEATURESTORE_ID}" +_TEST_FEATURESTORE_INVALID = f"{_TEST_PARENT}/featurestore/{_TEST_FEATURESTORE_ID}" + +# featurestore online +_TEST_ONLINE_SERVING_CONFIG = 1 +_TEST_ONLINE_SERVING_CONFIG_UPDATE = 2 + +# entity_type +_TEST_ENTITY_TYPE_ID = "entity_type_id" +_TEST_ENTITY_TYPE_NAME = f"{_TEST_FEATURESTORE_NAME}/entityTypes/{_TEST_ENTITY_TYPE_ID}" +_TEST_ENTITY_TYPE_INVALID = ( + f"{_TEST_FEATURESTORE_NAME}/entityType/{_TEST_ENTITY_TYPE_ID}" +) + +# feature +_TEST_FEATURE_ID = "feature_id" +_TEST_FEATURE_NAME = f"{_TEST_ENTITY_TYPE_NAME}/features/{_TEST_FEATURE_ID}" +_TEST_FEATURE_INVALID = f"{_TEST_ENTITY_TYPE_NAME}/feature/{_TEST_FEATURE_ID}" + +# misc +_TEST_DESCRIPTION = "my description" +_TEST_LABELS = {"my_key": "my_value"} +_TEST_DESCRIPTION_UPDATE = "my description update" +_TEST_LABELS_UPDATE = {"my_key_update": "my_value_update"} + +# request_metadata +_TEST_REQUEST_METADATA = () + +# CMEK encryption +_TEST_ENCRYPTION_KEY_NAME = "key_1234" +_TEST_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec( + kms_key_name=_TEST_ENCRYPTION_KEY_NAME +) + +# Lists +_TEST_FEATURESTORE_LIST = [ + gca_featurestore.Featurestore( + name=_TEST_FEATURESTORE_NAME, + online_serving_config=gca_featurestore.Featurestore.OnlineServingConfig( + fixed_node_count=_TEST_ONLINE_SERVING_CONFIG + ), + encryption_spec=_TEST_ENCRYPTION_SPEC, + ), + gca_featurestore.Featurestore( + name=_TEST_FEATURESTORE_NAME, + online_serving_config=gca_featurestore.Featurestore.OnlineServingConfig( + fixed_node_count=_TEST_ONLINE_SERVING_CONFIG + ), + ), + gca_featurestore.Featurestore( + name=_TEST_FEATURESTORE_NAME, + online_serving_config=gca_featurestore.Featurestore.OnlineServingConfig( + fixed_node_count=_TEST_ONLINE_SERVING_CONFIG + ), + encryption_spec=_TEST_ENCRYPTION_SPEC, + ), +] + +_TEST_ENTITY_TYPE_LIST = [ + gca_entity_type.EntityType(name=_TEST_ENTITY_TYPE_NAME,), + gca_entity_type.EntityType(name=_TEST_ENTITY_TYPE_NAME,), + gca_entity_type.EntityType(name=_TEST_ENTITY_TYPE_NAME,), +] + +_TEST_FEATURE_LIST = [ + gca_feature.Feature(name=_TEST_FEATURE_NAME,), + gca_feature.Feature(name=_TEST_FEATURE_NAME,), + gca_feature.Feature(name=_TEST_FEATURE_NAME,), +] + + +# All Featurestore Mocks +@pytest.fixture +def get_featurestore_mock(): + with patch.object( + featurestore_service_client.FeaturestoreServiceClient, "get_featurestore" + ) as get_featurestore_mock: + get_featurestore_mock.return_value = gca_featurestore.Featurestore( + name=_TEST_FEATURESTORE_NAME, + online_serving_config=gca_featurestore.Featurestore.OnlineServingConfig( + fixed_node_count=_TEST_ONLINE_SERVING_CONFIG + ), + encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + yield get_featurestore_mock + + +@pytest.fixture +def update_featurestore_mock(): + with patch.object( + featurestore_service_client.FeaturestoreServiceClient, "update_featurestore" + ) as update_featurestore_mock: + update_featurestore_lro_mock = mock.Mock(operation.Operation) + update_featurestore_mock.return_value = update_featurestore_lro_mock + yield update_featurestore_mock + + +@pytest.fixture +def list_featurestores_mock(): + with patch.object( + featurestore_service_client.FeaturestoreServiceClient, "list_featurestores" + ) as list_featurestores_mock: + list_featurestores_mock.return_value = _TEST_FEATURESTORE_LIST + yield list_featurestores_mock + + +@pytest.fixture +def delete_featurestore_mock(): + with mock.patch.object( + featurestore_service_client.FeaturestoreServiceClient, "delete_featurestore" + ) as delete_featurestore_mock: + delete_featurestore_lro_mock = mock.Mock(operation.Operation) + delete_featurestore_mock.return_value = delete_featurestore_lro_mock + yield delete_featurestore_mock + + +@pytest.fixture +def search_features_mock(): + with patch.object( + featurestore_service_client.FeaturestoreServiceClient, "search_features" + ) as search_features_mock: + search_features_mock.return_value = _TEST_FEATURE_LIST + yield search_features_mock + + +# ALL EntityType Mocks +@pytest.fixture +def get_entity_type_mock(): + with patch.object( + featurestore_service_client.FeaturestoreServiceClient, "get_entity_type" + ) as get_entity_type_mock: + get_entity_type_mock.return_value = gca_entity_type.EntityType( + name=_TEST_ENTITY_TYPE_NAME, + ) + yield get_entity_type_mock + + +@pytest.fixture +def update_entity_type_mock(): + with patch.object( + featurestore_service_client.FeaturestoreServiceClient, "update_entity_type" + ) as update_entity_type_mock: + update_entity_type_lro_mock = mock.Mock(operation.Operation) + update_entity_type_mock.return_value = update_entity_type_lro_mock + yield update_entity_type_mock + + +@pytest.fixture +def list_entity_types_mock(): + with patch.object( + featurestore_service_client.FeaturestoreServiceClient, "list_entity_types" + ) as list_entity_types_mock: + list_entity_types_mock.return_value = _TEST_ENTITY_TYPE_LIST + yield list_entity_types_mock + + +@pytest.fixture +def delete_entity_type_mock(): + with mock.patch.object( + featurestore_service_client.FeaturestoreServiceClient, "delete_entity_type" + ) as delete_entity_type_mock: + delete_entity_type_lro_mock = mock.Mock(operation.Operation) + delete_entity_type_mock.return_value = delete_entity_type_lro_mock + yield delete_entity_type_mock + + +# ALL Feature Mocks +@pytest.fixture +def get_feature_mock(): + with patch.object( + featurestore_service_client.FeaturestoreServiceClient, "get_feature" + ) as get_feature_mock: + get_feature_mock.return_value = gca_feature.Feature(name=_TEST_FEATURE_NAME,) + yield get_feature_mock + + +@pytest.fixture +def update_feature_mock(): + with patch.object( + featurestore_service_client.FeaturestoreServiceClient, "update_feature" + ) as update_feature_mock: + update_feature_lro_mock = mock.Mock(operation.Operation) + update_feature_mock.return_value = update_feature_lro_mock + yield update_feature_mock + + +@pytest.fixture +def list_features_mock(): + with patch.object( + featurestore_service_client.FeaturestoreServiceClient, "list_features" + ) as list_features_mock: + list_features_mock.return_value = _TEST_FEATURE_LIST + yield list_features_mock + + +@pytest.fixture +def delete_feature_mock(): + with mock.patch.object( + featurestore_service_client.FeaturestoreServiceClient, "delete_feature" + ) as delete_feature_mock: + delete_feature_lro_mock = mock.Mock(operation.Operation) + delete_feature_mock.return_value = delete_feature_lro_mock + yield delete_feature_mock + + +class TestFeaturestoreUtils: + @pytest.mark.parametrize( + "resource_id, expected", + [ + ("resource_id", True), + ("resource_id12345", True), + ("12345resource_id", False), + ("_resource_id", True), + ("resource_id/1234", False), + ("_resource_id/1234", False), + ("resource-id-1234", False), + ("123456", False), + ("c" * 61, False), + ("_123456", True), + ], + ) + def test_validate_resource_id(self, resource_id: str, expected: bool): + assert expected == featurestore_utils.validate_id(resource_id) + + @pytest.mark.parametrize( + "feature_name, featurestore_id, entity_type_id", + [ + (_TEST_FEATURE_NAME, None, None,), + (_TEST_FEATURE_ID, _TEST_FEATURESTORE_ID, _TEST_ENTITY_TYPE_ID,), + ], + ) + def test_validate_and_get_feature_resource_ids( + self, feature_name: str, featurestore_id: str, entity_type_id: str, + ): + assert ( + _TEST_FEATURESTORE_ID, + _TEST_ENTITY_TYPE_ID, + _TEST_FEATURE_ID, + ) == featurestore_utils.validate_and_get_feature_resource_ids( + feature_name=feature_name, + featurestore_id=featurestore_id, + entity_type_id=entity_type_id, + ) + + @pytest.mark.parametrize( + "feature_name, featurestore_id, entity_type_id", + [ + (_TEST_FEATURE_INVALID, None, None,), + (_TEST_FEATURE_ID, None, _TEST_ENTITY_TYPE_ID,), + (_TEST_FEATURE_ID, None, None,), + (_TEST_FEATURE_ID, _TEST_FEATURESTORE_NAME, None,), + ], + ) + def test_validate_and_get_feature_resource_ids_with_raise( + self, feature_name: str, featurestore_id: str, entity_type_id: str, + ): + with pytest.raises(ValueError): + featurestore_utils.validate_and_get_feature_resource_ids( + feature_name=feature_name, + featurestore_id=featurestore_id, + entity_type_id=entity_type_id, + ) + + @pytest.mark.parametrize( + "entity_type_name, featurestore_id", + [ + (_TEST_ENTITY_TYPE_NAME, None,), + (_TEST_ENTITY_TYPE_ID, _TEST_FEATURESTORE_ID,), + ], + ) + def test_validate_and_get_entity_type_resource_ids( + self, entity_type_name: str, featurestore_id: str + ): + assert ( + _TEST_FEATURESTORE_ID, + _TEST_ENTITY_TYPE_ID, + ) == featurestore_utils.validate_and_get_entity_type_resource_ids( + entity_type_name=entity_type_name, featurestore_id=featurestore_id + ) + + @pytest.mark.parametrize( + "entity_type_name, featurestore_id", + [ + (_TEST_ENTITY_TYPE_INVALID, None,), + (_TEST_ENTITY_TYPE_ID, None,), + (_TEST_ENTITY_TYPE_ID, _TEST_FEATURESTORE_NAME,), + ], + ) + def test_validate_and_get_entity_type_resource_ids_with_raise( + self, entity_type_name: str, featurestore_id: str, + ): + with pytest.raises(ValueError): + featurestore_utils.validate_and_get_entity_type_resource_ids( + entity_type_name=entity_type_name, featurestore_id=featurestore_id + ) + + +class TestFeaturestore: + def setup_method(self): + reload(initializer) + reload(aiplatform) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + @pytest.mark.parametrize( + "featurestore_name", [_TEST_FEATURESTORE_ID, _TEST_FEATURESTORE_NAME] + ) + def test_init_featurestore(self, featurestore_name, get_featurestore_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_featurestore = aiplatform.Featurestore(featurestore_name=featurestore_name) + + get_featurestore_mock.assert_called_once_with( + name=my_featurestore.resource_name, retry=base._DEFAULT_RETRY + ) + + @pytest.mark.usefixtures("get_featurestore_mock") + def test_get_entity_type(self, get_entity_type_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_featurestore = aiplatform.Featurestore( + featurestore_name=_TEST_FEATURESTORE_ID + ) + my_entity_type = my_featurestore.get_entity_type( + entity_type_id=_TEST_ENTITY_TYPE_ID + ) + + get_entity_type_mock.assert_called_once_with( + name=_TEST_ENTITY_TYPE_NAME, retry=base._DEFAULT_RETRY + ) + assert type(my_entity_type) == aiplatform.EntityType + + @pytest.mark.usefixtures("get_featurestore_mock") + def test_update_featurestore(self, update_featurestore_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_featurestore = aiplatform.Featurestore( + featurestore_name=_TEST_FEATURESTORE_ID + ) + my_featurestore.update(labels=_TEST_LABELS_UPDATE) + + expected_featurestore = gca_featurestore.Featurestore( + name=_TEST_FEATURESTORE_NAME, + labels=_TEST_LABELS_UPDATE, + online_serving_config=gca_featurestore.Featurestore.OnlineServingConfig(), + ) + update_featurestore_mock.assert_called_once_with( + featurestore=expected_featurestore, + update_mask=field_mask_pb2.FieldMask(paths=["labels"]), + metadata=_TEST_REQUEST_METADATA, + ) + + @pytest.mark.usefixtures("get_featurestore_mock") + def test_update_featurestore_online(self, update_featurestore_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_featurestore = aiplatform.Featurestore( + featurestore_name=_TEST_FEATURESTORE_ID + ) + my_featurestore.update_online_store( + fixed_node_count=_TEST_ONLINE_SERVING_CONFIG_UPDATE + ) + + expected_featurestore = gca_featurestore.Featurestore( + name=_TEST_FEATURESTORE_NAME, + online_serving_config=gca_featurestore.Featurestore.OnlineServingConfig( + fixed_node_count=_TEST_ONLINE_SERVING_CONFIG_UPDATE + ), + ) + update_featurestore_mock.assert_called_once_with( + featurestore=expected_featurestore, + update_mask=field_mask_pb2.FieldMask( + paths=["online_serving_config.fixed_node_count"] + ), + metadata=_TEST_REQUEST_METADATA, + ) + + def test_list_featurestores(self, list_featurestores_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_featurestore_list = aiplatform.Featurestore.list() + + list_featurestores_mock.assert_called_once_with( + request={"parent": _TEST_PARENT, "filter": None} + ) + assert len(my_featurestore_list) == len(_TEST_FEATURESTORE_LIST) + for my_featurestore in my_featurestore_list: + assert type(my_featurestore) == aiplatform.Featurestore + + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_featurestore_mock") + def test_delete_featurestore(self, delete_featurestore_mock, sync): + aiplatform.init(project=_TEST_PROJECT) + + my_featurestore = aiplatform.Featurestore( + featurestore_name=_TEST_FEATURESTORE_ID + ) + my_featurestore.delete(sync=sync) + + if not sync: + my_featurestore.wait() + + delete_featurestore_mock.assert_called_once_with( + name=my_featurestore.resource_name + ) + + @pytest.mark.usefixtures("get_featurestore_mock") + def test_list_entity_types(self, list_entity_types_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_featurestore = aiplatform.Featurestore( + featurestore_name=_TEST_FEATURESTORE_ID + ) + my_entity_type_list = my_featurestore.list_entity_types() + + list_entity_types_mock.assert_called_once_with( + request={"parent": _TEST_FEATURESTORE_NAME, "filter": None} + ) + assert len(my_entity_type_list) == len(_TEST_ENTITY_TYPE_LIST) + for my_entity_type in my_entity_type_list: + assert type(my_entity_type) == aiplatform.EntityType + + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_featurestore_mock", "get_entity_type_mock") + def test_delete_entity_types(self, delete_entity_type_mock, sync): + aiplatform.init(project=_TEST_PROJECT) + + my_featurestore = aiplatform.Featurestore( + featurestore_name=_TEST_FEATURESTORE_ID + ) + my_featurestore.delete_entity_types( + entity_type_ids=[_TEST_ENTITY_TYPE_ID, _TEST_ENTITY_TYPE_ID], sync=sync + ) + + if not sync: + my_featurestore.wait() + + delete_entity_type_mock.assert_has_calls( + calls=[ + mock.call(name=_TEST_ENTITY_TYPE_NAME), + mock.call(name=_TEST_ENTITY_TYPE_NAME), + ], + any_order=True, + ) + + +class TestEntityType: + def setup_method(self): + reload(initializer) + reload(aiplatform) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + @pytest.mark.parametrize( + "entity_type_name, featurestore_id", + [ + (_TEST_ENTITY_TYPE_NAME, None), + (_TEST_ENTITY_TYPE_ID, _TEST_FEATURESTORE_ID), + ], + ) + def test_init_entity_type( + self, entity_type_name, featurestore_id, get_entity_type_mock + ): + aiplatform.init(project=_TEST_PROJECT) + + aiplatform.EntityType( + entity_type_name=entity_type_name, featurestore_id=featurestore_id + ) + + get_entity_type_mock.assert_called_once_with( + name=_TEST_ENTITY_TYPE_NAME, retry=base._DEFAULT_RETRY + ) + + @pytest.mark.usefixtures("get_entity_type_mock") + def test_get_featurestore(self, get_featurestore_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_entity_type = aiplatform.EntityType(entity_type_name=_TEST_ENTITY_TYPE_NAME) + my_featurestore = my_entity_type.get_featurestore() + + get_featurestore_mock.assert_called_once_with( + name=my_featurestore.resource_name, retry=base._DEFAULT_RETRY + ) + assert type(my_featurestore) == aiplatform.Featurestore + + @pytest.mark.usefixtures("get_entity_type_mock") + def test_get_feature(self, get_feature_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_entity_type = aiplatform.EntityType(entity_type_name=_TEST_ENTITY_TYPE_NAME) + my_feature = my_entity_type.get_feature(feature_id=_TEST_FEATURE_ID) + + get_feature_mock.assert_called_once_with( + name=my_feature.resource_name, retry=base._DEFAULT_RETRY + ) + assert type(my_feature) == aiplatform.Feature + + @pytest.mark.usefixtures("get_entity_type_mock") + def test_update_entity_type(self, update_entity_type_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_entity_type = aiplatform.EntityType(entity_type_name=_TEST_ENTITY_TYPE_NAME) + my_entity_type.update(labels=_TEST_LABELS_UPDATE) + + expected_entity_type = gca_entity_type.EntityType( + name=_TEST_ENTITY_TYPE_NAME, labels=_TEST_LABELS_UPDATE, + ) + update_entity_type_mock.assert_called_once_with( + entity_type=expected_entity_type, + update_mask=field_mask_pb2.FieldMask(paths=["labels"]), + metadata=_TEST_REQUEST_METADATA, + ) + + @pytest.mark.parametrize( + "featurestore_name", [_TEST_FEATURESTORE_NAME, _TEST_FEATURESTORE_ID] + ) + def test_list_entity_types(self, featurestore_name, list_entity_types_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_entity_type_list = aiplatform.EntityType.list( + featurestore_name=featurestore_name + ) + + list_entity_types_mock.assert_called_once_with( + request={"parent": _TEST_FEATURESTORE_NAME, "filter": None} + ) + assert len(my_entity_type_list) == len(_TEST_ENTITY_TYPE_LIST) + for my_entity_type in my_entity_type_list: + assert type(my_entity_type) == aiplatform.EntityType + + @pytest.mark.usefixtures("get_entity_type_mock") + def test_list_features(self, list_features_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_entity_type = aiplatform.EntityType(entity_type_name=_TEST_ENTITY_TYPE_NAME) + my_feature_list = my_entity_type.list_features() + + list_features_mock.assert_called_once_with( + request={"parent": _TEST_ENTITY_TYPE_NAME, "filter": None} + ) + assert len(my_feature_list) == len(_TEST_FEATURE_LIST) + for my_feature in my_feature_list: + assert type(my_feature) == aiplatform.Feature + + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_entity_type_mock", "get_feature_mock") + def test_delete_features(self, delete_feature_mock, sync): + aiplatform.init(project=_TEST_PROJECT) + + my_entity_type = aiplatform.EntityType(entity_type_name=_TEST_ENTITY_TYPE_NAME) + my_entity_type.delete_features( + feature_ids=[_TEST_FEATURE_ID, _TEST_FEATURE_ID], sync=sync + ) + + if not sync: + my_entity_type.wait() + + delete_feature_mock.assert_has_calls( + calls=[ + mock.call(name=_TEST_FEATURE_NAME), + mock.call(name=_TEST_FEATURE_NAME), + ], + any_order=True, + ) + + +class TestFeature: + def setup_method(self): + reload(initializer) + reload(aiplatform) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + @pytest.mark.parametrize( + "feature_name, entity_type_id, featurestore_id", + [ + (_TEST_FEATURE_NAME, None, None), + (_TEST_FEATURE_ID, _TEST_ENTITY_TYPE_ID, _TEST_FEATURESTORE_ID), + ], + ) + def test_init_feature( + self, feature_name, entity_type_id, featurestore_id, get_feature_mock + ): + aiplatform.init(project=_TEST_PROJECT) + aiplatform.Feature( + feature_name=feature_name, + entity_type_id=entity_type_id, + featurestore_id=featurestore_id, + ) + get_feature_mock.assert_called_once_with( + name=_TEST_FEATURE_NAME, retry=base._DEFAULT_RETRY + ) + + @pytest.mark.usefixtures("get_feature_mock") + def test_get_featurestore(self, get_featurestore_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_feature = aiplatform.Feature(feature_name=_TEST_FEATURE_NAME) + my_featurestore = my_feature.get_featurestore() + + get_featurestore_mock.assert_called_once_with( + name=my_featurestore.resource_name, retry=base._DEFAULT_RETRY + ) + assert type(my_featurestore) == aiplatform.Featurestore + + @pytest.mark.usefixtures("get_feature_mock") + def test_get_entity_type(self, get_entity_type_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_feature = aiplatform.Feature(feature_name=_TEST_FEATURE_NAME) + my_entity_type = my_feature.get_entity_type() + + get_entity_type_mock.assert_called_once_with( + name=my_entity_type.resource_name, retry=base._DEFAULT_RETRY + ) + assert type(my_entity_type) == aiplatform.EntityType + + @pytest.mark.usefixtures("get_feature_mock") + def test_update_feature(self, update_feature_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_feature = aiplatform.Feature(feature_name=_TEST_FEATURE_NAME) + my_feature.update(labels=_TEST_LABELS_UPDATE) + + expected_feature = gca_feature.Feature( + name=_TEST_FEATURE_NAME, labels=_TEST_LABELS_UPDATE, + ) + update_feature_mock.assert_called_once_with( + feature=expected_feature, + update_mask=field_mask_pb2.FieldMask(paths=["labels"]), + metadata=_TEST_REQUEST_METADATA, + ) + + @pytest.mark.parametrize( + "entity_type_name, featurestore_id", + [ + (_TEST_ENTITY_TYPE_NAME, None), + (_TEST_ENTITY_TYPE_ID, _TEST_FEATURESTORE_ID), + ], + ) + def test_list_features(self, entity_type_name, featurestore_id, list_features_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_feature_list = aiplatform.Feature.list( + entity_type_name=entity_type_name, featurestore_id=featurestore_id + ) + + list_features_mock.assert_called_once_with( + request={"parent": _TEST_ENTITY_TYPE_NAME, "filter": None} + ) + assert len(my_feature_list) == len(_TEST_FEATURE_LIST) + for my_feature in my_feature_list: + assert type(my_feature) == aiplatform.Feature + + @pytest.mark.usefixtures("get_feature_mock") + def test_search_features(self, search_features_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_feature_list = aiplatform.Feature.search() + + search_features_mock.assert_called_once_with( + request={"location": _TEST_PARENT, "query": None} + ) + assert len(my_feature_list) == len(_TEST_FEATURE_LIST) + for my_feature in my_feature_list: + assert type(my_feature) == aiplatform.Feature