Bug 2059875 - Expose isAvailable() for native ONNX runtime r=valentinp,webidl,ai-platform-reviewers,saschanaz

This commit exposes an isAvailable() API for the native
ONNX runtime so that callers can more easily make an
informed decisions.

Many of the current code solutions get around this by
attempting to use ONNX native, and falling back to
an alternative only if ONNX native fails.

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

dom/onnx/InferenceSession.cpp

@@ -15,6 +15,7 @@
 #include "mozilla/Logging.h"
 #include "mozilla/RefPtr.h"
 #include "mozilla/ScopeExit.h"
+#include "mozilla/StaticMutex.h"
 #include "mozilla/dom/BindingDeclarations.h"
 #include "mozilla/dom/ContentChild.h"
 #include "mozilla/dom/ONNXBinding.h"
@@ -38,6 +39,8 @@ namespace mozilla::dom {
 static OrtEnv* sEnv = nullptr;
 static OrtApi* sAPI = nullptr;
 
+static StaticMutex sOrtAPIMutex;
+
 // RAII wrapper over OrtStatus.
 // Takes ownership of a externally allocated OrtStatus* passed at construction.
 // Move-only. OrtStatus released through OrtApi::ReleaseStatus.
@@ -166,6 +169,11 @@ OrtSessionOptions* ToOrtSessionOption(
 }  // namespace mozilla::dom
 
 OrtApi* GetOrtAPI() {
+  StaticMutexAutoLock lock(sOrtAPIMutex);
+  if (sAPI) {
+    return sAPI;
+  }
+
 #ifdef XP_WIN
   PathString path = GetLibraryFilePathname(LXUL_DLL, (PRFuncPtr)&GetOrtAPI);
 #else
@@ -224,6 +232,7 @@ OrtApi* GetOrtAPI() {
     return nullptr;
   }
 
+  sAPI = ortAPI;
   return ortAPI;
 }
 
@@ -235,6 +244,10 @@ bool InferenceSession::InInferenceProcess(JSContext*, JSObject*) {
       INFERENCE_REMOTE_TYPE);
 }
 
+bool InferenceSession::IsAvailable(const GlobalObject&) {
+  return GetOrtAPI() != nullptr;
+}
+
 nsCString InferenceSessionSessionOptionsToString(
     const InferenceSessionSessionOptions& aOptions) {
   return nsFmtCString(
@@ -310,8 +323,7 @@ void InferenceSession::Init(const RefPtr<Promise>& aPromise,
        aUriOrBuffer.IsUTF8String() ? "string" : "buffer");
 
   if (!sEnv) {
-    sAPI = GetOrtAPI();
-    if (!sAPI) {
+    if (!GetOrtAPI()) {
       LOGD("Couldn't get ahold of ORT API");
       // Use a distinguishable error so JS callers can recognize that the
       // native runtime is unavailable on this machine and fall back to the

dom/onnx/InferenceSession.h

@@ -35,6 +35,7 @@ class InferenceSession final : public nsISupports, public nsWrapperCache {
   }
 
   static bool InInferenceProcess(JSContext*, JSObject*);
+  static bool IsAvailable(const GlobalObject&);
 
  protected:
   virtual ~InferenceSession() { Destroy(); }

dom/webidl/ONNX.webidl

@@ -65,6 +65,7 @@ dictionary InferenceSessionSessionOptions {
 
 [Func="InferenceSession::InInferenceProcess", Exposed=(DedicatedWorker,Window)]
 interface InferenceSession {
+  static boolean isAvailable();
   [NewObject]
   Promise<InferenceSessionReturnType> run(InferenceSessionFeedsType feeds, optional InferenceSessionRunOptions options = {});
   [NewObject] static Promise<InferenceSession> create((UTF8String or Uint8Array) uriOrBuffer, optional InferenceSessionSessionOptions options = {});

toolkit/components/ml/actors/MLEngineChild.sys.mjs

@@ -93,6 +93,9 @@ export class MLEngineChild extends JSProcessActorChild {
       case "MLEngine:GetStatusByEngineId": {
         return this.getStatusByEngineId();
       }
+      case "MLEngine:RequestIsNativeOnnxRuntimeAvailable": {
+        return this.requestIsNativeOnnxRuntimeAvailable();
+      }
       case "MLEngine:ForceShutdown": {
         for (const engineDispatcher of this.#engineDispatchers.values()) {
           await engineDispatcher.terminate(
@@ -202,6 +205,25 @@ export class MLEngineChild extends JSProcessActorChild {
     return this.sendQuery("MLEngine:GetWorkerConfig");
   }
 
+  /**
+   * Resolves to true if the native ONNX runtime is available, otherwise false.
+   *
+   * @returns {Promise<boolean>}
+   */
+  async requestIsNativeOnnxRuntimeAvailable() {
+    const workerConfig = await this.getWorkerConfig();
+    const worker = new lazy.BasePromiseWorker(
+      workerConfig.url,
+      workerConfig.options
+    );
+
+    try {
+      return await worker.post("isNativeOnnxRuntimeAvailable", []);
+    } finally {
+      worker.terminate();
+    }
+  }
+
   /**
    * Resolves a requested backend to a concrete backend identifier. "best-onnx"
    * returns the cached choice if one exists, otherwise optimistically tries

toolkit/components/ml/actors/MLEngineParent.sys.mjs

@@ -632,6 +632,8 @@ export class MLEngineParent extends JSProcessActorParent {
 
   /**
    * Gets the configuration of the worker
+   *
+   * @returns {{ url: string, options: WorkerOptions }}
    */
   static getWorkerConfig() {
     return {
@@ -878,6 +880,15 @@ export class MLEngineParent extends JSProcessActorParent {
     return this.sendQuery("MLEngine:GetStatusByEngineId");
   }
 
+  /**
+   * Resolves to true if the native ONNX runtime is available, otherwise false.
+   *
+   * @returns {Promise<boolean>}
+   */
+  requestIsNativeOnnxRuntimeAvailable() {
+    return this.sendQuery("MLEngine:RequestIsNativeOnnxRuntimeAvailable");
+  }
+
   /**
    * Send a message to gracefully shutdown all of the ML engines in the engine process.
    * This mostly exists for testing the shutdown paths of the code.

toolkit/components/ml/content/EngineProcess.sys.mjs

@@ -1139,6 +1139,13 @@ export class PipelineOptions {
  * Translations engine and the MLEngine component.
  */
 export class EngineProcess {
+  /**
+   * The cached native ONNX runtime availability request.
+   *
+   * @type {Promise<boolean> | null}
+   */
+  static #nativeOnnxRuntimeAvailabilityPromise = null;
+
   /**
    * Get a reference to all running "inference" processes.
    *
@@ -1179,6 +1186,69 @@ export class EngineProcess {
     return EngineProcess.#getEngineActor({ actorName: "MLEngine" });
   }
 
+  /**
+   * Resolves to true if the native ONNX runtime is available, otherwise false.
+   *
+   * @returns {Promise<boolean>}
+   */
+  static requestIsNativeOnnxRuntimeAvailable() {
+    if (!Services.prefs.getBoolPref("browser.ml.enable")) {
+      return Promise.resolve(false);
+    }
+
+    if (!EngineProcess.#nativeOnnxRuntimeAvailabilityPromise) {
+      EngineProcess.#nativeOnnxRuntimeAvailabilityPromise =
+        EngineProcess.#requestNativeOnnxRuntimeAvailability();
+    }
+
+    const availabilityPromise =
+      EngineProcess.#nativeOnnxRuntimeAvailabilityPromise;
+    return availabilityPromise.catch(() => {
+      if (
+        EngineProcess.#nativeOnnxRuntimeAvailabilityPromise ===
+        availabilityPromise
+      ) {
+        // We weren't able to determine the availability definitively,
+        // so we shouldn't block future retry attempts.
+        EngineProcess.#nativeOnnxRuntimeAvailabilityPromise = null;
+      }
+
+      return false;
+    });
+  }
+
+  /**
+   * Clears the cached native ONNX runtime availability for tests.
+   */
+  static resetNativeOnnxRuntimeAvailabilityForTests() {
+    if (!Cu.isInAutomation) {
+      throw new Error("This function is only available in automation.");
+    }
+    EngineProcess.#nativeOnnxRuntimeAvailabilityPromise = null;
+  }
+
+  /**
+   * Requests native ONNX runtime availability from the inference process.
+   *
+   * @returns {Promise<boolean>}
+   */
+  static async #requestNativeOnnxRuntimeAvailability() {
+    const keepAlive =
+      await ChromeUtils.ensureHeadlessContentProcess("inference");
+
+    if (!keepAlive?.domProcess?.canSend) {
+      keepAlive?.invalidateKeepAlive();
+      throw new Error("Could not start the MLEngine inference process.");
+    }
+
+    try {
+      const actor = keepAlive.domProcess.getActor("MLEngine");
+      return await actor.requestIsNativeOnnxRuntimeAvailable();
+    } finally {
+      keepAlive.invalidateKeepAlive();
+    }
+  }
+
   /**
    * @returns {Promise<JSProcessActorParent>}
    */

toolkit/components/ml/content/MLEngine.worker.mjs

@@ -78,6 +78,15 @@ export class MLEngineWorker {
     throw new Error("Method not implemented.");
   }
 
+  /**
+   * Returns true if the native ONNX runtime is available, otherwise false.
+   *
+   * @returns {boolean}
+   */
+  isNativeOnnxRuntimeAvailable() {
+    return globalThis.InferenceSession?.isAvailable() ?? false;
+  }
+
   /**
    * @param {ArrayBuffer} wasm
    * @param {object} options received as an object, converted to a PipelineOptions instance

toolkit/components/ml/tests/browser/browser.toml

@@ -57,6 +57,13 @@ skip-if = [
   "verify",
 ]
 
+["browser_ml_native_ort_availability.js"]
+support-files = [
+  "ml_native_ort_available_stub.worker.mjs",
+  "ml_native_ort_error_stub.worker.mjs",
+  "ml_native_ort_unavailable_stub.worker.mjs",
+]
+
 ["browser_ml_nlp_utils.js"]
 
 ["browser_ml_openai.js"]

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

@@ -0,0 +1,110 @@
+/* Any copyright is dedicated to the Public Domain.
+   http://creativecommons.org/publicdomain/zero/1.0/ */
+
+"use strict";
+
+const { sinon } = ChromeUtils.importESModule(
+  "resource://testing-common/Sinon.sys.mjs"
+);
+
+const AVAILABLE_WORKER_URL =
+  "chrome://mochitests/content/browser/toolkit/components/ml/tests/browser/ml_native_ort_available_stub.worker.mjs";
+const ERROR_WORKER_URL =
+  "chrome://mochitests/content/browser/toolkit/components/ml/tests/browser/ml_native_ort_error_stub.worker.mjs";
+const UNAVAILABLE_WORKER_URL =
+  "chrome://mochitests/content/browser/toolkit/components/ml/tests/browser/ml_native_ort_unavailable_stub.worker.mjs";
+
+add_setup(async function () {
+  await SpecialPowers.pushPrefEnv({
+    set: [["browser.ml.enable", true]],
+  });
+  registerCleanupFunction(async () => {
+    EngineProcess.resetNativeOnnxRuntimeAvailabilityForTests();
+    await SpecialPowers.popPrefEnv();
+  });
+});
+
+async function testStubbedAvailability({
+  workerUrl,
+  expectedIsAvailable,
+  expectedProbeCount = 1,
+}) {
+  const workerConfigStub = sinon
+    .stub(MLEngineParent, "getWorkerConfig")
+    .returns({ url: workerUrl, options: { type: "module" } });
+
+  try {
+    const first = await EngineProcess.requestIsNativeOnnxRuntimeAvailable();
+    Assert.equal(
+      first,
+      expectedIsAvailable,
+      `The probe returned ${expectedIsAvailable}`
+    );
+
+    await TestUtils.waitForCondition(
+      () => EngineProcess.areAllEnginesTerminated(),
+      "The availability probe did not keep an inference process alive"
+    );
+
+    const second = await EngineProcess.requestIsNativeOnnxRuntimeAvailable();
+    Assert.equal(
+      second,
+      expectedIsAvailable,
+      `The second request returned ${expectedIsAvailable}`
+    );
+    Assert.equal(
+      workerConfigStub.callCount,
+      expectedProbeCount,
+      `The availability probe ran ${expectedProbeCount} time(s)`
+    );
+
+    await TestUtils.waitForCondition(
+      () => EngineProcess.areAllEnginesTerminated(),
+      "The availability probe did not keep an inference process alive"
+    );
+  } finally {
+    workerConfigStub.restore();
+    EngineProcess.resetNativeOnnxRuntimeAvailabilityForTests();
+  }
+}
+
+add_task(async function test_native_ort_unavailable() {
+  await testStubbedAvailability({
+    workerUrl: UNAVAILABLE_WORKER_URL,
+    expectedIsAvailable: false,
+  });
+});
+
+add_task(async function test_native_ort_available() {
+  await testStubbedAvailability({
+    workerUrl: AVAILABLE_WORKER_URL,
+    expectedIsAvailable: true,
+  });
+});
+
+add_task(async function test_native_ort_probe_error() {
+  await testStubbedAvailability({
+    workerUrl: ERROR_WORKER_URL,
+    expectedIsAvailable: false,
+    expectedProbeCount: 2,
+  });
+});
+
+add_task(async function test_native_ort_integration() {
+  const first = await EngineProcess.requestIsNativeOnnxRuntimeAvailable();
+
+  Assert.equal(typeof first, "boolean", "The real probe returns a boolean");
+
+  await TestUtils.waitForCondition(
+    () => EngineProcess.areAllEnginesTerminated(),
+    "The real availability probe did not keep an inference process alive"
+  );
+
+  const second = await EngineProcess.requestIsNativeOnnxRuntimeAvailable();
+
+  Assert.equal(second, first, "The real probe result is cached");
+  Assert.ok(
+    EngineProcess.areAllEnginesTerminated(),
+    "The real cached result did not recreate the inference process"
+  );
+});

toolkit/components/ml/tests/browser/ml_native_ort_available_stub.worker.mjs

@@ -0,0 +1,11 @@
+/* Any copyright is dedicated to the Public Domain.
+   https://creativecommons.org/publicdomain/zero/1.0/ */
+
+// eslint-disable-next-line import/no-unassigned-import
+import "chrome://global/content/ml/MLEngine.worker.mjs";
+
+if (typeof globalThis.InferenceSession?.isAvailable !== "function") {
+  throw new Error("InferenceSession.isAvailable is not defined");
+}
+
+globalThis.InferenceSession.isAvailable = () => true;

toolkit/components/ml/tests/browser/ml_native_ort_error_stub.worker.mjs

@@ -0,0 +1,9 @@
+/* Any copyright is dedicated to the Public Domain.
+   https://creativecommons.org/publicdomain/zero/1.0/ */
+
+// eslint-disable-next-line import/no-unassigned-import
+import "chrome://global/content/ml/MLEngine.worker.mjs";
+
+globalThis.InferenceSession.isAvailable = () => {
+  throw new Error("Native ORT availability probe failed");
+};

toolkit/components/ml/tests/browser/ml_native_ort_unavailable_stub.worker.mjs

@@ -0,0 +1,7 @@
+/* Any copyright is dedicated to the Public Domain.
+   https://creativecommons.org/publicdomain/zero/1.0/ */
+
+// eslint-disable-next-line import/no-unassigned-import
+import "chrome://global/content/ml/MLEngine.worker.mjs";
+
+globalThis.InferenceSession.isAvailable = () => false;