perf: update dump/load paths

perf: use pickle.dumps

Author: andhreljaKern

controller.py

@@ -489,8 +489,8 @@ def delete_embedding(project_id: str, embedding_id: str) -> int:
     org_id = organization.get_id_by_project_id(project_id)
     s3.delete_object(org_id, f"{project_id}/{object_name}")
     request_util.delete_embedding_from_neural_search(embedding_id)
-    json_path = util.INFERENCE_DIR / project_id / embedding_id / "embedder.json"
-    shutil.rmtree(json_path.parent)
+    json_path = util.INFERENCE_DIR / project_id / f"embedder-{embedding_id}.json"
+    json_path.unlink(missing_ok=True)
     return status.HTTP_200_OK
 
 
@@ -627,7 +627,7 @@ def re_embed_records(project_id: str, changes: Dict[str, List[Dict[str, str]]]):
 
 
 def __setup_tmp_embedder(project_id: str, embedder_id: str) -> Transformer:
-    embedder_path = util.INFERENCE_DIR / project_id / embedder_id / "embedder.json"
+    embedder_path = util.INFERENCE_DIR / project_id / f"embedder-{embedder_id}.json"
     if not embedder_path.exists():
         raise Exception(f"Embedder {embedder_id} not found")
     return __load_embedder_by_path(embedder_path)

src/embedders/__init__.py

@@ -55,8 +55,7 @@ def load(self, embedder: dict) -> None:
         """Loads the model configuration and weights from disk.
 
         Args:
-            project_id (str): The ID of the project.
-            embedding_id (str): The ID of the embedding.
+            embedder (dict): The dumped model configuration.
         """
         raise NotImplementedError
 

src/embedders/classification/contextual.py

@@ -49,7 +49,7 @@ def to_json(self) -> dict:
         }
 
     def dump(self, project_id: str, embedding_id: str) -> None:
-        export_file = util.INFERENCE_DIR / project_id / embedding_id / "embedder.json"
+        export_file = util.INFERENCE_DIR / project_id / f"embedder-{embedding_id}.json"
         export_file.parent.mkdir(parents=True, exist_ok=True)
         util.write_json(self.to_json(), export_file, indent=2)
 
@@ -176,6 +176,7 @@ def load(embedder: dict) -> "OpenAISentenceEmbedder":
             model_name=embedder["model_name"],
             batch_size=embedder["batch_size"],
             openai_api_key=embedder["openai_api_key"],
+            # only set for Azure
             api_base=embedder["api_base"],
             api_type=embedder["api_type"],
             api_version=embedder["api_version"],
@@ -187,13 +188,14 @@ def to_json(self) -> dict:
             "model_name": self.model_name,
             "batch_size": self.batch_size,
             "openai_api_key": self.openai_api_key,
+            # only set for Azure
             "api_base": self.api_base,
             "api_type": self.api_type,
             "api_version": self.api_version,
             "use_azure": self.use_azure,
         }
 
     def dump(self, project_id: str, embedding_id: str) -> None:
-        export_file = util.INFERENCE_DIR / project_id / embedding_id / "embedder.json"
+        export_file = util.INFERENCE_DIR / project_id / f"embedder-{embedding_id}.json"
         export_file.parent.mkdir(parents=True, exist_ok=True)
         util.write_json(self.to_json(), export_file, indent=2)

src/embedders/classification/reduce.py

@@ -1,7 +1,10 @@
 from spacy.tokens.doc import Doc
 from typing import Union, List, Generator
 import numpy as np
+import pickle
 from src.embedders import PCAReducer, util
+
+# Embedder imports are used by eval(Embedder) in load methods
 from src.embedders.classification.contextual import (
     OpenAISentenceEmbedder,
     HuggingFaceSentenceEmbedder,
@@ -57,7 +60,9 @@ def _reduce(
 
     @staticmethod
     def load(embedder: dict) -> "PCASentenceReducer":
-        reducer = util.read_pickle(embedder["reducer_pkl"])
+        reducer = pickle.loads(
+            embedder["reducer_pkl_bytes"].encode("latin-1")
+        )  # Decode to latin1 to avoid binary issues in JSON
         Embedder = eval(embedder["embedder"]["cls"])
         return PCASentenceReducer(
             embedder=Embedder.load(embedder["embedder"]),
@@ -68,14 +73,12 @@ def to_json(self) -> dict:
         return {
             "cls": "PCASentenceReducer",
             "embedder": self.embedder.to_json(),
+            "reducer_pkl_bytes": pickle.dumps(self.reducer).decode(
+                "latin-1"
+            ),  # Encode to latin1 to avoid binary issues in JSON
         }
 
     def dump(self, project_id: str, embedding_id: str) -> None:
-        export_file = util.INFERENCE_DIR / project_id / embedding_id / "embedder.json"
+        export_file = util.INFERENCE_DIR / project_id / f"embedder-{embedding_id}.json"
         export_file.parent.mkdir(parents=True, exist_ok=True)
-        pkl_file = util.INFERENCE_DIR / project_id / embedding_id / "reducer.pkl"
-        util.write_pickle(self.reducer, pkl_file)
-
-        json_obj = self.to_json()
-        json_obj["reducer_pkl"] = str(pkl_file)
-        util.write_json(json_obj, export_file, indent=2)
+        util.write_json(self.to_json(), export_file, indent=2)

src/embedders/util.py

@@ -1,11 +1,10 @@
 from typing import Any, Generator, List
 from pathlib import Path
 import numpy as np
-import os
 import json
 import pickle
 
-INFERENCE_DIR = Path(os.getenv("INFERENCE_DIR", "/inference"))
+INFERENCE_DIR = Path("/inference")
 
 
 def batch(documents: List[Any], batch_size: int) -> Generator[List[Any], None, None]:
@@ -28,7 +27,7 @@ def write_pickle(obj: Any, file_path: str, **kwargs) -> None:
         pickle.dump(obj, f, **kwargs)
 
 
-def read_json(file_path: str) -> Any:
+def read_json(file_path: str) -> dict[str, Any]:
     with open(file_path, "r") as f:
         return json.load(f)
 

start

@@ -68,7 +68,6 @@ docker run -d --rm \
 -e MODEL_PROVIDER=http://refinery-model-provider:80 \
 -e WS_NOTIFY_ENDPOINT="http://refinery-websocket:8080" \
 -e NEURAL_SEARCH=http://refinery-neural-search:80 \
--e INFERENCE_DIR=/inference \
 --mount type=bind,source="$(pwd)"/,target=/app \
 -v /var/run/docker.sock:/var/run/docker.sock \
 -v "$MODEL_DIR":/models \