Add the support for downloading the model before checking it (#452)

Also put the right configuration in the check pipeline

infra/check-pipeline.yml

@@ -9,6 +9,9 @@ tasks:
     payload:
       maxRunTime: 43200
       image: mozilla/bugbug-base
+      env:
+        MODEL_DOWNLOAD_URL: https://index.taskcluster.net/v1/task/project.relman.bugbug.train_component.latest/artifacts/public/componentmodel.xz
+        SHOULD_DOWNLOAD_URL: 1
       command:
         - bugbug-check
         - component

scripts/check.py

@@ -1,8 +1,12 @@
 # -*- coding: utf-8 -*-
 
 import argparse
+import lzma
+import os
+import shutil
 import sys
 from logging import INFO, basicConfig, getLogger
+from urllib.request import urlretrieve
 
 from bugbug.models import load_model
 
@@ -10,8 +14,24 @@ basicConfig(level=INFO)
 logger = getLogger(__name__)
 
 
+def download_model(model_url, file_path):
+    logger.info(f"Downloading model from {model_url!r} and save it in {file_path!r}")
+    urlretrieve(model_url, f"{file_path}.xz")
+
+    with lzma.open(f"{file_path}.xz", "rb") as input_f:
+        with open(file_path, "wb") as output_f:
+            shutil.copyfileobj(input_f, output_f)
+            logger.info(f"Written model in {file_path}")
+
+
 class ModelChecker:
     def go(self, model_name):
+        should_download_model = bool(os.getenv("SHOULD_DOWNLOAD_MODEL"))
+        download_url = os.getenv("MODEL_DOWNLOAD_URL")
+
+        if should_download_model and download_url:
+            download_url = download_model(download_url, f"{model_name}model")
+
         # Load the model
         model = load_model(model_name)