Source code
Revision control
Copy as Markdown
Other Tools
Test Info:
/* Any copyright is dedicated to the Public Domain.
"use strict";
const { sinon } = ChromeUtils.importESModule(
);
const { BACKENDS } = ChromeUtils.importESModule(
"chrome://global/content/ml/EngineProcess.sys.mjs"
);
const { MLUtils } = ChromeUtils.importESModule(
"chrome://global/content/ml/Utils.sys.mjs"
);
/**
* End to End test that the engine is indeed initialized with wllama when it is the
* best-llama.
*/
add_task(async function test_e2e_choose_backend_best_wllama() {
// Allow any url
Services.env.set("MOZ_ALLOW_EXTERNAL_ML_HUB", "true");
const backendData = new Uint8Array([10, 20, 30]);
const expectedBackendData = JSON.stringify(backendData);
// Mocking function used in the workers or child doesn't work.
// So we are stubbing the code run by the worker.
const workerCode = `
// Inject the MLEngine.worker.mjs code
${await getMLEngineWorkerCode()}
// Stub
ChromeUtils.defineESModuleGetters(
lazy,
{
createFileUrl: "chrome://global/content/ml/Utils.sys.mjs",
},
{ global: "current" }
);
// Change the getBackend to a mocked version that doesn't actually do inference
// but does initiate model downloads and engine initialization
lazy.getBackend = async function (
mlEngineWorker,
backendData,
{
modelHubUrlTemplate,
modelHubRootUrl,
modelId,
modelRevision,
modelFile,
engineId,
} = {}
) {
const receivedBackendData = JSON.stringify(backendData);
if (receivedBackendData !== '${expectedBackendData}'){
throw new Error("BackendData not equal Received: ".concat(receivedBackendData, " Expected: ", '${expectedBackendData}'));
}
return {
run: () => {},
};
};
`;
const blob = new Blob([workerCode], { type: "application/javascript" });
const blobURL = URL.createObjectURL(blob);
await EngineProcess.destroyMLEngine();
await IndexedDBCache.init({ reset: true });
let promiseStub = sinon
.stub(MLEngineParent, "getWorkerConfig")
.callsFake(function () {
return { url: blobURL, options: { type: "module" } };
});
let wasmBufferStub = sinon
.stub(MLEngineParent, "getWasmArrayBuffer")
.returns(backendData);
let chooseBestBackendStub = sinon
.stub(MLEngineParent, "chooseBestBackend")
.returns(BACKENDS.wllama);
try {
await createEngine({
engineId: "main",
taskName: "real-wllama-text-generation",
featureId: "link-preview",
backend: BACKENDS.bestLlama,
modelId: "acme/bert",
modelHubUrlTemplate: "{model}/resolve/{revision}",
modelRevision: "v0.4",
modelHubRootUrl:
modelFile: "onnx/config.json",
});
} finally {
await EngineProcess.destroyMLEngine();
await IndexedDBCache.init({ reset: true });
wasmBufferStub.restore();
promiseStub.restore();
chooseBestBackendStub.restore();
}
});
/**
* End to End test that the engine can indeed fail if it doesn't use best-llama.
*/
add_task(async function test_e2e_choose_backend_can_detect_failure() {
// Allow any url
Services.env.set("MOZ_ALLOW_EXTERNAL_ML_HUB", "true");
const backendData = new Uint8Array([10, 20, 30]);
const expectedBackendData = JSON.stringify("data so no matches");
// Mocking function used in the workers or child doesn't work.
// So we are stubbing the code run by the worker.
const workerCode = `
// Inject the MLEngine.worker.mjs code
${await getMLEngineWorkerCode()}
// Stub
ChromeUtils.defineESModuleGetters(
lazy,
{
createFileUrl: "chrome://global/content/ml/Utils.sys.mjs",
},
{ global: "current" }
);
// Change the getBackend to a mocked version that doesn't actually do inference
// but does initiate model downloads and engine initialization
lazy.getBackend = async function (
mlEngineWorker,
backendData,
{
modelHubUrlTemplate,
modelHubRootUrl,
modelId,
modelRevision,
modelFile,
engineId,
} = {}
) {
const receivedBackendData = JSON.stringify(backendData);
if (receivedBackendData !== '${expectedBackendData}'){
throw new Error("BackendData not equal Received: ".concat(receivedBackendData, " Expected: ", '${expectedBackendData}'));
}
return {
run: () => {},
};
};
`;
const blob = new Blob([workerCode], { type: "application/javascript" });
const blobURL = URL.createObjectURL(blob);
await EngineProcess.destroyMLEngine();
await IndexedDBCache.init({ reset: true });
let promiseStub = sinon
.stub(MLEngineParent, "getWorkerConfig")
.callsFake(function () {
return { url: blobURL, options: { type: "module" } };
});
let wasmBufferStub = sinon
.stub(MLEngineParent, "getWasmArrayBuffer")
.returns(backendData);
let chooseBestBackendStub = sinon
.stub(MLEngineParent, "chooseBestBackend")
.returns(BACKENDS.wllama);
try {
await Assert.rejects(
createEngine({
engineId: "main",
taskName: "real-wllama-text-generation",
featureId: "link-preview",
backend: BACKENDS.bestLlama,
modelId: "acme/bert",
modelHubUrlTemplate: "{model}/resolve/{revision}",
modelRevision: "v0.4",
modelHubRootUrl:
modelFile: "onnx/config.json",
}),
/BackendData not equal Received:/,
"The call should be rejected because it used the wrong backend"
);
} finally {
await EngineProcess.destroyMLEngine();
await IndexedDBCache.init({ reset: true });
wasmBufferStub.restore();
promiseStub.restore();
chooseBestBackendStub.restore();
}
});
/**
* End to End test that the engine is indeed initialized with llama.cpp when it is the
* best-llama.
*/
add_task(async function test_e2e_choose_backend_best_llamma_cpp() {
// Allow any url
Services.env.set("MOZ_ALLOW_EXTERNAL_ML_HUB", "true");
const backendData = new Uint8Array([10, 20, 30]);
const expectedBackendData = JSON.stringify(null);
// Mocking function used in the workers or child doesn't work.
// So we are stubbing the code run by the worker.
const workerCode = `
// Inject the MLEngine.worker.mjs code
${await getMLEngineWorkerCode()}
// Stub
ChromeUtils.defineESModuleGetters(
lazy,
{
createFileUrl: "chrome://global/content/ml/Utils.sys.mjs",
},
{ global: "current" }
);
// Change the getBackend to a mocked version that doesn't actually do inference
// but does initiate model downloads and engine initialization
lazy.getBackend = async function (
mlEngineWorker,
backendData,
{
modelHubUrlTemplate,
modelHubRootUrl,
modelId,
modelRevision,
modelFile,
engineId,
} = {}
) {
const receivedBackendData = JSON.stringify(backendData);
if (receivedBackendData !== '${expectedBackendData}'){
throw new Error("BackendData not equal Received: ".concat(receivedBackendData, " Expected: ", '${expectedBackendData}'));
}
return {
run: () => {},
};
};
`;
const blob = new Blob([workerCode], { type: "application/javascript" });
const blobURL = URL.createObjectURL(blob);
await EngineProcess.destroyMLEngine();
await IndexedDBCache.init({ reset: true });
let promiseStub = sinon
.stub(MLEngineParent, "getWorkerConfig")
.callsFake(function () {
return { url: blobURL, options: { type: "module" } };
});
let wasmBufferStub = sinon
.stub(MLEngineParent, "getWasmArrayBuffer")
.returns(backendData);
let chooseBestBackendStub = sinon
.stub(MLEngineParent, "chooseBestBackend")
.returns(BACKENDS.llamaCpp);
try {
await createEngine({
engineId: "main",
taskName: "real-wllama-text-generation",
featureId: "link-preview",
backend: BACKENDS.bestLlama,
modelId: "acme/bert",
modelHubUrlTemplate: "{model}/resolve/{revision}",
modelRevision: "v0.4",
modelHubRootUrl:
modelFile: "onnx/config.json",
});
} finally {
await EngineProcess.destroyMLEngine();
await IndexedDBCache.init({ reset: true });
wasmBufferStub.restore();
promiseStub.restore();
chooseBestBackendStub.restore();
}
});
/**
* End to End test that the engine can be cancelled.
*/
add_task(async function test_e2e_engine_can_be_cancelled() {
// Allow any url
Services.env.set("MOZ_ALLOW_EXTERNAL_ML_HUB", "true");
const backendData = new Uint8Array([10, 20, 30]);
// Mocking function used in the workers or child doesn't work.
// So we are stubbing the code run by the worker.
const workerCode = `
// Inject the MLEngine.worker.mjs code
${await getMLEngineWorkerCode()}
// Stub
ChromeUtils.defineESModuleGetters(
lazy,
{
createFileUrl: "chrome://global/content/ml/Utils.sys.mjs",
},
{ global: "current" }
);
// Change the getBackend to a mocked version that doesn't actually do inference
// but does initiate model downloads and engine initialization
lazy.getBackend = async function (
mlEngineWorker,
backendData,
{
modelHubUrlTemplate,
modelHubRootUrl,
modelId,
modelRevision,
modelFile,
engineId,
} = {}
) {
const url = lazy.createFileUrl({
model: modelId,
revision: modelRevision,
file: modelFile,
urlTemplate: modelHubUrlTemplate,
rootUrl: modelHubRootUrl,
});
await mlEngineWorker.getModelFile({url});
return {
run: () => {},
};
};
`;
const blob = new Blob([workerCode], { type: "application/javascript" });
const blobURL = URL.createObjectURL(blob);
await EngineProcess.destroyMLEngine();
await IndexedDBCache.init({ reset: true });
let promiseStub = sinon
.stub(MLEngineParent, "getWorkerConfig")
.callsFake(function () {
return { url: blobURL, options: { type: "module" } };
});
let wasmBufferStub = sinon
.stub(MLEngineParent, "getWasmArrayBuffer")
.returns(backendData);
const controller = new AbortController();
const { signal } = controller;
controller.abort();
try {
await Assert.rejects(
createEngine(
{
engineId: "main5",
taskName: "real-wllama-text-generation",
featureId: "link-preview",
backend: BACKENDS.llamaCpp,
modelId: "acme/bert",
modelHubUrlTemplate: "{model}/resolve/{revision}",
modelRevision: "v0.1",
modelHubRootUrl:
modelFile: "onnx/config.json",
},
null,
signal
),
/AbortError:/,
"The call should be cancelled"
);
} catch (err) {
Assert.ok(false, `Expected AbortError. Got ${err}`);
} finally {
await EngineProcess.destroyMLEngine();
await IndexedDBCache.init({ reset: true });
wasmBufferStub.restore();
promiseStub.restore();
}
});
/**
* End to End test that the engine can be cancelled after fetch success.
*/
add_task(async function test_e2e_engine_can_be_cancelled_after_fetch() {
// Allow any url
Services.env.set("MOZ_ALLOW_EXTERNAL_ML_HUB", "true");
const backendData = new Uint8Array([10, 20, 30]);
// Mocking function used in the workers or child doesn't work.
// So we are stubbing the code run by the worker.
const workerCode = `
// Inject the MLEngine.worker.mjs code
${await getMLEngineWorkerCode()}
// Stub
ChromeUtils.defineESModuleGetters(
lazy,
{
createFileUrl: "chrome://global/content/ml/Utils.sys.mjs",
},
{ global: "current" }
);
// Change the getBackend to a mocked version that doesn't actually do inference
// but does initiate model downloads and engine initialization
lazy.getBackend = async function (
mlEngineWorker,
backendData,
{
modelHubUrlTemplate,
modelHubRootUrl,
modelId,
modelRevision,
modelFile,
engineId,
} = {}
) {
const url = lazy.createFileUrl({
model: modelId,
revision: modelRevision,
file: modelFile,
urlTemplate: modelHubUrlTemplate,
rootUrl: modelHubRootUrl,
});
await mlEngineWorker.getModelFile({url});
return {
run: () => {},
};
};
`;
const blob = new Blob([workerCode], { type: "application/javascript" });
const blobURL = URL.createObjectURL(blob);
await EngineProcess.destroyMLEngine();
await IndexedDBCache.init({ reset: true });
let promiseStub = sinon
.stub(MLEngineParent, "getWorkerConfig")
.callsFake(function () {
return { url: blobURL, options: { type: "module" } };
});
let wasmBufferStub = sinon
.stub(MLEngineParent, "getWasmArrayBuffer")
.returns(backendData);
const controller = new AbortController();
const { signal } = controller;
const fetchUrlStub = sinon
.stub(MLUtils, "fetchUrl")
.callsFake((url, { signal: _, ...rest } = {}) => {
const p = fetch(url, rest);
controller.abort();
return p;
});
try {
await Assert.rejects(
createEngine(
{
engineId: "main5",
taskName: "real-wllama-text-generation",
featureId: "link-preview",
backend: BACKENDS.llamaCpp,
modelId: "acme/bert",
modelHubUrlTemplate: "{model}/resolve/{revision}",
modelRevision: "v0.1",
modelHubRootUrl:
modelFile: "onnx/config.json",
},
null,
signal
),
/AbortError:/,
"The call should be cancelled"
);
} catch (err) {
Assert.ok(false, `Expected AbortError. Got ${err}`);
} finally {
await EngineProcess.destroyMLEngine();
await IndexedDBCache.init({ reset: true });
wasmBufferStub.restore();
promiseStub.restore();
fetchUrlStub.restore();
}
});