feat: add cachesize for inference cache

Signed-off-by: s0nicboOm <i.kushalbatra@gmail.com>
numaproj · Aug 31, 2023 · 0243142 · 0243142
1 parent 3581a82
commit 0243142
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/numalogic/udfs/inference.py b/numalogic/udfs/inference.py
@@ -21,6 +21,7 @@
 
 # TODO: move to config
 LOCAL_CACHE_TTL = int(os.getenv("LOCAL_CACHE_TTL", "3600"))
+LOCAL_CACHE_SIZE = int(os.getenv("LOCAL_CACHE_SIZE", "10000"))
 
 
 class InferenceUDF(NumalogicUDF):
@@ -35,7 +36,7 @@ class InferenceUDF(NumalogicUDF):
     def __init__(self, r_client: redis_client_t, pl_conf: Optional[PipelineConf] = None):
         super().__init__(is_async=False)
         self.model_registry = RedisRegistry(
-            client=r_client, cache_registry=LocalLRUCache(ttl=LOCAL_CACHE_TTL)
+            client=r_client, cache_registry=LocalLRUCache(ttl=LOCAL_CACHE_TTL, cachesize=LOCAL_CACHE_SIZE)
         )
         self.pl_conf = pl_conf or PipelineConf()