wakeup update

embed_utils.py

@@ -2,6 +2,7 @@ import json
 import asyncio
 import logging
 import time
+import requests
 
 from tqdm.asyncio import tqdm_asyncio
 from huggingface_hub import get_inference_endpoint
@@ -46,9 +47,23 @@ async def embed_wrapper(input_ds, temp_file):
 
 
 def wake_up_endpoint():
-    endpoint.fetch()
-    if endpoint.status != 'running':
-        logger.info("Starting up TEI endpoint")
-        endpoint.resume().wait()
+    logger.info("Starting up TEI endpoint")
+    n_loop = 0
+    while requests.get(
+            url=endpoint.url,
+            headers={"Authorization": f"Bearer {env_config.hf_token}"}
+    ).status_code != 200:
+        time.sleep(2)
+        n_loop += 1
+        if n_loop > 40:
+            raise TimeoutError("TEI endpoint is unavailable")
     logger.info("TEI endpoint is up")
+
     return
+
+    # endpoint.fetch()
+    # if endpoint.status != 'running':
+    #     logger.info("Starting up TEI endpoint")
+    #     endpoint.resume().wait()
+    # logger.info("TEI endpoint is up")
+    # return