Browse Source

fix(ml): clear cache if bad zip file (#3600)

Mert 2 years ago
parent
commit
fa57853bd2
1 changed files with 2 additions and 1 deletions
  1. 2 1
      machine-learning/app/models/base.py

+ 2 - 1
machine-learning/app/models/base.py

@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
 from pathlib import Path
 from shutil import rmtree
 from typing import Any
+from zipfile import BadZipFile
 
 from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf  # type: ignore
 
@@ -23,7 +24,7 @@ class InferenceModel(ABC):
         loader = self.load if eager else self.download
         try:
             loader(**model_kwargs)
-        except (OSError, InvalidProtobuf):
+        except (OSError, InvalidProtobuf, BadZipFile):
             self.clear_cache()
             loader(**model_kwargs)