Adding basic handler cache tests

Signed-off-by: Sachin Varghese <[email protected]>

mlserver/cache/cache.py

@@ -16,3 +16,12 @@ async def lookup(self, key: str) -> str:
         **This method should be overriden to implement your custom cache logic.**
         """
         raise NotImplementedError("lookup() method not implemented")
+
+    async def size(self) -> int:
+        """
+        Method responsible for returning the size of the cache.
+
+
+        **This method should be overriden to implement your custom cache logic.**
+        """
+        raise NotImplementedError("size() method not implemented")

mlserver/cache/local/local.py

@@ -9,7 +9,7 @@ def __init__(self, size=100):
 
     async def insert(self, key: str, value: str):
         self.cache[key] = value
-        if len(self.cache) > self.size_limit:
+        if self.size() > self.size_limit:
             # The cache removes the first entry if it overflows (i.e. in FIFO order)
             self.cache.popitem(last=False)
         return None
@@ -19,3 +19,6 @@ async def lookup(self, key: str) -> str:
             return self.cache[key]
         else:
             return ""
+
+    async def size(self) -> int:
+        return len(self.cache)

mlserver/handlers/dataplane.py

@@ -29,7 +29,7 @@ class DataPlane:
     def __init__(self, settings: Settings, model_registry: MultiModelRegistry):
         self._settings = settings
         self._model_registry = model_registry
-
+        self._response_cache = None
         if settings.cache_enabled:
             self._response_cache = self._create_response_cache()
         self._inference_middleware = InferenceMiddlewares(
@@ -132,3 +132,6 @@ async def infer(
 
     def _create_response_cache(self) -> ResponseCache:
         return LocalCache(size=self._settings.cache_size)
+
+    def _get_response_cache(self) -> Optional[ResponseCache]:
+        return self._response_cache

tests/handlers/test_dataplane.py

@@ -3,7 +3,7 @@
 
 from mlserver.errors import ModelNotReady
 from mlserver.settings import ModelSettings, ModelParameters
-from mlserver.types import MetadataTensor
+from mlserver.types import MetadataTensor, InferenceResponse
 
 from ..fixtures import SumModel
 
@@ -114,3 +114,18 @@ async def test_infer_generates_uuid(data_plane, sum_model, inference_request):
 
     assert prediction.id is not None
     assert prediction.id == str(uuid.UUID(prediction.id))
+
+
+async def test_infer_response_cache(data_plane, sum_model, inference_request):
+    prediction = await data_plane.infer(
+        payload=inference_request, name=sum_model.name, version=sum_model.version
+    )
+    response_cache = data_plane._get_response_cache()
+
+    assert response_cache is not None
+    assert response_cache.size() == 1
+    assert len(
+        InferenceResponse.parse_raw(
+            response_cache.lookup(inference_request.json())
+        ).outputs
+    ) == len(prediction.outputs)

tests/testdata/model-settings.json

@@ -28,5 +28,7 @@
 
   "parameters": {
     "version": "v1.2.3"
-  }
+  },
+
+  "cache_enabled": true
 }

tests/testdata/settings.json

@@ -4,5 +4,6 @@
   "parallel_workers": 2,
   "cors_settings": {
     "allow_origins": ["*"]
-  }
+  },
+  "cache_enabled": true
 }