From d17ab9d4013debb880260121761f7af1bd30352f Mon Sep 17 00:00:00 2001 From: marcsello Date: Wed, 4 Aug 2021 15:50:36 +0200 Subject: [PATCH] Fixed model cache loader --- cnn_classification_service/classifier_cache.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/cnn_classification_service/classifier_cache.py b/cnn_classification_service/classifier_cache.py index 0b7af93..887af21 100644 --- a/cnn_classification_service/classifier_cache.py +++ b/cnn_classification_service/classifier_cache.py @@ -78,14 +78,25 @@ class ClassifierCache: 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 ('id' not in model_details) or not model_details['id']: + raise KeyError("Model info is invalid!") + + current_model_id = None + if self._current_model_details: + current_model_id = self._current_model_details['id'] + + if current_model_id != model_details['id']: + logging.info( + f"Model needs to be loaded (local: {current_model_id}; modelsvc def: {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']})") + logging.debug( + f"Currently loaded model seems up to date ({current_model_id} == {model_details['id']})" + ) else: logging.debug("Cache is still valid. Not fetching model info")