Tweak error handling for pickled models

api_server.py

@@ -212,21 +212,21 @@ def on_get(self, req, resp, client, start_slot, end_slot=None):
 if not DISABLE_CLASSIFIER:
     if MODEL_PATH != "":
         if MODEL_PATH.endswith(".pkl"):
-            classifier = import_classifier(MODEL_PATH)
+            try:
+                classifier = import_classifier(MODEL_PATH)
+            except Exception as e:
+                print(f"Failed to persist classifier due to {e}")
+                exit(1)
 
         else:
             print("model path must end with .pkl")
-            exit(0)
+            exit(1)
 
     else:
         print("Initialising classifier, this could take a moment...")
         classifier = MultiClassifier(DATA_DIR) if not DISABLE_CLASSIFIER else None
         print("Done")
 
-if classifier is None:
-    print("The classifier was not loaded")
-    exit(0)
-
 block_db = open_block_db(BLOCK_DB)
 
 app.add_route("/classify/no_store", ClassifyNoStore(classifier))

classifier.py

@@ -291,15 +291,14 @@ def persist_classifier(classifier: Classifier, name: str) -> None:
 
 
 def import_classifier(model_path: str) -> Classifier:
-    print(f"""Loading classifier from {model_path}""")
-
-    try:
-        classifier = pickle.load(open(model_path, "rb"))
-        print("Loaded classifier into memory")
-        return classifier
+    """Load a pickled classifier.
 
-    except Exception as e:
-        print(f"Failed to import classifier due to {e}")
+    This function may throw an exception if the data is corrupt or the file does not exist.
+    """
+    print(f"""Loading classifier from {model_path}""")
+    classifier = pickle.load(open(model_path, "rb"))
+    print("Loaded classifier into memory")
+    return classifier
 
 
 def main():