Split models download and check scripts (#279)

This commit is contained in:
Bastien Abadie 2019-04-12 15:26:06 +02:00 коммит произвёл Marco
Родитель aa578b00fe
Коммит 0ee0cd60bf
3 изменённых файлов: 61 добавлений и 47 удалений

Просмотреть файл

@ -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"]
CMD ["gunicorn", "-b", "0.0.0.0:8000", "bugbug_http_service.app", "--preload", "--timeout", "30", "-w", "3"]

Просмотреть файл

@ -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)

Просмотреть файл

@ -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()