Make Docker Images for VectorDB Bench

Dockerfile

@@ -0,0 +1,18 @@
+FROM python:3.11-buster as builder-image
+
+RUN apt-get update
+
+COPY install/requirements_py3.11.txt .
+RUN pip3 install -U pip
+RUN pip3 install --no-cache-dir -r requirements_py3.11.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
+
+FROM python:3.11-slim-buster
+
+COPY --from=builder-image /usr/local/bin /usr/local/bin
+COPY --from=builder-image /usr/local/lib/python3.11/site-packages /usr/local/lib/python3.11/site-packages
+
+WORKDIR /opt/code
+COPY . .
+ENV PYTHONPATH /opt/code
+
+ENTRYPOINT ["python3", "-m", "vectordb_bench"]

install.py

@@ -0,0 +1,72 @@
+import os
+import argparse
+import subprocess
+
+def docker_tag_base():
+    return 'vdbbench'
+
+def dockerfile_path_base():
+    return os.path.join('vectordb_bench/', '../Dockerfile')
+
+def docker_tag(track, algo):
+    return docker_tag_base() + '-' + track + '-' + algo
+
+
+def build(tag, args, dockerfile):
+    print('Building %s...' % tag)
+    if args is not None and len(args) != 0:
+        q = " ".join(["--build-arg " + x.replace(" ", "\\ ") for x in args])
+    else:
+        q = ""
+
+    try:
+        command = 'docker build %s --rm -t %s -f' \
+                   % (q, tag)
+        command += ' %s .' % dockerfile
+        print(command)
+        subprocess.check_call(command, shell=True)
+        return {tag: 'success'}
+    except subprocess.CalledProcessError:
+        return {tag: 'fail'}
+
+def build_multiprocess(args):
+    return build(*args)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        "--proc",
+        default=1,
+        type=int,
+        help="the number of process to build docker images")
+    parser.add_argument(
+        '--track',
+        choices=['none'],
+        default='none'
+    )
+    parser.add_argument(
+        '--algorithm',
+        metavar='NAME',
+        help='build only the named algorithm image',
+        default=None)
+    parser.add_argument(
+        '--dockerfile',
+        metavar='PATH',
+        help='build only the image from a Dockerfile path',
+        default=None)
+    parser.add_argument(
+        '--build-arg',
+        help='pass given args to all docker builds',
+        nargs="+")
+    args = parser.parse_args()
+
+    print('Building base image...')
+
+    subprocess.check_call(
+        'docker build \
+        --rm -t %s -f %s .' % (docker_tag_base(), dockerfile_path_base()), shell=True)
+
+    print('Building end.')
+

install/requirements_py3.11.txt

@@ -0,0 +1,23 @@
+grpcio==1.53.0
+grpcio-tools==1.53.0
+qdrant-client
+pinecone-client
+weaviate-client
+elasticsearch
+pgvector
+sqlalchemy
+redis
+chromadb
+pytz
+streamlit-autorefresh
+streamlit>=1.23.0
+streamlit_extras
+tqdm
+s3fs
+psutil
+polars
+plotly
+environs
+pydantic<v2
+scikit-learn
+pymilvus