1385 lines
36 KiB
JavaScript
1385 lines
36 KiB
JavaScript
/* Any copyright is dedicated to the Public Domain.
|
|
http://creativecommons.org/publicdomain/zero/1.0/ */
|
|
|
|
"use strict";
|
|
|
|
/// <reference path="head.js" />
|
|
|
|
requestLongerTimeout(2);
|
|
|
|
const RAW_PIPELINE_OPTIONS = { taskName: "moz-echo", timeoutMS: -1 };
|
|
const PIPELINE_OPTIONS = new PipelineOptions({
|
|
taskName: "moz-echo",
|
|
timeoutMS: -1,
|
|
});
|
|
|
|
async function checkForRemoteType(remoteType) {
|
|
let procinfo3 = await ChromeUtils.requestProcInfo();
|
|
for (const child of procinfo3.children) {
|
|
if (child.type === remoteType) {
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
add_task(async function test_ml_engine_basics() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine(RAW_PIPELINE_OPTIONS);
|
|
|
|
info("Check the inference process is running");
|
|
Assert.equal(await checkForRemoteType("inference"), true);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
const res = await inferencePromise;
|
|
Assert.equal(
|
|
res.output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.equal(res.output.dtype, "q8", "The config was enriched by RS");
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
|
|
await cleanup();
|
|
});
|
|
|
|
add_task(async function test_ml_engine_pick_feature_id() {
|
|
// one record sent back from RS contains featureId
|
|
const records = [
|
|
{
|
|
taskName: "moz-echo",
|
|
modelId: "mozilla/distilvit",
|
|
processorId: "mozilla/distilvit",
|
|
tokenizerId: "mozilla/distilvit",
|
|
modelRevision: "main",
|
|
processorRevision: "main",
|
|
tokenizerRevision: "main",
|
|
dtype: "q8",
|
|
id: "74a71cfd-1734-44e6-85c0-69cf3e874138",
|
|
},
|
|
{
|
|
featureId: "pdfjs-alt-text",
|
|
taskName: "moz-echo",
|
|
modelId: "mozilla/distilvit",
|
|
processorId: "mozilla/distilvit",
|
|
tokenizerId: "mozilla/distilvit",
|
|
modelRevision: "v1.0",
|
|
processorRevision: "v1.0",
|
|
tokenizerRevision: "v1.0",
|
|
dtype: "fp16",
|
|
id: "74a71cfd-1734-44e6-85c0-69cf3e874138",
|
|
},
|
|
];
|
|
|
|
const { cleanup, remoteClients } = await setup({ records });
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine({
|
|
featureId: "pdfjs-alt-text",
|
|
taskName: "moz-echo",
|
|
});
|
|
|
|
info("Check the inference process is running");
|
|
Assert.equal(await checkForRemoteType("inference"), true);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
const res = await inferencePromise;
|
|
Assert.equal(
|
|
res.output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.equal(
|
|
res.output.dtype,
|
|
"fp16",
|
|
"The config was enriched by RS - using a feature Id"
|
|
);
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
|
|
await cleanup();
|
|
});
|
|
|
|
add_task(async function test_ml_engine_wasm_rejection() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine(RAW_PIPELINE_OPTIONS);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].rejectPendingDownloads(1);
|
|
//await remoteClients.models.resolvePendingDownloads(1);
|
|
|
|
let error;
|
|
try {
|
|
await inferencePromise;
|
|
} catch (e) {
|
|
error = e;
|
|
}
|
|
|
|
is(
|
|
error?.message,
|
|
"Intentionally rejecting downloads.",
|
|
"The error is correctly surfaced."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
/**
|
|
* Tests that the engineInstanceModel's internal errors are correctly surfaced.
|
|
*/
|
|
add_task(async function test_ml_engine_model_error() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine(RAW_PIPELINE_OPTIONS);
|
|
|
|
info("Run the inference with a throwing example.");
|
|
const inferencePromise = engineInstance.run("throw");
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
//await remoteClients.models.resolvePendingDownloads(1);
|
|
|
|
let error;
|
|
try {
|
|
await inferencePromise;
|
|
} catch (e) {
|
|
error = e;
|
|
}
|
|
is(
|
|
error?.message,
|
|
'Error: Received the message "throw", so intentionally throwing an error.',
|
|
"The error is correctly surfaced."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
/**
|
|
* This test is really similar to the "basic" test, but tests manually destroying
|
|
* the engineInstance.
|
|
*/
|
|
add_task(async function test_ml_engine_destruction() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
info("Get the engine process");
|
|
const mlEngineParent = await EngineProcess.getMLEngineParent();
|
|
|
|
info("Get engineInstance");
|
|
const engineInstance = await mlEngineParent.getEngine(PIPELINE_OPTIONS);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await engineInstance.terminate(
|
|
/* shutDownIfEmpty */ true,
|
|
/* replacement */ false
|
|
);
|
|
|
|
info(
|
|
"The engineInstance is manually destroyed. The cleanup function should wait for the engine process to be destroyed."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
/**
|
|
* Tests that we display a nice error message when the pref is off
|
|
*/
|
|
add_task(async function test_pref_is_off() {
|
|
await SpecialPowers.pushPrefEnv({
|
|
set: [["browser.ml.enable", false]],
|
|
});
|
|
|
|
info("Get the engine process");
|
|
let error;
|
|
|
|
try {
|
|
await EngineProcess.getMLEngineParent();
|
|
} catch (e) {
|
|
error = e;
|
|
}
|
|
is(
|
|
error?.message,
|
|
"MLEngine is disabled. Check the browser.ml prefs.",
|
|
"The error is correctly surfaced."
|
|
);
|
|
|
|
await SpecialPowers.pushPrefEnv({
|
|
set: [["browser.ml.enable", true]],
|
|
});
|
|
});
|
|
|
|
/**
|
|
* Tests the generic pipeline API
|
|
*/
|
|
add_task(async function test_ml_generic_pipeline() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
info("Get the engine process");
|
|
const mlEngineParent = await EngineProcess.getMLEngineParent();
|
|
|
|
info("Get engineInstance");
|
|
|
|
const options = new PipelineOptions({
|
|
taskName: "summarization",
|
|
modelId: "test-echo",
|
|
modelRevision: "main",
|
|
});
|
|
|
|
const engineInstance = await mlEngineParent.getEngine(options);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({
|
|
args: ["This gets echoed."],
|
|
});
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
/**
|
|
* Tests that the engine is reused.
|
|
*/
|
|
add_task(async function test_ml_engine_reuse_same() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
const options = { taskName: "moz-echo", engineId: "echo" };
|
|
const engineInstance = await createEngine(options);
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
let engineInstance2 = await createEngine(options);
|
|
is(engineInstance2.engineId, "echo", "The engine ID matches");
|
|
is(engineInstance, engineInstance2, "The engine is reused.");
|
|
const inferencePromise2 = engineInstance2.run({ data: "This gets echoed." });
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise2).output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
/**
|
|
* Tests that we can have two competing engines
|
|
*/
|
|
add_task(async function test_ml_two_engines() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
const engineInstance = await createEngine({
|
|
taskName: "moz-echo",
|
|
engineId: "engine1",
|
|
});
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
let engineInstance2 = await createEngine({
|
|
taskName: "moz-echo",
|
|
engineId: "engine2",
|
|
});
|
|
|
|
const inferencePromise2 = engineInstance2.run({ data: "This gets echoed." });
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise2).output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.notEqual(
|
|
engineInstance.engineId,
|
|
engineInstance2.engineId,
|
|
"Should be different engines"
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
/**
|
|
* Tests that we can have the same engine reinitialized
|
|
*/
|
|
add_task(async function test_ml_dupe_engines() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
const engineInstance = await createEngine({
|
|
taskName: "moz-echo",
|
|
engineId: "engine1",
|
|
});
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
let engineInstance2 = await createEngine({
|
|
taskName: "moz-echo",
|
|
engineId: "engine1",
|
|
timeoutMS: 2000, // that makes the options different
|
|
});
|
|
const inferencePromise2 = engineInstance2.run({ data: "This gets echoed." });
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise2).output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.notEqual(
|
|
engineInstance,
|
|
engineInstance2,
|
|
"Should be different engines"
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
add_task(async function test_ml_engine_override_options() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine({
|
|
taskName: "moz-echo",
|
|
modelRevision: "v1",
|
|
});
|
|
|
|
info("Check the inference process is running");
|
|
Assert.equal(await checkForRemoteType("inference"), true);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.modelRevision,
|
|
"v1",
|
|
"The config options goes through and overrides."
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
/**
|
|
* Tests a custom model hub
|
|
*/
|
|
add_task(async function test_ml_custom_hub() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
info("Get the engine process");
|
|
const mlEngineParent = await EngineProcess.getMLEngineParent();
|
|
|
|
info("Get engineInstance");
|
|
|
|
const options = new PipelineOptions({
|
|
taskName: "summarization",
|
|
modelId: "test-echo",
|
|
modelRevision: "main",
|
|
modelHubRootUrl: "https://example.com",
|
|
modelHubUrlTemplate: "models/{model}/{revision}",
|
|
});
|
|
|
|
const engineInstance = await mlEngineParent.getEngine(options);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({
|
|
args: ["This gets echoed."],
|
|
});
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
let res = await inferencePromise;
|
|
|
|
Assert.equal(
|
|
res.output,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.equal(
|
|
res.config.modelHubRootUrl,
|
|
"https://example.com",
|
|
"The pipeline used the custom hub"
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
/**
|
|
* Make sure we don't get race conditions when running several inference runs in parallel
|
|
*
|
|
*/
|
|
|
|
add_task(async function test_ml_engine_parallel() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
// We're doing 10 calls and each echo call will take from 0 to 1000ms
|
|
// So we're sure we're mixing runs.
|
|
let sleepTimes = [300, 1000, 700, 0, 500, 900, 400, 800, 600, 100];
|
|
let numCalls = 10;
|
|
|
|
async function run(x) {
|
|
const engineInstance = await createEngine(RAW_PIPELINE_OPTIONS);
|
|
|
|
let msg = `${x} - This gets echoed.`;
|
|
let res = engineInstance.run({
|
|
data: msg,
|
|
sleepTime: sleepTimes[x],
|
|
});
|
|
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
res = await res;
|
|
|
|
return res;
|
|
}
|
|
|
|
info(`Run ${numCalls} inferences in parallel`);
|
|
let runs = [];
|
|
for (let x = 0; x < numCalls; x++) {
|
|
runs.push(run(x));
|
|
}
|
|
|
|
// await all runs
|
|
const results = await Promise.all(runs);
|
|
Assert.equal(results.length, numCalls, `All ${numCalls} were successful`);
|
|
|
|
// check that each one got their own stuff
|
|
for (let y = 0; y < numCalls; y++) {
|
|
Assert.equal(
|
|
results[y].output.echo,
|
|
`${y} - This gets echoed.`,
|
|
`Result ${y} is correct`
|
|
);
|
|
}
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
|
|
await cleanup();
|
|
});
|
|
|
|
/**
|
|
* Test threading support
|
|
*/
|
|
add_task(async function test_ml_threading_support() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
info("Get the engine process");
|
|
const mlEngineParent = await EngineProcess.getMLEngineParent();
|
|
|
|
info("Get engineInstance");
|
|
|
|
const options = new PipelineOptions({
|
|
taskName: "summarization",
|
|
modelId: "test-echo",
|
|
modelRevision: "main",
|
|
});
|
|
|
|
const engineInstance = await mlEngineParent.getEngine(options);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({
|
|
args: ["This gets echoed."],
|
|
});
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
let res = await inferencePromise;
|
|
|
|
ok(res.multiThreadSupported, "Multi-thread should be supported");
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
add_task(async function test_ml_engine_get_status() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine({ taskName: "moz-echo" });
|
|
|
|
info("Check the inference process is running");
|
|
Assert.equal(await checkForRemoteType("inference"), true);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
const res = await inferencePromise;
|
|
Assert.equal(
|
|
res.output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
const expected = {
|
|
"default-engine": {
|
|
status: "IDLING",
|
|
options: {
|
|
useExternalDataFormat: false,
|
|
engineId: "default-engine",
|
|
featureId: null,
|
|
taskName: "moz-echo",
|
|
timeoutMS: 1000,
|
|
modelHubRootUrl:
|
|
"chrome://mochitests/content/browser/toolkit/components/ml/tests/browser/data",
|
|
modelHubUrlTemplate: "{model}/{revision}",
|
|
modelId: "mozilla/distilvit",
|
|
modelRevision: "main",
|
|
tokenizerId: "mozilla/distilvit",
|
|
tokenizerRevision: "main",
|
|
processorId: "mozilla/distilvit",
|
|
processorRevision: "main",
|
|
logLevel: "All",
|
|
runtimeFilename: "ort-wasm-simd-threaded.jsep.wasm",
|
|
device: null,
|
|
dtype: "q8",
|
|
numThreads: "NOT_COMPARED",
|
|
executionPriority: null,
|
|
kvCacheDtype: null,
|
|
numContext: 1024,
|
|
numBatch: 1024,
|
|
numUbatch: 1024,
|
|
flashAttn: false,
|
|
useMmap: false,
|
|
useMlock: true,
|
|
numThreadsDecoding: null,
|
|
modelFile: null,
|
|
backend: null,
|
|
modelHub: null,
|
|
},
|
|
engineId: "default-engine",
|
|
},
|
|
};
|
|
|
|
let status = await engineInstance.mlEngineParent.getStatus();
|
|
status = JSON.parse(JSON.stringify(Object.fromEntries(status)));
|
|
|
|
status["default-engine"].options.numThreads = "NOT_COMPARED";
|
|
Assert.deepEqual(status, expected);
|
|
|
|
await ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
|
|
await cleanup();
|
|
});
|
|
|
|
add_task(async function test_ml_engine_not_enough_memory() {
|
|
const { cleanup } = await setup({
|
|
prefs: [
|
|
["browser.ml.checkForMemory", true],
|
|
["browser.ml.minimumPhysicalMemory", 99999],
|
|
],
|
|
});
|
|
|
|
info("Get the greedy engine");
|
|
|
|
await Assert.rejects(
|
|
createEngine({
|
|
modelId: "testing/greedy",
|
|
taskName: "moz-echo",
|
|
dtype: "q8",
|
|
numThreads: 1,
|
|
device: "wasm",
|
|
}),
|
|
/Not enough physical memory/,
|
|
"The call should be rejected because of a lack of memory"
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
/**
|
|
* Helper function to create a basic set of valid options
|
|
*/
|
|
function getValidOptions(overrides = {}) {
|
|
return Object.assign(
|
|
{
|
|
engineId: "validEngine1",
|
|
featureId: "pdfjs-alt-text",
|
|
taskName: "valid_task",
|
|
modelHubRootUrl: "https://example.com",
|
|
modelHubUrlTemplate: "https://example.com/{modelId}",
|
|
timeoutMS: 5000,
|
|
modelId: "validModel",
|
|
modelRevision: "v1",
|
|
tokenizerId: "validTokenizer",
|
|
tokenizerRevision: "v1",
|
|
processorId: "validProcessor",
|
|
processorRevision: "v1",
|
|
logLevel: null,
|
|
runtimeFilename: "runtime.wasm",
|
|
device: InferenceDevice.GPU,
|
|
numThreads: 4,
|
|
executionPriority: ExecutionPriority.NORMAL,
|
|
},
|
|
overrides
|
|
);
|
|
}
|
|
|
|
/**
|
|
* A collection of test cases for invalid and valid values.
|
|
*/
|
|
const commonInvalidCases = [
|
|
{ description: "Invalid value (special characters)", value: "org1/my!value" },
|
|
{
|
|
description: "Invalid value (special characters in organization)",
|
|
value: "org@1/my-value",
|
|
},
|
|
{ description: "Invalid value (missing name part)", value: "org1/" },
|
|
{
|
|
description: "Invalid value (invalid characters in name)",
|
|
value: "my$value",
|
|
},
|
|
];
|
|
|
|
const commonValidCases = [
|
|
{ description: "Valid organization/name", value: "org1/my-value" },
|
|
{ description: "Valid name only", value: "my-value" },
|
|
{
|
|
description: "Valid name with underscores and dashes",
|
|
value: "my_value-123",
|
|
},
|
|
{
|
|
description: "Valid organization with underscores and dashes",
|
|
value: "org_123/my-value",
|
|
},
|
|
];
|
|
|
|
const pipelineOptionsCases = [
|
|
// Invalid cases for various fields
|
|
...commonInvalidCases.map(test => ({
|
|
description: `Invalid processorId (${test.description})`,
|
|
options: { processorId: test.value },
|
|
expectedError: /Invalid value/,
|
|
})),
|
|
...commonInvalidCases.map(test => ({
|
|
description: `Invalid tokenizerId (${test.description})`,
|
|
options: { tokenizerId: test.value },
|
|
expectedError: /Invalid value/,
|
|
})),
|
|
...commonInvalidCases.map(test => ({
|
|
description: `Invalid modelId (${test.description})`,
|
|
options: { modelId: test.value },
|
|
expectedError: /Invalid value/,
|
|
})),
|
|
|
|
// Valid cases for various fields
|
|
...commonValidCases.map(test => ({
|
|
description: `Valid processorId (${test.description})`,
|
|
options: { processorId: test.value },
|
|
expected: { processorId: test.value },
|
|
})),
|
|
...commonValidCases.map(test => ({
|
|
description: `Valid tokenizerId (${test.description})`,
|
|
options: { tokenizerId: test.value },
|
|
expected: { tokenizerId: test.value },
|
|
})),
|
|
...commonValidCases.map(test => ({
|
|
description: `Valid modelId (${test.description})`,
|
|
options: { modelId: test.value },
|
|
expected: { modelId: test.value },
|
|
})),
|
|
|
|
// Invalid values
|
|
{
|
|
description: "Invalid hub",
|
|
options: { modelHub: "rogue" },
|
|
expectedError: /Invalid value/,
|
|
},
|
|
{
|
|
description: "Invalid timeoutMS",
|
|
options: { timeoutMS: -3 },
|
|
expectedError: /Invalid value/,
|
|
},
|
|
{
|
|
description: "Invalid timeoutMS",
|
|
options: { timeoutMS: 40000000 },
|
|
expectedError: /Invalid value/,
|
|
},
|
|
{
|
|
description: "Invalid featureId",
|
|
options: { featureId: "unknown" },
|
|
expectedError: /Invalid value/,
|
|
},
|
|
{
|
|
description: "Invalid dtype",
|
|
options: { dtype: "invalid_dtype" },
|
|
expectedError: /Invalid value/,
|
|
},
|
|
{
|
|
description: "Invalid device",
|
|
options: { device: "invalid_device" },
|
|
expectedError: /Invalid value/,
|
|
},
|
|
{
|
|
description: "Invalid executionPriority",
|
|
options: { executionPriority: "invalid_priority" },
|
|
expectedError: /Invalid value/,
|
|
},
|
|
{
|
|
description: "Invalid logLevel",
|
|
options: { logLevel: "invalid_log_level" },
|
|
expectedError: /Invalid value/,
|
|
},
|
|
|
|
// Valid values
|
|
{
|
|
description: "valid hub",
|
|
options: { modelHub: "huggingface" },
|
|
expected: { modelHub: "huggingface" },
|
|
},
|
|
{
|
|
description: "valid hub",
|
|
options: { modelHub: "mozilla" },
|
|
expected: { modelHub: "mozilla" },
|
|
},
|
|
{
|
|
description: "valid timeoutMS",
|
|
options: { timeoutMS: 12345 },
|
|
expected: { timeoutMS: 12345 },
|
|
},
|
|
{
|
|
description: "valid timeoutMS",
|
|
options: { timeoutMS: -1 },
|
|
expected: { timeoutMS: -1 },
|
|
},
|
|
|
|
{
|
|
description: "Valid dtype",
|
|
options: { dtype: QuantizationLevel.FP16 },
|
|
expected: { dtype: QuantizationLevel.FP16 },
|
|
},
|
|
{
|
|
description: "Valid device",
|
|
options: { device: InferenceDevice.WASM },
|
|
expected: { device: InferenceDevice.WASM },
|
|
},
|
|
{
|
|
description: "Valid executionPriority",
|
|
options: { executionPriority: ExecutionPriority.HIGH },
|
|
expected: { executionPriority: ExecutionPriority.HIGH },
|
|
},
|
|
{
|
|
description: "Valid logLevel (Info)",
|
|
options: { logLevel: LogLevel.INFO },
|
|
expected: { logLevel: LogLevel.INFO },
|
|
},
|
|
{
|
|
description: "Valid logLevel (Critical)",
|
|
options: { logLevel: LogLevel.CRITICAL },
|
|
expected: { logLevel: LogLevel.CRITICAL },
|
|
},
|
|
{
|
|
description: "Valid logLevel (All)",
|
|
options: { logLevel: LogLevel.ALL },
|
|
expected: { logLevel: LogLevel.ALL },
|
|
},
|
|
{
|
|
description: "Valid modelId",
|
|
options: { modelId: "Qwen2.5-0.5B-Instruct" },
|
|
expected: { modelId: "Qwen2.5-0.5B-Instruct" },
|
|
},
|
|
|
|
// Invalid revision cases
|
|
{
|
|
description: "Invalid revision (random string)",
|
|
options: { modelRevision: "invalid_revision" },
|
|
expectedError: /Invalid value/,
|
|
},
|
|
{
|
|
description: "Invalid revision (too many version numbers)",
|
|
options: { tokenizerRevision: "v1.0.3.4.5" },
|
|
expectedError: /Invalid value/,
|
|
},
|
|
{
|
|
description: "Invalid revision (unknown suffix)",
|
|
options: { processorRevision: "v1.0.0-unknown" },
|
|
expectedError: /Invalid value/,
|
|
},
|
|
|
|
// Valid revision cases with new format
|
|
{
|
|
description: "Valid revision (main)",
|
|
options: { modelRevision: "main" },
|
|
expected: { modelRevision: "main" },
|
|
},
|
|
{
|
|
description: "Valid revision (v-prefixed version with alpha)",
|
|
options: { tokenizerRevision: "v1.2.3-alpha1" },
|
|
expected: { tokenizerRevision: "v1.2.3-alpha1" },
|
|
},
|
|
{
|
|
description:
|
|
"Valid revision (v-prefixed version with beta and dot separator)",
|
|
options: { tokenizerRevision: "v1.2.3.beta2" },
|
|
expected: { tokenizerRevision: "v1.2.3.beta2" },
|
|
},
|
|
{
|
|
description:
|
|
"Valid revision (non-prefixed version with rc and dash separator)",
|
|
options: { processorRevision: "1.0.0-rc3" },
|
|
expected: { processorRevision: "1.0.0-rc3" },
|
|
},
|
|
{
|
|
description:
|
|
"Valid revision (non-prefixed version with pre and dot separator)",
|
|
options: { processorRevision: "1.0.0.pre4" },
|
|
expected: { processorRevision: "1.0.0.pre4" },
|
|
},
|
|
{
|
|
description: "Valid revision (version without suffix)",
|
|
options: { modelRevision: "1.0.0" },
|
|
expected: { modelRevision: "1.0.0" },
|
|
},
|
|
|
|
// Valid engineID cases
|
|
{
|
|
description: "Valid engineID (qwen)",
|
|
options: { engineId: "SUM-ONNX-COMMUNITY_QWEN2_5-0_5B-INSTRUCT_BIG" },
|
|
expected: { engineId: "SUM-ONNX-COMMUNITY_QWEN2_5-0_5B-INSTRUCT_BIG" },
|
|
},
|
|
];
|
|
|
|
/**
|
|
* Testing PipelineOption validation
|
|
*/
|
|
add_task(async function test_pipeline_options_validation() {
|
|
pipelineOptionsCases.forEach(testCase => {
|
|
if (testCase.expectedError) {
|
|
Assert.throws(
|
|
() => new PipelineOptions(getValidOptions(testCase.options)),
|
|
testCase.expectedError,
|
|
`${testCase.description} throws the expected error`
|
|
);
|
|
} else {
|
|
const pipelineOptions = new PipelineOptions(
|
|
getValidOptions(testCase.options)
|
|
);
|
|
Object.keys(testCase.expected).forEach(key => {
|
|
is(
|
|
pipelineOptions[key],
|
|
testCase.expected[key],
|
|
`${testCase.description} sets ${key} correctly`
|
|
);
|
|
});
|
|
}
|
|
});
|
|
});
|
|
|
|
add_task(async function test_ml_engine_infinite_worker() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
const options = { taskName: "moz-echo", timeoutMS: -1 };
|
|
const engineInstance = await createEngine(options);
|
|
|
|
info("Check the inference process is running");
|
|
Assert.equal(await checkForRemoteType("inference"), true);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
const res = await inferencePromise;
|
|
Assert.equal(
|
|
res.output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.equal(res.output.timeoutMS, -1, "This should be an infinite worker.");
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
|
|
await cleanup();
|
|
});
|
|
|
|
add_task(async function test_ml_engine_model_hub_applied() {
|
|
const options = {
|
|
taskName: "moz-echo",
|
|
timeoutMS: -1,
|
|
modelHub: "huggingface",
|
|
};
|
|
const parsedOptions = new PipelineOptions(options);
|
|
|
|
Assert.equal(
|
|
parsedOptions.modelHubRootUrl,
|
|
"https://huggingface.co/",
|
|
"modelHubRootUrl is set"
|
|
);
|
|
|
|
Assert.equal(
|
|
parsedOptions.modelHubUrlTemplate,
|
|
"{model}/resolve/{revision}",
|
|
"modelHubUrlTemplate is set"
|
|
);
|
|
});
|
|
|
|
add_task(async function test_ml_engine_blessed_model() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
const options = { taskName: "test-echo" };
|
|
const engineInstance = await createEngine(options);
|
|
|
|
info("Check the inference process is running");
|
|
Assert.equal(await checkForRemoteType("inference"), true);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
const res = await inferencePromise;
|
|
|
|
Assert.equal(
|
|
res.config.modelId,
|
|
"test-echo",
|
|
"The blessed model was picked."
|
|
);
|
|
|
|
Assert.equal(res.config.dtype, "q8", "With the right quantization level");
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
|
|
await cleanup();
|
|
});
|
|
|
|
add_task(async function test_ml_engine_two_tasknames_in_rs() {
|
|
// RS has two records with the same taskName
|
|
// we should use the modelId match in that case
|
|
const records = [
|
|
{
|
|
taskName: "moz-echo",
|
|
modelId: "mozilla/anothermodel",
|
|
processorId: "mozilla/distilvit",
|
|
tokenizerId: "mozilla/distilvit",
|
|
modelRevision: "main",
|
|
processorRevision: "main",
|
|
tokenizerRevision: "main",
|
|
dtype: "q8",
|
|
id: "74a71cfd-1734-44e6-85c0-69cf3e874138",
|
|
},
|
|
{
|
|
taskName: "moz-echo",
|
|
modelId: "mozilla/distilvit",
|
|
processorId: "mozilla/distilvit",
|
|
tokenizerId: "mozilla/distilvit",
|
|
modelRevision: "v1.0",
|
|
processorRevision: "v1.0",
|
|
tokenizerRevision: "v1.0",
|
|
dtype: "fp16",
|
|
id: "74a71cfd-1734-44e6-85c0-69cf3e874138",
|
|
},
|
|
];
|
|
|
|
const { cleanup, remoteClients } = await setup({ records });
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine({
|
|
featureId: "pdfjs-alt-text",
|
|
taskName: "moz-echo",
|
|
});
|
|
|
|
info("Check the inference process is running");
|
|
Assert.equal(await checkForRemoteType("inference"), true);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
const res = await inferencePromise;
|
|
Assert.equal(
|
|
res.output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.equal(
|
|
res.output.dtype,
|
|
"fp16",
|
|
"The config was enriched by RS - using a feature Id"
|
|
);
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
|
|
await cleanup();
|
|
});
|
|
|
|
add_task(
|
|
async function test_override_ml_engine_pipeline_options_in_allow_list() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
await SpecialPowers.pushPrefEnv({
|
|
set: [
|
|
[
|
|
"browser.ml.overridePipelineOptions",
|
|
'{"about-inference": {"modelRevision": "v0.2.0"}}',
|
|
],
|
|
],
|
|
});
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine({
|
|
taskName: "moz-echo",
|
|
featureId: "about-inference",
|
|
});
|
|
|
|
info("Check the inference process is running");
|
|
Assert.equal(await checkForRemoteType("inference"), true);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.modelRevision,
|
|
"v0.2.0",
|
|
"The config options goes through and overrides."
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
}
|
|
);
|
|
|
|
add_task(async function test_override_ml_pipeline_options_not_in_allow_list() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
await SpecialPowers.pushPrefEnv({
|
|
set: [
|
|
[
|
|
"browser.ml.overridePipelineOptions",
|
|
'{"about-inferences": {"modelRevision": "v0.2.0"}}',
|
|
],
|
|
],
|
|
});
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine({
|
|
taskName: "moz-echo",
|
|
featureId: "about-inference",
|
|
});
|
|
info("Check the inference process is running");
|
|
Assert.equal(await checkForRemoteType("inference"), true);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.modelRevision,
|
|
"main",
|
|
"The config options goes through and overrides."
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
add_task(async function test_override_ml_pipeline_options_unsafe_options() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
await SpecialPowers.pushPrefEnv({
|
|
set: [
|
|
[
|
|
"browser.ml.overridePipelineOptions",
|
|
'{"about-inference": {"modelRevision": "v0.2.0", "modelId": "unsafe-model-id"}}',
|
|
],
|
|
],
|
|
});
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine({
|
|
taskName: "moz-echo",
|
|
featureId: "about-inference",
|
|
});
|
|
|
|
info("Check the inference process is running");
|
|
Assert.equal(await checkForRemoteType("inference"), true);
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.echo,
|
|
"This gets echoed.",
|
|
"The text get echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.modelRevision,
|
|
"v0.2.0",
|
|
"The config options goes through and overrides."
|
|
);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.modelId,
|
|
"mozilla/distilvit",
|
|
"The config should not override."
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
add_task(async function test_q8_by_default() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine({
|
|
taskName: "moz-echo",
|
|
modelId: "Xenova/distilbart-cnn-6-6",
|
|
modelHub: "huggingface",
|
|
});
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.echo,
|
|
"This gets echoed.",
|
|
"The text gets echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.dtype,
|
|
"q8",
|
|
"dtype should be set to q8"
|
|
);
|
|
|
|
// the model hub sets the revision
|
|
Assert.equal(
|
|
(await inferencePromise).output.modelRevision,
|
|
"main",
|
|
"modelRevision should be main"
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|
|
|
|
add_task(async function test_hub_by_default() {
|
|
const { cleanup, remoteClients } = await setup();
|
|
|
|
info("Get the engine");
|
|
const engineInstance = await createEngine({
|
|
taskName: "moz-echo",
|
|
});
|
|
|
|
info("Run the inference");
|
|
const inferencePromise = engineInstance.run({ data: "This gets echoed." });
|
|
|
|
info("Wait for the pending downloads.");
|
|
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.echo,
|
|
"This gets echoed.",
|
|
"The text gets echoed exercising the whole flow."
|
|
);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.modelHubUrlTemplate,
|
|
"{model}/{revision}",
|
|
"Default template should be model/revision"
|
|
);
|
|
|
|
Assert.equal(
|
|
(await inferencePromise).output.modelRevision,
|
|
"main",
|
|
"modelRevision should be main"
|
|
);
|
|
|
|
ok(
|
|
!EngineProcess.areAllEnginesTerminated(),
|
|
"The engine process is still active."
|
|
);
|
|
|
|
await EngineProcess.destroyMLEngine();
|
|
await cleanup();
|
|
});
|