Support base64-encoded image in inference cache (#1285)

Author: papa99do

requirements.txt

@@ -9,4 +9,5 @@
 opentelemetry-api==1.33.1
 opentelemetry-sdk==1.33.1
 
-cachetools==6.1.0
+cachetools==6.1.0
+blake3==1.0.5

src/marqo/inference/inference_cache/caching_inference.py

@@ -1,11 +1,13 @@
 import hashlib
-from typing import Tuple, List
+from typing import Tuple, List, Optional
 
+import blake3
 import numpy as np
 import orjson
 
 from marqo.core.inference.api import Inference, InferenceRequest, InferenceResult, Modality, \
     InferenceErrorModel
+from marqo.core.inference.modality_utils import is_base64_image
 from marqo.inference.inference_cache.marqo_inference_cache import MarqoInferenceCache
 
 
@@ -24,7 +26,12 @@ def vectorise(self, request: InferenceRequest) -> InferenceResult:
         contents_to_vectorise: List[str] = []
 
         for index, content in enumerate(request.contents):
-            embedding = self.inference_cache.get(model_cache_key, content)
+            content_cache_key = self.content_cache_key(content, request.modality)
+            if not content_cache_key:
+                contents_to_vectorise.append(content)
+                continue
+
+            embedding = self.inference_cache.get(model_cache_key, content_cache_key)
             if embedding is not None:
                 cached_result.append((index, content, embedding))
             else:
@@ -43,7 +50,9 @@ def vectorise(self, request: InferenceRequest) -> InferenceResult:
                                        f"Preprocessing config: "
                                        f"{orjson.dumps(dict(new_request.preprocessing_config)).decode('utf-8')}")
                 content, embedding = r[0]
-                self.inference_cache.set(model_cache_key, content, embedding)
+                content_cache_key = self.content_cache_key(content, request.modality)
+                if content_cache_key:
+                    self.inference_cache.set(model_cache_key, content_cache_key, embedding)
 
         # Merge result
         if cached_result:
@@ -61,10 +70,37 @@ def model_cache_key(self, model_properties) -> str:
         data = orjson.dumps(model_properties, option=orjson.OPT_SORT_KEYS)
         return hashlib.md5(data).hexdigest()
 
+    def content_cache_key(self, content: str, modality: Modality) -> Optional[str]:
+        """
+        Generate appropriate cache key for content based on modality.
+        
+        For TEXT modality: use content directly
+        For IMAGE modality: 
+            - if base64 image: use blake3 hash with prefix
+            - otherwise: use content directly (will be skipped in caching logic)
+        
+        Args:
+            content: The content string
+            modality: The modality type
+            
+        Returns:
+            Cache key string, None if it should not be cached
+        """
+        if modality == Modality.TEXT:
+            # Use original content for text and non-base64 images
+            return content
+        elif modality == Modality.IMAGE and is_base64_image(content):
+            # Use blake3 hash for base64 images to save memory
+            hash_digest = blake3.blake3(content.encode()).hexdigest()
+            return f"blake3:{hash_digest}"
+        else:
+            # should not cache non-base64-encoded images
+            return None
+
     def should_skip_cache(self, request):
         return (
             not request.use_inference_cache
             or request.device  # device is only specified to debug embedding, skip caching
-            or request.modality != Modality.TEXT  # we only support text modality for now
+            or request.modality not in [Modality.TEXT, Modality.IMAGE]  # we support text and image modalities
             or request.preprocessing_config.should_chunk  # we do not support caching chunks
         )

tests/api_tests/v1/tests/application_tests/test_env_var_changes.py

@@ -18,17 +18,13 @@
  this test suite's runtime from growing too large.
 """
 import json
-import unittest
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Callable, Optional, List
 
-import math
-
+from marqo import Client
 from tests import marqo_test
 from tests import utilities
 
-from marqo import Client
-
 
 class TestEnvVarChanges(marqo_test.MarqoTestCase):
 
@@ -83,7 +79,7 @@ def test_inference_cache(self):
         """
 
         # Restart marqo with new max values
-        new_models = ["hf/e5-large-v2"]
+        new_models = ["open_clip/ViT-B-32/laion2b_s34b_b79k"]
         index_name = "test_multiple_env_vars"
         utilities.rerun_marqo_with_env_vars(
             env_vars=[
@@ -111,10 +107,12 @@ def test_inference_cache(self):
         telemetry_client = Client(**self.client_settings, return_telemetry=True)
 
         min_inference_time_ms = 8      # inference usually takes at least 8ms
-        cache_reading_time_ms = 2      # if it hits cache, it's usually less than 2ms
+        cache_reading_time_ms = 3      # if it hits cache, the pipeline should take less than 3ms
 
         # Test search query's embedding is cached when inference cache is enabled
-        for query in ["test", {"random": 1, "query": 2}]:
+        base64_image = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
+        image_url = marqo_test.TestImageUrls.HIPPO_STATUE.value
+        for query in ["test", {"random": 1, "query": 2}, base64_image, image_url]:
             with self.subTest(f"Search query: {query}"):
                 # Single query
                 # First search that misses cache should take longer
@@ -129,7 +127,13 @@ def test_inference_cache(self):
                 inference_latency = self._run_in_threads(
                     lambda client: client.index(index_name).search(q=query),
                     max_workers=1, count=10, telemetry_name="search.vector_inference_full_pipeline")
-                self.assertTrue(sum(inference_latency) / 10 < cache_reading_time_ms, inference_latency)
+
+                if query == image_url:
+                    # image url is not cached, so avg latency will usually be > min_inference_time_ms
+                    self.assertTrue(sum(inference_latency) / 10 > min_inference_time_ms, inference_latency)
+                else:
+                    # other queries are all cached, so avg latency should be < cache_reading_time_ms
+                    self.assertTrue(sum(inference_latency) / 10 < cache_reading_time_ms, inference_latency)
 
         # Test to ensure inference cache is not working for add_documents:
         with self.subTest("Add document"):

tests/integ_tests/inference/inference_cache/test_inference_cache.py

@@ -13,7 +13,7 @@
 from orjson import orjson
 
 from marqo.core.inference.api import InferenceRequest, Modality, ModelConfig, TextPreprocessingConfig, Inference, \
-    InferenceResult, InferenceErrorModel
+    InferenceResult, InferenceErrorModel, ImagePreprocessingConfig
 from marqo.inference.inference_cache.caching_inference import CachingInference
 
 
@@ -199,6 +199,82 @@ def test_caching_inference_should_capture_key_metrics(self):
 
                 provider.shutdown()
 
+    def test_base64_image_selective_caching(self):
+        """Test selective caching: base64 images cached, URL images processed normally."""
+        caching_inference = CachingInference(self.inference_local, 10, "LRU")
+
+        base64_png = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
+        base64_jpeg = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQEAYABgAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/"
+        url_image = "https://example.com/image.jpg"
+
+        # Mixed request: 2 base64 images + 1 URL
+        mixed_request = InferenceRequest(
+            modality=Modality.IMAGE,
+            contents=[base64_png, url_image, base64_jpeg],
+            model_config=ModelConfig(
+                model_name="test/clip-model",
+                model_properties={
+                    "name": "test-clip-model",
+                    "dimensions": 512,
+                    "type": "clip"
+                }
+            ),
+            preprocessing_config=ImagePreprocessingConfig(should_chunk=False),
+            use_inference_cache=True
+        )
+
+        # First call - base64 images should be cached, URL processed normally
+        result1 = caching_inference.vectorise(mixed_request)
+
+        # Verify cache contains blake3 keys for both base64 images
+        import blake3
+        model_key = caching_inference.model_cache_key(mixed_request.model_config.model_properties)
+
+        hash1 = blake3.blake3(base64_png.encode()).hexdigest()
+        hash2 = blake3.blake3(base64_jpeg.encode()).hexdigest()
+        cache_key1 = f"blake3:{hash1}"
+        cache_key2 = f"blake3:{hash2}"
+
+        cached_embedding1 = caching_inference.inference_cache.get(model_key, cache_key1)
+        cached_embedding2 = caching_inference.inference_cache.get(model_key, cache_key2)
+
+        self.assertIsNotNone(cached_embedding1, "First base64 image should be cached")
+        self.assertIsNotNone(cached_embedding2, "Second base64 image should be cached")
+
+        # Verify URL image is NOT cached
+        url_cached_embedding = caching_inference.inference_cache.get(model_key, url_image)
+        self.assertIsNone(url_cached_embedding, "URL image should not be cached")
+
+        # Verify cache size (only 2 base64 images cached)
+        self.assertEqual(caching_inference.inference_cache._cache.currsize, 2)
+
+        # Second call with same mixed content
+        result2 = caching_inference.vectorise(mixed_request)
+
+        # Base64 results should return original base64 content (not blake3 keys)
+        png_content1, png_embedding1 = result1.result[0][0]
+        png_content2, png_embedding2 = result2.result[0][0]
+
+        # Content should be original base64, embeddings should be identical (from cache)
+        self.assertEqual(png_content1, base64_png)
+        self.assertEqual(png_content2, base64_png)
+        self.assertTrue(np.array_equal(png_embedding1, png_embedding2))
+
+        jpeg_content1, jpeg_embedding1 = result1.result[2][0]
+        jpeg_content2, jpeg_embedding2 = result2.result[2][0]
+
+        # Content should be original base64, embeddings should be identical (from cache)
+        self.assertEqual(jpeg_content1, base64_jpeg)
+        self.assertEqual(jpeg_content2, base64_jpeg)
+        self.assertTrue(np.array_equal(jpeg_embedding1, jpeg_embedding2))
+
+        # URL results should be unchanged (original URL returned)
+        url_content1, url_embedding1 = result1.result[1][0]
+        url_content2, url_embedding2 = result2.result[1][0]
+        self.assertEqual(url_content1, url_image)  # Original URL unchanged
+        self.assertEqual(url_content2, url_image)  # Original URL unchanged
+        self.assertTrue(np.array_equal(url_embedding1, url_embedding2))
+
     def _assert_metric_value(self, metric_data: MetricsData, name: str, expected_value: Any):
         cache_metrics = metric_data.resource_metrics[0].scope_metrics[0].metrics
         metric = next((metric for metric in cache_metrics if metric.name == name), None)

tests/integ_tests/core/inference/test_cache.py renamed to tests/unit_tests/marqo/inference/inference_cache/test_cache.py

@@ -1,17 +1,10 @@
-import random
-import time
 import unittest
-from concurrent.futures import ThreadPoolExecutor
-from queue import Queue
 
-import numpy as np
-
-from marqo.api.exceptions import EnvVarError
 from marqo.inference.inference_cache.marqo_lfu_cache import MarqoLFUCache
 from marqo.inference.inference_cache.marqo_lru_cache import MarqoLRUCache
 
 
-class TestLFUCache(unittest.TestCase):
+class TestCache(unittest.TestCase):
     """This class tests the LRU and LFU cache implementations."""
 
     def setUp(self):