Fix cache (#278)

Add inference cache --------- Signed-off-by: s0nicboOm <i.kushalbatra@gmail.com>
numaproj · Aug 31, 2023 · ac08953 · ac08953
1 parent 3581a82
commit ac08953
Showing 1 changed file with 3 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,8 @@ 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()