Repository: mozilla/bugbug

Issue:
Pre-load model in HTTP service background worker
Due to RQ worker design, we cannot cache models without implementing our own worker script.

https://github.com/mozilla/bugbug/pull/321/files/225b0a03969aeb5c63c064dec0215e18a632b59f#r291106686


Discussion:

Resolving patch:
diff --git a/http_service/check_models.py b/http_service/check_models.py
--- a/http_service/check_models.py
+++ b/http_service/check_models.py
@@ -8,7 +8,7 @@
 import sys
 
 # Non-relative imports might be brittle
-from models import MODELS_NAMES, load_model
+from models import MODELS_NAMES, get_model
 
 logging.basicConfig(level=logging.INFO)
 LOGGER = logging.getLogger()
@@ -17,7 +17,7 @@
 def check_models():
     for model_name in MODELS_NAMES:
         # Try loading the model
-        load_model(model_name)
+        get_model(model_name)
 
 
 if __name__ == "__main__":
diff --git a/http_service/models.py b/http_service/models.py
--- a/http_service/models.py
+++ b/http_service/models.py
@@ -14,7 +14,7 @@
 from redis import Redis
 
 from bugbug import bugzilla, get_bugbug_version
-from bugbug.models import load_model as bugbug_load_model
+from bugbug.models import load_model
 
 logging.basicConfig(level=logging.INFO)
 LOGGER = logging.getLogger()
@@ -25,9 +25,23 @@
 DEFAULT_EXPIRATION_TTL = 7 * 24 * 3600  # A week
 
 
-def load_model(model):
-    # TODO: Do not crash when the asked model is not one of the trained models
-    return bugbug_load_model(model, MODELS_DIR)
+MODEL_CACHE = {}
+
+
+def get_model(model_name):
+    if model_name not in MODEL_CACHE:
+        print("Recreating the model in cache")
+        model = load_model(model_name, MODELS_DIR)
+
+        MODEL_CACHE[model_name] = model
+        return model
+
+    return MODEL_CACHE[model_name]
+
+
+def preload_models():
+    for model in MODELS_NAMES:
+        get_model(model)
 
 
 def retrieve_model(name):
@@ -93,9 +107,7 @@ def classify_bug(
     if not bugs:
         return "NOK"
 
-    # TODO: Cache the model in the process memory, it's quite hard as the RQ
-    # worker is forking before starting
-    model = load_model(model_name)
+    model = get_model(model_name)
 
     model_extra_data = model.get_extra_data()
 
diff --git a/http_service/worker.py b/http_service/worker.py
new file mode 100644
--- /dev/null
+++ b/http_service/worker.py
@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+# -*- 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 os
+import sys
+from os.path import abspath, dirname, join
+
+from redis import Redis
+from rq import Connection, Worker
+
+sys.path.insert(0, abspath(join(dirname(__file__), "..")))
+
+# We need to make sure that the models.py file is imported with the same name
+# than in the HTTP docker container or the cache will not match. If we import
+# it as models, the cache will be located at
+# `sys.modules['models'].MODEL_CACHE` while the job will look the models in
+# `sys.modules['http_service.models'].MODEL_CACHE`
+import http_service.models  # noqa: E402 isort:skip
+
+# Preload libraries
+http_service.models.preload_models()
+
+# Provide queue names to listen to as arguments to this script,
+# similar to rq worker
+redis_url = os.environ.get("REDIS_URL", "redis://localhost/0")
+redis_conn = Redis.from_url(redis_url)
+with Connection(connection=redis_conn):
+    qs = sys.argv[1:] or ["default"]
+
+    w = Worker(qs)
+    w.work()


Test patch: