Source code
Revision control
Copy as Markdown
Other Tools
Test Info:
-
CI results over the last 22 days
- Manifest: toolkit/components/ml/tests/browser/browser.toml
/* Any copyright is dedicated to the Public Domain.
"use strict";
const { LlamaCppPipeline } = ChromeUtils.importESModule(
"chrome://global/content/ml/backends/LlamaCppPipeline.mjs"
);
/**
* A generator standing in for LlamaRunner. Its stream delivers one chunk and
* then leaves a read outstanding, and its cancel algorithm only settles once
* the test releases it, the way the generation thread holds the backend after
* a cancellation.
*
* @returns {{generator: object, control: object}} The stub and the handles the
* test drives it with: `streams` counts generations started,
* `releases` holds a resolver per cancelled stream, and `firstChunk`
* settles once a generation has produced output.
*/
function createGeneratorStub() {
let deliverFirstChunk;
const control = {
streams: 0,
releases: [],
firstChunk: new Promise(resolve => {
deliverFirstChunk = resolve;
}),
};
const generator = {
createGenerationStream() {
control.streams++;
let pulls = 0;
return new ReadableStream({
pull(controller) {
pulls++;
if (pulls === 1) {
controller.enqueue({
phase: "generation",
tokens: [1],
piece: "x",
});
deliverFirstChunk();
return undefined;
}
return new Promise(() => {});
},
cancel() {
return new Promise(resolve => control.releases.push(resolve));
},
});
},
};
return { generator, control };
}
async function settle() {
await TestUtils.waitForTick();
await TestUtils.waitForTick();
}
// Cancelling a stream closes it, so the run ends and cancels the same stream
// again on its way out. That second cancel resolves immediately, and taking it
// as the drain would let the next generation start while the backend, which is
// not thread-safe, is still held by the cancelled one.
add_task(async function test_next_run_waits_for_a_cancelled_generation() {
const { generator, control } = createGeneratorStub();
const pipeline = new LlamaCppPipeline(generator, {}, error => error);
const request = { prompt: "Once upon a time there was", nPredict: 200 };
const cancelledRun = pipeline.run(request, "run-1");
await control.firstChunk;
await settle();
// Not awaited: it settles only once the backend is released below.
const cancelled = pipeline.cancel("run-1");
// Closing the stream ends the read loop, so the run itself returns here.
await cancelledRun;
Assert.equal(control.streams, 1, "The cancelled run started one generation");
Assert.equal(control.releases.length, 1, "The stream's cancel algorithm ran");
const nextRun = pipeline.run(request, "run-2");
await settle();
Assert.equal(
control.streams,
1,
"The next run waits for the cancelled generation to release the backend"
);
control.releases.shift()();
await cancelled;
await settle();
Assert.equal(
control.streams,
2,
"The next run starts once the backend is free"
);
const nextCancelled = pipeline.cancel("run-2");
await settle();
control.releases.forEach(release => release());
await nextCancelled;
await nextRun;
});
// A cancellation can also land before the run has opened its stream, in which
// case the run cancels on its own behalf once it has a reader. That drain has
// to be waited on too, otherwise the same overlap happens through a narrower
// window.
add_task(async function test_next_run_waits_for_a_cancel_before_the_reader() {
const { generator, control } = createGeneratorStub();
let enterFormatChat;
let releaseFormatChat;
const formatChatEntered = new Promise(resolve => {
enterFormatChat = resolve;
});
const formatChatReleased = new Promise(resolve => {
releaseFormatChat = resolve;
});
// formatChat is only awaited for a chat prompt, which parks the run after it
// has registered but before it opens a stream.
generator.formatChat = async () => {
enterFormatChat();
await formatChatReleased;
return "Once upon a time there was";
};
const pipeline = new LlamaCppPipeline(generator, {}, error => error);
const request = {
prompt: [{ role: "user", content: "Once upon a time there was" }],
nPredict: 200,
};
const cancelledRun = pipeline.run(request, "run-1");
await formatChatEntered;
await pipeline.cancel("run-1");
releaseFormatChat();
await settle();
Assert.equal(control.streams, 1, "The cancelled run started one generation");
Assert.equal(control.releases.length, 1, "The stream's cancel algorithm ran");
const nextRun = pipeline.run(request, "run-2");
await settle();
Assert.equal(
control.streams,
1,
"The next run waits even though the cancel preceded the reader"
);
control.releases.shift()();
await cancelledRun;
await settle();
Assert.equal(
control.streams,
2,
"The next run starts once the backend is free"
);
const nextCancelled = pipeline.cancel("run-2");
await settle();
control.releases.forEach(release => release());
await nextCancelled;
await nextRun;
});