fea: upgrade to v1 API in feature store llm grounding tutorial. (#3009)

* Upgrade to v1 API in feature store llm grounding tutorial.

* Sleep for 5min before starting serving to wait for DNS to be ready.

* use data_key in fetch request
This commit is contained in:
Tianrui Yang
2024-05-30 12:44:30 +00:00
committed by GitHub
parent c306cfaae1
commit 8bed394caf
@@ -347,18 +347,17 @@
"import uuid\n",
"\n",
"from google.cloud import aiplatform, bigquery\n",
"from google.cloud.aiplatform_v1beta1 import (\n",
" FeatureOnlineStoreAdminServiceClient, FeatureOnlineStoreServiceClient)\n",
"from google.cloud.aiplatform_v1beta1.types import NearestNeighborQuery\n",
"from google.cloud.aiplatform_v1beta1.types import \\\n",
"from google.cloud.aiplatform_v1 import (FeatureOnlineStoreAdminServiceClient,\n",
" FeatureOnlineStoreServiceClient)\n",
"from google.cloud.aiplatform_v1.types import NearestNeighborQuery\n",
"from google.cloud.aiplatform_v1.types import \\\n",
" feature_online_store as feature_online_store_pb2\n",
"from google.cloud.aiplatform_v1beta1.types import \\\n",
"from google.cloud.aiplatform_v1.types import \\\n",
" feature_online_store_admin_service as \\\n",
" feature_online_store_admin_service_pb2\n",
"from google.cloud.aiplatform_v1beta1.types import \\\n",
"from google.cloud.aiplatform_v1.types import \\\n",
" feature_online_store_service as feature_online_store_service_pb2\n",
"from google.cloud.aiplatform_v1beta1.types import \\\n",
" feature_view as feature_view_pb2"
"from google.cloud.aiplatform_v1.types import feature_view as feature_view_pb2"
]
},
{
@@ -707,14 +706,7 @@
"outputs": [],
"source": [
"online_store_config = feature_online_store_pb2.FeatureOnlineStore(\n",
" bigtable=feature_online_store_pb2.FeatureOnlineStore.Bigtable(\n",
" auto_scaling=feature_online_store_pb2.FeatureOnlineStore.Bigtable.AutoScaling(\n",
" min_node_count=1, max_node_count=3, cpu_utilization_target=50\n",
" )\n",
" ),\n",
" embedding_management=feature_online_store_pb2.FeatureOnlineStore.EmbeddingManagement(\n",
" enabled=True\n",
" ),\n",
" optimized=feature_online_store_pb2.FeatureOnlineStore.Optimized(),\n",
")\n",
"\n",
"create_store_lro = admin_client.create_feature_online_store(\n",
@@ -810,7 +802,7 @@
"* A data source (BigQuery table or view URI or `FeatureGroup/features`) synced to the `FeatureOnlineStore` instance for serving.\n",
"* The [cron](https://en.wikipedia.org/wiki/Cron) schedule to run the sync pipeline.\n",
"\n",
"During feature view creation, a sync job will be scheduled, and either started immediately or following the cron schedule. In the sync job, data is exported to Cloud Bigtable, a index is built and deployed to GKE cluster."
"During feature view creation, a sync job will be scheduled, and either started immediately or following the cron schedule. In the sync job, data is exported, a index is built and deployed to GKE cluster."
]
},
{
@@ -835,7 +827,7 @@
},
"outputs": [],
"source": [
"# Vector search configs\n",
"# Index building configs\n",
"DIMENSIONS = 768 # @param {type: \"number\"}\n",
"EMBEDDING_COLUMN = \"embedding\" # @param {type: \"string\"}\n",
"# Optional\n",
@@ -871,15 +863,15 @@
"\n",
"sync_config = feature_view_pb2.FeatureView.SyncConfig(cron=CRON_SCHEDULE)\n",
"\n",
"vector_search_config = feature_view_pb2.FeatureView.VectorSearchConfig(\n",
"index_config = feature_view_pb2.FeatureView.IndexConfig(\n",
" embedding_column=EMBEDDING_COLUMN,\n",
" # filter_columns=FILTER_COLUMNS,\n",
" # crowding_column=CROWDING_COLUMN,\n",
" embedding_dimension=DIMENSIONS,\n",
" tree_ah_config=feature_view_pb2.FeatureView.VectorSearchConfig.TreeAHConfig(),\n",
" tree_ah_config=feature_view_pb2.FeatureView.IndexConfig.TreeAHConfig(),\n",
")\n",
"\n",
"print(f\"vector_search_config: {vector_search_config}\")\n",
"print(f\"index_config: {index_config}\")\n",
"\n",
"create_view_lro = admin_client.create_feature_view(\n",
" feature_online_store_admin_service_pb2.CreateFeatureViewRequest(\n",
@@ -888,7 +880,7 @@
" feature_view=feature_view_pb2.FeatureView(\n",
" big_query_source=big_query_source,\n",
" sync_config=sync_config,\n",
" vector_search_config=vector_search_config,\n",
" index_config=index_config,\n",
" ),\n",
" )\n",
")"
@@ -1109,6 +1101,9 @@
},
"outputs": [],
"source": [
"# It will take some time for the DNS to be fully ready\n",
"time.sleep(300)\n",
"\n",
"data_client = FeatureOnlineStoreServiceClient(\n",
" client_options={\"api_endpoint\": PUBLIC_ENDPOINT}\n",
")"
@@ -1222,7 +1217,7 @@
"data_client.fetch_feature_values(\n",
" request=feature_online_store_service_pb2.FetchFeatureValuesRequest(\n",
" feature_view=f\"projects/{PROJECT_ID}/locations/{REGION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}\",\n",
" id=ENTITY_ID,\n",
" data_key=feature_online_store_service_pb2.FeatureViewDataKey(key=ENTITY_ID),\n",
" )\n",
")"
]