diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-19 01:13:33 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-19 01:13:33 +0000 |
commit | 086c044dc34dfc0f74fbe41f4ecb402b2cd34884 (patch) | |
tree | a4f824bd33cb075dd5aa3eb5a0a94af221bbe83a /toolkit/components/ml/content | |
parent | Adding debian version 124.0.1-1. (diff) | |
download | firefox-086c044dc34dfc0f74fbe41f4ecb402b2cd34884.tar.xz firefox-086c044dc34dfc0f74fbe41f4ecb402b2cd34884.zip |
Merging upstream version 125.0.1.
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'toolkit/components/ml/content')
-rw-r--r-- | toolkit/components/ml/content/EngineProcess.sys.mjs | 241 | ||||
-rw-r--r-- | toolkit/components/ml/content/MLEngine.html | 16 | ||||
-rw-r--r-- | toolkit/components/ml/content/MLEngine.worker.mjs | 91 | ||||
-rw-r--r-- | toolkit/components/ml/content/SummarizerModel.sys.mjs | 160 |
4 files changed, 508 insertions, 0 deletions
diff --git a/toolkit/components/ml/content/EngineProcess.sys.mjs b/toolkit/components/ml/content/EngineProcess.sys.mjs new file mode 100644 index 0000000000..36a9381192 --- /dev/null +++ b/toolkit/components/ml/content/EngineProcess.sys.mjs @@ -0,0 +1,241 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +const lazy = {}; +ChromeUtils.defineESModuleGetters(lazy, { + HiddenFrame: "resource://gre/modules/HiddenFrame.sys.mjs", +}); + +/** + * @typedef {import("../actors/MLEngineParent.sys.mjs").MLEngineParent} MLEngineParent + */ + +/** + * @typedef {import("../../translations/actors/TranslationsEngineParent.sys.mjs").TranslationsEngineParent} TranslationsEngineParent + */ + +/** + * This class controls the life cycle of the engine process used both in the + * Translations engine and the MLEngine component. + */ +export class EngineProcess { + /** + * @type {Promise<{ hiddenFrame: HiddenFrame, actor: TranslationsEngineParent }> | null} + */ + + /** @type {Promise<HiddenFrame> | null} */ + static #hiddenFrame = null; + /** @type {Promise<TranslationsEngineParent> | null} */ + static translationsEngineParent = null; + /** @type {Promise<MLEngineParent> | null} */ + static mlEngineParent = null; + + /** @type {((actor: TranslationsEngineParent) => void) | null} */ + resolveTranslationsEngineParent = null; + + /** @type {((actor: MLEngineParent) => void) | null} */ + resolveMLEngineParent = null; + + /** + * See if all engines are terminated. This is useful for testing. + * + * @returns {boolean} + */ + static areAllEnginesTerminated() { + return ( + !EngineProcess.#hiddenFrame && + !EngineProcess.translationsEngineParent && + !EngineProcess.mlEngineParent + ); + } + + /** + * @returns {Promise<TranslationsEngineParent>} + */ + static async getTranslationsEngineParent() { + if (!this.translationsEngineParent) { + this.translationsEngineParent = this.#attachBrowser({ + id: "translations-engine-browser", + url: "chrome://global/content/translations/translations-engine.html", + resolverName: "resolveTranslationsEngineParent", + }); + } + return this.translationsEngineParent; + } + + /** + * @returns {Promise<MLEngineParent>} + */ + static async getMLEngineParent() { + if (!this.mlEngineParent) { + this.mlEngineParent = this.#attachBrowser({ + id: "ml-engine-browser", + url: "chrome://global/content/ml/MLEngine.html", + resolverName: "resolveMLEngineParent", + }); + } + return this.mlEngineParent; + } + + /** + * @param {object} config + * @param {string} config.url + * @param {string} config.id + * @param {string} config.resolverName + * @returns {Promise<TranslationsEngineParent>} + */ + static async #attachBrowser({ url, id, resolverName }) { + const hiddenFrame = await this.#getHiddenFrame(); + const chromeWindow = await hiddenFrame.get(); + const doc = chromeWindow.document; + + if (doc.getElementById(id)) { + throw new Error( + "Attempting to append the translations-engine.html <browser> when one " + + "already exists." + ); + } + + const browser = doc.createXULElement("browser"); + browser.setAttribute("id", id); + browser.setAttribute("remote", "true"); + browser.setAttribute("remoteType", "web"); + browser.setAttribute("disableglobalhistory", "true"); + browser.setAttribute("type", "content"); + browser.setAttribute("src", url); + + ChromeUtils.addProfilerMarker( + "EngineProcess", + {}, + `Creating the "${id}" process` + ); + doc.documentElement.appendChild(browser); + + const { promise, resolve } = Promise.withResolvers(); + + // The engine parents must resolve themselves when they are ready. + this[resolverName] = resolve; + + return promise; + } + + /** + * @returns {HiddenFrame} + */ + static async #getHiddenFrame() { + if (!EngineProcess.#hiddenFrame) { + EngineProcess.#hiddenFrame = new lazy.HiddenFrame(); + } + return EngineProcess.#hiddenFrame; + } + + /** + * Destroy the translations engine, and remove the hidden frame if no other + * engines exist. + */ + static destroyTranslationsEngine() { + return this.#destroyEngine({ + id: "translations-engine-browser", + keyName: "translationsEngineParent", + }); + } + + /** + * Destroy the ML engine, and remove the hidden frame if no other engines exist. + */ + static destroyMLEngine() { + return this.#destroyEngine({ + id: "ml-engine-browser", + keyName: "mlEngineParent", + }); + } + + /** + * Destroy the specified engine and maybe the entire hidden frame as well if no engines + * are remaining. + */ + static #destroyEngine({ id, keyName }) { + ChromeUtils.addProfilerMarker( + "EngineProcess", + {}, + `Destroying the "${id}" engine` + ); + + const actorShutdown = this.forceActorShutdown(id, keyName).catch( + error => void console.error(error) + ); + + this[keyName] = null; + + const hiddenFrame = EngineProcess.#hiddenFrame; + if (hiddenFrame && !this.translationsEngineParent && !this.mlEngineParent) { + EngineProcess.#hiddenFrame = null; + + // Both actors are destroyed, also destroy the hidden frame. + actorShutdown.then(() => { + // Double check a race condition that no new actors have been created during + // shutdown. + if (this.translationsEngineParent && this.mlEngineParent) { + return; + } + if (!hiddenFrame) { + return; + } + hiddenFrame.destroy(); + ChromeUtils.addProfilerMarker( + "EngineProcess", + {}, + `Removing the hidden frame` + ); + }); + } + + // Infallibly resolve the promise even if there are errors. + return Promise.resolve(); + } + + /** + * Shut down an actor and remove its <browser> element. + * + * @param {string} id + * @param {string} keyName + */ + static async forceActorShutdown(id, keyName) { + const actorPromise = this[keyName]; + if (!actorPromise) { + return; + } + + let actor; + try { + actor = await actorPromise; + } catch { + // The actor failed to initialize, so it doesn't need to be shut down. + return; + } + + // Shut down the actor. + try { + await actor.forceShutdown(); + } catch (error) { + console.error("Failed to shut down the actor " + id, error); + return; + } + + if (!EngineProcess.#hiddenFrame) { + // The hidden frame was already removed. + return; + } + + // Remove the <brower> element. + const chromeWindow = EngineProcess.#hiddenFrame.getWindow(); + const doc = chromeWindow.document; + const element = doc.getElementById(id); + if (!element) { + console.error("Could not find the <browser> element for " + id); + return; + } + element.remove(); + } +} diff --git a/toolkit/components/ml/content/MLEngine.html b/toolkit/components/ml/content/MLEngine.html new file mode 100644 index 0000000000..8763995102 --- /dev/null +++ b/toolkit/components/ml/content/MLEngine.html @@ -0,0 +1,16 @@ +<!-- This Source Code Form is subject to the terms of the Mozilla Public + - License, v. 2.0. If a copy of the MPL was not distributed with this + - file, You can obtain one at http://mozilla.org/MPL/2.0/. --> + +<!DOCTYPE html> +<html> + <head> + <meta charset="utf-8" /> + <meta + http-equiv="Content-Security-Policy" + content="default-src chrome: resource:; object-src 'none'" + /> + <!-- Run the machine learning inference engine in its own singleton content process. --> + </head> + <body></body> +</html> diff --git a/toolkit/components/ml/content/MLEngine.worker.mjs b/toolkit/components/ml/content/MLEngine.worker.mjs new file mode 100644 index 0000000000..1013977e07 --- /dev/null +++ b/toolkit/components/ml/content/MLEngine.worker.mjs @@ -0,0 +1,91 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +import { PromiseWorker } from "resource://gre/modules/workers/PromiseWorker.mjs"; + +// Respect the preference "browser.ml.logLevel". +let _loggingLevel = "Error"; +function log(...args) { + if (_loggingLevel !== "Error" && _loggingLevel !== "Warn") { + console.log("ML:", ...args); + } +} +function trace(...args) { + if (_loggingLevel === "Trace" || _loggingLevel === "All") { + console.log("ML:", ...args); + } +} + +/** + * The actual MLEngine lives here in a worker. + */ +class MLEngineWorker { + /** @type {ArrayBuffer} */ + #wasm; + /** @type {ArrayBuffer} */ + #model; + + constructor() { + // Connect the provider to the worker. + this.#connectToPromiseWorker(); + } + + /** + * @param {ArrayBuffer} wasm + * @param {ArrayBuffer} model + * @param {string} loggingLevel + */ + initializeEngine(wasm, model, loggingLevel) { + this.#wasm = wasm; + this.#model = model; + _loggingLevel = loggingLevel; + // TODO - Initialize the engine for real here. + log("MLEngineWorker is initalized"); + } + + /** + * Run the worker. + * + * @param {string} request + */ + run(request) { + if (!this.#wasm) { + throw new Error("Expected the wasm to exist."); + } + if (!this.#model) { + throw new Error("Expected the model to exist"); + } + if (request === "throw") { + throw new Error( + 'Received the message "throw", so intentionally throwing an error.' + ); + } + trace("inference run requested with:", request); + return request.slice(0, Math.floor(request.length / 2)); + } + + /** + * Glue code to connect the `MLEngineWorker` to the PromiseWorker interface. + */ + #connectToPromiseWorker() { + const worker = new PromiseWorker.AbstractWorker(); + worker.dispatch = (method, args = []) => { + if (!this[method]) { + throw new Error("Method does not exist: " + method); + } + return this[method](...args); + }; + worker.close = () => self.close(); + worker.postMessage = (message, ...transfers) => { + self.postMessage(message, ...transfers); + }; + + self.addEventListener("message", msg => worker.handleMessage(msg)); + self.addEventListener("unhandledrejection", function (error) { + throw error.reason; + }); + } +} + +new MLEngineWorker(); diff --git a/toolkit/components/ml/content/SummarizerModel.sys.mjs b/toolkit/components/ml/content/SummarizerModel.sys.mjs new file mode 100644 index 0000000000..7cac55d92f --- /dev/null +++ b/toolkit/components/ml/content/SummarizerModel.sys.mjs @@ -0,0 +1,160 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +/** + * @typedef {object} LazyImports + * @property {typeof import("../actors/MLEngineParent.sys.mjs").MLEngineParent} MLEngineParent + */ + +/** @type {LazyImports} */ +const lazy = {}; + +ChromeUtils.defineESModuleGetters(lazy, { + RemoteSettings: "resource://services-settings/remote-settings.sys.mjs", + TranslationsParent: "resource://gre/actors/TranslationsParent.sys.mjs", +}); + +ChromeUtils.defineLazyGetter(lazy, "console", () => { + return console.createInstance({ + maxLogLevelPref: "browser.ml.logLevel", + prefix: "ML", + }); +}); + +export class SummarizerModel { + /** + * The RemoteSettingsClient that downloads the summarizer model. + * + * @type {RemoteSettingsClient | null} + */ + static #remoteClient = null; + + /** @type {Promise<WasmRecord> | null} */ + static #modelRecord = null; + + /** + * The following constant controls the major version for wasm downloaded from + * Remote Settings. When a breaking change is introduced, Nightly will have these + * numbers incremented by one, but Beta and Release will still be on the previous + * version. Remote Settings will ship both versions of the records, and the latest + * asset released in that version will be used. For instance, with a major version + * of "1", assets can be downloaded for "1.0", "1.2", "1.3beta", but assets marked + * as "2.0", "2.1", etc will not be downloaded. + */ + static MODEL_MAJOR_VERSION = 1; + + /** + * Remote settings isn't available in tests, so provide mocked responses. + */ + static mockRemoteSettings(remoteClient) { + lazy.console.log("Mocking remote client in SummarizerModel."); + SummarizerModel.#remoteClient = remoteClient; + SummarizerModel.#modelRecord = null; + } + + /** + * Remove anything that could have been mocked. + */ + static removeMocks() { + lazy.console.log("Removing mocked remote client in SummarizerModel."); + SummarizerModel.#remoteClient = null; + SummarizerModel.#modelRecord = null; + } + /** + * Download or load the model from remote settings. + * + * @returns {Promise<ArrayBuffer>} + */ + static async getModel() { + const client = SummarizerModel.#getRemoteClient(); + + if (!SummarizerModel.#modelRecord) { + // Place the records into a promise to prevent any races. + SummarizerModel.#modelRecord = (async () => { + // Load the wasm binary from remote settings, if it hasn't been already. + lazy.console.log(`Getting the summarizer model record.`); + + // TODO - The getMaxVersionRecords should eventually migrated to some kind of + // shared utility. + const { getMaxVersionRecords } = lazy.TranslationsParent; + + /** @type {WasmRecord[]} */ + const wasmRecords = await getMaxVersionRecords(client, { + // TODO - This record needs to be created with the engine wasm payload. + filters: { name: "summarizer-model" }, + majorVersion: SummarizerModel.MODEL_MAJOR_VERSION, + }); + + if (wasmRecords.length === 0) { + // The remote settings client provides an empty list of records when there is + // an error. + throw new Error("Unable to get the models from Remote Settings."); + } + + if (wasmRecords.length > 1) { + SummarizerModel.reportError( + new Error("Expected the ml engine to only have 1 record."), + wasmRecords + ); + } + const [record] = wasmRecords; + lazy.console.log( + `Using ${record.name}@${record.release} release version ${record.version} first released on Fx${record.fx_release}`, + record + ); + return record; + })(); + } + + try { + /** @type {{buffer: ArrayBuffer}} */ + const { buffer } = await client.attachments.download( + await SummarizerModel.#modelRecord + ); + + return buffer; + } catch (error) { + SummarizerModel.#modelRecord = null; + throw error; + } + } + + /** + * Lazily initializes the RemoteSettingsClient. + * + * @returns {RemoteSettingsClient} + */ + static #getRemoteClient() { + if (SummarizerModel.#remoteClient) { + return SummarizerModel.#remoteClient; + } + + /** @type {RemoteSettingsClient} */ + const client = lazy.RemoteSettings("ml-model"); + + SummarizerModel.#remoteClient = client; + + client.on("sync", async ({ data: { created, updated, deleted } }) => { + lazy.console.log(`"sync" event for ml-model`, { + created, + updated, + deleted, + }); + + // Remove all the deleted records. + for (const record of deleted) { + await client.attachments.deleteDownloaded(record); + } + + // Remove any updated records, and download the new ones. + for (const { old: oldRecord } of updated) { + await client.attachments.deleteDownloaded(oldRecord); + } + + // Do nothing for the created records. + }); + + return client; + } +} |