From a9b70b0d522a34d09efb4be24b2a720dbd767a1d Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Fri, 2 Aug 2024 11:36:02 +1000
Subject: [PATCH] Tweak error handling

---
 api_server.py | 12 ++++++------
 classifier.py | 15 +++++++--------
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/api_server.py b/api_server.py
index 1bf06d1..49f30c7 100644
--- a/api_server.py
+++ b/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))
diff --git a/classifier.py b/classifier.py
index 072d2a4..4252da9 100755
--- a/classifier.py
+++ b/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():