From 8e06ced83ed2cc480d869318c4debef9c28ad214 Mon Sep 17 00:00:00 2001 From: Cheng Li Date: Wed, 22 Sep 2021 15:41:59 -0700 Subject: [PATCH] feat: code snippets for feature store control plane (#709) * Add code snippets for feature store service * Add resource creation and deletion including featurestore, entity type and feature --- samples/snippets/conftest.py | 34 ++++++++++++- .../create_entity_type_sample.py | 45 +++++++++++++++++ .../create_entity_type_sample_test.py | 42 ++++++++++++++++ .../create_feature_sample.py | 47 ++++++++++++++++++ .../create_feature_sample_test.py | 46 ++++++++++++++++++ .../create_featurestore_sample.py | 48 +++++++++++++++++++ .../create_featurestore_sample_test.py | 39 +++++++++++++++ .../delete_featurestore_sample.py | 40 ++++++++++++++++ .../delete_featurestore_sample_test.py | 45 +++++++++++++++++ samples/snippets/helpers.py | 7 +++ 10 files changed, 392 insertions(+), 1 deletion(-) create mode 100644 samples/snippets/feature_store_service/create_entity_type_sample.py create mode 100644 samples/snippets/feature_store_service/create_entity_type_sample_test.py create mode 100644 samples/snippets/feature_store_service/create_feature_sample.py create mode 100644 samples/snippets/feature_store_service/create_feature_sample_test.py create mode 100644 samples/snippets/feature_store_service/create_featurestore_sample.py create mode 100644 samples/snippets/feature_store_service/create_featurestore_sample_test.py create mode 100644 samples/snippets/feature_store_service/delete_featurestore_sample.py create mode 100644 samples/snippets/feature_store_service/delete_featurestore_sample_test.py diff --git a/samples/snippets/conftest.py b/samples/snippets/conftest.py index 2d2dd2fa78..b33aebfd16 100644 --- a/samples/snippets/conftest.py +++ b/samples/snippets/conftest.py @@ -15,7 +15,7 @@ import os from uuid import uuid4 -from google.cloud import aiplatform +from google.cloud import aiplatform, aiplatform_v1beta1 from google.cloud import storage import pytest @@ -83,6 +83,14 @@ def dataset_client(): yield dataset_client +@pytest.fixture +def featurestore_client(): + featurestore_client = aiplatform_v1beta1.FeaturestoreServiceClient( + client_options={"api_endpoint": "us-central1-aiplatform.googleapis.com"} + ) + yield featurestore_client + + # Shared setup/teardown. @pytest.fixture() def teardown_batch_prediction_job(shared_state, job_client): @@ -201,6 +209,30 @@ def teardown_dataset(shared_state, dataset_client): dataset_client.delete_dataset(name=shared_state["dataset_name"]) +@pytest.fixture() +def teardown_featurestore(shared_state, featurestore_client): + yield + + # Delete the created featurestore + featurestore_client.delete_featurestore(name=shared_state["featurestore_name"]) + + +@pytest.fixture() +def teardown_entity_type(shared_state, featurestore_client): + yield + + # Delete the created entity type + featurestore_client.delete_entity_type(name=shared_state["entity_type_name"]) + + +@pytest.fixture() +def teardown_feature(shared_state, featurestore_client): + yield + + # Delete the created feature + featurestore_client.delete_feature(name=shared_state["feature_name"]) + + @pytest.fixture() def create_endpoint(shared_state, endpoint_client): def create(project, location, test_name="temp_deploy_model_test"): diff --git a/samples/snippets/feature_store_service/create_entity_type_sample.py b/samples/snippets/feature_store_service/create_entity_type_sample.py new file mode 100644 index 0000000000..81e938c6bd --- /dev/null +++ b/samples/snippets/feature_store_service/create_entity_type_sample.py @@ -0,0 +1,45 @@ +# 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 +# +# https://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. + +# [START aiplatform_create_entity_type_sample] +from google.cloud import aiplatform_v1beta1 as aiplatform + + +def create_entity_type_sample( + project: str, + featurestore_id: str, + entity_type_id: str, + description: str = "sample entity type", + location: str = "us-central1", + api_endpoint: str = "us-central1-aiplatform.googleapis.com", + timeout: int = 300, +): + # The AI Platform services require regional API endpoints. + client_options = {"api_endpoint": api_endpoint} + # Initialize client that will be used to create and send requests. + # This client only needs to be created once, and can be reused for multiple requests. + client = aiplatform.FeaturestoreServiceClient(client_options=client_options) + parent = f"projects/{project}/locations/{location}/featurestores/{featurestore_id}" + create_entity_type_request = aiplatform.CreateEntityTypeRequest( + parent=parent, + entity_type_id=entity_type_id, + entity_type=aiplatform.EntityType(description=description), + ) + lro_response = client.create_entity_type(request=create_entity_type_request) + print("Long running operation:", lro_response.operation.name) + create_entity_type_response = lro_response.result(timeout=timeout) + print("create_entity_type_response:", create_entity_type_response) + + +# [END aiplatform_create_entity_type_sample] diff --git a/samples/snippets/feature_store_service/create_entity_type_sample_test.py b/samples/snippets/feature_store_service/create_entity_type_sample_test.py new file mode 100644 index 0000000000..a2311cfe28 --- /dev/null +++ b/samples/snippets/feature_store_service/create_entity_type_sample_test.py @@ -0,0 +1,42 @@ +# 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 os +from uuid import uuid4 + +import create_entity_type_sample +import pytest + +import helpers + +PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT") + + +@pytest.fixture(scope="function", autouse=True) +def teardown(teardown_entity_type): + yield + + +def test_ucaip_generated_create_entity_type_sample_vision(capsys, shared_state): + featurestore_id = "perm_sample_featurestore" + entity_type_id = f"temp_create_entity_type_test_{uuid4()}".replace("-", "_")[:60] + create_entity_type_sample.create_entity_type_sample( + project=PROJECT_ID, + featurestore_id=featurestore_id, + entity_type_id=entity_type_id, + ) + out, _ = capsys.readouterr() + assert "create_entity_type_response" in out + + shared_state["entity_type_name"] = helpers.get_featurestore_resource_name(out) diff --git a/samples/snippets/feature_store_service/create_feature_sample.py b/samples/snippets/feature_store_service/create_feature_sample.py new file mode 100644 index 0000000000..e3ac2b1933 --- /dev/null +++ b/samples/snippets/feature_store_service/create_feature_sample.py @@ -0,0 +1,47 @@ +# 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 +# +# https://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. + +# [START aiplatform_create_feature_sample] +from google.cloud import aiplatform_v1beta1 as aiplatform + + +def create_feature_sample( + project: str, + featurestore_id: str, + entity_type_id: str, + feature_id: str, + value_type: aiplatform.Feature.ValueType, + description: str = "sample feature", + location: str = "us-central1", + api_endpoint: str = "us-central1-aiplatform.googleapis.com", + timeout: int = 300, +): + # The AI Platform services require regional API endpoints. + client_options = {"api_endpoint": api_endpoint} + # Initialize client that will be used to create and send requests. + # This client only needs to be created once, and can be reused for multiple requests. + client = aiplatform.FeaturestoreServiceClient(client_options=client_options) + parent = f"projects/{project}/locations/{location}/featurestores/{featurestore_id}/entityTypes/{entity_type_id}" + create_feature_request = aiplatform.CreateFeatureRequest( + parent=parent, + feature=aiplatform.Feature(value_type=value_type, description=description), + feature_id=feature_id, + ) + lro_response = client.create_feature(request=create_feature_request) + print("Long running operation:", lro_response.operation.name) + create_feature_response = lro_response.result(timeout=timeout) + print("create_feature_response:", create_feature_response) + + +# [END aiplatform_create_feature_sample] diff --git a/samples/snippets/feature_store_service/create_feature_sample_test.py b/samples/snippets/feature_store_service/create_feature_sample_test.py new file mode 100644 index 0000000000..d1af2f34a4 --- /dev/null +++ b/samples/snippets/feature_store_service/create_feature_sample_test.py @@ -0,0 +1,46 @@ +# 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 os +from uuid import uuid4 + +import create_feature_sample +from google.cloud import aiplatform_v1beta1 as aiplatform +import pytest + +import helpers + +PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT") + + +@pytest.fixture(scope="function", autouse=True) +def teardown(teardown_feature): + yield + + +def test_ucaip_generated_create_feature_sample_vision(capsys, shared_state): + featurestore_id = "perm_sample_featurestore" + entity_type_id = "perm_sample_entity_type" + feature_id = f"temp_create_feature_test_{uuid4()}".replace("-", "_")[:60] + create_feature_sample.create_feature_sample( + project=PROJECT_ID, + featurestore_id=featurestore_id, + entity_type_id=entity_type_id, + feature_id=feature_id, + value_type=aiplatform.Feature.ValueType.INT64, + ) + out, _ = capsys.readouterr() + assert "create_feature_response" in out + + shared_state["feature_name"] = helpers.get_featurestore_resource_name(out) diff --git a/samples/snippets/feature_store_service/create_featurestore_sample.py b/samples/snippets/feature_store_service/create_featurestore_sample.py new file mode 100644 index 0000000000..aa2936e8d4 --- /dev/null +++ b/samples/snippets/feature_store_service/create_featurestore_sample.py @@ -0,0 +1,48 @@ +# 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 +# +# https://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. + +# [START aiplatform_create_featurestore_sample] +from google.cloud import aiplatform_v1beta1 as aiplatform + + +def create_featurestore_sample( + project: str, + featurestore_id: str, + fixed_node_count: int = 1, + location: str = "us-central1", + api_endpoint: str = "us-central1-aiplatform.googleapis.com", + timeout: int = 300, +): + # The AI Platform services require regional API endpoints. + client_options = {"api_endpoint": api_endpoint} + # Initialize client that will be used to create and send requests. + # This client only needs to be created once, and can be reused for multiple requests. + client = aiplatform.FeaturestoreServiceClient(client_options=client_options) + parent = f"projects/{project}/locations/{location}" + create_featurestore_request = aiplatform.CreateFeaturestoreRequest( + parent=parent, + featurestore_id=featurestore_id, + featurestore=aiplatform.Featurestore( + online_serving_config=aiplatform.Featurestore.OnlineServingConfig( + fixed_node_count=fixed_node_count, + ), + ), + ) + lro_response = client.create_featurestore(request=create_featurestore_request) + print("Long running operation:", lro_response.operation.name) + create_featurestore_response = lro_response.result(timeout=timeout) + print("create_featurestore_response:", create_featurestore_response) + + +# [END aiplatform_create_featurestore_sample] diff --git a/samples/snippets/feature_store_service/create_featurestore_sample_test.py b/samples/snippets/feature_store_service/create_featurestore_sample_test.py new file mode 100644 index 0000000000..0de26762e1 --- /dev/null +++ b/samples/snippets/feature_store_service/create_featurestore_sample_test.py @@ -0,0 +1,39 @@ +# 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 os +from uuid import uuid4 + +import create_featurestore_sample +import pytest + +import helpers + +PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT") + + +@pytest.fixture(scope="function", autouse=True) +def teardown(teardown_featurestore): + yield + + +def test_ucaip_generated_create_featurestore_sample_vision(capsys, shared_state): + featurestore_id = f"temp_create_featurestore_test_{uuid4()}".replace("-", "_")[:60] + create_featurestore_sample.create_featurestore_sample( + project=PROJECT_ID, featurestore_id=featurestore_id, fixed_node_count=1 + ) + out, _ = capsys.readouterr() + assert "create_featurestore_response" in out + + shared_state["featurestore_name"] = helpers.get_featurestore_resource_name(out) diff --git a/samples/snippets/feature_store_service/delete_featurestore_sample.py b/samples/snippets/feature_store_service/delete_featurestore_sample.py new file mode 100644 index 0000000000..7b16e835b9 --- /dev/null +++ b/samples/snippets/feature_store_service/delete_featurestore_sample.py @@ -0,0 +1,40 @@ +# 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 +# +# https://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. + +# [START aiplatform_delete_featurestore_sample] +from google.cloud import aiplatform_v1beta1 as aiplatform + + +def delete_featurestore_sample( + project: str, + featurestore_id: str, + location: str = "us-central1", + api_endpoint: str = "us-central1-aiplatform.googleapis.com", + timeout: int = 300, +): + # The AI Platform services require regional API endpoints. + client_options = {"api_endpoint": api_endpoint} + # Initialize client that will be used to create and send requests. + # This client only needs to be created once, and can be reused for multiple requests. + client = aiplatform.FeaturestoreServiceClient(client_options=client_options) + name = client.featurestore_path( + project=project, location=location, featurestore=featurestore_id + ) + response = client.delete_featurestore(name=name) + print("Long running operation:", response.operation.name) + delete_featurestore_response = response.result(timeout=timeout) + print("delete_featurestore_response:", delete_featurestore_response) + + +# [END aiplatform_delete_featurestore_sample] diff --git a/samples/snippets/feature_store_service/delete_featurestore_sample_test.py b/samples/snippets/feature_store_service/delete_featurestore_sample_test.py new file mode 100644 index 0000000000..ef09095d9a --- /dev/null +++ b/samples/snippets/feature_store_service/delete_featurestore_sample_test.py @@ -0,0 +1,45 @@ +# 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 os +from uuid import uuid4 + +import create_featurestore_sample +import delete_featurestore_sample +import pytest + +PROJECT_ID = os.getenv("BUILD_SPECIFIC_GCLOUD_PROJECT") + + +@pytest.fixture(scope="function", autouse=True) +def teardown(): + yield + + +def setup_test(capsys, featurestore_id): + create_featurestore_sample.create_featurestore_sample( + project=PROJECT_ID, featurestore_id=featurestore_id, fixed_node_count=1 + ) + out, _ = capsys.readouterr() + assert "create_featurestore_response" in out + + +def test_ucaip_generated_delete_featurestore_sample_vision(capsys, shared_state): + featurestore_id = f"temp_create_featurestore_test_{uuid4()}".replace("-", "_")[:60] + setup_test(capsys, featurestore_id) + delete_featurestore_sample.delete_featurestore_sample( + project=PROJECT_ID, featurestore_id=featurestore_id + ) + out, _ = capsys.readouterr() + assert "delete_featurestore_response" in out diff --git a/samples/snippets/helpers.py b/samples/snippets/helpers.py index d86da0238c..17baa0a34d 100644 --- a/samples/snippets/helpers.py +++ b/samples/snippets/helpers.py @@ -18,6 +18,13 @@ def get_state(out): return state +def get_featurestore_resource_name(out, key="name"): + pattern = re.compile(fr'{key}:\s*"([\_\-a-zA-Z0-9/]+)"') + name = re.search(pattern, out).group(1) + + return name + + def wait_for_job_state( get_job_method: Callable[[str], "proto.Message"], # noqa: F821 name: str,