Source code
Revision control
Copy as Markdown
Other Tools
Test Info:
/* Any copyright is dedicated to the Public Domain.
"use strict";
/// <reference path="head.js" />
const { sinon } = ChromeUtils.importESModule(
);
const BEST_ONNX_OPTIONS = {
taskName: "text-classification",
modelId: "acme/bert",
dtype: "q8",
backend: "best-onnx",
modelHubUrlTemplate: "{model}/resolve/{revision}",
};
/**
* Stubs the worker's getBackend so any onnx-native engine creation rejects
* with the same message InferenceSession.cpp uses when libonnxruntime is
* missing. Returns a cleanup function the caller must invoke in `finally`.
*/
async function stubNativeUnavailable() {
const workerCode = `
${await getMLEngineWorkerCode()}
lazy.getBackend = async function (mlEngineWorker, wasm, options) {
if (options.backend === "onnx-native") {
throw new Error("onnxruntime shared library could not be loaded");
}
return { run: () => ({}) };
};
`;
const blob = new Blob([workerCode], { type: "application/javascript" });
const blobURL = URL.createObjectURL(blob);
const workerConfigStub = sinon
.stub(MLEngineParent, "getWorkerConfig")
.callsFake(() => ({ url: blobURL, options: { type: "module" } }));
return () => {
workerConfigStub.restore();
URL.revokeObjectURL(blobURL);
};
}
/**
* Demonstrates the lazy-init contract for `gBestOnnxBackend` (see
* MLEngineChild.sys.mjs). On the very first best-onnx call in an inference
* child process the cache is null. `chooseBestBackend` is required to
* tolerate that state — it returns "onnx-native" optimistically and lets
* the engine-creation try/catch resolve the actual outcome, caching either
* "onnx-native" (success) or "onnx" (ORT-unavailable). No priming step is
* needed for correctness; if `chooseBestBackend` instead threw on null,
* the very first best-onnx call in any inference child would fail.
*
* This test exercises the failure half of that contract by stubbing the
* native attempt so the outcome is deterministic across platforms. Placed
* first in this file so it runs against the freshest module state.
*/
add_task(async function test_best_onnx_lazy_init_handles_null_cache() {
const { cleanup, remoteClients } = await setup();
const restoreStub = await stubNativeUnavailable();
try {
// First best-onnx call: cache may be null. The lazy-init path inside
// chooseBestBackend returns "onnx-native" optimistically; the stub
// forces the catch in initializeInferenceEngine, which caches "onnx"
// and retries with wasm. If chooseBestBackend threw on null, this
// createEngine would reject before any of that.
const enginePromise = createEngine(BEST_ONNX_OPTIONS);
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
const engine = await enginePromise;
Assert.equal(
engine.pipelineOptions.backend,
"onnx",
"First best-onnx call resolves correctly without a primed cache."
);
} finally {
restoreStub();
await EngineProcess.destroyMLEngine();
await cleanup();
}
});
/**
* Verifies best-onnx falls back to the wasm onnx backend when the native
* runtime fails to load. The dispatcher's catch + retry path runs with
* backend "onnx" and we observe the resolved backend on the parent-side
* MLEngine.
*
* A non-mocked task (text-classification) is required so the child awaits
* dispatcher.isReady() and runs initializeInferenceEngine; moz-echo would
* short-circuit that path and the stub would never be reached.
*/
add_task(async function test_best_onnx_falls_back_to_wasm() {
const { cleanup, remoteClients } = await setup();
const restoreStub = await stubNativeUnavailable();
try {
const enginePromise = createEngine(BEST_ONNX_OPTIONS);
// Only the wasm retry pulls the runtime; the native attempt is stubbed
// to throw before getWasmArrayBuffer is reached.
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
const engine = await enginePromise;
Assert.equal(
engine.pipelineOptions.backend,
"onnx",
"best-onnx falls back to wasm onnx when the native runtime is missing."
);
} finally {
restoreStub();
await EngineProcess.destroyMLEngine();
await cleanup();
}
});
/**
* Repeated best-onnx requests with the same options must reuse the existing
* engine even when "best-onnx" resolved via the wasm fallback path. Before
* PipelineOptions.equals understood the sentinel, the cached engine's
* backend ("onnx") never equalled the new request's backend ("best-onnx").
*/
add_task(async function test_best_onnx_engine_is_reused_after_fallback() {
const { cleanup, remoteClients } = await setup();
const restoreStub = await stubNativeUnavailable();
try {
const enginePromise1 = createEngine(BEST_ONNX_OPTIONS);
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
const engine1 = await enginePromise1;
const engine2 = await createEngine(BEST_ONNX_OPTIONS);
Assert.strictEqual(
engine1,
engine2,
"Repeated best-onnx createEngine returns the cached engine instance."
);
} finally {
restoreStub();
await EngineProcess.destroyMLEngine();
await cleanup();
}
});