From 0ee0cd60bf2db2c9d1dbb0de8c6626c88c23e8b0 Mon Sep 17 00:00:00 2001
From: Bastien Abadie <bastien@nextcairn.com>
Date: Fri, 12 Apr 2019 15:26:06 +0200
Subject: [PATCH] Split models download and check scripts (#279)

---
 http_service/Dockerfile                       |  4 +-
 http_service/check_models.py                  | 51 ++++++++++++++++++
 .../{models.py => download_models.py}         | 53 +++----------------
 3 files changed, 61 insertions(+), 47 deletions(-)
 create mode 100644 http_service/check_models.py
 rename http_service/{models.py => download_models.py} (57%)

diff --git a/http_service/Dockerfile b/http_service/Dockerfile
index 062a0e11..97a4a17d 100644
--- a/http_service/Dockerfile
+++ b/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"]
\ No newline at end of file
+CMD ["gunicorn", "-b", "0.0.0.0:8000", "bugbug_http_service.app", "--preload", "--timeout", "30", "-w", "3"]
diff --git a/http_service/check_models.py b/http_service/check_models.py
new file mode 100644
index 00000000..3a5c274d
--- /dev/null
+++ b/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)
diff --git a/http_service/models.py b/http_service/download_models.py
similarity index 57%
rename from http_service/models.py
rename to http_service/download_models.py
index f20ab650..db4d2d6a 100644
--- a/http_service/models.py
+++ b/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]}")
+    preload_models()