Split models download and check scripts (#279)

http_service/Dockerfile

@@ -11,6 +11,6 @@ WORKDIR /code/
 
 COPY models/ /code/models/
 
-RUN python bugbug_http_service/models.py check
+RUN python bugbug_http_service/check_models.py
 
 CMD ["gunicorn", "-b", "0.0.0.0:8000", "bugbug_http_service.app", "--preload", "--timeout", "30", "-w", "3"]

http_service/check_models.py

@@ -0,0 +1,51 @@
+# -*- coding: utf-8 -*-
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this file,
+# You can obtain one at http://mozilla.org/MPL/2.0/.
+
+import logging
+import os
+import sys
+
+from bugbug.models.component import ComponentModel
+from bugbug.models.defect_enhancement_task import DefectEnhancementTaskModel
+from bugbug.models.regression import RegressionModel
+
+logging.basicConfig(level=logging.INFO)
+LOGGER = logging.getLogger()
+
+MODELS = {
+    "defectenhancementtask": DefectEnhancementTaskModel,
+    "component": ComponentModel,
+    "regression": RegressionModel,
+}
+
+
+def get_model_path(name):
+    file_name = f"{name}model"
+    file_path = os.path.join("models", file_name)
+
+    return file_path
+
+
+def load_model(model):
+    model_file_path = get_model_path(model)
+    model = MODELS[model].load(model_file_path)
+    return model
+
+
+def check_models():
+    for model_name in MODELS.keys():
+        # Try loading the model
+        load_model(model_name)
+
+
+if __name__ == "__main__":
+    try:
+        check_models()
+    except Exception:
+        LOGGER.warning(
+            "Failed to validate the models, please run `python models.py download`",
+            exc_info=True,
+        )
+        sys.exit(1)

http_service/download_models.py

@@ -7,27 +7,20 @@ import logging
 import lzma
 import os
 import shutil
-import sys
 from urllib.request import urlretrieve
 
 import requests
 
-from bugbug.models.component import ComponentModel
-from bugbug.models.defect_enhancement_task import DefectEnhancementTaskModel
-from bugbug.models.regression import RegressionModel
-
 logging.basicConfig(level=logging.INFO)
 LOGGER = logging.getLogger()
 
-
-MODELS = {
-    "defectenhancementtask": DefectEnhancementTaskModel,
-    "component": ComponentModel,
-    "regression": RegressionModel,
-}
-
 BASE_URL = "https://index.taskcluster.net/v1/task/project.releng.services.project.testing.bugbug_train.latest/artifacts/public"
 
+MODELS_NAMES = (
+    "defectenhancementtask",
+    "component",
+    "regression",
+)
 
 def retrieve_model(name):
     os.makedirs("models", exist_ok=True)
@@ -38,6 +31,7 @@ def retrieve_model(name):
     model_url = f"{BASE_URL}/{file_name}.xz"
     LOGGER.info(f"Checking ETAG of {model_url}")
     r = requests.head(model_url, allow_redirects=True)
+    r.raise_for_status()
     new_etag = r.headers["ETag"]
 
     try:
@@ -62,41 +56,10 @@ def retrieve_model(name):
     return file_path
 
 
-def get_model_path(name):
-    file_name = f"{name}model"
-    file_path = os.path.join("models", file_name)
-
-    return file_path
-
-
 def preload_models():
-    for model_name in MODELS.keys():
+    for model_name in MODELS_NAMES:
         retrieve_model(model_name)
 
 
-def load_model(model):
-    model_file_path = get_model_path(model)
-    model = MODELS[model].load(model_file_path)
-    return model
-
-
-def check_models():
-    for model_name in MODELS.keys():
-        # Try loading the model
-        load_model(model_name)
-
-
 if __name__ == "__main__":
-    if sys.argv[1] == "download":
     preload_models()
-    elif sys.argv[1] == "check":
-        try:
-            check_models()
-        except Exception:
-            LOGGER.warning(
-                "Failed to validate the models, please run `python models.py download`",
-                exc_info=True,
-            )
-            sys.exit(1)
-    else:
-        raise ValueError(f"Unknown command {sys.argv[1]}")