From 8bed394caf4ae67d31455dae4658e85c8eb14d7c Mon Sep 17 00:00:00 2001 From: Tianrui Yang Date: Thu, 30 May 2024 05:44:30 -0700 Subject: [PATCH] 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 --- ...e_store_based_llm_grounding_tutorial.ipynb | 41 ++++++++----------- 1 file changed, 18 insertions(+), 23 deletions(-) diff --git a/notebooks/official/feature_store/vertex_ai_feature_store_based_llm_grounding_tutorial.ipynb b/notebooks/official/feature_store/vertex_ai_feature_store_based_llm_grounding_tutorial.ipynb index 4d59eaec0..078369169 100644 --- a/notebooks/official/feature_store/vertex_ai_feature_store_based_llm_grounding_tutorial.ipynb +++ b/notebooks/official/feature_store/vertex_ai_feature_store_based_llm_grounding_tutorial.ipynb @@ -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", ")" ]