wakeup update

embed_utils.py

@@ -13,7 +13,9 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 endpoint = get_inference_endpoint(env_config.tei_name, token=env_config.hf_token)
-
+endpoint.resume().wait().fetch()
+url = endpoint.url
+a = 5
 
 async def embed_chunk(sentence, semaphore, tmp_file):
     async with semaphore:
@@ -47,16 +49,18 @@ async def embed_wrapper(input_ds, temp_file):
 
 
 def wake_up_endpoint():
-    logger.info("Starting up TEI endpoint")
-    n_loop = 0
-    while requests.get(
+    if endpoint.status != 'running':
+        logger.info("Starting up TEI endpoint")
+        endpoint.resume().wait().fetch()
+        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")
+        ).status_code != 200:
+            time.sleep(2)
+            n_loop += 1
+            if n_loop > 20:
+                raise TimeoutError("TEI endpoint is unavailable")
     logger.info("TEI endpoint is up")
 
     return