Inference.

app/api/api_v1/endpoints/openai.py

@@ -26,6 +26,8 @@ import glob
 import papermill as pm
 import shutil
 import nbformat as nbf
+import random
+import time
 
 from langchain.vectorstores.pgvector import (
     CollectionStore,
@@ -41,6 +43,25 @@ PIL.Image.MAX_IMAGE_PIXELS = None
 router = APIRouter()
 embeddings = OpenAIEmbeddings()
 
+def retry_with_backoff(retries = 5, backoff_in_seconds = 1):
+    def rwb(f):
+        def wrapper(*args, **kwargs):
+          x = 0
+          while True:
+            try:
+              return f(*args, **kwargs)
+            except:
+              if x == retries:
+                raise
+
+              sleep = (backoff_in_seconds * 2 ** x +
+                       random.uniform(0, 1))
+              time.sleep(sleep)
+              x += 1
+                
+        return wrapper
+    return rwb
+
 def convert_from_ls(result):
     if 'original_width' not in result or 'original_height' not in result:
         return None
@@ -65,6 +86,14 @@ def embed_documents(collection: CollectionStore):
         pre_delete_embeddings=True
     )
 
+@retry_with_backoff(retries=6)
+def export_tasks(task_id, tmp_archive):
+    label_studio_project.export_tasks("COCO", False, True, [task_id], tmp_archive)
+
+@retry_with_backoff(retries=6)
+def import_tasks(tmp_file):
+    return label_studio_project.import_tasks(tmp_file)
+
 def embed(db: Session, payload: schemas.LabelStudio):
 
     task_id = payload.task["id"]
@@ -83,7 +112,7 @@ def embed(db: Session, payload: schemas.LabelStudio):
 
     notebook_dir = "{}/notebooks".format(tmp_dir)
 
-    label_studio_project.export_tasks("COCO", False, True, [task_id], tmp_archive)
+    export_tasks(task_id, tmp_archive)
 
     with zipfile.ZipFile(tmp_archive, 'r') as zip_ref:
         zip_ref.extractall(tmp_dir)
@@ -178,7 +207,7 @@ async def create_collection(
     target_height = round(height - (height * 0.75))
     cover = cover.resize((target_width, target_height), PIL.Image.Resampling.LANCZOS)
     cover.save(tmp_file, optimize=True, quality=95)
-    tasks = label_studio_project.import_tasks(tmp_file)
+    tasks = import_tasks(tmp_file)
     os.remove(tmp_file)