fix(api): bypass model cache entirely when limit is 0

ssube · Feb 15, 2023 · df5fb93 · df5fb93
1 parent 2b29b09
commit df5fb93
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/api/logging.yaml b/api/logging.yaml
@@ -5,14 +5,14 @@ formatters:
 handlers:
   console:
     class: logging.StreamHandler
-    level: INFO
+    level: DEBUG
     formatter: simple
     stream: ext://sys.stdout
 loggers:
   '':
-    level: INFO
+    level: DEBUG
     handlers: [console]
     propagate: True
 root:
-  level: INFO
-  handlers: [console]
+  level: DEBUG
+  handlers: [console]
diff --git a/api/onnx_web/diffusion/load.py b/api/onnx_web/diffusion/load.py
@@ -96,7 +96,7 @@ def load_pipeline(
     lpw: bool,
 ):
     pipe_key = (pipeline, model, device.device, device.provider, lpw)
-    scheduler_key = (scheduler_type,)
+    scheduler_key = (scheduler_type, model)
 
     cache_pipe = server.cache.get("diffusion", pipe_key)
 

diff --git a/api/onnx_web/server/model_cache.py b/api/onnx_web/server/model_cache.py
@@ -25,6 +25,10 @@ def get(self, tag: str, key: Any) -> Any:
         return None
 
     def set(self, tag: str, key: Any, value: Any) -> None:
+        if self.limit == 0:
+            logger.debug("Cache limit set to 0, not caching model: %s", tag)
+            return
+
         for i in range(len(self.cache)):
             t, k, v = self.cache[i]
             if tag == t: