Bug 1956003 - use MOZ_ML_LOCAL_DIR instead of MOZ_FETCHES_DIR for local ML test files r=padenot

Differential Revision: https://phabricator.services.mozilla.com/D242747

toolkit/components/ml/docs/perf.rst

@@ -85,7 +85,7 @@ We provide a script to automate this.
   (Use `--model <key>` to clone one of these repositories.)
 
 You can then use `--model` to download locally models, by specifying the local
-`MOZ_FETCHES_DIR` directory, via the env var or command line argument :
+`MOZ_ML_LOCAL_DIR` directory, via the env var or command line argument :
 
 .. code-block:: bash
 
@@ -102,9 +102,9 @@ Once done, you should then be able to run it locally with :
 
 .. code-block:: bash
 
-   $ MOZ_FETCHES_DIR=~/ml-fetches ./mach perftest toolkit/components/ml/tests/browser/browser_ml_engine_perf.js --mochitest-extra-args=headless
+   $ MOZ_ML_LOCAL_DIR=~/ml-fetches ./mach perftest toolkit/components/ml/tests/browser/browser_ml_engine_perf.js --mochitest-extra-args=headless
 
-Notice that `MOZ_FETCHES_DIR` is an absolute path to the `root` directory.
+Notice that `MOZ_ML_LOCAL_DIR` is an absolute path to the `root` directory.
 
 Add in the CI
 -------------

toolkit/components/ml/tests/browser/browser_ml_llama_summarizer_perf.js

@@ -103,7 +103,7 @@ requestLongerTimeout(120);
 
 // To run locally
 // pip install huggingface-hub
-// huggingface-cli download {model_id} --local-dir MOZ_FETCHES_DIR/onnx-models/{model_id}/{revision}
+// huggingface-cli download {model_id} --local-dir MOZ_ML_LOCAL_DIR/onnx-models/{model_id}/{revision}
 
 // Update your test in
 // Then run:  ./mach lint -l perfdocs --fix .

toolkit/components/ml/tests/browser/browser_ml_summarizer_perf.js

@@ -91,7 +91,7 @@ requestLongerTimeout(120);
 
 // To run locally
 // pip install huggingface-hub
-// huggingface-cli download {model_id} --local-dir MOZ_FETCHES_DIR/onnx-models/{model_id}/{revision}
+// huggingface-cli download {model_id} --local-dir MOZ_ML_LOCAL_DIR/onnx-models/{model_id}/{revision}
 
 // Update your test in
 // Then run:  ./mach lint -l perfdocs --fix .

toolkit/components/ml/tests/browser/head.js

@@ -394,7 +394,7 @@ async function perfSetup({ disabled = false, prefs = [], backend } = {}) {
     set: finalPrefs,
   });
 
-  const artifactDirectory = normalizePathForOS(
+  let artifactDirectory = normalizePathForOS(
     `${Services.env.get("MOZ_FETCHES_DIR")}`
   );
 
@@ -407,18 +407,24 @@ async function perfSetup({ disabled = false, prefs = [], backend } = {}) {
   }
 
   // Stop immediately if this fails.
+  if (!artifactDirectory) {
+    artifactDirectory = normalizePathForOS(
+      `${Services.env.get("MOZ_ML_LOCAL_DIR")}`
+    );
+  }
+
   if (!artifactDirectory) {
     throw new Error(
       `The wasm artifact directory is not set. This usually happens when running locally. " +
       "Please download all the files from taskcluster/kinds/fetch/onnxruntime-web-fetch.yml. " +
-      "Place them in a directory and rerun the test with the environment variable 'MOZ_FETCHES_DIR' " +
-      "set such that all the files are directly inside 'MOZ_FETCHES_DIR'`
+      "Place them in a directory and rerun the test with the environment variable 'MOZ_ML_LOCAL_DIR' " +
+      "set such that all the files are directly inside 'MOZ_ML_LOCAL_DIR'`
     );
   }
 
   if (!PathUtils.isAbsolute(artifactDirectory)) {
     throw new Error(
-      "Please provide an absolute path for 'MOZ_FETCHES_DIR and not a relative path"
+      "Please provide an absolute path for 'MOZ_ML_LOCAL_DIR and not a relative path"
     );
   }
 

toolkit/components/ml/tests/tools/create_local_hub.py

@@ -205,7 +205,7 @@ def main():
         description="Download ort.jsep.wasm and optionally clone specified models."
     )
 
-    default_dir = os.getenv("MOZ_FETCHES_DIR", None)
+    default_dir = os.getenv("MOZ_ML_LOCAL_DIR", None)
 
     parser.add_argument(
         "--fetches-dir",
@@ -235,7 +235,7 @@ def main():
 
     if args.fetches_dir is None:
         raise ValueError(
-            "Missing --fetches-dir argument or MOZ_FETCHES_DIR env var. Please specify a directory to store the downloaded files"
+            "Missing --fetches-dir argument or MOZ_ML_LOCAL_DIR env var. Please specify a directory to store the downloaded files"
         )
 
     fetches_dir = Path(args.fetches_dir).resolve()

toolkit/components/ml/tests/tools/hooks_local_hub.py

@@ -8,7 +8,7 @@ import socketserver
 import threading
 from pathlib import Path
 
-_THREAD = None
+THREADS = []
 
 
 class CustomHTTPRequestHandler(http.server.SimpleHTTPRequestHandler):
@@ -70,23 +70,26 @@ def start_hub(root_directory):
 def before_runs(env):
     """Runs before all performance tests.
 
-    We grab MOZ_FETCHES_DIR. If set we serve MOZ_FETCHES_DIR/onnx-models as our local hub.
+    We grab MOZ_ML_LOCAL_DIR. If set we serve MOZ_ML_LOCAL_DIR/onnx-models as our local hub.
+
+    MOZ_FETCHES_DIR is used in the CI as an alternate localtion.
     """
-    global _THREAD
-    fetches_dir = os.environ.get("MOZ_FETCHES_DIR")
+    fetches_dir = os.environ.get("MOZ_ML_LOCAL_DIR")
+    if fetches_dir is None:
+        fetches_dir = os.environ.get("MOZ_FETCHES_DIR")
     if fetches_dir is None:
         return
+
     hub_dir = Path(fetches_dir) / "onnx-models"
     if not hub_dir.is_dir():
         return
     port, server_thread = start_hub(hub_dir)
     os.environ["MOZ_MODELS_HUB"] = f"http://localhost:{port}"
-    _THREAD = server_thread
+    THREADS.append(server_thread)
 
 
 def after_runs(env):
-    global _THREAD
-    if _THREAD is not None:
+    if len(THREADS) > 0:
         print("Shutting down")
-        _THREAD.join(timeout=0)
-        _THREAD = None
+        THREADS[0].join(timeout=0)
+        THREADS.clear()