diff --git a/controller.py b/controller.py index 3e5c9e1..cc81153 100644 --- a/controller.py +++ b/controller.py @@ -10,6 +10,7 @@ notification, organization, ) +import pickle import torch import traceback from requests.exceptions import HTTPError @@ -449,6 +450,16 @@ def run_encoding( request_util.post_embedding_to_neural_search( request.project_id, embedding_id ) + + if get_config_value("is_managed"): + pickle_path = os.path.join( + "/inference", request.project_id, f"embedder-{embedding_id}.pkl" + ) + if not os.path.exists(pickle_path): + os.makedirs(os.path.dirname(pickle_path), exist_ok=True) + with open(pickle_path, "wb") as f: + pickle.dump(embedder, f) + upload_embedding_as_file(request.project_id, embedding_id) embedding.update_embedding_state_finished( request.project_id, @@ -485,6 +496,9 @@ 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, project_id + "/" + object_name) request_util.delete_embedding_from_neural_search(embedding_id) + pickle_path = os.path.join("/inference", project_id, f"embedder-{embedding_id}.pkl") + if os.path.exists(pickle_path): + os.remove(pickle_path) return 200 diff --git a/start b/start index 3aedfd3..e0dcd48 100755 --- a/start +++ b/start @@ -27,6 +27,20 @@ then fi fi +INFERENCE_DIR=${PWD%/*}/dev-setup/inference/ +if [ ! -d "$_DIR" ] +then + INFERENCE_DIR=${PWD%/*/*}/dev-setup/inference/ + if [ ! -d "$INFERENCE_DIR" ] + then + # to include volume for local development, use the dev-setup inference folder: + # alternative use manual logic with + # -v /path/to/dev-setup/inference:/models \ + echo "Can't find model data directory: $INFERENCE_DIR -> stopping" + exit 1 + fi +fi + echo -ne 'starting...' docker run -d --rm \ --name refinery-embedder \ @@ -42,6 +56,7 @@ docker run -d --rm \ --mount type=bind,source="$(pwd)"/,target=/app \ -v /var/run/docker.sock:/var/run/docker.sock \ -v "$MODEL_DIR":/models \ +-v "$INFERENCE_DIR":/inference \ --network dev-setup_default \ refinery-embedder-dev > /dev/null 2>&1 echo -ne '\t\t\t [done]\n'