From 188a31508e66dce47d655d6c815f36dd4bebfa88 Mon Sep 17 00:00:00 2001
From: marcsello <punkosdmarcell@rocketmail.com>
Date: Wed, 4 Aug 2021 14:49:40 +0200
Subject: [PATCH] Added a little caching to model info

---
 .../classifier_cache.py                       | 33 ++++++++++++++-----
 cnn_classification_service/config.py          |  2 ++
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/cnn_classification_service/classifier_cache.py b/cnn_classification_service/classifier_cache.py
index 0a7cf3f..0b7af93 100644
--- a/cnn_classification_service/classifier_cache.py
+++ b/cnn_classification_service/classifier_cache.py
@@ -9,6 +9,8 @@ from typing import Tuple
 from urllib.parse import urljoin
 
 from cnn_classifier import Classifier
+from config import Config
+import time
 
 
 class ClassifierCache:
@@ -20,6 +22,8 @@ class ClassifierCache:
         self._current_classifier = None  # Latest classifier is a classifier that uses the $default model
         self._downloaded_files = []
 
+        self._last_fetch_time = 0
+
         self._session = SessionTracing(propagate=True)
 
     def _cleanup(self):
@@ -63,16 +67,27 @@ class ClassifierCache:
         self._current_classifier = Classifier(model_file_path, weights_file_path)
 
     def get_default_classifier(self) -> Tuple[dict, Classifier]:
-        logging.debug("Fetching model info...")
-        r = self._session.get(self._model_info_url)
-        r.raise_for_status()
 
-        model_details = r.json()
+        if ((time.time() - self._last_fetch_time) > Config.MODEL_CACHE_LIFETIME_SEC) or \
+                (not self._current_model_details):
 
-        if (not self._current_model_details) or (self._current_model_details['id'] != model_details['id']):
-            # If the currently loaded model is not the default... then load it
-            self._cleanup()  # delete/unload everything
-            self._download_and_load_model(model_details['files']['model'], model_details['files']['weights'])
-            self._current_model_details = model_details
+            logging.debug("Fetching model info...")
+            r = self._session.get(self._model_info_url)
+            r.raise_for_status()
+            self._last_fetch_time = time.time()
+
+            model_details = r.json()
+
+            if (not self._current_model_details) or (self._current_model_details['id'] != model_details['id']):
+                logging.info(f"Model needs to be loaded (local: {self._current_model_details['id']} model service default: {model_details['id']})")
+                # If the currently loaded model is not the default... then load it
+                self._cleanup()  # delete/unload everything
+                self._download_and_load_model(model_details['files']['model'], model_details['files']['weights'])
+                self._current_model_details = model_details
+            else:
+                logging.debug(f"Currently loaded model seems up to date ({self._current_model_details['id']} == {model_details['id']})")
+
+        else:
+            logging.debug("Cache is still valid. Not fetching model info")
 
         return self._current_model_details, self._current_classifier
diff --git a/cnn_classification_service/config.py b/cnn_classification_service/config.py
index 12ffae3..7baf953 100644
--- a/cnn_classification_service/config.py
+++ b/cnn_classification_service/config.py
@@ -16,6 +16,8 @@ class Config:
     RELEASE_ID = os.environ.get('RELEASE_ID', 'test')
     RELEASEMODE = os.environ.get('RELEASEMODE', 'dev')
 
+    MODEL_CACHE_LIFETIME_SEC = int(os.environ.get("MODEL_CACHE_LIFETIME_SEC", 15))
+
     LOG_LEVEL = logging.DEBUG if (
                                          '--debug' in sys.argv
                                  ) or (