Merge pull request #16 from ittia-research/dev

add /health endpoint to wiki_dpr server

datasets/wiki_dpr/Dockerfile.serve.cpu

@@ -2,11 +2,17 @@ FROM python:3.11-bookworm
 RUN apt update && apt install python3 python3-pip git -y
 RUN pip install --no-cache-dir colbert-ai[torch,faiss-cpu]
 RUN pip install --no-cache-dir "numpy<2"
+RUN pip install --no-cache-dir tenacity
 
 WORKDIR /app
 
-CP prepare_files.py .
-CP server.py .
-CP start.sh .
+COPY prepare_files.py .
+COPY server.py .
+COPY start.sh .
 
-CMD ['/app/start.sh']
+RUN chmod +x /app/start.sh
+
+ENV INDEX_ROOT=/data/indexes
+ENV INDEX_NAME=wiki
+
+CMD ["/app/start.sh"]

datasets/wiki_dpr/Dockerfile.serve.gpu

@@ -2,11 +2,17 @@ FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
 RUN apt update && apt install python3 python3-pip git -y
 RUN pip install --no-cache-dir colbert-ai[torch,faiss-gpu]
 RUN pip install --no-cache-dir "numpy<2"
+RUN pip install --no-cache-dir tenacity
 
 WORKDIR /app
 
-CP prepare_files.py .
-CP server.py .
-CP start.sh .
+COPY prepare_files.py .
+COPY server.py .
+COPY start.sh .
 
-CMD ['/app/start.sh']
+RUN chmod +x /app/start.sh
+
+ENV INDEX_ROOT=/data/indexes
+ENV INDEX_NAME=wiki
+
+CMD ["/app/start.sh"]

datasets/wiki_dpr/server.py

@@ -33,6 +33,10 @@ def api_search_query(query, k):
     topk = list(sorted(topk, key=lambda p: (-1 * p['score'], p['pid'])))
     return {"query" : query, "topk": topk}
 
+@app.route("/health", methods=["GET"])
+def health():
+    return ('ok', 200)
+
 @app.route("/api/search", methods=["GET"])
 def api_search():
     if request.method == "GET":

datasets/wiki_dpr/start.sh

@@ -1,4 +1,4 @@
-#/bin/bash
+#!/bin/bash
 
 python ./prepare_files.py
 python ./server.py