Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Setting max size to local cache implementation
Browse files Browse the repository at this point in the history
Signed-off-by: Sachin Varghese <[email protected]>
SachinVarghese committed Oct 15, 2023
1 parent bce5d36 commit 4c552de
Showing 6 changed files with 15 additions and 7 deletions.
2 changes: 2 additions & 0 deletions mlserver/cache/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
from .cache import ResponseCache
from .local import LocalCache
1 change: 1 addition & 0 deletions mlserver/cache/local/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from .local import LocalCache
9 changes: 6 additions & 3 deletions mlserver/cache/local/local.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,15 @@
from collections import OrderedDict
from ..cache import ResponseCache


class LocalCache(ResponseCache):
def __init__(self):
self.cache = {}
def __init__(self,size=100):
self.cache = OrderedDict()
self.size_limit = size

async def insert(self, key: str, value: str):
self.cache[key] = value
if len(self.cache) > self.size_limit:
self.cache.popitem(last=False)
return None

async def lookup(self, key: str) -> str:
2 changes: 1 addition & 1 deletion mlserver/handlers/dataplane.py
Original file line number Diff line number Diff line change
@@ -18,7 +18,7 @@
from ..middleware import InferenceMiddlewares
from ..cloudevents import CloudEventsMiddleware
from ..utils import generate_uuid
from ..cache.cache import ResponseCache
from ..cache import ResponseCache


class DataPlane:
5 changes: 2 additions & 3 deletions mlserver/server.py
Original file line number Diff line number Diff line change
@@ -18,8 +18,7 @@
from .metrics import MetricsServer
from .kafka import KafkaServer
from .utils import logger
from .cache.cache import ResponseCache
from .cache.local.local import LocalCache
from .cache import ResponseCache, LocalCache

HANDLED_SIGNALS = [signal.SIGINT, signal.SIGTERM, signal.SIGQUIT]

@@ -63,7 +62,7 @@ def __init__(self, settings: Settings):

def _create_response_cache(self) -> ResponseCache:
if self._settings.cache_enabled:
return LocalCache()
return LocalCache(size=self._settings.cache_size)
else:
return None

3 changes: 3 additions & 0 deletions mlserver/settings.py
Original file line number Diff line number Diff line change
@@ -244,6 +244,9 @@ class Config:
cache_enabled: Optional[bool] = False
"""Enable Caching for the model predictions."""

cache_size: Optional[int] = 100
"""Cache size to be used if caching is enabled."""


class ModelParameters(BaseSettings):
"""

0 comments on commit 4c552de

Please sign in to comment.