summaryrefslogtreecommitdiffstats
path: root/toolkit/components/ml/content
diff options
context:
space:
mode:
authorDaniel Baumann <daniel.baumann@progress-linux.org>2024-04-19 01:13:33 +0000
committerDaniel Baumann <daniel.baumann@progress-linux.org>2024-04-19 01:13:33 +0000
commit086c044dc34dfc0f74fbe41f4ecb402b2cd34884 (patch)
treea4f824bd33cb075dd5aa3eb5a0a94af221bbe83a /toolkit/components/ml/content
parentAdding debian version 124.0.1-1. (diff)
downloadfirefox-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.mjs241
-rw-r--r--toolkit/components/ml/content/MLEngine.html16
-rw-r--r--toolkit/components/ml/content/MLEngine.worker.mjs91
-rw-r--r--toolkit/components/ml/content/SummarizerModel.sys.mjs160
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;
+ }
+}