Source code

Revision control

Copy as Markdown

Other Tools

Test Info:

/* Any copyright is dedicated to the Public Domain.
"use strict";
/* global TextGenerator */
const { ProfilerTestUtils } = ChromeUtils.importESModule(
);
async function hwInferencePid() {
const procInfo = await ChromeUtils.requestProcInfo();
for (const child of procInfo.children) {
if (
child.type.startsWith("utility") &&
child.utilityActors.some(actor => actor.actorName === "hwInference")
) {
return child.pid;
}
}
return 0;
}
add_task(async function test_hwinference_crash_and_respawn() {
const { cleanup } = await setup({
prefs: [["browser.ml.llama.hwInference", true]],
});
try {
// A generation long enough to still be running when the kill lands:
// one-token chunks, no stop on end of text, and a context to fill.
const engine = await createEngine({
...TINYSTORIES_ENGINE_OPTIONS,
engineId: "hwi-crash",
numContext: 4096,
});
const pid = await hwInferencePid();
Assert.greater(pid, 0, "The HWInference utility process is running");
const processTools = Cc["@mozilla.org/processtools-service;1"].getService(
Ci.nsIProcessToolsService
);
const gone = TestUtils.topicObserved(
"ipc:utility-shutdown",
(subject, data) => parseInt(data, 10) === pid
);
let killed = false;
await Assert.rejects(
(async () => {
for await (const chunk of engine.runWithGenerator({
prompt: TINYSTORIES_STORYTELLER_PROMPT,
samplers: TINYSTORIES_GREEDY_SAMPLERS,
nPredict: 4096,
minOutputBufferSize: 1,
stopOnEndOfGenerationTokens: false,
})) {
info(`chunk before kill: ${chunk.text}`);
// Chunks already in flight arrive after the kill; a second kill of
// a dying process fails on Windows.
if (!killed) {
killed = true;
processTools.kill(pid);
}
}
})(),
/went away|abort/i,
"A generation interrupted by a process crash rejects with the " +
"generator-teardown error"
);
await gone;
noteIntentionalUtilityCrash(pid);
await engine.terminate();
const respawned = await createEngine({
...TINYSTORIES_ENGINE_OPTIONS,
engineId: "hwi-crash-respawn",
});
const text = await collectGeneratedText(
respawned.runWithGenerator({
prompt: TINYSTORIES_STORYTELLER_PROMPT,
samplers: TINYSTORIES_GREEDY_SAMPLERS,
nPredict: 8,
})
);
info(`Respawned engine generated: ${text}`);
Assert.greater(
text.length,
0,
"A fresh engine after the crash generates again"
);
const newPid = await hwInferencePid();
Assert.greater(newPid, 0, "A HWInference utility process is running again");
Assert.notEqual(newPid, pid, "The utility process was respawned");
await respawned.terminate();
} finally {
await EngineProcess.destroyMLEngine();
await cleanup();
}
});
function getPayloads(profile, type) {
return ProfilerTestUtils.getPayloadsOfTypeFromAllThreads(profile, type);
}
add_task(async function test_process_lifecycle_markers() {
// The suite pins the idle grace to 0; reuse needs a real idle window. A
// deferred drop cannot be recalled by lowering the pref afterwards, so it
// is short enough to wait out.
const graceMs = 2000;
await SpecialPowers.pushPrefEnv({
set: [["browser.ml.hwinference.browser_idle_shutdown_grace_ms", graceMs]],
});
await ProfilerTestUtils.startProfiler({
features: ["stackwalk", "js"],
threads: ["GeckoMain", "TextGenerator"],
});
const cold = await createTinyStoriesGenerator();
const pid = await hwInferencePid();
Assert.greater(pid, 0, "The HWInference utility process is running");
const gone = TestUtils.topicObserved(
"ipc:utility-shutdown",
(subject, data) => parseInt(data, 10) === pid
);
cold.terminate();
const warm = await createTinyStoriesGenerator();
warm.terminate();
await gone;
const profile = await ProfilerTestUtils.stopNowAndGetProfile();
const spawn = getPayloads(profile, "MLProcessSpawn");
const releases = getPayloads(profile, "MLProcessRelease");
const creates = getPayloads(profile, "MLGeneratorCreate");
Assert.equal(spawn.length, 1, "The process was launched exactly once");
Assert.equal(
getPayloads(profile, "MLFailed").length,
0,
"A successful launch reports no failed phase"
);
Assert.equal(creates.length, 2, "Both creates are marked");
const acquires = getPayloads(profile, "MLProcessAcquire");
Assert.equal(acquires.length, 2, "Both process acquisitions are marked");
Assert.ok(
!acquires[0].processReused,
"The first acquisition paid for a spawn"
);
Assert.ok(acquires[1].processReused, "The second one reused the process");
Assert.equal(
creates.filter(payload => payload.spawnMs !== undefined).length,
0,
"Generator create does not fold in the process acquisition"
);
Assert.equal(releases.length, 2, "Both keep-alive releases are marked");
Assert.ok(
releases.every(payload => payload.graceMs === graceMs),
"Each release carries the grace it waited"
);
Assert.ok(
!releases[0].retired,
"The first release found the process still held by the second generator"
);
Assert.ok(releases[1].retired, "The last release retired the process");
await SpecialPowers.popPrefEnv();
});