From 8dd16259287f58f9273002717ec4d27e97127719 Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Wed, 12 Jun 2024 07:43:14 +0200 Subject: Merging upstream version 127.0. Signed-off-by: Daniel Baumann --- toolkit/components/ml/actors/MLEngineChild.sys.mjs | 156 +++++++++++++++++---- 1 file changed, 125 insertions(+), 31 deletions(-) (limited to 'toolkit/components/ml/actors/MLEngineChild.sys.mjs') diff --git a/toolkit/components/ml/actors/MLEngineChild.sys.mjs b/toolkit/components/ml/actors/MLEngineChild.sys.mjs index 925ce59266..17a8b3511a 100644 --- a/toolkit/components/ml/actors/MLEngineChild.sys.mjs +++ b/toolkit/components/ml/actors/MLEngineChild.sys.mjs @@ -21,6 +21,8 @@ ChromeUtils.defineESModuleGetters(lazy, { BasePromiseWorker: "resource://gre/modules/PromiseWorker.sys.mjs", setTimeout: "resource://gre/modules/Timer.sys.mjs", clearTimeout: "resource://gre/modules/Timer.sys.mjs", + ModelHub: "chrome://global/content/ml/ModelHub.sys.mjs", + PipelineOptions: "chrome://global/content/ml/EngineProcess.sys.mjs", }); ChromeUtils.defineLazyGetter(lazy, "console", () => { @@ -32,9 +34,20 @@ ChromeUtils.defineLazyGetter(lazy, "console", () => { XPCOMUtils.defineLazyPreferenceGetter( lazy, - "loggingLevel", - "browser.ml.logLevel" + "CACHE_TIMEOUT_MS", + "browser.ml.modelCacheTimeout" ); +XPCOMUtils.defineLazyPreferenceGetter( + lazy, + "MODEL_HUB_ROOT_URL", + "browser.ml.modelHubRootUrl" +); +XPCOMUtils.defineLazyPreferenceGetter( + lazy, + "MODEL_HUB_URL_TEMPLATE", + "browser.ml.modelHubUrlTemplate" +); +XPCOMUtils.defineLazyPreferenceGetter(lazy, "LOG_LEVEL", "browser.ml.logLevel"); /** * The engine child is responsible for the life cycle and instantiation of the local @@ -52,10 +65,21 @@ export class MLEngineChild extends JSWindowActorChild { async receiveMessage({ name, data }) { switch (name) { case "MLEngine:NewPort": { - const { engineName, port, timeoutMS } = data; + const { port, pipelineOptions } = data; + + // Override some options using prefs + let options = new lazy.PipelineOptions(pipelineOptions); + + options.updateOptions({ + modelHubRootUrl: lazy.MODEL_HUB_ROOT_URL, + modelHubUrlTemplate: lazy.MODEL_HUB_URL_TEMPLATE, + timeoutMS: lazy.CACHE_TIMEOUT_MS, + logLevel: lazy.LOG_LEVEL, + }); + this.#engineDispatchers.set( - engineName, - new EngineDispatcher(this, port, engineName, timeoutMS) + options.taskName, + new EngineDispatcher(this, port, options) ); break; } @@ -78,12 +102,23 @@ export class MLEngineChild extends JSWindowActorChild { } /** - * @returns {ArrayBuffer} + * Gets the wasm array buffer from RemoteSettings. + * + * @returns {Promise} */ getWasmArrayBuffer() { return this.sendQuery("MLEngine:GetWasmArrayBuffer"); } + /** + * Gets the inference options from RemoteSettings. + * + * @returns {Promise} + */ + getInferenceOptions(taskName) { + return this.sendQuery(`MLEngine:GetInferenceOptions:${taskName}`); + } + /** * @param {string} engineName */ @@ -113,28 +148,45 @@ class EngineDispatcher { #engine = null; /** @type {string} */ - #engineName; + #taskName; + + /** Creates the inference engine given the wasm runtime and the run options. + * + * The initialization is done in three steps: + * 1. The wasm runtime is fetched from RS + * 2. The inference options are fetched from RS and augmented with the pipeline options. + * 3. The inference engine is created with the wasm runtime and the options. + * + * Any exception here will be bubbled up for the constructor to log. + * + * @param {PipelineOptions} pipelineOptions + * @returns {Promise} + */ + async initializeInferenceEngine(pipelineOptions) { + // Create the inference engine given the wasm runtime and the options. + const wasm = await this.mlEngineChild.getWasmArrayBuffer(); + const inferenceOptions = await this.mlEngineChild.getInferenceOptions( + this.#taskName + ); + lazy.console.debug("Inference engine options:", inferenceOptions); + pipelineOptions.updateOptions(inferenceOptions); + + return InferenceEngine.create(wasm, pipelineOptions); + } /** * @param {MLEngineChild} mlEngineChild * @param {MessagePort} port - * @param {string} engineName - * @param {number} timeoutMS + * @param {PipelineOptions} pipelineOptions */ - constructor(mlEngineChild, port, engineName, timeoutMS) { - /** @type {MLEngineChild} */ + constructor(mlEngineChild, port, pipelineOptions) { this.mlEngineChild = mlEngineChild; + this.#taskName = pipelineOptions.taskName; + this.timeoutMS = pipelineOptions.timeoutMS; - /** @type {number} */ - this.timeoutMS = timeoutMS; - - this.#engineName = engineName; - - this.#engine = Promise.all([ - this.mlEngineChild.getWasmArrayBuffer(), - this.getModel(port), - ]).then(([wasm, model]) => FakeEngine.create(wasm, model)); + this.#engine = this.initializeInferenceEngine(pipelineOptions); + // Trigger the keep alive timer. this.#engine .then(() => void this.keepAlive()) .catch(error => { @@ -265,20 +317,61 @@ class EngineDispatcher { port.close(); } this.#ports = new Set(); - this.mlEngineChild.removeEngine(this.#engineName); + this.mlEngineChild.removeEngine(this.#taskName); try { const engine = await this.#engine; engine.terminate(); } catch (error) { - console.error("Failed to get the engine", error); + lazy.console.error("Failed to get the engine", error); } } } +let modelHub = null; // This will hold the ModelHub instance to reuse it. + +/** + * Retrieves a model file as an ArrayBuffer from the specified URL. + * This function normalizes the URL, extracts the organization, model name, and file path, + * then fetches the model file using the ModelHub API. The `modelHub` instance is created + * only once and reused for subsequent calls to optimize performance. + * + * @param {string} url - The URL of the model file to fetch. Can be a path relative to + * the model hub root or an absolute URL. + * @returns {Promise} A promise that resolves to a Meta object containing the URL, response headers, + * and data as an ArrayBuffer. The data is marked for transfer to avoid cloning. + */ +async function getModelFile(url) { + // Create the model hub instance if needed + if (!modelHub) { + lazy.console.debug("Creating model hub instance"); + modelHub = new lazy.ModelHub({ + rootUrl: lazy.MODEL_HUB_ROOT_URL, + urlTemplate: lazy.MODEL_HUB_URL_TEMPLATE, + }); + } + + if (url.startsWith(lazy.MODEL_HUB_ROOT_URL)) { + url = url.slice(lazy.MODEL_HUB_ROOT_URL.length); + // Make sure we get a front slash + if (!url.startsWith("/")) { + url = `/${url}`; + } + } + + // Parsing url to get model name, and file path. + // if this errors out, it will be caught in the worker + const parsedUrl = modelHub.parseUrl(url); + + let [data, headers] = await modelHub.getModelFileAsArrayBuffer(parsedUrl); + return new lazy.BasePromiseWorker.Meta([url, headers, data], { + transfers: [data], + }); +} + /** - * Fake the engine by slicing the text in half. + * Wrapper around the ChromeWorker that runs the inference. */ -class FakeEngine { +class InferenceEngine { /** @type {BasePromiseWorker} */ #worker; @@ -286,21 +379,22 @@ class FakeEngine { * Initialize the worker. * * @param {ArrayBuffer} wasm - * @param {ArrayBuffer} model - * @returns {FakeEngine} + * @param {PipelineOptions} pipelineOptions + * @returns {InferenceEngine} */ - static async create(wasm, model) { + static async create(wasm, pipelineOptions) { /** @type {BasePromiseWorker} */ const worker = new lazy.BasePromiseWorker( "chrome://global/content/ml/MLEngine.worker.mjs", - { type: "module" } + { type: "module" }, + { getModelFile } ); - const args = [wasm, model, lazy.loggingLevel]; + const args = [wasm, pipelineOptions]; const closure = {}; - const transferables = [wasm, model]; + const transferables = [wasm]; await worker.post("initializeEngine", args, closure, transferables); - return new FakeEngine(worker); + return new InferenceEngine(worker); } /** -- cgit v1.2.3