Merge branch 'master' of https://github.com/datacommonsorg/website in…

…to api

deploy/nl/embeddings.yaml

@@ -11,6 +11,42 @@
 # 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.
+#
+#
+# This is the main config file for NL models and embeddings,
+# with two sections.
+# 
+# models:
+# - type: what type of model serving infra? (LOCAL, VERTEXAI)
+# - usage: what is the model used for?  (EMBEDDINGS, RERANKING)
+# - score_threshold: For embeddings model, what is the cutoff threshold
+#                    below which we drop matches? (default: 0.5)
+#
+# indexes:
+# - store: what type of embeddings store?  (MEMORY, LANCEDB, VERTEXAI)
+# - model: the name of the associated model from `models` section
+# - embeddings: For MEMORY/LANCEDB, the path to the index files.
+#               Can be a local absolute path or GCS (gs://) path.
+# - healthcheck_query: if this index were the default index, what is
+#                      the query to use for health-checking the index?
+# - Additional params specific to VERTEXAI:
+#   - project_id
+#   - location
+#   - index_endpoint_root
+#   - index_endpoint
+#   - index_id
+# TODO: Move the VERTEXAI index params to per-env yamls.
+#
+# Additionally, every environment (prod, etc) may define the following
+# parameters if NL is enabled:
+# - defaultIndex: Which index in `indexes` section is the default?
+# - enabledIndexes: What are the indexes to load on server start?
+# - enableRanking: Should the RANKING models be loaded on server start?
+# - vertexAIModels: Additional endpoint configuration for Vertex AI models
+#
+# NOTE: Ensure every VERTEXAI model below has parameters in the
+# env yaml files for all the environments you're enabling it for.
+#
 
 version: 1
 
@@ -19,26 +55,32 @@ indexes:
     store: MEMORY
     embeddings: gs://datcom-nl-models/embeddings_medium_2024_05_09_18_01_32.ft_final_v20230717230459.all-MiniLM-L6-v2.csv
     model: ft-final-v20230717230459-all-MiniLM-L6-v2
+    healthcheck_query: "Life expectancy"
   sdg_ft:
     store: MEMORY
     embeddings: gs://datcom-nl-models/embeddings_sdg_2023_12_26_10_03_03.ft_final_v20230717230459.all-MiniLM-L6-v2.csv
     model: ft-final-v20230717230459-all-MiniLM-L6-v2
+    healthcheck_query: "Hunger"
   undata_ft:
     store: MEMORY
     embeddings: gs://datcom-nl-models/embeddings_undata_2024_03_20_11_01_12.ft_final_v20230717230459.all-MiniLM-L6-v2.csv
     model: ft-final-v20230717230459-all-MiniLM-L6-v2
+    healthcheck_query: "Hunger"
   undata_ilo_ft:
     store: MEMORY
     embeddings: gs://datcom-nl-models/embeddings_undata_ilo_2024_05_15_11_18_05.ft_final_v20230717230459.all-MiniLM-L6-v2.csv
     model: ft-final-v20230717230459-all-MiniLM-L6-v2
+    healthcheck_query: "Employment"
   bio_ft:
     store: MEMORY
     embeddings: gs://datcom-nl-models/embeddings_bio_2024_03_19_16_39_03.ft_final_v20230717230459.all-MiniLM-L6-v2.csv
     model: ft-final-v20230717230459-all-MiniLM-L6-v2
+    healthcheck_query: "Gene"
   medium_lance_ft:
     store: LANCEDB
     embeddings: gs://datcom-nl-models-dev/lancedb/lancedb_embeddings_medium_2024_04_09_07_22_21.ft_final_v20230717230459.all-MiniLM-L6-v2
     model: ft-final-v20230717230459-all-MiniLM-L6-v2
+    healthcheck_query: "Life expectancy"
   medium_vertex_ft:
     store: VERTEXAI
     project_id: datcom-website-dev
@@ -47,10 +89,12 @@ indexes:
     index_endpoint: projects/496370955550/locations/us-central1/indexEndpoints/8500794985312944128
     index_id: dc_all_minilm_l6_v2_ft_1709655496660
     model: dc-all-minilm-l6-v2-model
+    healthcheck_query: "Life expectancy"
   base_uae_mem:
     store: MEMORY
     embeddings: gs://datcom-nl-models/embeddings_medium_2024_05_16_13_45_32.8110162693219942400.csv
     model: uae-large-v1-model
+    healthcheck_query: "Life expectancy"
   medium_vertex_mistral:
     store: VERTEXAI
     project_id: datcom-website-dev
@@ -59,6 +103,7 @@ indexes:
     index_endpoint: projects/496370955550/locations/us-central1/indexEndpoints/8500794985312944128
     index_id: sfr_embedding_mistral_1709682235285
     model: sfr-embedding-mistral-model
+    healthcheck_query: "Life expectancy"
 
 models:
   dc-all-minilm-l6-v2-model:

nl_server/config.py

@@ -28,6 +28,9 @@
 NL_EMBEDDINGS_VERSION_KEY: str = 'NL_EMBEDDINGS_VERSION_MAP'
 EMBEDDINGS_SPEC_KEY: str = 'EMBEDDINGS_SPEC'
 
+# Query to use to check index health if this is the default index.
+_HEALTHCHECK_QUERY = 'health'
+
 
 class StoreType(str, Enum):
   MEMORY = 'MEMORY'
@@ -68,6 +71,7 @@ class LocalModelConfig(ModelConfig):
 class IndexConfig(ABC):
   store_type: str
   model: str
+  healthcheck_query: str
 
 
 @dataclass
@@ -132,16 +136,19 @@ def parse_v1(embeddings_map: Dict[str, any], vertex_ai_model_info: Dict[str,
   for index_name, index_info in embeddings_map.get('indexes', {}).items():
     store_type = index_info['store']
     used_models.add(index_info['model'])
+    healthcheck_query = index_info.get('healthcheck_query', _HEALTHCHECK_QUERY)
     if store_type == StoreType.MEMORY:
       indexes[index_name] = MemoryIndexConfig(
           store_type=store_type,
           model=index_info['model'],
-          embeddings_path=index_info['embeddings'])
+          embeddings_path=index_info['embeddings'],
+          healthcheck_query=healthcheck_query)
     elif store_type == StoreType.LANCEDB:
       indexes[index_name] = LanceDBIndexConfig(
           store_type=store_type,
           model=index_info['model'],
-          embeddings_path=index_info['embeddings'])
+          embeddings_path=index_info['embeddings'],
+          healthcheck_query=healthcheck_query)
     elif store_type == StoreType.VERTEXAI:
       indexes[index_name] = VertexAIIndexConfig(
           store_type=store_type,
@@ -150,7 +157,8 @@ def parse_v1(embeddings_map: Dict[str, any], vertex_ai_model_info: Dict[str,
           location=index_info['location'],
           index_endpoint_root=index_info['index_endpoint_root'],
           index_endpoint=index_info['index_endpoint'],
-          index_id=index_info['index_id'])
+          index_id=index_info['index_id'],
+          healthcheck_query=healthcheck_query)
     else:
       raise AssertionError(
           'Error parsing information for index {index_name}: unsupported store type {store_type}'

nl_server/embeddings.py

@@ -58,7 +58,8 @@ def encode(self, queries: List[str]) -> List[List[float]] | torch.Tensor:
 #
 class EmbeddingsStore(ABC):
 
-  def __init__(self, needs_tensor=False):
+  def __init__(self, healthcheck_query: str, needs_tensor=False):
+    self.healthcheck_query = healthcheck_query
     self.needs_tensor = needs_tensor
 
   @abstractmethod

nl_server/routes.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 
 import json
-from typing import List
+import logging
+from typing import Dict, List
 
 from flask import Blueprint
 from flask import current_app
@@ -38,14 +39,20 @@ def healthz():
   default_index_type = current_app.config[
       config.EMBEDDINGS_SPEC_KEY].default_index
   if not default_index_type:
+    logging.warning('Health Check Failed: Default index name empty!')
     return 'Service Unavailable', 500
-  nl_embeddings = current_app.config[config.NL_EMBEDDINGS_KEY].get_index(
-      default_index_type)
+  nl_embeddings: Embeddings = current_app.config[
+      config.NL_EMBEDDINGS_KEY].get_index(default_index_type)
   if nl_embeddings:
-    result: VarCandidates = search.search_vars(
-        [nl_embeddings], ['life expectancy'])['life expectancy']
-    if result.svs and 'Expectancy' in result.svs[0]:
+    query = nl_embeddings.store.healthcheck_query
+    result: VarCandidates = search.search_vars([nl_embeddings],
+                                               [query]).get(query)
+    if result and result.svs:
       return 'OK', 200
+    else:
+      logging.warning(f'Health Check Failed: query "{query}" failed!')
+  else:
+    logging.warning('Health Check Failed: Default index not yet loaded!')
   return 'Service Unavailable', 500
 
 

nl_server/store/lancedb.py

@@ -38,7 +38,8 @@ class LanceDBStore(EmbeddingsStore):
   """Manages the embeddings."""
 
   def __init__(self, idx_info: LanceDBIndexConfig) -> None:
-    super().__init__(needs_tensor=False)
+    super().__init__(healthcheck_query=idx_info.healthcheck_query,
+                     needs_tensor=False)
 
     if idx_info.embeddings_path.startswith('/'):
       lance_db_dir = idx_info.embeddings_path

nl_server/store/memory.py

@@ -33,7 +33,8 @@ class MemoryEmbeddingsStore(EmbeddingsStore):
   """Manages the embeddings."""
 
   def __init__(self, idx_info: MemoryIndexConfig) -> None:
-    super().__init__(needs_tensor=True)
+    super().__init__(healthcheck_query=idx_info.healthcheck_query,
+                     needs_tensor=True)
 
     if idx_info.embeddings_path.startswith('/'):
       embeddings_path = idx_info.embeddings_path

nl_server/store/vertexai.py

@@ -28,7 +28,8 @@ class VertexAIStore(EmbeddingsStore):
   """Manages the embeddings."""
 
   def __init__(self, index_info: VertexAIIndexConfig) -> None:
-    super().__init__()
+    super().__init__(healthcheck_query=index_info.healthcheck_query,
+                     needs_tensor=False)
 
     self.index_endpoint = index_info.index_endpoint
     self.index_id = index_info.index_id