From 26849e7f20d9bd9238ac64aaafb072cd9f43caa0 Mon Sep 17 00:00:00 2001 From: Tianrui Yang Date: Tue, 22 Oct 2024 07:50:02 -0700 Subject: [PATCH] fea: add PSC example code in Feature Store embedding notebook (#3663) --- ...val_bigquery_data_with_feature_store.ipynb | 230 +++++++++++++++++- 1 file changed, 227 insertions(+), 3 deletions(-) diff --git a/notebooks/official/feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.ipynb b/notebooks/official/feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.ipynb index d53ac58cd..63bdd6281 100644 --- a/notebooks/official/feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.ipynb +++ b/notebooks/official/feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.ipynb @@ -279,6 +279,7 @@ "outputs": [], "source": [ "from google.cloud import bigquery\n", + "from google.cloud.aiplatform_v1 import FeatureOnlineStoreAdminServiceClient\n", "from google.cloud.aiplatform_v1.types import NearestNeighborQuery\n", "from vertexai.resources.preview import (FeatureOnlineStore, FeatureView,\n", " FeatureViewBigQuerySource)\n", @@ -475,6 +476,28 @@ "2. Define the data (FeatureView) to be served by the newly-created instance." ] }, + { + "cell_type": "markdown", + "metadata": { + "id": "8b6888369eaa" + }, + "source": [ + "### Initialize Service Client\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "379e5135591d" + }, + "outputs": [], + "source": [ + "admin_client = FeatureOnlineStoreAdminServiceClient(\n", + " client_options={\"api_endpoint\": API_ENDPOINT}\n", + ")" + ] + }, { "cell_type": "markdown", "metadata": { @@ -856,13 +879,22 @@ ")" ] }, + { + "cell_type": "markdown", + "metadata": { + "id": "2e35bea02946" + }, + "source": [ + "#### Option 1: Search with public endpoint" + ] + }, { "cell_type": "markdown", "metadata": { "id": "v5gYpORiBW1n" }, "source": [ - "#### Search with `ENTITY_ID`" + "##### Search with `ENTITY_ID`" ] }, { @@ -890,7 +922,7 @@ "id": "OQEpt08GBX-b" }, "source": [ - "#### Search with `Embedding`" + "##### Search with `Embedding`" ] }, { @@ -926,7 +958,7 @@ "id": "MKALOxbsZfce" }, "source": [ - "#### Use the `FetchFeatureValues` API to retrieve the full data without search\n" + "##### Use the `FetchFeatureValues` API to retrieve the full data without search\n" ] }, { @@ -940,6 +972,198 @@ "my_fv.read(key=[ENTITY_ID])" ] }, + { + "cell_type": "markdown", + "metadata": { + "id": "808d95976090" + }, + "source": [ + "#### Option 2: Search with private endpoint\n", + "\n", + "You need to connect to private endpoint over gRPC. Follow these instructions to set up [Private Service Connect](https://cloud.google.com/vpc/docs/private-service-connect)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "7883e04698c5" + }, + "source": [ + "**Get Private Service Connect information**\n", + "\n", + "Retrieve the information to set up Private Service Connect from your `FeatureOnlineStore` instance." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "dfd03c754b88" + }, + "outputs": [], + "source": [ + "# Get Optimized online store\n", + "admin_client.get_feature_online_store(\n", + " name=f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}\"\n", + ")\n", + "\n", + "FEATURE_VIEW = f\"projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "44cd5a3bee94" + }, + "source": [ + "You should see output similar to the following:\n", + "\n", + "```\n", + "dedicated_serving_endpoint {\n", + " private_service_connect_config {\n", + " enable_private_service_connect: true\n", + " project_allowlist: \"your_allowlisted_project\"\n", + " }\n", + " service_attachment: \"service_attachment_string\"\n", + "}\n", + "optimized {}\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "7221d68485e3" + }, + "source": [ + "* `your_allowlisted_project` represents the name of your allowlisted project\n", + "\n", + "---\n", + "\n", + "where you created your `FeatureOnlineStore` instance.\n", + "\n", + "* `service_attachment_string` represents the target service that you need to specify while [adding Private Service Connect to your network configuration]()." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "e3e80f0b30ac" + }, + "source": [ + "**Add Private Service Connect to your network configuration:**\n", + "\n", + "1. In the Google Cloud console, select the project represented by `your_allowlisted_project`. This is the project you allowlisted while creating the `FeatureOnlineStore` instance.\n", + "1. On the [**Private Service Connect** page](https://console.cloud.google.com/net-services/psc/list/consumers) of the Google Cloud console, in the **Connected endpoints** tab, click **Connect endpoint**.\n", + "1. Under **Target**, click **Published service**\n", + "1. In the **Target service** field, specify the value of `service_attachment_string`.\n", + "1. Enter your **endpoint name**.\n", + "1. In the **Network** field, select **default**.\n", + "1. In the **Subetwork** field, select **default**.\n", + "1. In the **IP address** list, click **Create IP address** to create an IP address that you use to connect the Feature Store API.\n", + "1. Select **Enable global access**.\n", + "1. Click **Add Endpoint**.\n", + "\n", + "After the connection is successfully added, it appears in the **Connected endpoints** tab on the **Private Service Connect** page.\n", + "\n", + "Retrieve the IP address of the new connection from the **IP addresses** column and replace `{endpoint_ip}` before running the following code." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "1967dfab5a95" + }, + "outputs": [], + "source": [ + "# # Uncomment the following code blocks after your PSC setup is complete. Replace {endpoint_ip} with the IP of the new connection.\n", + "\n", + "# from google.cloud.aiplatform_v1.services.feature_online_store_service.transports.grpc import FeatureOnlineStoreServiceGrpcTransport\n", + "# from google.cloud.aiplatform_v1 import FeatureOnlineStoreServiceClient\n", + "# import grpc\n", + "\n", + "# data_client = FeatureOnlineStoreServiceClient(\n", + "# transport = FeatureOnlineStoreServiceGrpcTransport(\n", + "# # Add the IP address of the Endpoint you just created.\n", + "# channel = grpc.insecure_channel(\"{endpoint_ip}:10002\")\n", + "# ))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "a29e552d8a6e" + }, + "source": [ + "##### Search with `ENTITY_ID`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "bfab8ee53cd6" + }, + "outputs": [], + "source": [ + "# Uncomment the following code block.\n", + "\n", + "# data_client.search_nearest_entities(\n", + "# request=feature_online_store_service_pb2.SearchNearestEntitiesRequest(\n", + "# feature_view=FEATURE_VIEW,\n", + "# query= NearestNeighborQuery(\n", + "# entity_id = ENTITY_ID,\n", + "# neighbor_count = 5,\n", + "# string_filters = [country_filter]\n", + "# ),\n", + "# return_full_entity=True, # returning entities with metadata\n", + "# ))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "7d62f53a1f62" + }, + "source": [ + "##### Search with `Embedding`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "3f2a4b149643" + }, + "outputs": [], + "source": [ + "EMBEDDINGS = [1] * DIMENSIONS" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "bba020ff509c" + }, + "outputs": [], + "source": [ + "# Uncomment the following code block.\n", + "\n", + "# data_client.search_nearest_entities(\n", + "# request=feature_online_store_service_pb2.SearchNearestEntitiesRequest(\n", + "# feature_view=FEATURE_VIEW,\n", + "# query= NearestNeighborQuery(\n", + "# embedding = NearestNeighborQuery.Embedding(\n", + "# value = EMBEDDINGS),\n", + "# neighbor_count = 5,\n", + "# string_filters = [country_filter]\n", + "# ),\n", + "# return_full_entity=True, # returning entities with metadata\n", + "# ))" + ] + }, { "cell_type": "markdown", "metadata": {