Source code
Revision control
Copy as Markdown
Other Tools
Test Info:
/* Any copyright is dedicated to the Public Domain.
"use strict";
const MOZ_ECHO_OPTIONS_RAW = { taskName: "moz-echo", timeoutMS: -1 };
const MOZ_ECHO_OPTIONS = new PipelineOptions({
taskName: "moz-echo",
timeoutMS: -1,
});
/**
* Performing a basic engine initialization and run.
*/
add_task(async function test_ml_engine_basics() {
const { cleanup, remoteClients } = await setup();
info("Get the engine");
const engineInstance = await createEngine(MOZ_ECHO_OPTIONS_RAW);
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();
});
/**
* Test the Wasm failing to download triggering a rejection.
*/
add_task(async function test_ml_engine_wasm_rejection() {
const { cleanup, remoteClients } = await setup();
info("Get the engine");
const engineInstance = await createEngine(MOZ_ECHO_OPTIONS_RAW);
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);
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();
});
/**
* 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;
const enginesSeen = new Set();
async function run(x) {
const engineInstance = await createEngine(MOZ_ECHO_OPTIONS_RAW);
enginesSeen.add(engineInstance);
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`
);
}
Assert.equal(enginesSeen.size, 1, "Only one engine was created.");
ok(
!EngineProcess.areAllEnginesTerminated(),
"The engine process is still active."
);
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(MOZ_ECHO_OPTIONS_RAW);
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);
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 engineInstance");
const engineInstance = await createEngine(MOZ_ECHO_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 creating an engine after an error.
*/
add_task(async function test_ml_engine_model_error() {
const { cleanup, remoteClients } = await setup();
info("Get the engine");
const engineInstance = await createEngine(MOZ_ECHO_OPTIONS_RAW);
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);
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();
});
/**
* Tests that we display a nice error message when the "browser.ml.enable" 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 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();
});
/**
* Tests that a worker can have an infinite timeout.
*/
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();
});
/**
* These status are visualized in about:inference, but aren't used for business
* logic.
*/
add_task(async function test_ml_engine_get_status_by_engine_id() {
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: "IDLE",
options: {
useExternalDataFormat: false,
engineId: "default-engine",
featureId: null,
taskName: "moz-echo",
timeoutMS: 1000,
modelHubRootUrl:
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",
staticEmbeddingsOptions: null,
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,
baseURL: null,
apiKey: null,
},
},
};
const statusByEngineId = Object.fromEntries(
await engineInstance.mlEngineParent.getStatusByEngineId()
);
statusByEngineId["default-engine"].options.numThreads = "NOT_COMPARED";
Assert.deepEqual(statusByEngineId, expected);
await ok(
!EngineProcess.areAllEnginesTerminated(),
"The engine process is still active."
);
await EngineProcess.destroyMLEngine();
await cleanup();
});