Add back the support of Languagebind Models (#1188)

We removed the support for Languagebind Models in 2.17. Now we add them back. Languagebind Models are working again.

This PR mainly includes the following changes:
- 1. Add back the support of the Languagebind Models.
- 2. Move the `MARQO_MAX_SEARCH_VIDEO_AUDIO_FILE_SIZE` and `MARQO_MAX_ADD_DOCS_VIDEO_AUDIO_FILE_SIZE` environment variables to the API part. The inference server will get this value from the request.
- 3. Add unit tests for media downloading and preprocessing in the inference server.

Author: wanliAlex

src/marqo/api/configs.py

@@ -64,18 +64,21 @@ def default_env_vars() -> dict:
         EnvVars.MARQO_INFERENCE_POOL_SIZE: 20,  # Please adjust this based on the throttling config
         EnvVars.MARQO_INFERENCE_TIMEOUT: 300,   # 300s to support inference of large batch of media files
 
+        # 370 megabytes in bytes, read in API and passed to inference server
+        EnvVars.MARQO_MAX_SEARCH_VIDEO_AUDIO_FILE_SIZE: 387973120,
+        # 370 megabytes in bytes, read in API and passed to inference server
+        EnvVars.MARQO_MAX_ADD_DOCS_VIDEO_AUDIO_FILE_SIZE: 387973120,
 
+        # Read in API and passed to inference server
+        EnvVars.MARQO_MEDIA_DOWNLOAD_THREAD_COUNT_PER_REQUEST: 5,
+        EnvVars.MARQO_IMAGE_DOWNLOAD_THREAD_COUNT_PER_REQUEST: 20,
+
+        ##########################################
         # Inference Server config (In Inference)
         EnvVars.MARQO_MODELS_TO_PRELOAD: [],
         EnvVars.MARQO_MAX_CPU_MODEL_MEMORY: 4,
         EnvVars.MARQO_MAX_CUDA_MODEL_MEMORY: 4,  # For multi-GPU, this is the max memory for each GPU.
 
-        EnvVars.MARQO_MEDIA_DOWNLOAD_THREAD_COUNT_PER_REQUEST: 5,
-        EnvVars.MARQO_IMAGE_DOWNLOAD_THREAD_COUNT_PER_REQUEST: 20,
-
-        EnvVars.MARQO_MAX_SEARCH_VIDEO_AUDIO_FILE_SIZE: 387973120,  # 370 megabytes in bytes
-        EnvVars.MARQO_MAX_ADD_DOCS_VIDEO_AUDIO_FILE_SIZE: 387973120,  # 370 megabytes in bytes
-
         EnvVars.MARQO_MAX_VECTORISE_BATCH_SIZE: 16,  # static inference batching
         EnvVars.MARQO_INFERENCE_CACHE_SIZE: 0,
         EnvVars.MARQO_INFERENCE_CACHE_TYPE: "LRU",

src/marqo/api/models/add_docs_objects.py

@@ -4,9 +4,7 @@
 from pydantic.v1 import BaseModel, root_validator
 from pydantic.v1 import Field
 
-from marqo.tensor_search.enums import EnvVars
 from marqo.tensor_search.models.private_models import ModelAuth
-from marqo.tensor_search.utils import read_env_vars_and_defaults_ints
 
 
 class AddDocsBodyParams(BaseModel):
@@ -24,18 +22,10 @@ class Config:
     modelAuth: Optional[ModelAuth] = None
     mappings: Optional[dict] = None
     documents: Sequence[Dict[str, Any]]
-    imageDownloadThreadCount: int = Field(default_factory=lambda: read_env_vars_and_defaults_ints(EnvVars.MARQO_IMAGE_DOWNLOAD_THREAD_COUNT_PER_REQUEST))
+    imageDownloadThreadCount: Optional[int] = None
     mediaDownloadThreadCount: Optional[int] = None
     textChunkPrefix: Optional[str] = None
 
-    @root_validator
-    def validate_thread_counts(cls, values):
-        image_count = values.get('imageDownloadThreadCount')
-        media_count = values.get('mediaDownloadThreadCount')
-        if media_count is not None and image_count != read_env_vars_and_defaults_ints(EnvVars.MARQO_IMAGE_DOWNLOAD_THREAD_COUNT_PER_REQUEST):
-            raise ValueError("Cannot set both imageDownloadThreadCount and mediaDownloadThreadCount")
-        return values
-
     @root_validator(skip_on_failure=True)
     def _validate_image_download_headers_and_media_download_headers(cls, values):
         """Validate imageDownloadHeaders and mediaDownloadHeaders. Raise an error if both are set.
@@ -53,4 +43,4 @@ def _validate_image_download_headers_and_media_download_headers(cls, values):
                              "Use mediaDownloadHeaders instead.")
         if image_download_headers:
             values['mediaDownloadHeaders'] = image_download_headers
-        return values
+        return values

src/marqo/core/inference/api/exceptions.py

@@ -24,3 +24,8 @@ class MediaDownloadError(InferenceError):
 class UnsupportedModalityError(InferenceError):
     """Raises if a modality is not supported by a specific model"""
     pass
+
+
+class MediaExceedsMaxSizeError(InferenceError):
+    """Raised when the media exceeds the maximum size limit"""
+    pass

src/marqo/core/inference/api/preprocessing_config.py

@@ -81,6 +81,7 @@ class AudioPreprocessingConfig(PreprocessingConfig):
     download_thread_count: Optional[int] = Field(default=None, alias='downloadThreadCount')
     download_header: Optional[Dict[str, str]] = Field(default=None, alias='downloadHeader')
     chunk_config: Optional[ChunkConfig] = Field(default=None, alias='chunkConfig')
+    max_media_size_bytes: int = Field(ge=1, default=387973120, alias='maxMediaSizeBytes')
 
     @root_validator
     def validate_chunk_config(cls, values):
@@ -99,6 +100,7 @@ class VideoPreprocessingConfig(PreprocessingConfig):
     download_thread_count: Optional[int] = Field(default=None, alias='downloadThreadCount')
     download_header: Optional[Dict[str, str]] = Field(default=None, alias='downloadHeader')
     chunk_config: Optional[ChunkConfig] = Field(default=None, alias='chunkConfig')
+    max_media_size_bytes: int = Field(ge=1, default=387973120, alias='maxMediaSizeBytes')
 
     @root_validator
     def validate_chunk_config(cls, values):

src/marqo/core/models/add_docs_params.py

@@ -22,7 +22,7 @@ class AddDocsParams(BaseModel):
         use_existing_tensors: Whether to use the vectors already in doc (for update docs)
         device: Device used to carry out the document update, if `None` is given, it will be determined inference
         image_download_thread_count: number of threads used to concurrently download images
-        media_download_headers: headers to authenticate media download requests
+        media_download_headers: headers to authenticate media download requests for audio and video
         mappings: a dictionary used to handle all the object field content in the doc,
             e.g., multimodal_combination field
         model_auth: an object used to authorise downloading an object from a datastore
@@ -39,9 +39,8 @@ class Config:
     index_name: str
     device: Optional[str]
     tensor_fields: Optional[List] = Field(default_factory=None)
-    image_download_thread_count: int = Field(default_factory=lambda: read_env_vars_and_defaults_ints(
-        EnvVars.MARQO_IMAGE_DOWNLOAD_THREAD_COUNT_PER_REQUEST))
-    media_download_thread_count: Optional[int]
+    image_download_thread_count: Optional[int] = None
+    media_download_thread_count: Optional[int] = None
     media_download_headers: Optional[dict] = None
     use_existing_tensors: bool = False
     mappings: Optional[dict] = None
@@ -51,12 +50,39 @@ class Config:
     def __init__(self, **data: Any):
         super().__init__(**data)
 
-    @root_validator
+    @root_validator(pre=True)
     def validate_thread_counts(cls, values):
+        """
+        Set the values for image_download_thread_count and media_download_thread_count.
+        There are 4 cases:
+            1. Both not given -> Both reads default values
+            2. Image given, media not given -> media reads default value, image uses given value
+            3. Media set, image not set -> media uses given value, image uses media value
+            4. Both set -> error:
+        Once set, media_download_thread_count is used for audio and video, image_download_thread_count is
+        used for images, when sending inference requests to the inference server.
+        """
         image_count = values.get('image_download_thread_count')
         media_count = values.get('media_download_thread_count')
-        if media_count is not None and image_count != read_env_vars_and_defaults_ints(EnvVars.MARQO_IMAGE_DOWNLOAD_THREAD_COUNT_PER_REQUEST):
-            raise ValueError("Cannot set both image_download_thread_count and media_download_thread_count")
+        if media_count and image_count:
+            raise ValueError("Cannot set both 'image_download_thread_count' and 'media_download_thread_count'.")
+        elif image_count is None and media_count is None:
+            # Set default values for both
+            values['image_download_thread_count'] = (
+                read_env_vars_and_defaults_ints(EnvVars.MARQO_IMAGE_DOWNLOAD_THREAD_COUNT_PER_REQUEST))
+            values['media_download_thread_count'] = (
+                read_env_vars_and_defaults_ints(EnvVars.MARQO_MEDIA_DOWNLOAD_THREAD_COUNT_PER_REQUEST)
+            )
+        elif image_count is None and media_count is not None:
+            values['media_download_thread_count'] = media_count
+            values['image_download_thread_count'] = media_count
+        elif image_count is not None and media_count is None:
+            values['media_download_thread_count'] = (
+                read_env_vars_and_defaults_ints(EnvVars.MARQO_MEDIA_DOWNLOAD_THREAD_COUNT_PER_REQUEST)
+            )
+            values['image_download_thread_count'] = image_count
+        else:
+            raise ValueError("Invalid combination of image_download_thread_count and media_download_thread_count.")
         return values
 
     @validator('docs')

src/marqo/core/vespa_index/add_documents_handler.py

@@ -15,7 +15,9 @@
 from marqo.core.models.marqo_add_documents_response import MarqoAddDocumentsItem, MarqoAddDocumentsResponse
 from marqo.logging import get_logger
 from marqo.tensor_search import validation
+from marqo.tensor_search.enums import EnvVars
 from marqo.tensor_search.telemetry import RequestMetricsStore
+from marqo.tensor_search.utils import read_env_vars_and_defaults_ints
 from marqo.vespa.models import VespaDocument, FeedBatchResponse
 from marqo.vespa.models.get_document_response import Document
 from marqo.vespa.vespa_client import VespaClient
@@ -391,7 +393,7 @@ def _get_preprocessing_config(self, modality: Modality, for_top_level_field: boo
                 should_chunk=for_top_level_field and patch_method is not None,
                 download_thread_count=self.add_docs_params.image_download_thread_count,
                 download_header=self.add_docs_params.media_download_headers,
-                patch_method=None if not for_top_level_field or not patch_method else patch_method.value
+                patch_method=None if not for_top_level_field or not patch_method else patch_method.value,
             )
         elif modality == Modality.AUDIO:
             return AudioPreprocessingConfig(
@@ -401,7 +403,8 @@ def _get_preprocessing_config(self, modality: Modality, for_top_level_field: boo
                 chunk_config=ChunkConfig(
                     split_length=self.marqo_index.audio_preprocessing.split_length,
                     split_overlap=self.marqo_index.audio_preprocessing.split_overlap,
-                )
+                ),
+                max_media_size_bytes=read_env_vars_and_defaults_ints(EnvVars.MARQO_MAX_ADD_DOCS_VIDEO_AUDIO_FILE_SIZE)
             )
         elif modality == Modality.VIDEO:
             return VideoPreprocessingConfig(
@@ -411,7 +414,8 @@ def _get_preprocessing_config(self, modality: Modality, for_top_level_field: boo
                 chunk_config=ChunkConfig(
                     split_length=self.marqo_index.video_preprocessing.split_length,
                     split_overlap=self.marqo_index.video_preprocessing.split_overlap,
-                )
+                ),
+                max_media_size_bytes=read_env_vars_and_defaults_ints(EnvVars.MARQO_MAX_ADD_DOCS_VIDEO_AUDIO_FILE_SIZE)
             )
         else:
             raise InternalError(f'The modality {modality} is not supported.')