From 6591f20fab2e97cd59c04821de5c10ef3339979d Mon Sep 17 00:00:00 2001 From: Joan Fontanals Date: Thu, 15 Jun 2023 17:53:59 +0200 Subject: [PATCH] feat: add deploy to jcloud journey (#20) * feat: add deploy to jcloud journey * feat: push to hubble * feat: complete deploy method code * feat: provide CLI in main Signed-off-by: Joan Fontanals Martinez * ci: add docker building steps * docs: adapt readme to new CLI --------- Signed-off-by: Joan Fontanals Martinez --- .github/workflows/force-docker-build.yml | 135 ++------------- .github/workflows/force-docs-build.yml | 113 ------------- .gitignore | 1 + Dockerfiles/vectordb.Dockerfile | 13 ++ README.md | 69 +++++++- requirements.txt | 1 + resources/jcloud_exec_template/Dockerfile | 6 + resources/jcloud_exec_template/__init__.py | 0 resources/jcloud_exec_template/config.yml | 6 + resources/jcloud_exec_template/executor.py | 5 + .../jcloud_exec_template/vectordb_app.py | 0 setup.py | 6 +- vectordb/__main__.py | 125 ++++++++++++++ vectordb/client/client.py | 5 +- vectordb/db/base.py | 160 ++++++++++++++---- .../db/executors/inmemory_exact_indexer.py | 3 + vectordb/db/executors/typed_executor.py | 23 ++- vectordb/db/service.py | 4 +- vectordb/utils/push_to_hubble.py | 113 +++++++++++++ 19 files changed, 502 insertions(+), 286 deletions(-) delete mode 100644 .github/workflows/force-docs-build.yml create mode 100644 Dockerfiles/vectordb.Dockerfile create mode 100644 resources/jcloud_exec_template/Dockerfile create mode 100644 resources/jcloud_exec_template/__init__.py create mode 100644 resources/jcloud_exec_template/config.yml create mode 100644 resources/jcloud_exec_template/executor.py create mode 100644 resources/jcloud_exec_template/vectordb_app.py create mode 100644 vectordb/__main__.py create mode 100644 vectordb/utils/push_to_hubble.py diff --git a/.github/workflows/force-docker-build.yml b/.github/workflows/force-docker-build.yml index 39e1bf1..b2bedf6 100644 --- a/.github/workflows/force-docker-build.yml +++ b/.github/workflows/force-docker-build.yml @@ -21,118 +21,28 @@ jobs: touch SUCCESS if: inputs.release_token == env.release_token env: - release_token: ${{ secrets.JINA_CORE_RELEASE_TOKEN }} + release_token: ${{ secrets.VECTORDB_RELEASE_TOKEN }} - name: Fail release token run: | [[ -f SUCCESS ]] - regular-release: needs: token-check runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - pip_tag: [ "", "perf", "standard", "devel"] # default: "" = core - py_version: [ "3.7", "3.8", "3.9" , "3.10", "3.11"] # default "" = 3.7 steps: - uses: actions/checkout@v2.5.0 with: fetch-depth: 100 - - name: Set envs and versions + - name: Set up Python 3.9 + uses: actions/setup-python@v1 + with: + python-version: 3.9 + + - name: Get vectordb version run: | - DEFAULT_PY_VERSION="3.8" - VCS_REF=${{ github.ref }} - echo "VCS_REF=$VCS_REF" >> $GITHUB_ENV - echo "Will build $VCS_REF" - echo "BUILD_DATE=$(date -u +'%Y-%m-%dT%H:%M:%SZ')" >> $GITHUB_ENV - echo "BUILD_TARGET=jina" >> $GITHUB_ENV - - if [[ "${{ matrix.pip_tag }}" == "perf" ]]; then - echo "JINA_PIP_INSTALL_PERF=1" >> $GITHUB_ENV - fi - - if [[ "${{ matrix.pip_tag }}" == "" ]]; then - echo "JINA_PIP_INSTALL_CORE=1" >> $GITHUB_ENV - fi - - JINA_VERSION=$(sed -n '/^__version__/p' ./jina/__init__.py | cut -d \' -f2) - V_JINA_VERSION=v${JINA_VERSION} - JINA_MINOR_VERSION=${JINA_VERSION%.*} - JINA_MAJOR_VERSION=${JINA_MINOR_VERSION%.*} - - PY_TAG=${{matrix.py_version}} - if [ -n "${PY_TAG}" ]; then - PY_TAG=-py${PY_TAG//./} - fi - - PIP_TAG=${{ matrix.pip_tag }} - if [ -n "${PIP_TAG}" ]; then - PIP_TAG=-${PIP_TAG} - fi - - git fetch --depth=1 origin +refs/tags/*:refs/tags/* - LAST_VER_TAG=$(git tag -l | sort -V | tail -n1) - PRE_VERSION=-dev$(git rev-list $LAST_VER_TAG..HEAD --count) - - if [[ "${{ github.event.inputs.triggered_by }}" == "CD" ]]; then - - if [[ "${{ matrix.py_version }}" == "$DEFAULT_PY_VERSION" ]]; then - echo "TAG_ALIAS=\ - jinaai/jina:master${PY_TAG}${PIP_TAG}, \ - jinaai/jina:master${PIP_TAG}, \ - jinaai/jina:${JINA_VERSION}${PRE_VERSION}${PIP_TAG}, \ - jinaai/jina:${JINA_VERSION}${PRE_VERSION}${PY_TAG}${PIP_TAG}" \ - >> $GITHUB_ENV - else - # on every CD - echo "TAG_ALIAS=\ - jinaai/jina:master${PY_TAG}${PIP_TAG}, \ - jinaai/jina:${JINA_VERSION}${PRE_VERSION}${PY_TAG}${PIP_TAG}" \ - >> $GITHUB_ENV - fi - - elif [[ "${{ github.event.inputs.triggered_by }}" == "TAG" ]]; then - # on every tag release - - if [[ "${{ matrix.py_version }}" == "$DEFAULT_PY_VERSION" ]]; then - echo "TAG_ALIAS=\ - jinaai/jina:latest${PY_TAG}${PIP_TAG}, \ - jinaai/jina:${JINA_VERSION}${PY_TAG}${PIP_TAG}, \ - jinaai/jina:${JINA_MINOR_VERSION}${PY_TAG}${PIP_TAG}, \ - jinaai/jina:${JINA_MAJOR_VERSION}${PY_TAG}${PIP_TAG}, \ - jinaai/jina:latest${PIP_TAG}, \ - jinaai/jina:${JINA_VERSION}${PIP_TAG}, \ - jinaai/jina:${JINA_MINOR_VERSION}${PIP_TAG}, \ - jinaai/jina:${JINA_MAJOR_VERSION}${PIP_TAG} \ - " >> $GITHUB_ENV - else - echo "TAG_ALIAS=\ - jinaai/jina:latest${PY_TAG}${PIP_TAG}, \ - jinaai/jina:${JINA_VERSION}${PY_TAG}${PIP_TAG}, \ - jinaai/jina:${JINA_MINOR_VERSION}${PY_TAG}${PIP_TAG}, \ - jinaai/jina:${JINA_MAJOR_VERSION}${PY_TAG}${PIP_TAG} \ - " >> $GITHUB_ENV - fi - elif [[ "${{ github.event.inputs.triggered_by }}" == "MANUAL" ]]; then - # on every manual release - if [[ "${{ matrix.py_version }}" == "$DEFAULT_PY_VERSION" ]]; then - echo "TAG_ALIAS=\ - jinaai/jina:${JINA_VERSION}${PIP_TAG}, \ - jinaai/jina:${JINA_VERSION}${PY_TAG}${PIP_TAG} \ - " >> $GITHUB_ENV - else - echo "TAG_ALIAS=\ - jinaai/jina:${JINA_VERSION}${PY_TAG}${PIP_TAG} \ - " >> $GITHUB_ENV - fi - else - echo "Bad triggered_by: ${{ github.event.inputs.triggered_by }}!" - exit 1 - fi - - echo "JINA_VERSION=${JINA_VERSION}" >> $GITHUB_ENV - + pip install -e . + echo "VECTORDB_VERSION=$(python -c 'import vectordb; print(vectordb.__version__)')" >> $GITHUB_ENV + - name: Set up Docker Buildx id: buildx uses: docker/setup-buildx-action@v1 @@ -141,27 +51,12 @@ jobs: - name: Login to DockerHub uses: docker/login-action@v1 with: - username: ${{ secrets.DOCKERHUB_DEVBOT_USER }} - password: ${{ secrets.DOCKERHUB_DEVBOT_TOKEN }} - - run: | - # https://github.com/docker/buildx/issues/464#issuecomment-741507760 - # https://github.com/kubernetes-sigs/azuredisk-csi-driver/pull/808/files - docker run --privileged --rm tonistiigi/binfmt --uninstall qemu-aarch64 - docker run --rm --privileged tonistiigi/binfmt --install all + username: ${{ secrets.DOCKERHUB_JINAVECTORDB_USER }} + password: ${{ secrets.DOCKERHUB_JINAVECTORDB_TOKEN }} - name: Build and push uses: docker/build-push-action@v2 with: - context: . - file: Dockerfiles/debianx.Dockerfile - platforms: linux/amd64,linux/arm64 push: true - tags: ${{env.TAG_ALIAS}} - build-args: | - BUILD_DATE=${{env.BUILD_DATE}} - JINA_VERSION=${{env.JINA_VERSION}} - VCS_REF=${{env.VCS_REF}} - PIP_INSTALL_CORE=${{env.JINA_PIP_INSTALL_CORE}} - PIP_INSTALL_PERF=${{env.JINA_PIP_INSTALL_PERF}} - PY_VERSION=${{matrix.py_version}} - PIP_TAG=${{matrix.pip_tag}} - target: ${{env.BUILD_TARGET}} + context: . + file: Dockerfiles/vectordb.Dockerfile + tags: jinaai/vectordb:latest, jinaai/vectordb:${{ env.VECTORDB_VERSION }} diff --git a/.github/workflows/force-docs-build.yml b/.github/workflows/force-docs-build.yml deleted file mode 100644 index a7ad000..0000000 --- a/.github/workflows/force-docs-build.yml +++ /dev/null @@ -1,113 +0,0 @@ -name: Manual Docs Build - -on: - workflow_dispatch: - inputs: - release_token: - description: 'Your release token' - required: true - triggered_by: - description: 'CD | TAG | MANUAL' - required: false - default: MANUAL - build_old_docs: - description: 'Whether to build old docs (TRUE | FALSE)' - type: string - default: 'FALSE' - package: - description: The name of the repo to build documentation for. - type: string - default: jina - repo_owner: - description: The owner of the repo to build documentation for. Defaults to 'jina-ai'. - type: string - default: jina-ai - pages_branch: - description: Branch that Github Pages observes - type: string - default: gh-pages - git_config_name: - type: string - default: Jina Dev Bot - git_config_email: - type: string - default: dev-bot@jina.ai - -jobs: - token-check: - runs-on: ubuntu-latest - steps: - - name: Check release token - id: token-check - run: | - touch SUCCESS - if: inputs.release_token == env.release_token - env: - release_token: ${{ secrets.JINA_CORE_RELEASE_TOKEN }} - - name: Fail release token - run: | - [[ -f SUCCESS ]] - - build-and-push-latest-docs: - needs: token-check - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v3 - with: - fetch-depth: 1 - - uses: actions/setup-python@v4 - with: - python-version: '3.7' - - name: Install Dependencies - run: | - pip install .[devel] - cd docs - pip install -r requirements.txt - pip install --pre -U furo - pip install sphinx-markdown-tables==0.0.17 - - name: Sphinx Build - run: | - cd docs - bash makedoc.sh local-only - mv ./_build/dirhtml /tmp/gen-html - cd .. - - name: Checkout to GH pages branch (${{ inputs.pages_branch }}) - run: | - git fetch origin ${{ inputs.pages_branch }}:${{ inputs.pages_branch }} --depth 1 - git checkout -f ${{ inputs.pages_branch }} - git reset --hard HEAD - - name: Small config stuff - run: | - touch /tmp/gen-html/.nojekyll - cp ./docs/_versions.json /tmp/gen-html/_versions.json - cp ./docs/CNAME /tmp/gen-html/CNAME - cp /tmp/gen-html/404/index.html /tmp/gen-html/404.html - sed -i 's/href="\.\./href="/' /tmp/gen-html/404.html # fix asset urls that needs to be updated in 404.html - - name: Moving old doc versions - run: | - cd docs - for i in $(cat _versions.json | jq '.[].version' | tr -d '"'); do if [ -d "$i" ]; then mv "$i" /tmp/gen-html; fi; done - - name: Swap in new docs - run: | - rm -rf ./docs - mv /tmp/gen-html ./docs - - name: Push it up! - run: | - git config --local user.email "${{ inputs.git_config_email }}" - git config --local user.name "${{ inputs.git_config_name }}" - git show --summary - git add ./docs && git commit -m "chore(docs): update docs due to ${{github.event_name}} on ${{github.repository}}" - git push origin ${{ inputs.pages_branch }} - - build-old-docs: - needs: build-and-push-latest-docs - runs-on: ubuntu-latest - if: inputs.build_old_docs == 'TRUE' - steps: - - uses: benc-uk/workflow-dispatch@v1 - with: - workflow: Build old docs - token: ${{ secrets.JINA_DEV_BOT }} - inputs: '{ "release_token": "${{ env.release_token }}", "triggered_by": "TAG"}' - env: - release_token: ${{ secrets.JINA_CORE_RELEASE_TOKEN }} diff --git a/.gitignore b/.gitignore index b757764..d87a60e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ # Byte-compiled / optimized / DLL files __pycache__/ +*__pycache__/ *.py[cod] *$py.class diff --git a/Dockerfiles/vectordb.Dockerfile b/Dockerfiles/vectordb.Dockerfile new file mode 100644 index 0000000..a16544b --- /dev/null +++ b/Dockerfiles/vectordb.Dockerfile @@ -0,0 +1,13 @@ +ARG PY_VERSION=3.9 + +FROM python:${PY_VERSION}-slim + +RUN apt-get update && apt-get install --no-install-recommends -y gcc libc6-dev pkg-config wget build-essential + +COPY . /vectordb/ + +RUN cd /vectordb && pip install -U pip && pip install . + +RUN pip install -U docarray[hnswlib]>=0.33 + +ENTRYPOINT ["vectordb"] \ No newline at end of file diff --git a/README.md b/README.md index a372aa5..1019b16 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,9 @@ # Vector Database for Python Developers + +Vector Databases are databases that store embeddings representing data to provide semantic similarity between objects. Vector databases +are used to perform similarity search between multimodal data, such as text, image, audio or videos and also are powering LLM applications +to provide context for LLMs to improve the results of the generation and prevent evaluations. + `vectordb` is a simple, user-friendly solution for Python developers looking to create their own vector database with CRUD support. Vector databases are a key component of the stack needed to use LLMs as they allow them to have access to context and memory. Many of the solutions out there require developers and users to use complex solutions that are often not needed. With `vectordb`, you can easily create your own vector database solution that can work locally and still be easily deployed and served with scalability features such as sharding and replication. Start with your solution as a local library and seamlessly transition into a served database with all the needed capability. No extra complexity than the needed one. @@ -122,6 +127,38 @@ When serving or deploying your Vector Databases you can set 2 scaling parameters ** When deployed to JCloud, the number of replicas will be set to 1. We are working to enable replication in the cloud +## 💻 `vectordb` CLI + +`vectordb` is a simple CLI that helps you to serve and deploy your `vectordb` db. + +First, you need to embed your database instance or class in a python file. + +```python +# example.py +from docarray import DocList, BaseDoc +from docarray.typing import NdArray +from vectordb import InMemoryExactNNVectorDB + + +class MyDoc(BaseDoc): + text: str + embedding: NdArray[128] + + +db = InMemoryExactNNVectorDB[MyDoc](workspace='./vectordb') # notice how `db` is the instance that we want to serve + +if __name__ == '__main__': + # make sure to protect this part of the code + with app.serve() as service: + service.block() +``` + + +| Description | Command | +| --- | ---: | +| Serve your app locally | `vectordb serve --db example:db` | +| Deploy your app on JCloud |`vectordb deploy --db example:db` | + ## :cloud: Deploy it to the cloud @@ -134,17 +171,43 @@ When serving or deploying your Vector Databases you can set 2 scaling parameters ```jc login``` 3. Deploy: + ```bash + vectordb deploy --db example:db + ``` + +
+ Show command output + + ```text + ╭──────────────┬────────────────────────────────────────────────────────────────────────────────────────────────────────╮ + │ App ID │ │ + ├──────────────┼────────────────────────────────────────────────────────────────────────────────────────────────────────┤ + │ Phase │ Serving │ + ├──────────────┼────────────────────────────────────────────────────────────────────────────────────────────────────────┤ + │ Endpoint │ grpc://.wolf.jina.ai │ + ├──────────────┼────────────────────────────────────────────────────────────────────────────────────────────────────────┤ + │ App logs │ dashboards.wolf.jina.ai │ + ╰──────────────┴────────────────────────────────────────────────────────────────────────────────────────────────────────╯ + ``` + +
+ +4. Connect from Client + +Once deployed, you can use `vectordb` Client to access the given endpoint. ```python -HNSWLibDB[MyTextDoc].deploy(config={'data_path'= './hnswlib_path'}, replicas=1, shards=1) +from vectordb import Client + +c = Client(address=' grpc://.wolf.jina.ai') ``` -You can then list and delete your deployed DBs with `jc`: +5. Manage your deployed instances using [jcloud](https://github.com/jina-ai/jcloud): +You can then list and delete your deployed DBs with `jc` command: ```jc list <>``` ```jc delete <>``` - ## 🛣️ Roadmap We have big plans for the future of Vector Database! Here are some of the features we have in the works: diff --git a/requirements.txt b/requirements.txt index eb1a995..2712635 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ jina>=3.17.0 +click #docarray[hnswlib]>=0.33.0 diff --git a/resources/jcloud_exec_template/Dockerfile b/resources/jcloud_exec_template/Dockerfile new file mode 100644 index 0000000..73969a0 --- /dev/null +++ b/resources/jcloud_exec_template/Dockerfile @@ -0,0 +1,6 @@ +FROM jinaai/vectordb:latest + +COPY . /workdir/ +WORKDIR /workdir + +ENTRYPOINT ["jina", "executor", "--uses", "config.yml"] \ No newline at end of file diff --git a/resources/jcloud_exec_template/__init__.py b/resources/jcloud_exec_template/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/resources/jcloud_exec_template/config.yml b/resources/jcloud_exec_template/config.yml new file mode 100644 index 0000000..b8a15f0 --- /dev/null +++ b/resources/jcloud_exec_template/config.yml @@ -0,0 +1,6 @@ +jtype: VectorDBExecutor +metas: + name: vectordb_executor +py_modules: + - vectordb_app.py + - executor.py diff --git a/resources/jcloud_exec_template/executor.py b/resources/jcloud_exec_template/executor.py new file mode 100644 index 0000000..8a79c8c --- /dev/null +++ b/resources/jcloud_exec_template/executor.py @@ -0,0 +1,5 @@ +from vectordb_app import db + + +class VectorDBExecutor(db._executor_cls): + pass diff --git a/resources/jcloud_exec_template/vectordb_app.py b/resources/jcloud_exec_template/vectordb_app.py new file mode 100644 index 0000000..e69de29 diff --git a/setup.py b/setup.py index bde054a..cdce30f 100644 --- a/setup.py +++ b/setup.py @@ -32,13 +32,17 @@ setup( 'License :: OSI Approved :: MIT License', 'Programming Language :: Python', 'Programming Language :: Python :: 3', - 'Programming Language :: Python :: 3.7', 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9', 'Programming Language :: Python :: 3.10', 'Programming Language :: Python :: 3.11', ], python_requires='>=3.7', + entry_points={ + 'console_scripts': [ + 'vectordb=vectordb.__main__:deploy', + ], + }, extras_require={ 'test': [ 'pytest', diff --git a/vectordb/__main__.py b/vectordb/__main__.py new file mode 100644 index 0000000..2db74fa --- /dev/null +++ b/vectordb/__main__.py @@ -0,0 +1,125 @@ +import click + +from vectordb.db.base import VectorDB +from vectordb import __version__ + + +@click.group() +@click.version_option(__version__, '-v', '--version', prog_name='vectordb') +@click.help_option('-h', '--help') +def vectordb(): + pass + + +@vectordb.command(help='Deploy a vectorDB app to Jina AI Cloud') +@click.option( + '--db', + '--app', + type=str, + required=True, + help='VectorDB to deploy, in the format ":"', +) +@click.option( + '--protocol', + '--protocols', + type=str, + default='grpc', + help='Protocol to use to communicate with the VectorDB. It can be a single one or a list of multiple protocols to use. Options are grpc, http and websocket', + required=False, + show_default=True, +) +@click.option( + '--shards', + '-s', + type=int, + default=1, + help='Number of shards to use for the deployed VectorDB', + required=False, + show_default=True, +) +def deploy(db, protocol, shards): + definition_file, _, obj_name = db.partition(":") + protocol = protocol.split(',') + VectorDB.deploy(protoocl=protocol, + shards=shards, + definition_file=definition_file, + obj_name=obj_name) + + +@vectordb.command(help='Deploy a vectorDB app to Jina AI Cloud') +@click.option( + '--db', + '--app', + type=str, + required=True, + help='VectorDB to serve, in the format ":"', +) +@click.option( + '--port', + '-p', + type=str, + default='8081', + help='Port to use to access the VectorDB. It can be a single one or a list corresponding to the protocols', + required=False, + show_default=True, +) +@click.option( + '--protocol', + '--protocols', + type=str, + default='grpc', + help='Protocol to use to communicate with the VectorDB. It can be a single one or a list of multiple protocols to use. Options are grpc, http and websocket', + required=False, + show_default=True, +) +@click.option( + '--replicas', + '-r', + type=int, + default=1, + help='Number of replicas to use for the serving VectorDB', + required=False, + show_default=True, +) +@click.option( + '--shards', + '-s', + type=int, + default=1, + help='Number of shards to use for the served VectorDB', + required=False, + show_default=True, +) +@click.option( + '--workspace', + '-w', + type=str, + default='.', + help='Workspace for the VectorDB to persist its data', + required=False, + show_default=True, +) +def serve(db, port, protocol, replicas, shards, workspace): + import importlib + definition_file, _, obj_name = db.partition(":") + port = port.split(',') + if len(port) == 1: + port = int(port[0]) + else: + port = [int(p) for p in port] + protocol = protocol.split(',') + if definition_file.endswith('.py'): + definition_file = definition_file[:-3] + module = importlib.import_module(definition_file) + db = getattr(module, obj_name) + service = db.serve(protocol=protocol, + port=port, + shards=shards, + replicas=replicas, + workspace=workspace) + with service: + service.block() + + +if __name__ == '__main__': + vectordb() diff --git a/vectordb/client/client.py b/vectordb/client/client.py index 248ee1d..7a0eff8 100644 --- a/vectordb/client/client.py +++ b/vectordb/client/client.py @@ -2,6 +2,7 @@ from typing import TypeVar, Generic, Type, Optional, TYPE_CHECKING from vectordb.utils.unify_input_output import unify_input_output from vectordb.utils.pass_parameters import pass_kwargs_as_params from vectordb.utils.create_doc_type import create_output_doc_type +from vectordb.utils.sort_matches_by_score import sort_matches_by_scores if TYPE_CHECKING: @@ -43,8 +44,9 @@ class Client(Generic[TSchema]): return ClientTyped - def __init__(self, address): + def __init__(self, address, reverse_order=False): from jina import Client as jClient + self.reverse_score_order = reverse_order self._client = jClient(host=address) @unify_input_output @@ -52,6 +54,7 @@ class Client(Generic[TSchema]): def index(self, *args, **kwargs): return self._client.index(*args, **kwargs) + @sort_matches_by_scores @unify_input_output @pass_kwargs_as_params def search(self, *args, **kwargs): diff --git a/vectordb/db/base.py b/vectordb/db/base.py index f9a7561..726eeef 100644 --- a/vectordb/db/base.py +++ b/vectordb/db/base.py @@ -1,10 +1,10 @@ from typing import TypeVar, Generic, Type, Optional, Union, List, Dict, TYPE_CHECKING from vectordb.db.executors.typed_executor import TypedExecutor from vectordb.db.service import Service -from vectordb.utils.create_doc_type import create_output_doc_type from vectordb.utils.unify_input_output import unify_input_output from vectordb.utils.pass_parameters import pass_kwargs_as_params from vectordb.utils.sort_matches_by_score import sort_matches_by_scores +from vectordb.utils.push_to_hubble import push_vectordb_to_hubble if TYPE_CHECKING: from jina import Deployment, Flow @@ -19,7 +19,6 @@ class VectorDB(Generic[TSchema]): # the BaseDoc that defines the schema of the store # for subclasses this is filled automatically _input_schema: Optional[Type['BaseDoc']] = None - _output_schema: Optional[Type['BaseDoc']] = None _executor_type: Optional[Type[TypedExecutor]] = None _executor_cls: Type[TypedExecutor] @@ -38,12 +37,9 @@ class VectorDB(Generic[TSchema]): f'{cls.__name__}[item] `item` should be a Document not a {item} ' ) - out_item = create_output_doc_type(item) - class VectorDBTyped(cls): # type: ignore _input_schema: Type[TSchema] = item - _output_schema: Type[TSchema] = out_item - _executor_cls: Type[TypedExecutor] = cls._executor_type[item, out_item] + _executor_cls: Type[TypedExecutor] = cls._executor_type[item] VectorDBTyped.__name__ = f'{cls.__name__}[{item.__name__}]' VectorDBTyped.__qualname__ = f'{cls.__qualname__}[{item.__name__}]' @@ -61,23 +57,43 @@ class VectorDB(Generic[TSchema]): self._executor = self._executor_cls(*args, **kwargs) @classmethod - def serve(cls, - *, - port: Optional[Union[str, List[str]]] = 8081, - workspace: Optional[str] = None, - protocol: Optional[Union[str, List[str]]] = None, - shards: Optional[int] = None, - replicas: Optional[int] = None, - peer_ports: Optional[Union[Dict[str, List], List]] = None, - **kwargs): + def _get_jina_object(cls, + *, + to_deploy: bool = False, + port: Optional[Union[str, List[str]]] = 8081, + protocol: Optional[Union[str, List[str]]] = None, + workspace: Optional[str] = None, + shards: Optional[int] = None, + replicas: Optional[int] = None, + peer_ports: Optional[Union[Dict[str, List], List]] = None, + definition_file: Optional[str] = None, + obj_name: Optional[str] = None, + **kwargs): from jina import Deployment, Flow + is_instance = False + if isinstance(cls, VectorDB): + is_instance = True + + if is_instance: + workspace = workspace or cls._workspace + replicas = replicas or 1 + shards = shards or 1 protocol = protocol or 'grpc' protocol_list = [p.lower() for p in protocol] if isinstance(protocol, list) else [protocol.lower()] stateful = replicas is not None and replicas > 1 - executor_cls_name = ''.join(cls._executor_cls.__name__.split('[')[0:2]) - ServedExecutor = type(f'{executor_cls_name.replace("[", "").replace("]", "")}', (cls._executor_cls,), - {}) + if not to_deploy: + executor_cls_name = ''.join(cls._executor_cls.__name__.split('[')[0:2]) + ServedExecutor = type(f'{executor_cls_name.replace("[", "").replace("]", "")}', (cls._executor_cls,), + {}) + uses = ServedExecutor polling = {'/index': 'ANY', '/search': 'ALL', '/update': 'ALL', '/delete': 'ALL'} + if to_deploy and replicas > 1: + import warnings + warnings.warn( + 'Deployment with replicas > 1 is not currently available. The deployment will have 1 replica per each ' + 'shard') + replicas = 1 + if 1 < replicas < 3: raise Exception(f'In order for consensus to properly work, at least 3 replicas need to be provided.') @@ -102,6 +118,12 @@ class VectorDB(Generic[TSchema]): kwargs.pop('stateful') use_deployment = True + if to_deploy: + # here we would need to push the EXECUTOR TO HUBBLE AND CHANGE THE USES + assert definition_file is not None, 'Trying to create a Jina Object for Deployment without the file where the vectordb object/class is defined' + assert obj_name is not None, 'Trying to create a Jina Object for Deployment without the name of the vectordb object/class to deploy' + uses = f'jinaai+docker://{push_vectordb_to_hubble(vectordb_name=obj_name, definition_file_path=definition_file)}' + use_deployment = False if 'websocket' in protocol_list: # websocket not supported for Deployment use_deployment = False @@ -110,26 +132,92 @@ class VectorDB(Generic[TSchema]): use_deployment = False if use_deployment: - ctxt_manager = Deployment(uses=ServedExecutor, - port=port, - protocol=protocol, - shards=shards, - replicas=replicas, - stateful=stateful, - peer_ports=peer_ports, - workspace=workspace, - polling=polling, - **kwargs) + jina_object = Deployment(name='indexer', + uses=uses, + port=port, + protocol=protocol, + shards=shards, + replicas=replicas, + stateful=stateful, + peer_ports=peer_ports, + workspace=workspace, + polling=polling, **kwargs) else: - ctxt_manager = Flow(port=port, protocol=protocol, **kwargs).add(uses=ServedExecutor, - shards=shards, - replicas=replicas, - stateful=stateful, - peer_ports=peer_ports, - polling=polling, - workspace=workspace) + jina_object = Flow(port=port, protocol=protocol, **kwargs).add(name='indexer', + uses=uses, + shards=shards, + replicas=replicas, + stateful=stateful, + peer_ports=peer_ports, + polling=polling, + workspace=workspace) - return Service(ctxt_manager, address=f'{protocol_list[0]}://0.0.0.0:{port}', schema=cls._input_schema, reverse_order=cls.reverse_score_order) + return jina_object + + @classmethod + def serve(cls, + *, + port: Optional[Union[str, List[str]]] = 8081, + protocol: Optional[Union[str, List[str]]] = None, + **kwargs): + protocol = protocol or 'grpc' + protocol_list = [p.lower() for p in protocol] if isinstance(protocol, list) else [protocol.lower()] + ctxt_manager = cls._get_jina_object(to_deploy=False, port=port, protocol=protocol, **kwargs) + port = port[0] if isinstance(port, list) else port + return Service(ctxt_manager, address=f'{protocol_list[0]}://0.0.0.0:{port}', schema=cls._input_schema, + reverse_order=cls.reverse_score_order) + + @classmethod + def deploy(cls, + **kwargs): + from tempfile import mkdtemp + import os + import yaml + from yaml.loader import SafeLoader + jina_obj = cls._get_jina_object(to_deploy=True, **kwargs) + + tmpdir = mkdtemp() + jina_obj.save_config(os.path.join(tmpdir, 'flow.yml')) + with open(os.path.join(tmpdir, 'flow.yml'), 'r') as f: + flow_dict = yaml.load(f, SafeLoader) + + executor_jcloud_config = {'resources': {'instance': 'C5', 'autoscale': {'min': 0, 'max': 1}, + 'storage': {'kind': 'ebs', 'size': '10G', 'retain': True}}} + for executor in flow_dict['executors']: + executor['jcloud'] = executor_jcloud_config + + global_jcloud_config = { + 'labels': { + 'app': 'vectordb', + }, + 'monitor': { + 'traces': { + 'enable': True, + }, + 'metrics': { + 'enable': True, + 'host': 'http://opentelemetry-collector.monitor.svc.cluster.local', + 'port': 4317, + }, + }, + } + flow_dict['jcloud'] = global_jcloud_config + import tempfile + from jcloud.flow import CloudFlow + + with tempfile.TemporaryDirectory() as tmpdir: + flow_path = os.path.join(tmpdir, 'flow.yml') + with open(flow_path, 'w') as f: + yaml.safe_dump(flow_dict, f, sort_keys=False) + + cloud_flow = CloudFlow(path=flow_path) + + async def _deploy(): + await cloud_flow.__aenter__() + + import asyncio + ret = asyncio.run(_deploy()) + return ret @pass_kwargs_as_params @unify_input_output diff --git a/vectordb/db/executors/inmemory_exact_indexer.py b/vectordb/db/executors/inmemory_exact_indexer.py index 01cb711..31b4266 100644 --- a/vectordb/db/executors/inmemory_exact_indexer.py +++ b/vectordb/db/executors/inmemory_exact_indexer.py @@ -85,3 +85,6 @@ class InMemoryExactNNIndexer(TypedExecutor): def close(self): if self._index_file_path is not None: self._indexer.persist(self._index_file_path) + + + diff --git a/vectordb/db/executors/typed_executor.py b/vectordb/db/executors/typed_executor.py index 8883305..051fb0d 100644 --- a/vectordb/db/executors/typed_executor.py +++ b/vectordb/db/executors/typed_executor.py @@ -1,7 +1,9 @@ from jina import Executor from jina.serve.executors import _FunctionWithSchema -from typing import TypeVar, Generic, Type, Optional, Tuple, TYPE_CHECKING +from typing import TypeVar, Generic, Type, Optional, TYPE_CHECKING +from vectordb.utils.create_doc_type import create_output_doc_type + if TYPE_CHECKING: from docarray import BaseDoc, DocList @@ -40,18 +42,21 @@ class TypedExecutor(Executor, Generic[InputSchema, OutputSchema]): # Behind-the-scenes magic # # Subclasses should not need to implement these # ################################################## - def __class_getitem__(cls, item: Tuple[Type[InputSchema], Type[OutputSchema]]): + def __class_getitem__(cls, item: Type[InputSchema]): + from docarray import BaseDoc - input_schema, output_schema = item - if not issubclass(input_schema, BaseDoc): + if not isinstance(item, type): + # do nothing + # enables use in static contexts with type vars, e.g. as type annotation + return Generic.__class_getitem__.__func__(cls, item) + if not issubclass(item, BaseDoc): raise ValueError( - f'{cls.__name__}[item, ...] `item` should be a Document not a {input_schema} ' - ) - if not issubclass(output_schema, BaseDoc): - raise ValueError( - f'{cls.__name__}[..., item] `item` should be a Document not a {output_schema} ' + f'{cls.__name__}[item] `item` should be a Document not a {item} ' ) + input_schema = item + output_schema = create_output_doc_type(input_schema) + class ExecutorTyped(cls): # type: ignore _input_schema: Type[InputSchema] = input_schema _output_schema: Type[OutputSchema] = output_schema diff --git a/vectordb/db/service.py b/vectordb/db/service.py index 2739506..7469bdb 100644 --- a/vectordb/db/service.py +++ b/vectordb/db/service.py @@ -1,7 +1,6 @@ from vectordb.client.client import Client from vectordb.utils.unify_input_output import unify_input_output from vectordb.utils.pass_parameters import pass_kwargs_as_params -from vectordb.utils.sort_matches_by_score import sort_matches_by_scores class Service: @@ -9,7 +8,7 @@ class Service: def __init__(self, ctxt_manager, schema, address, reverse_order=False): self.ctxt_manager = ctxt_manager self._reverse_order = reverse_order - self._client = Client[schema](address) + self._client = Client[schema](address, reverse_order=reverse_order) def __enter__(self): self.ctxt_manager.__enter__() @@ -33,7 +32,6 @@ class Service: def index(self, *args, **kwargs): return self._client.index(*args, **kwargs) - @sort_matches_by_scores @pass_kwargs_as_params @unify_input_output def search(self, *args, **kwargs): diff --git a/vectordb/utils/push_to_hubble.py b/vectordb/utils/push_to_hubble.py new file mode 100644 index 0000000..2f3e676 --- /dev/null +++ b/vectordb/utils/push_to_hubble.py @@ -0,0 +1,113 @@ +import uuid +import os +from tempfile import mktemp +import shutil +import sys +import secrets +from shutil import copytree + +from pathlib import Path + +__resources_path__ = os.path.join( + Path(os.path.dirname(sys.modules['vectordb'].__file__)).parent.absolute(), 'resources' +) + + +class EnvironmentVarCtxtManager: + """a class to wrap env vars""" + + def __init__(self, envs): + """ + :param envs: a dictionary of environment variables + """ + self._env_keys_added = envs + self._env_keys_old = {} + + def __enter__(self): + for key, val in self._env_keys_added.items(): + # Store the old value, if it exists + if key in os.environ: + self._env_keys_old[key] = os.environ[key] + # Update the environment variable with the new value + os.environ[key] = str(val) + + def __exit__(self, exc_type, exc_val, exc_tb): + # Restore the old values of updated environment variables + for key, val in self._env_keys_old.items(): + os.environ[key] = str(val) + # Remove any newly added environment variables + for key in self._env_keys_added.keys(): + os.unsetenv(key) + + +def get_random_tag(): + return 't-' + uuid.uuid4().hex[:5] + + +def get_random_name(): + return 'n-' + uuid.uuid4().hex[:5] + + +def _push_to_hubble( + tmpdir: str, name: str, tag: str, verbose: bool, public: bool +) -> str: + from hubble.executor.hubio import HubIO + from hubble.executor.parsers import set_hub_push_parser + + secret = secrets.token_hex(8) + args_list = [ + tmpdir, + '--tag', + tag, + '--no-usage', + '--no-cache', + ] + if verbose: + args_list.remove('--no-usage') + args_list.append('--verbose') + if not public: + args_list.append('--secret') + args_list.append(secret) + args_list.append('--private') + + args = set_hub_push_parser().parse_args(args_list) + + push_envs = ( + {'JINA_HUBBLE_HIDE_EXECUTOR_PUSH_SUCCESS_MSG': 'true'} if not verbose else {} + ) + # return 'a' + ':' + tag + + with EnvironmentVarCtxtManager(push_envs): + executor_id = HubIO(args).push().get('id') + return executor_id + ':' + tag + + +def push_vectordb_to_hubble( + vectordb_name, + definition_file_path, +) -> str: + tmpdir = mktemp() + image_name = get_random_name() + tag = get_random_name() + copytree(os.path.join(__resources_path__, 'jcloud_exec_template'), tmpdir) + shutil.copy(definition_file_path, os.path.join(tmpdir, 'vectordb_app.py')) + # replace `vectordb_name` in `vectordb_app` + with open(os.path.join(tmpdir, 'executor.py'), encoding='utf-8') as f: + content = f.read() + + content = content.replace('from vectordb_app import db', f'from vectordb_app import {vectordb_name}') + content = content.replace('class VectorDBExecutor(db._executor_cls):', + f'class VectorDBExecutor({vectordb_name}._executor_cls):') + + with open(os.path.join(tmpdir, 'executor.py'), mode='w', encoding='utf-8') as f: + f.write(content) + + with open(os.path.join(tmpdir, 'config.yml'), encoding='utf-8') as f: + content = f.read() + + content = content.replace('vectordb_executor', f'vectordb_executor-{image_name}') + + with open(os.path.join(tmpdir, 'config.yml'), mode='w', encoding='utf-8') as f: + f.write(content) + + return _push_to_hubble(tmpdir, image_name, tag, True, False)