commit 4903e6275b2f9d65ad17d1dbb9b367ac2415955e
parent 0974467706f2067587546c466e799d484eed8dd9
Author: Greg Tatum <tatum.creative@gmail.com>
Date: Thu, 30 Oct 2025 19:45:48 +0000
Bug 1996291 - Rework getStatus to getStatusByEngineId and fix status reporting in about:inference r=firefox-ai-ml-reviewers,tarek
When going through the tests I could not figure out what this code did.
After some sleuthing I realized that it was gathering information for
the about:inference. However, the code was badly broken, and incorrect
in what it was doing. I simplified the implementation a bit, and fixed
the breakage in about:inference.
There are no tests for this behavior in about:inference, and there as
far as I can tell there are no live models in about:inference that could
be used to test this behavior. Since it's a developer-only facing
feature, I'm not planning on adding a test for it.
Also please note that I'm adding an ml.d.ts file here to centralize some
shared types. See the in-flight TypeScript documentation here:
https://phabricator.services.mozilla.com/D268571
Differential Revision: https://phabricator.services.mozilla.com/D269982
Diffstat:
5 files changed, 138 insertions(+), 56 deletions(-)
diff --git a/toolkit/components/aboutinference/content/aboutInference.js b/toolkit/components/aboutinference/content/aboutInference.js
@@ -6,6 +6,11 @@
"use strict";
/**
+ * @import { MLEngineParent } from "resource://gre/actors/MLEngineParent.sys.mjs"
+ * @import { StatusByEngineId } from "../../ml/ml.d.ts"
+ */
+
+/**
* Imports necessary modules from ChromeUtils.
*/
const lazy = {};
@@ -75,6 +80,7 @@ function getNumThreadsArray() {
);
}
+/** @type {MLEngineParent | null} */
let engineParent = null;
const TINY_ARTICLE =
@@ -460,25 +466,26 @@ function ts2str(ts) {
*/
async function updateStatus() {
- if (!engineParent) {
- return;
- }
+ const engineParent = await getEngineParent();
- let info;
+ /**
+ * @type {StatusByEngineId}
+ */
+ let statusByEngineId;
// Fetch the engine status info
try {
- info = await engineParent.getStatus();
- } catch (e) {
- engineParent = null; // let's re-create it on errors.
- info = new Map();
+ statusByEngineId = await engineParent.getStatusByEngineId();
+ } catch (error) {
+ console.error("Failed to get the engine status", error);
+ statusByEngineId = new Map();
}
// Get the container where the table will be displayed
let tableContainer = document.getElementById("statusTableContainer");
// Clear the container if the map is empty
- if (info.size === 0) {
+ if (statusByEngineId.size === 0) {
tableContainer.innerHTML = ""; // Clear any existing table
if (updateStatusInterval) {
clearInterval(updateStatusInterval); // Clear the interval if it exists
@@ -520,7 +527,7 @@ async function updateStatus() {
let tbody = document.createElement("tbody");
// Iterate over the info map
- for (let [engineId, engineInfo] of info.entries()) {
+ for (let [engineId, { status, options }] of statusByEngineId.entries()) {
let row = document.createElement("tr");
// Create a cell for each piece of data
@@ -529,23 +536,23 @@ async function updateStatus() {
row.appendChild(engineIdCell);
let statusCell = document.createElement("td");
- statusCell.textContent = engineInfo.status;
+ statusCell.textContent = status;
row.appendChild(statusCell);
let modelIdCell = document.createElement("td");
- modelIdCell.textContent = engineInfo.options?.modelId || "N/A";
+ modelIdCell.textContent = options?.modelId || "N/A";
row.appendChild(modelIdCell);
let dtypeCell = document.createElement("td");
- dtypeCell.textContent = engineInfo.options?.dtype || "N/A";
+ dtypeCell.textContent = options?.dtype || "N/A";
row.appendChild(dtypeCell);
let deviceCell = document.createElement("td");
- deviceCell.textContent = engineInfo.options?.device || "N/A";
+ deviceCell.textContent = options?.device || "N/A";
row.appendChild(deviceCell);
let timeoutCell = document.createElement("td");
- timeoutCell.textContent = engineInfo.options?.timeoutMS || "N/A";
+ timeoutCell.textContent = options?.timeoutMS || "N/A";
row.appendChild(timeoutCell);
// Append the row to the table body
@@ -1133,6 +1140,9 @@ function showTab(button) {
button.setAttribute("selected", "true");
}
+/**
+ * @returns {Promise<MLEngineParent>}
+ */
async function getEngineParent() {
if (!engineParent) {
engineParent = await EngineProcess.getMLEngineParent();
diff --git a/toolkit/components/ml/actors/MLEngineChild.sys.mjs b/toolkit/components/ml/actors/MLEngineChild.sys.mjs
@@ -9,6 +9,7 @@ import { XPCOMUtils } from "resource://gre/modules/XPCOMUtils.sys.mjs";
/**
* @import { BasePromiseWorker } from "resource://gre/modules/PromiseWorker.sys.mjs"
* @import { PipelineOptions } from "chrome://global/content/ml/EngineProcess.sys.mjs"
+ * @import { EngineStatus, EngineId, StatusByEngineId } from "../ml.d.ts"
* @import { ProgressAndStatusCallbackParams } from "chrome://global/content/ml/Utils.sys.mjs"
*/
@@ -62,11 +63,11 @@ export class MLEngineChild extends JSProcessActorChild {
#engineDispatchers = new Map();
/**
- * Engine statuses
+ * Tracks that an engine is present, even if the dispatcher is not present yet.
*
- * @type {Map<string, string>}
+ * @type {Map<EngineId, PipelineOptions>}
*/
- #engineStatuses = new Map();
+ #enginesPresent = new Map();
// eslint-disable-next-line consistent-return
async receiveMessage({ name, data }) {
@@ -75,8 +76,8 @@ export class MLEngineChild extends JSProcessActorChild {
await this.#onNewPortCreated(data);
break;
}
- case "MLEngine:GetStatus": {
- return this.getStatus();
+ case "MLEngine:GetStatusByEngineId": {
+ return this.getStatusByEngineId();
}
case "MLEngine:ForceShutdown": {
for (const engineDispatcher of this.#engineDispatchers.values()) {
@@ -114,7 +115,7 @@ export class MLEngineChild extends JSProcessActorChild {
this.getUpdatedPipelineOptions(pipelineOptions);
options.updateOptions(updatedPipelineOptions);
const engineId = options.engineId;
- this.#engineStatuses.set(engineId, "INITIALIZING");
+ this.#enginesPresent.set(engineId, options);
// Check if we already have an engine under this id.
if (this.#engineDispatchers.has(engineId)) {
@@ -126,8 +127,6 @@ export class MLEngineChild extends JSProcessActorChild {
type: "EnginePort:EngineReady",
error: null,
});
- this.#engineStatuses.set(engineId, "READY");
-
return;
}
@@ -139,8 +138,6 @@ export class MLEngineChild extends JSProcessActorChild {
this.#engineDispatchers.delete(engineId);
}
- this.#engineStatuses.set(engineId, "CREATING");
-
const dispatcher = new EngineDispatcher(this, port, options);
this.#engineDispatchers.set(engineId, dispatcher);
@@ -154,7 +151,6 @@ export class MLEngineChild extends JSProcessActorChild {
await dispatcher.isReady();
}
- this.#engineStatuses.set(engineId, "READY");
port.postMessage({
type: "EnginePort:EngineReady",
error: null,
@@ -238,7 +234,7 @@ export class MLEngineChild extends JSProcessActorChild {
*/
removeEngine(engineId, shutDownIfEmpty, replacement) {
this.#engineDispatchers.delete(engineId);
- this.#engineStatuses.delete(engineId);
+ this.#enginesPresent.delete(engineId);
try {
this.sendAsyncMessage("MLEngine:Removed", {
@@ -262,19 +258,26 @@ export class MLEngineChild extends JSProcessActorChild {
}
}
- /*
+ /**
* Collects information about the current status.
+ *
+ * @returns {StatusByEngineId}
*/
- async getStatus() {
+ getStatusByEngineId() {
+ /** @type {StatusByEngineId} */
const statusMap = new Map();
- for (const [key, value] of this.#engineStatuses) {
- if (this.#engineDispatchers.has(key)) {
- statusMap.set(key, this.#engineDispatchers.get(key).getStatus());
- } else {
- // The engine is probably being created
- statusMap.set(key, { status: value });
+ for (let [engineId, options] of this.#enginesPresent) {
+ const dispatcher = this.#engineDispatchers.get(engineId);
+ let status = dispatcher.getStatus();
+ if (!status) {
+ // This engine doesn't have a dispatcher yet.
+ status = {
+ status: "SHUTTING_DOWN_PREVIOUS_ENGINE",
+ options,
+ };
}
+ statusMap.set(engineId, status);
}
return statusMap;
}
@@ -329,7 +332,7 @@ class EngineDispatcher {
/** @type {PipelineOptions | null} */
pipelineOptions = null;
- /** @type {string} */
+ /** @type {EngineStatus} */
#status;
/**
@@ -414,7 +417,7 @@ class EngineDispatcher {
* @param {PipelineOptions} pipelineOptions
*/
constructor(mlEngineChild, port, pipelineOptions) {
- this.#status = "CREATED";
+ this.#status = "INITIALIZING";
/** @type {MLEngineChild} */
this.mlEngineChild = mlEngineChild;
this.#featureId = pipelineOptions.featureId;
@@ -431,7 +434,7 @@ class EngineDispatcher {
this.#engine
.then(() => {
- this.#status = "READY";
+ this.#status = "IDLE";
// Trigger the keep alive timer.
void this.keepAlive();
})
@@ -454,7 +457,6 @@ class EngineDispatcher {
return {
status: this.#status,
options: this.pipelineOptions,
- engineId: this.#engineId,
};
}
@@ -583,7 +585,7 @@ class EngineDispatcher {
error,
});
}
- this.#status = "IDLING";
+ this.#status = "IDLE";
break;
}
default:
diff --git a/toolkit/components/ml/actors/MLEngineParent.sys.mjs b/toolkit/components/ml/actors/MLEngineParent.sys.mjs
@@ -3,6 +3,10 @@
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */
import { XPCOMUtils } from "resource://gre/modules/XPCOMUtils.sys.mjs";
+/**
+ * @import { MLEngineChild } from "./MLEngineChild.sys.mjs"
+ */
+
const lazy = XPCOMUtils.declareLazy({
RemoteSettings: "resource://services-settings/remote-settings.sys.mjs",
Utils: "resource://services-settings/Utils.sys.mjs",
@@ -227,10 +231,8 @@ export class MLEngineParent extends JSProcessActorParent {
// Wait for the existing lock to resolve
await MLEngineParent.engineLocks.get(engineId);
}
- let resolveLock;
- const lockPromise = new Promise(resolve => {
- resolveLock = resolve;
- });
+ const { promise: lockPromise, resolve: resolveLock } =
+ Promise.withResolvers();
MLEngineParent.engineLocks.set(engineId, lockPromise);
MLEngineParent.engineCreationAbortSignal.set(engineId, abortSignal);
try {
@@ -783,10 +785,15 @@ export class MLEngineParent extends JSProcessActorParent {
}
/**
- * Gets a status
+ * Goes through the engines and determines their status. This is used by about:inference
+ * to display debug information about the engines.
+ *
+ * @see MLEngineChild#getStatusByEngineId
+ *
+ * @returns {Promise<StatusByEngineId>}
*/
- getStatus() {
- return this.sendQuery("MLEngine:GetStatus");
+ getStatusByEngineId() {
+ return this.sendQuery("MLEngine:GetStatusByEngineId");
}
/**
@@ -941,7 +948,7 @@ export class MLEngine {
#requests = new Map();
/**
- * @type {"uninitialized" | "ready" | "error" | "closed"}
+ * @type {"uninitialized" | "ready" | "error" | "closed" | "crashed"}
*/
engineStatus = "uninitialized";
diff --git a/toolkit/components/ml/ml.d.ts b/toolkit/components/ml/ml.d.ts
@@ -0,0 +1,60 @@
+/* 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/. */
+
+/**
+ * This file contains the shared types for the machine learning component. The intended
+ * use is for defining types to be used in JSDoc. They are used in a form that the
+ * TypeScript language server can read them, and provide code hints.
+ *
+ * @see https://firefox-source-docs.mozilla.org/code-quality/typescript/
+ */
+
+import { type PipelineOptions } from "chrome://global/content/ml/EngineProcess.sys.mjs";
+
+export type EngineStatus =
+ // The engine is waiting for a previous one to shut down.
+ | "SHUTTING_DOWN_PREVIOUS_ENGINE"
+ // The engine dispatcher has been created, and the engine is still initializing.
+ | "INITIALIZING"
+ // The engine is fully ready and idle.
+ | "IDLE"
+ // The engine is currently processing a run request.
+ | "RUNNING"
+ // The engine is in the process of terminating, but hasn't fully shut down.
+ | "TERMINATING"
+ // The engine has been fully terminated and removed.
+ | "TERMINATED";
+
+/**
+ * The EngineId is used to identify a unique engine that can be shared across multiple
+ * consumers. This way a single model can be loaded into memory and used in different
+ * locations, assuming the other parameters match as well.
+ */
+export type EngineId = string;
+
+/**
+ * Utility type to extract the data fields from a class. It removes all of the
+ * functions.
+ */
+type DataFields<T> = {
+ [K in keyof T as T[K] extends Function ? never : K]: T[K];
+};
+
+/**
+ * The PipelineOptions are a nominal class that validates the options. The
+ * PipelineOptionsRaw are the raw subset of those.
+ */
+type PipelineOptionsRaw = Partial<DataFields<PipelineOptions>>;
+
+/**
+ * Tracks the current status of the engines for about:inference. It's not used
+ * for deciding any business logic of the engines, only for debug info.
+ */
+export type StatusByEngineId = Map<
+ EngineId,
+ {
+ status: EngineStatus;
+ options: PipelineOptions | PipelineOptionsRaw;
+ }
+>;
diff --git a/toolkit/components/ml/tests/browser/browser_ml_engine_lifetime.js b/toolkit/components/ml/tests/browser/browser_ml_engine_lifetime.js
@@ -435,7 +435,11 @@ add_task(async function test_ml_engine_infinite_worker() {
await cleanup();
});
-add_task(async function test_ml_engine_get_status() {
+/**
+ * These status are visualized in about:inference, but aren't used for business
+ * logic.
+ */
+add_task(async function test_ml_engine_get_status_by_engine_id() {
const { cleanup, remoteClients } = await setup();
info("Get the engine");
@@ -459,7 +463,7 @@ add_task(async function test_ml_engine_get_status() {
const expected = {
"default-engine": {
- status: "IDLING",
+ status: "IDLE",
options: {
useExternalDataFormat: false,
engineId: "default-engine",
@@ -496,15 +500,14 @@ add_task(async function test_ml_engine_get_status() {
baseURL: null,
apiKey: null,
},
- engineId: "default-engine",
},
};
- let status = await engineInstance.mlEngineParent.getStatus();
- status = JSON.parse(JSON.stringify(Object.fromEntries(status)));
-
- status["default-engine"].options.numThreads = "NOT_COMPARED";
- Assert.deepEqual(status, expected);
+ const statusByEngineId = Object.fromEntries(
+ await engineInstance.mlEngineParent.getStatusByEngineId()
+ );
+ statusByEngineId["default-engine"].options.numThreads = "NOT_COMPARED";
+ Assert.deepEqual(statusByEngineId, expected);
await ok(
!EngineProcess.areAllEnginesTerminated(),
