Skip to content

Commit 9422acd

Browse files
authored
Implementing inference api (#21)
* save embedder as pickle * write out embedder pickle for embedding type if is managed * removes inference pickle on embedder removal * pr comments
1 parent 94b8c55 commit 9422acd

File tree

2 files changed

+29
-0
lines changed

2 files changed

+29
-0
lines changed

controller.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
notification,
1111
organization,
1212
)
13+
import pickle
1314
import torch
1415
import traceback
1516
from requests.exceptions import HTTPError
@@ -449,6 +450,16 @@ def run_encoding(
449450
request_util.post_embedding_to_neural_search(
450451
request.project_id, embedding_id
451452
)
453+
454+
if get_config_value("is_managed"):
455+
pickle_path = os.path.join(
456+
"/inference", request.project_id, f"embedder-{embedding_id}.pkl"
457+
)
458+
if not os.path.exists(pickle_path):
459+
os.makedirs(os.path.dirname(pickle_path), exist_ok=True)
460+
with open(pickle_path, "wb") as f:
461+
pickle.dump(embedder, f)
462+
452463
upload_embedding_as_file(request.project_id, embedding_id)
453464
embedding.update_embedding_state_finished(
454465
request.project_id,
@@ -485,6 +496,9 @@ def delete_embedding(project_id: str, embedding_id: str) -> int:
485496
org_id = organization.get_id_by_project_id(project_id)
486497
s3.delete_object(org_id, project_id + "/" + object_name)
487498
request_util.delete_embedding_from_neural_search(embedding_id)
499+
pickle_path = os.path.join("/inference", project_id, f"embedder-{embedding_id}.pkl")
500+
if os.path.exists(pickle_path):
501+
os.remove(pickle_path)
488502
return 200
489503

490504

start

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,20 @@ then
2727
fi
2828
fi
2929

30+
INFERENCE_DIR=${PWD%/*}/dev-setup/inference/
31+
if [ ! -d "$_DIR" ]
32+
then
33+
INFERENCE_DIR=${PWD%/*/*}/dev-setup/inference/
34+
if [ ! -d "$INFERENCE_DIR" ]
35+
then
36+
# to include volume for local development, use the dev-setup inference folder:
37+
# alternative use manual logic with
38+
# -v /path/to/dev-setup/inference:/models \
39+
echo "Can't find model data directory: $INFERENCE_DIR -> stopping"
40+
exit 1
41+
fi
42+
fi
43+
3044
echo -ne 'starting...'
3145
docker run -d --rm \
3246
--name refinery-embedder \
@@ -42,6 +56,7 @@ docker run -d --rm \
4256
--mount type=bind,source="$(pwd)"/,target=/app \
4357
-v /var/run/docker.sock:/var/run/docker.sock \
4458
-v "$MODEL_DIR":/models \
59+
-v "$INFERENCE_DIR":/inference \
4560
--network dev-setup_default \
4661
refinery-embedder-dev > /dev/null 2>&1
4762
echo -ne '\t\t\t [done]\n'

0 commit comments

Comments
 (0)