Merge pull request #4 from lklic/bulk_add

Add batch indexing API

Author: lklic

api/commands.py

@@ -11,7 +11,7 @@ def format_url_list(urls):
 
 
 def insert_images(
-    model_name, urls, metadatas, image_embeddings=None, replace_existing=True
+    model_name, urls, metadatas, image_embeddings=None, replace_existing=True, fail_on_error=True
 ):
     existing_urls = []
     if not replace_existing:
@@ -24,30 +24,51 @@ def insert_images(
         ]
 
     new_urls = [url for url in urls if url not in existing_urls]
-    image_embeddings = (
-        [
-            embeddings[model_name].get_image_embedding(load_image_from_url(url))
-            for url in new_urls
-        ]
-        if image_embeddings is None
-        else [
-            embedding
-            for url, embedding in zip(urls, image_embeddings)
-            if url not in existing_urls
-        ]
-    )
-    metadatas = [
-        json.dumps(metadata)
-        for url, metadata in zip(urls, metadatas)
-        if url not in existing_urls
-    ]
-
-    if len(new_urls) > 0:
-        collections[model_name].insert([new_urls, image_embeddings, metadatas])
+    
+    # Handle individual image failures
+    successful_urls = []
+    failed_urls = []
+    computed_embeddings = []
+    processed_metadatas = []
+    
+    if image_embeddings is None:
+        # Process each image individually
+        for i, url in enumerate(new_urls):
+            try:
+                embedding = embeddings[model_name].get_image_embedding(load_image_from_url(url))
+                computed_embeddings.append(embedding)
+                processed_metadatas.append(json.dumps(metadatas[urls.index(url)]))
+                successful_urls.append(url)
+            except Exception as e:
+                if fail_on_error:
+                    # Original behavior: propagate the exception
+                    raise
+                else:
+                    # New behavior: collect the error and continue
+                    failed_urls.append({"url": url, "error": str(e)})
+    else:
+        # Use provided embeddings
+        for url, embedding in zip(urls, image_embeddings):
+            if url not in existing_urls:
+                try:
+                    computed_embeddings.append(embedding)
+                    processed_metadatas.append(json.dumps(metadatas[urls.index(url)]))
+                    successful_urls.append(url)
+                except Exception as e:
+                    if fail_on_error:
+                        # Original behavior: propagate the exception
+                        raise
+                    else:
+                        # New behavior: collect the error and continue
+                        failed_urls.append({"url": url, "error": str(e)})
+
+    if len(successful_urls) > 0:
+        collections[model_name].insert([successful_urls, computed_embeddings, processed_metadatas])
 
     return {
-        "added": new_urls,
+        "added": successful_urls,
         "found": existing_urls,
+        "failed": failed_urls,
     }
 
 

api/embeddings.py

@@ -41,7 +41,10 @@ def get_text_embedding(self, text):
 
     @torch.no_grad()
     def get_image_embedding(self, images):
-        inputs = self.processor(images=images, return_tensors="pt")
+        # Ensure image is a list since processor expects a batch    
+        if not isinstance(images, list):         
+            images= [images] 
+        inputs = self.processor(images=images, return_tensors="pt", padding=True)
         inputs = inputs.to(self.device)
         image_embedding = self.model.get_image_features(**inputs)
         image_embedding /= image_embedding.norm(dim=-1, keepdim=True)

api/main.py

@@ -232,6 +232,38 @@ async def restore(model_name: ModelName, images: list[DatabaseEntry]):
     )
 
 
+class BulkInsertResult(BaseModel):
+    added: list[ImageUrl]
+    found: list[ImageUrl]
+    failed: list[dict]
+
+
+@app.post(
+    "/models/{model_name}/add_bulk",
+    status_code=status.HTTP_200_OK,
+    tags=["model"],
+    summary="""
+Adds multiple images to the index at once.
+Existing urls will have their metadata replaced with the provided ones.
+Returns lists of successfully added URLs, existing URLs, and failed URLs with error messages.
+    """.strip(),
+    response_model=BulkInsertResult
+)
+async def add_bulk(model_name: ModelName, images: list[ImageAndMetada]):
+    result = try_rpc(
+        "insert_images",
+        [
+            model_name.value,
+            [image.url for image in images],
+            [check_json_string_length(image.metadata) for image in images],
+            None,  # No pre-computed embeddings
+            True,   # Replace existing
+            False,  # Don't fail on error, continue processing
+        ],
+    )
+    return result
+
+
 @app.post(
     "/models/{model_name}/search",
     tags=["model"],

api/requirements-worker.txt

@@ -7,3 +7,4 @@ pymilvus==2.2.5
 torch==2.0.0
 transformers==4.28.1
 Pillow==9.5.0
+numpy<2.0

api/rpc_client.py

@@ -48,7 +48,7 @@ def __call__(self, command, args):
             ),
             body=json.dumps([command, args]),
         )
-        self.connection.process_data_events(time_limit=60)  # seconds
+        self.connection.process_data_events(time_limit=200)  # seconds
         if self.response_data is None:
             raise RuntimeError("No response (timeout?)")
         response = json.loads(self.response_data)

docker/docker-compose.dev.yml

@@ -69,16 +69,19 @@ services:
     depends_on:
       rabbitmq:
         condition: service_healthy
-      worker:
+      worker_1:
         condition: service_healthy
+      worker_2:
+        condition: service_healthy
+
     environment:
-      WEB_CONCURRENCY: 16
+      WEB_CONCURRENCY: 3
       RABBITMQ_URL: amqp://guest:guest@rabbitmq:5672
     volumes:
       - ../api:/app
     ports:
       - 4213:4213
-    command: "uvicorn main:app --reload --host 0.0.0.0 --port 4213"
+    command: "uvicorn main:app --workers 5 --host 0.0.0.0 --port 4213"
     logging:
       options:
         max-size: "10M"
@@ -103,7 +106,7 @@ services:
         max-size: "10M"
         max-file: "10"
 
-  worker:
+  worker_1: &worker_1
     build:
       context: ../api
       dockerfile: ../docker/Dockerfile
@@ -122,6 +125,9 @@ services:
         max-size: "10M"
         max-file: "10"
 
+  worker_2:
+    <<: *worker_1
+
   dev:
     build:
       context: ../api