Source code
Revision control
Copy as Markdown
Other Tools
Test Info:
/* Any copyright is dedicated to the Public Domain.
"use strict";
const { EngineProcess: PerfTestEngineProcess } = ChromeUtils.importESModule(
"chrome://global/content/ml/EngineProcess.sys.mjs"
);
const { MLEngine: PerfTestMLEngine, MLEngineParent: PerfTestMLEngineParent } =
ChromeUtils.importESModule(
"moz-src:///toolkit/components/ml/actors/MLEngineParent.sys.mjs"
);
const { MLPerfTestUtils } = ChromeUtils.importESModule(
);
const { sinon: perfTestSinon } = ChromeUtils.importESModule(
);
/**
* This test case ensures every measured lifecycle reports its individual
* replicates and median without reporting warmups.
*/
add_task(async function test_run_perf_scenario_lifecycles() {
const contexts = [];
let message;
await MLPerfTestUtils.runPerfScenario({
Assert,
info(value) {
message = value;
},
metricPrefix: "TEST",
metricSuffix: "NATIVE",
async scenario(context) {
contexts.push(context);
return { duration: contexts.length };
},
coldIterations: 2,
warmIterations: 2,
memoryIterations: 0,
});
Assert.deepEqual(
contexts,
[
{ lifecycle: "first-use", sampleKind: "latency", iteration: 0 },
{ lifecycle: "cold", sampleKind: "latency", iteration: 0 },
{ lifecycle: "cold", sampleKind: "latency", iteration: 1 },
{ lifecycle: "warm", sampleKind: "warmup", iteration: 0 },
{ lifecycle: "warm", sampleKind: "latency", iteration: 0 },
{ lifecycle: "warm", sampleKind: "latency", iteration: 1 },
],
"The scenario receives each lifecycle and sample kind"
);
Assert.deepEqual(
JSON.parse(message.replace("perfMetrics | ", "")),
[
{ name: "TEST-duration-first-use-NATIVE", values: [1], value: 1 },
{
name: "TEST-duration-cold-NATIVE",
values: [2, 3],
value: 2.5,
},
{ name: "TEST-duration-warm-NATIVE", values: [5, 6], value: 5.5 },
],
"Only measured scenario invocations are reported"
);
});
/**
* This test case ensures an unreported first use still populates caches before
* the first cold measurement.
*/
add_task(async function test_run_perf_scenario_unreported_first_use() {
const contexts = [];
let message;
await MLPerfTestUtils.runPerfScenario({
Assert,
info(value) {
message = value;
},
metricPrefix: "TEST",
async scenario(context) {
contexts.push(context);
return { duration: contexts.length };
},
measureFirstUse: false,
coldIterations: 1,
warmIterations: 0,
memoryIterations: 0,
});
Assert.deepEqual(
contexts,
[
{ lifecycle: "first-use", sampleKind: "warmup", iteration: 0 },
{ lifecycle: "cold", sampleKind: "latency", iteration: 0 },
],
"First use runs without being measured"
);
Assert.deepEqual(
JSON.parse(message.replace("perfMetrics | ", "")),
[{ name: "TEST-duration-cold", values: [2], value: 2 }],
"Only the cold invocation is reported"
);
});
/**
* This test case ensures the runner reports shared and peak memory while
* reusing one engine within each warm measurement phase.
*/
add_task(async function test_run_perf_scenario_reports_engine_metrics() {
const { cleanup, remoteClients } = await setup();
const warmEngines = [];
let engine;
let message;
try {
await MLPerfTestUtils.runPerfScenario({
Assert,
info(value) {
message = value;
},
metricPrefix: "TEST",
metricSuffix: "NATIVE",
engines: [
{
featureId: "formfill-classification",
overrides: {
numThreads: { expectValue: 2, replaceWith: 1 },
},
},
],
async scenario(context) {
let created = false;
if (!engine || ["closed", "error"].includes(engine.engineStatus)) {
engine = await createEngine({
backend: "onnx",
featureId: "formfill-classification",
numThreads: 2,
taskName: "moz-echo",
timeoutMS: -1,
});
created = true;
}
const run = engine.run({ data: "Shared metrics" });
if (created) {
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
}
const result = await run;
if (context.lifecycle === "warm") {
warmEngines.push(engine);
}
Assert.equal(
result.output.echo,
"Shared metrics",
"The production engine completes inference"
);
return { duration: 5 };
},
coldIterations: 0,
warmIterations: 1,
memoryIterations: 1,
});
const metrics = JSON.parse(message.replace("perfMetrics | ", ""));
const metricsByName = new Map(metrics.map(metric => [metric.name, metric]));
Assert.equal(
metrics.length,
10,
"Only feature-owned and shared engine measurements are reported"
);
Assert.deepEqual(
metricsByName.get("TEST-duration-first-use-NATIVE"),
{ name: "TEST-duration-first-use-NATIVE", values: [5], value: 5 },
"The feature-owned measurement is preserved"
);
Assert.deepEqual(
metricsByName.get("TEST-duration-warm-NATIVE"),
{ name: "TEST-duration-warm-NATIVE", values: [5], value: 5 },
"The warm feature measurement is preserved"
);
Assert.equal(
warmEngines.length,
3,
"Warmup, warm and memory runs happened"
);
Assert.equal(
new Set(warmEngines).size,
1,
"Warm latency and memory reuse the warmup engine"
);
const peakMemory = metricsByName.get("TEST-peak-memory-NATIVE");
Assert.equal(peakMemory.values.length, 1, "Peak memory has one sample");
Assert.greater(peakMemory.value, 0, "Peak memory is positive");
for (const name of [
"engine-creation-time",
"engine-run-time",
"memory-before-run",
"memory-after-run",
]) {
const lifecycles =
name === "engine-creation-time" ? ["first-use"] : ["first-use", "warm"];
for (const lifecycle of lifecycles) {
const metric = metricsByName.get(`TEST-${name}-${lifecycle}-NATIVE`);
Assert.ok(metric, `${name} is reported for ${lifecycle}`);
Assert.equal(
metric.values.length,
1,
`${name} has one ${lifecycle} replicate`
);
Assert.ok(
Number.isFinite(metric.value),
`${name} has a finite ${lifecycle} median`
);
}
}
} finally {
await PerfTestEngineProcess.destroyMLEngine();
await cleanup();
}
});
/**
* This test case ensures final engine teardown still runs when a performance
* scenario rejects.
*/
add_task(async function test_run_perf_scenario_cleans_up_after_failure() {
const destroySpy = perfTestSinon.spy(
PerfTestEngineProcess,
"destroyMLEngine"
);
try {
await Assert.rejects(
MLPerfTestUtils.runPerfScenario({
Assert,
info() {},
metricPrefix: "TEST",
async scenario() {
throw new Error("Expected scenario failure");
},
coldIterations: 0,
warmIterations: 0,
memoryIterations: 0,
}),
/Expected scenario failure/,
"The scenario failure is propagated"
);
Assert.equal(
destroySpy.callCount,
2,
"Engines are destroyed before first use and again after failure"
);
} finally {
destroySpy.restore();
}
});
/**
* This test case ensures the runner observes configured engines, passes
* unrelated engines through, and restores production methods.
*/
add_task(async function test_run_perf_scenario_observes_real_engines() {
const { cleanup, remoteClients } = await setup();
const originalGetEngine = PerfTestMLEngineParent.prototype.getEngine;
const originalRun = PerfTestMLEngine.prototype.run;
let message;
try {
await MLPerfTestUtils.runPerfScenario({
Assert,
info(value) {
message = value;
},
metricPrefix: "TEST",
engines: [
{
featureId: "formfill-encoder",
metricName: "encoder",
overrides: {
numThreads: { expectValue: 2, replaceWith: 1 },
},
},
{
featureId: "formfill-head",
metricName: "head",
overrides: {
numThreads: { expectValue: 2, replaceWith: 1 },
},
},
],
async scenario() {
const featureIds = [
"formfill-encoder",
"formfill-head",
"about-inference",
];
const engines = await Promise.all(
featureIds.map(featureId =>
createEngine({
backend: "onnx",
featureId,
numThreads: 2,
taskName: "moz-echo",
})
)
);
const runs = engines.map((engine, index) =>
engine.run({ data: `Engine ${index}` })
);
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(
engines.length
);
const results = await Promise.all(runs);
Assert.deepEqual(
results.map(result => result.output.echo),
["Engine 0", "Engine 1", "Engine 2"],
"Every production engine completes inference"
);
Assert.deepEqual(
engines.map(engine => engine.pipelineOptions.numThreads),
[1, 1, 2],
"Configured engines receive overrides and the unrelated engine does not"
);
return { duration: 5 };
},
coldIterations: 0,
warmIterations: 0,
memoryIterations: 0,
});
const metrics = JSON.parse(message.replace("perfMetrics | ", ""));
const metricsByName = new Map(metrics.map(metric => [metric.name, metric]));
Assert.equal(
metrics.length,
9,
"Only the feature and configured engines report measurements"
);
Assert.deepEqual(
metricsByName.get("TEST-duration-first-use"),
{ name: "TEST-duration-first-use", values: [5], value: 5 },
"The feature measurement is reported"
);
for (const engineName of ["encoder", "head"]) {
for (const metricName of [
"engine-creation-time",
"engine-run-time",
"memory-before-run",
"memory-after-run",
]) {
const metric = metricsByName.get(
`TEST-${engineName}-${metricName}-first-use`
);
Assert.ok(metric, `${engineName} reports ${metricName}`);
Assert.equal(metric.values.length, 1, "The metric has one replicate");
}
}
Assert.equal(
PerfTestMLEngineParent.prototype.getEngine,
originalGetEngine,
"Engine creation interception is restored"
);
Assert.equal(
PerfTestMLEngine.prototype.run,
originalRun,
"Engine run observation is restored"
);
} finally {
await PerfTestEngineProcess.destroyMLEngine();
await cleanup();
}
});
/**
* This test case ensures cleanup restores the production run method before
* rejecting unfinished inference.
*/
add_task(async function test_run_perf_scenario_rejects_an_active_engine_run() {
const { cleanup, remoteClients } = await setup();
let finishRun;
const runStub = perfTestSinon
.stub(PerfTestMLEngine.prototype, "run")
.callsFake(
() =>
new Promise(resolve => {
finishRun = resolve;
})
);
let runPromise;
try {
await Assert.rejects(
MLPerfTestUtils.runPerfScenario({
Assert,
info() {},
metricPrefix: "TEST",
engines: [{ featureId: "formfill-classification" }],
async scenario() {
const engine = await createEngine({
backend: "onnx",
featureId: "formfill-classification",
taskName: "moz-echo",
});
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
runPromise = engine.run({});
return { duration: 1 };
},
coldIterations: 0,
warmIterations: 0,
memoryIterations: 0,
}),
/1 engine run was still active during cleanup/,
"Cleanup rejects an unfinished engine run"
);
Assert.equal(
PerfTestMLEngine.prototype.run,
runStub,
"The production run method is restored before cleanup rejects"
);
} finally {
finishRun?.({});
await runPromise;
runStub.restore();
await cleanup();
}
});
/**
* This test case ensures completed and intentionally abandoned generators
* report generation metrics without fabricating final resources.
*/
add_task(async function test_run_perf_scenario_observes_generator_lifecycles() {
const { cleanup, remoteClients } = await setup();
const runWithGeneratorStub = perfTestSinon
.stub(PerfTestMLEngine.prototype, "runWithGenerator")
.callsFake(async function* () {
yield { text: "", tokens: [10, 11], isPrompt: true };
await TestUtils.waitForTick();
yield { text: "A", tokens: [1], isPrompt: false };
await TestUtils.waitForTick();
yield { text: "B", tokens: [2], isPrompt: false };
return {
resourcesBefore: { cpuTime: 1, memory: 2 * 1024 * 1024 },
resourcesAfter: { cpuTime: 2, memory: 3 * 1024 * 1024 },
metrics: {
decodingTime: 10,
inputTokens: 2,
outputTokens: 2,
},
};
});
let message;
try {
await MLPerfTestUtils.runPerfScenario({
Assert,
info(value) {
message = value;
},
metricPrefix: "TEST",
engines: [{ featureId: "link-preview", expectedRuns: 2 }],
async scenario() {
const engine = await createEngine({
backend: "onnx",
featureId: "link-preview",
taskName: "moz-echo",
});
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
const chunks = [];
for await (const chunk of engine.runWithGenerator({})) {
if (!chunk.isPrompt) {
chunks.push(chunk.text);
}
}
Assert.deepEqual(
chunks,
["A", "B"],
"The production chunks pass through"
);
for await (const chunk of engine.runWithGenerator({})) {
if (chunk.text === "B") {
break;
}
}
return { duration: 5 };
},
coldIterations: 0,
warmIterations: 0,
memoryIterations: 0,
});
const metrics = JSON.parse(message.replace("perfMetrics | ", ""));
const metricsByName = new Map(metrics.map(metric => [metric.name, metric]));
for (const name of [
"engine-run-time",
"time-to-first-token",
"tokens-per-second",
"output-tokens",
]) {
Assert.equal(
metricsByName.get(`TEST-${name}-first-use`).values.length,
2,
`${name} includes the completed and abandoned runs`
);
}
Assert.deepEqual(
metricsByName.get("TEST-output-tokens-first-use").values,
[2, 2],
"output-tokens counts the completed and abandoned streams"
);
for (const [name, expectedValue] of [
["memory-before-run", 2],
["memory-after-run", 3],
["decoding-time", 10],
["input-tokens", 2],
]) {
Assert.deepEqual(
metricsByName.get(`TEST-${name}-first-use`).values,
[expectedValue],
`${name} includes only the completed run`
);
}
Assert.equal(
PerfTestMLEngine.prototype.runWithGenerator,
runWithGeneratorStub,
"The production generator method is restored"
);
} finally {
runWithGeneratorStub.restore();
await cleanup();
}
});