fea: add PSC example code in Feature Store embedding notebook (#3663)

This commit is contained in:
Tianrui Yang
2024-10-22 14:50:02 +00:00
committed by GitHub
parent a96c2926af
commit 26849e7f20
@@ -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": {