Source code
Revision control
Copy as Markdown
Other Tools
Test Info: Warnings
-
CI results over the last 22 days2,868Runs99.93%Pass1Failures1Timeouts0Crashes0Skips
Runs with issues per day (up to 0.8%) Sep 14Oct 5- Failures
- Timeouts
- Manifest: toolkit/components/ml/tests/browser/browser.toml
/* Any copyright is dedicated to the Public Domain.
"use strict";
/* global TextGenerator */
const { ProfilerTestUtils } = ChromeUtils.importESModule(
);
const EXPECTED_MARKERS = [
"MLProcessAcquire",
"MLGeneratorCreate",
"MLGeneratorRun",
"MLChunkSend",
"MLBackendInit",
"MLModelPrefill",
"MLModelDecode",
];
function getPayloads(profile, type) {
return ProfilerTestUtils.getPayloadsOfTypeFromAllThreads(profile, type);
}
add_task(async function test_ml_pipeline_markers() {
await ProfilerTestUtils.startProfiler({
features: ["stackwalk", "js"],
threads: ["GeckoMain", "TextGenerator", "llama.cpp"],
});
const generator = await createTinyStoriesGenerator();
// Captured before the teardown: the process goes away with its last
// generator, and its markers with it.
let profile;
try {
await generator.generate(
{
messages: [{ role: "user", content: "Once upon a time" }],
maxTokens: 16,
bufferLength: 4,
samplers: TINYSTORIES_GREEDY_SAMPLERS,
},
() => {}
);
profile = await ProfilerTestUtils.stopNowAndGetProfile();
} finally {
generator.terminate();
}
for (const type of EXPECTED_MARKERS) {
Assert.greater(getPayloads(profile, type).length, 0, `${type} was emitted`);
Assert.ok(
ProfilerTestUtils.getSchema(profile, type),
`${type} has a marker schema`
);
}
const threadsFor = type => {
const found = new Set();
(function walk(proc) {
for (const thread of proc.threads) {
const { markers } = thread;
for (const t of markers.data) {
if (t[markers.schema.data]?.type === type) {
found.add(thread.name);
}
}
}
proc.processes.forEach(walk);
})(profile);
return found;
};
Assert.ok(
threadsFor("MLGeneratorRun").has("GeckoMain"),
"The parent half of a generation is on GeckoMain"
);
Assert.ok(
[...threadsFor("MLModelPrefill")].some(name =>
name.startsWith("TextGenerator")
),
"The model half is on the generator thread in HWInference"
);
const generate = getPayloads(profile, "MLGeneratorRun")[0];
Assert.greater(generate.computeMs, 0, "The model compute is reported");
Assert.greater(
generate.prefillTokensPerSecond,
0,
"The prefill throughput is reported alongside the decode one, so the " +
"label cannot imply the whole interval ran at decode speed"
);
Assert.greaterOrEqual(
generate.overheadMs,
0,
"The overhead around the model is reported"
);
Assert.greaterOrEqual(
generate.deliverMs,
0,
"What the caller's callback cost is reported"
);
Assert.greater(
generate.chunkTokens,
0,
"The chunk size the ttfc number refers to is reported"
);
const prefill = getPayloads(profile, "MLModelPrefill")[0];
Assert.greater(prefill.promptTokens, 0, "The prompt tokens are counted");
const decode = getPayloads(profile, "MLModelDecode")[0];
Assert.ok(
decode,
"The prompt phase boundary was observed, so a decode phase is reported"
);
Assert.equal(
getPayloads(profile, "MLFailed").length,
0,
"A clean run reports no failed phase"
);
Assert.greater(
getPayloads(profile, "MLChunkSend").length,
0,
"The child marked the chunks it streamed out"
);
});
add_task(async function test_teardown_mid_generation_is_not_a_failure() {
await ProfilerTestUtils.startProfiler({
features: ["stackwalk", "js"],
threads: ["GeckoMain", "TextGenerator"],
});
const generator = await createTinyStoriesGenerator();
// Terminate only once a chunk arrived, so it lands mid-generation.
let sawChunk;
const firstChunk = new Promise(resolve => {
sawChunk = resolve;
});
const generation = generator.generate(
{
messages: [{ role: "user", content: "Once upon a time" }],
maxTokens: 4096,
bufferLength: 1,
samplers: TINYSTORIES_GREEDY_SAMPLERS,
},
() => sawChunk()
);
await firstChunk;
generator.terminate();
await Assert.rejects(
generation,
/terminated/,
"The abandoned generation rejects"
);
const profile = await ProfilerTestUtils.stopNowAndGetProfile();
Assert.equal(
getPayloads(profile, "MLFailed").length,
0,
"A teardown the caller asked for is not marked as a failed phase"
);
const stages = getPayloads(profile, "MLCancel").map(p => p.stage);
Assert.ok(stages.includes("teardown"), "The teardown is marked as a cancel");
});
// The MLEngine path adds no markers of its own: what it contributes is the
// feature name on the generator markers.
add_task(async function test_engine_path_names_the_feature() {
const { cleanup } = await setup({
prefs: [["browser.ml.llama.hwInference", true]],
});
await ProfilerTestUtils.startProfiler({
features: ["stackwalk", "js"],
threads: ["GeckoMain", "TextGenerator"],
});
try {
const engine = await createEngine({
...TINYSTORIES_ENGINE_OPTIONS,
engineId: "marker-engine",
featureId: "link-preview",
});
for await (const chunk of engine.runWithGenerator({
prompt: [{ role: "user", content: "Once upon a time" }],
samplers: TINYSTORIES_GREEDY_SAMPLERS,
nPredict: 16,
})) {
void chunk;
}
const profile = await ProfilerTestUtils.stopNowAndGetProfile();
await engine.terminate();
const create = getPayloads(profile, "MLGeneratorCreate")[0];
Assert.ok(create, "The generator create was emitted");
Assert.equal(
create.featureId,
"link-preview",
"The create names the feature that asked for it"
);
const run = getPayloads(profile, "MLGeneratorRun")[0];
Assert.ok(run, "The generator run was emitted");
Assert.equal(run.featureId, "link-preview", "So does the run");
Assert.greater(run.generatedTokens, 0, "The run generated tokens");
} finally {
await cleanup();
}
});
add_task(async function test_abandoned_stream_is_marked_cancelled() {
const { cleanup } = await setup({
prefs: [["browser.ml.llama.hwInference", true]],
});
await ProfilerTestUtils.startProfiler({
features: ["stackwalk", "js"],
threads: ["GeckoMain", "TextGenerator"],
});
try {
const engine = await createEngine({
...TINYSTORIES_ENGINE_OPTIONS,
engineId: "cancelled-engine",
featureId: "link-preview",
numContext: 8192,
});
for await (const chunk of engine.runWithGenerator({
prompt: [{ role: "user", content: "Once upon a time" }],
samplers: TINYSTORIES_GREEDY_SAMPLERS,
nPredict: 8192,
minOutputBufferSize: 1,
stopOnEndOfGenerationTokens: false,
})) {
void chunk;
break;
}
await engine.terminate();
const profile = await ProfilerTestUtils.stopNowAndGetProfile();
const stages = getPayloads(profile, "MLCancel").map(p => p.stage);
Assert.ok(
stages.includes("requested"),
"Abandoning the loop asked the generator to stop"
);
const generatorRun = getPayloads(profile, "MLGeneratorRun")[0];
Assert.ok(
generatorRun,
"The cancelled generation is still marked as a generation"
);
Assert.equal(
generatorRun.reason,
"cancelled",
"A cancel is a finish reason on the run, not a run of its own"
);
Assert.greater(
generatorRun.generatedTokens,
0,
"The cancelled generation reports the tokens it produced, so it " +
"cannot read as a run that generated nothing"
);
Assert.ok(
!stages.includes("teardown"),
"Waiting for the cancelled run to settle means the terminate no " +
"longer races its reply"
);
} finally {
await cleanup();
}
});