Bug 1944691 - Local model managment - Display local models r=Mardak,extension-reviewers,rpl
basic hooking up model hub on init calculating stuff from files Differential Revision: https://phabricator.services.mozilla.com/D240550
This commit is contained in:
Nick Grato
@@ -388,7 +388,7 @@ async function updateModels() {
|
||||
for (const { name: model, revision } of models) {
|
||||
const icon = await hub.getOwnerIcon(model);
|
||||
|
||||
let files = await hub.listFiles({ model, revision });
|
||||
let { files } = await hub.listFiles({ model, revision });
|
||||
|
||||
// Create a new table for the current model
|
||||
let table = document.createElement("table");
|
||||
|
||||
@@ -1090,69 +1090,112 @@ class IndexedDBCache {
|
||||
}
|
||||
|
||||
/**
|
||||
* Lists all files for a given model and revision stored in the cache.
|
||||
* Lists all files for a given model and revision stored in the cache,
|
||||
* and aggregates metadata from the file headers.
|
||||
*
|
||||
* @param {object} config
|
||||
* @param {?string} config.model - The model name (organization/name).
|
||||
* When a `taskName` is provided, the method retrieves all model/revision
|
||||
* pairs associated with that task; otherwise, it uses the provided `model`
|
||||
* and `revision`. It then queries the store to retrieve file information (path
|
||||
* and headers) and aggregates metadata (totalSize, lastUsed, updateDate, engineIds)
|
||||
* across all files.
|
||||
*
|
||||
* @param {object} config - The configuration for querying the files.
|
||||
* @param {?string} config.model - The model name (in "organization/name" format).
|
||||
* @param {?string} config.revision - The model version.
|
||||
* @param {?string} config.taskName - name of the inference :wtask.
|
||||
* @returns {Promise<Array<{path:string, headers: object}>>} An array of file identifiers.
|
||||
* @param {?string} config.taskName - The name of the inference task.
|
||||
* @returns {Promise<{
|
||||
* files: Array<{ path: string, headers: object, engineIds: Array<string> }>,
|
||||
* metadata: { totalSize: number, lastUsed: number, updateDate: number, engineIds: Array<string> }
|
||||
* }>} An object containing:
|
||||
* - files: an array of file records with their path, headers, and engine IDs.
|
||||
* - metadata: aggregated metadata computed from all the files.
|
||||
*/
|
||||
async listFiles({ taskName, model, revision }) {
|
||||
// When not providing taskName, we want model and revision
|
||||
// When not providing taskName, both model and revision must be defined.
|
||||
if (!taskName && (!model || !revision)) {
|
||||
throw new Error("Both model and revision must be defined");
|
||||
}
|
||||
|
||||
// Determine which model/revision pairs we want files for.
|
||||
let modelRevisions = [{ model, revision }];
|
||||
if (taskName) {
|
||||
// Get all model/revision associated to this task.
|
||||
const data = await this.#getKeys({
|
||||
// Get all model/revision pairs associated with this task.
|
||||
const keysData = await this.#getKeys({
|
||||
storeName: this.taskStoreName,
|
||||
...this.#getFileQuery({ taskName, model, revision }),
|
||||
});
|
||||
|
||||
modelRevisions = [];
|
||||
for (const { key } of data) {
|
||||
modelRevisions.push({ model: key[1], revision: key[2] });
|
||||
}
|
||||
modelRevisions = keysData.map(({ key }) => ({
|
||||
model: key[1],
|
||||
revision: key[2],
|
||||
}));
|
||||
}
|
||||
|
||||
const filePromises = [];
|
||||
|
||||
for (const task of modelRevisions) {
|
||||
filePromises.push(
|
||||
// For each model/revision, query for headers data.
|
||||
const fileDataPromises = modelRevisions.map(task =>
|
||||
this.#getData({
|
||||
storeName: this.headersStoreName,
|
||||
indexName: this.#indices.modelRevisionIndex.name,
|
||||
key: [task.model, task.revision],
|
||||
})
|
||||
);
|
||||
}
|
||||
const fileData = (await Promise.all(fileDataPromises)).flat();
|
||||
|
||||
const data = (await Promise.all(filePromises)).flat();
|
||||
// Initialize aggregated metadata.
|
||||
let totalFileSize = 0;
|
||||
let aggregatedLastUsed = 0;
|
||||
let aggregatedUpdateDate = 0;
|
||||
let aggregatedEngineIds = [];
|
||||
|
||||
// Process each file entry.
|
||||
const files = [];
|
||||
for (const { file: path, headers } of data) {
|
||||
files.push({ path, headers });
|
||||
for (const { file: path, headers } of fileData) {
|
||||
// Aggregate metadata.
|
||||
totalFileSize += headers.fileSize;
|
||||
aggregatedLastUsed = Math.max(aggregatedLastUsed, headers.lastUsed);
|
||||
aggregatedUpdateDate = Math.max(
|
||||
aggregatedUpdateDate,
|
||||
headers.lastUpdated
|
||||
);
|
||||
if (headers.engineIds && headers.engineIds.length) {
|
||||
aggregatedEngineIds = headers.engineIds;
|
||||
}
|
||||
files.push({ path, headers, engineIds: headers.engineIds || [] });
|
||||
}
|
||||
|
||||
return files;
|
||||
return {
|
||||
files,
|
||||
metadata: {
|
||||
totalSize: totalFileSize,
|
||||
lastUsed: aggregatedLastUsed,
|
||||
updateDate: aggregatedUpdateDate,
|
||||
engineIds: aggregatedEngineIds,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Lists all models stored in the cache.
|
||||
*
|
||||
* @returns {Promise<Array<{name:string, revision:string}>>} An array of model identifiers.
|
||||
* @returns {Promise<Array<{name: string, revision: string}>>}
|
||||
* An array of model identifiers.
|
||||
*/
|
||||
async listModels() {
|
||||
// Get all keys (model/revision pairs) from the underlying store.
|
||||
const modelRevisions = await this.#getKeys({
|
||||
storeName: this.taskStoreName,
|
||||
indexName: this.#indices.modelRevisionIndex.name,
|
||||
});
|
||||
|
||||
const models = [];
|
||||
// Process each key entry.
|
||||
for (const { key } of modelRevisions) {
|
||||
models.push({ name: key[0], revision: key[1] });
|
||||
const model = key[0];
|
||||
const revision = key[1];
|
||||
|
||||
models.push({
|
||||
name: model,
|
||||
revision,
|
||||
});
|
||||
}
|
||||
return models;
|
||||
}
|
||||
|
||||
@@ -1093,7 +1093,10 @@ add_task(async function test_listFiles() {
|
||||
headers,
|
||||
});
|
||||
|
||||
const files = await cache.listFiles({ model: "org/model", revision: "v1" });
|
||||
const { files } = await cache.listFiles({
|
||||
model: "org/model",
|
||||
revision: "v1",
|
||||
});
|
||||
const expected = [
|
||||
{
|
||||
path: "file.txt",
|
||||
@@ -1104,6 +1107,7 @@ add_task(async function test_listFiles() {
|
||||
lastUsed: when1,
|
||||
lastUpdated: when1,
|
||||
},
|
||||
engineIds: [],
|
||||
},
|
||||
{
|
||||
path: "file2.txt",
|
||||
@@ -1114,6 +1118,7 @@ add_task(async function test_listFiles() {
|
||||
lastUsed: when2,
|
||||
lastUpdated: when2,
|
||||
},
|
||||
engineIds: [],
|
||||
},
|
||||
{
|
||||
path: "sub/file3.txt",
|
||||
@@ -1125,6 +1130,7 @@ add_task(async function test_listFiles() {
|
||||
lastUsed: when3,
|
||||
lastUpdated: when3,
|
||||
},
|
||||
engineIds: [],
|
||||
},
|
||||
];
|
||||
|
||||
@@ -1172,7 +1178,7 @@ add_task(async function test_listFilesUsingTaskName() {
|
||||
headers,
|
||||
});
|
||||
|
||||
const files = await cache.listFiles({ taskName });
|
||||
const { files } = await cache.listFiles({ taskName });
|
||||
const expected = [
|
||||
{
|
||||
path: "file.txt",
|
||||
@@ -1183,6 +1189,7 @@ add_task(async function test_listFilesUsingTaskName() {
|
||||
lastUsed: when1,
|
||||
lastUpdated: when1,
|
||||
},
|
||||
engineIds: [],
|
||||
},
|
||||
{
|
||||
path: "file2.txt",
|
||||
@@ -1193,6 +1200,7 @@ add_task(async function test_listFilesUsingTaskName() {
|
||||
lastUsed: when2,
|
||||
lastUpdated: when2,
|
||||
},
|
||||
engineIds: [],
|
||||
},
|
||||
{
|
||||
path: "sub/file3.txt",
|
||||
@@ -1204,6 +1212,7 @@ add_task(async function test_listFilesUsingTaskName() {
|
||||
lastUsed: when3,
|
||||
lastUpdated: when3,
|
||||
},
|
||||
engineIds: [],
|
||||
},
|
||||
];
|
||||
|
||||
@@ -1252,7 +1261,7 @@ add_task(async function test_listFilesUsingNonExistingTaskName() {
|
||||
}),
|
||||
]);
|
||||
|
||||
const files = await cache.listFiles({ taskName: "non-existing-task" });
|
||||
const { files } = await cache.listFiles({ taskName: "non-existing-task" });
|
||||
|
||||
Assert.deepEqual(files, []);
|
||||
|
||||
@@ -1322,11 +1331,12 @@ add_task(async function test_initDbFromExistingEmpty() {
|
||||
lastUsed: when,
|
||||
lastUpdated: when,
|
||||
},
|
||||
engineIds: [],
|
||||
},
|
||||
];
|
||||
|
||||
// Ensure every table & indices is on so that we can list files
|
||||
const files = await cache.listFiles({ taskName });
|
||||
const { files } = await cache.listFiles({ taskName });
|
||||
Assert.deepEqual(files, expected);
|
||||
|
||||
await deleteCache(cache);
|
||||
@@ -1369,7 +1379,7 @@ add_task(async function test_initDbFromExistingNoChange() {
|
||||
Assert.equal(cache.db.version, 2);
|
||||
|
||||
// Ensure tables are all empty.
|
||||
const files = await cache.listFiles({ taskName });
|
||||
const { files } = await cache.listFiles({ taskName });
|
||||
|
||||
Assert.deepEqual(files, []);
|
||||
|
||||
@@ -1422,11 +1432,12 @@ add_task(async function test_initDbFromExistingElseWhereStoreChanges() {
|
||||
lastUpdated: when,
|
||||
lastUsed: when,
|
||||
},
|
||||
engineIds: [],
|
||||
},
|
||||
];
|
||||
|
||||
// Ensure every table & indices is on so that we can list files
|
||||
const files = await cache2.listFiles({ taskName });
|
||||
const { files } = await cache2.listFiles({ taskName });
|
||||
Assert.deepEqual(files, expected);
|
||||
|
||||
await deleteCache(cache2);
|
||||
@@ -1750,25 +1761,17 @@ add_task(async function test_migrateStore_modelsDeleted() {
|
||||
cache = await IndexedDBCache.init({ dbName, version: 5 });
|
||||
|
||||
// Verify all unknown model data is deleted
|
||||
let remainingFiles = await cache.listFiles({
|
||||
const { files: random } = await cache.listFiles({
|
||||
model: "random/model",
|
||||
revision: "v1",
|
||||
});
|
||||
Assert.deepEqual(
|
||||
remainingFiles,
|
||||
[],
|
||||
"All unknown model files should be deleted."
|
||||
);
|
||||
Assert.deepEqual(random, [], "All unknown model files should be deleted.");
|
||||
|
||||
remainingFiles = await cache.listFiles({
|
||||
const { files: unknown } = await cache.listFiles({
|
||||
model: "unknown/model",
|
||||
revision: "v2",
|
||||
});
|
||||
Assert.deepEqual(
|
||||
remainingFiles,
|
||||
[],
|
||||
"All unknown model files should be deleted."
|
||||
);
|
||||
Assert.deepEqual(unknown, [], "All unknown model files should be deleted.");
|
||||
|
||||
await deleteCache(cache);
|
||||
});
|
||||
|
||||
@@ -15,4 +15,8 @@ mlmodel-description =
|
||||
|
||||
They process your queries on your device, meaning they are more private than online AI providers, but sometimes slower depending on the performance of your device. <a data-l10n-name="learn-more">Learn more</a>
|
||||
|
||||
addon-detail-file-size-label = File size
|
||||
# Label for button that when clicked removed local model
|
||||
mlmodel-remove-addon-button =
|
||||
.aria-label = Remove
|
||||
# Label for the aggregated value of all files for a model
|
||||
mlmodel-addon-detail-totalsize-label = File size
|
||||
|
||||
@@ -70,6 +70,10 @@
|
||||
src="chrome://mozapps/content/extensions/components/mlmodel-list-intro.mjs"
|
||||
type="module"
|
||||
></script>
|
||||
<script
|
||||
type="module"
|
||||
src="chrome://mozapps/content/extensions/components/mlmodel-card-header-additions.mjs"
|
||||
></script>
|
||||
<script
|
||||
type="module"
|
||||
src="chrome://global/content/elements/moz-message-bar.mjs"
|
||||
@@ -441,6 +445,7 @@
|
||||
data-l10n-id="extension-enable-addon-button-label"
|
||||
hidden
|
||||
></moz-toggle>
|
||||
<mlmodel-card-header-additions></mlmodel-card-header-additions>
|
||||
<button
|
||||
class="more-options-button"
|
||||
action="more-options"
|
||||
|
||||
@@ -2877,6 +2877,15 @@ class AddonCard extends HTMLElement {
|
||||
this.details.update();
|
||||
}
|
||||
|
||||
if (addon.type == "mlmodel") {
|
||||
this.optionsButton.hidden = this.expanded;
|
||||
const mlmodelHeaderAdditions = this.card.querySelector(
|
||||
"mlmodel-card-header-additions"
|
||||
);
|
||||
mlmodelHeaderAdditions.setAddon(addon);
|
||||
mlmodelHeaderAdditions.expanded = this.expanded;
|
||||
}
|
||||
|
||||
this.sendEvent("update");
|
||||
}
|
||||
|
||||
@@ -2969,7 +2978,6 @@ class AddonCard extends HTMLElement {
|
||||
if (addon.type != "extension" && addon.type != "sitepermission") {
|
||||
this.card.querySelector(".extension-enable-button").remove();
|
||||
}
|
||||
|
||||
let nameContainer = this.card.querySelector(".addon-name-container");
|
||||
let headingLevel = this.expanded ? "h1" : "h3";
|
||||
let nameHeading = document.createElement(headingLevel);
|
||||
|
||||
@@ -5,6 +5,11 @@
|
||||
import { html } from "chrome://global/content/vendor/lit.all.mjs";
|
||||
import { MozLitElement } from "chrome://global/content/lit-utils.mjs";
|
||||
|
||||
const lazy = {};
|
||||
ChromeUtils.defineESModuleGetters(lazy, {
|
||||
DownloadUtils: "resource://gre/modules/DownloadUtils.sys.mjs",
|
||||
});
|
||||
|
||||
export class AddonMLModelDetails extends MozLitElement {
|
||||
static properties = {
|
||||
addon: {
|
||||
@@ -28,9 +33,11 @@ export class AddonMLModelDetails extends MozLitElement {
|
||||
|
||||
get template() {
|
||||
return html`
|
||||
<div class="addon-detail-row addon-detail-row-mlmodel-filesize">
|
||||
<label data-l10n-id="addon-detail-file-size-label"></label>
|
||||
${this.addon?.fileSize}
|
||||
<div class="addon-detail-row addon-detail-row-mlmodel-totalsize">
|
||||
<label data-l10n-id="mlmodel-addon-detail-totalsize-label"></label>
|
||||
<span>
|
||||
${lazy.DownloadUtils.getTransferTotal(this.addon?.totalSize ?? 0)}
|
||||
</span>
|
||||
</div>
|
||||
`;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
/* 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 https://mozilla.org/MPL/2.0/. */
|
||||
|
||||
/* mlmodel-card-header-additions LitElement is not rendering into the
|
||||
isolated shadow DOM and so this stylesheet is not restricted to the
|
||||
content of the LitElement, and these rules are nested to make sure
|
||||
they don't apply to elements outside of the LitElement */
|
||||
|
||||
mlmodel-card-header-additions {
|
||||
.mlmodel-total-size-bubble {
|
||||
border: var(--border-width) solid var(--border-color-deemphasized);
|
||||
border-radius: var(--border-radius-circle);
|
||||
padding: var(--space-xsmall);
|
||||
font-size: var(--font-size-small);
|
||||
}
|
||||
|
||||
button.mlmodel-remove-addon-button {
|
||||
min-width: auto;
|
||||
min-height: auto;
|
||||
width: 24px;
|
||||
height: 24px;
|
||||
-moz-context-properties: fill;
|
||||
fill: var(--button-text-color-destructive);
|
||||
background-image: url("chrome://global/skin/icons/delete.svg");
|
||||
background-repeat: no-repeat;
|
||||
background-position: center center;
|
||||
padding: var(--button-padding-icon);
|
||||
display: flex;
|
||||
justify-content: flex-end;
|
||||
background-color: var(--button-background-color-destructive);
|
||||
}
|
||||
|
||||
button.mlmodel-remove-addon-button:active,
|
||||
button.mlmodel-remove-addon-button:hover,
|
||||
button.mlmodel-remove-addon-button:enabled:hover:active {
|
||||
background-color: var(--button-background-color-destructive-hover);
|
||||
color: var(--button-text-color-destructive-hover);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
/* 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 https://mozilla.org/MPL/2.0/. */
|
||||
|
||||
import { html } from "chrome://global/content/vendor/lit.all.mjs";
|
||||
import { MozLitElement } from "chrome://global/content/lit-utils.mjs";
|
||||
|
||||
const lazy = {};
|
||||
ChromeUtils.defineESModuleGetters(lazy, {
|
||||
DownloadUtils: "resource://gre/modules/DownloadUtils.sys.mjs",
|
||||
});
|
||||
|
||||
class MLModelCardHeaderAdditions extends MozLitElement {
|
||||
static properties = {
|
||||
addon: { type: Object, reflect: false },
|
||||
expanded: { type: Boolean, reflect: true, attribute: true },
|
||||
};
|
||||
|
||||
// NOTE: opt-out from using the shadow dom as render root (and use the element
|
||||
// itself to host the custom element content instead).
|
||||
createRenderRoot() {
|
||||
return this;
|
||||
}
|
||||
|
||||
setAddon(addon) {
|
||||
this.addon = addon;
|
||||
}
|
||||
|
||||
render() {
|
||||
if (this.addon?.type !== "mlmodel") {
|
||||
return null;
|
||||
}
|
||||
return html`
|
||||
<link
|
||||
href="chrome://mozapps/content/extensions/components/mlmodel-card-header-additions.css"
|
||||
rel="stylesheet"
|
||||
/>
|
||||
${this.expanded
|
||||
? html`<button
|
||||
class="mlmodel-remove-addon-button"
|
||||
action="remove"
|
||||
data-l10n-id="mlmodel-remove-addon-button"
|
||||
></button>`
|
||||
: html`<span class="mlmodel-total-size-bubble"
|
||||
>${this.getModelSizeString(this.addon)}</span
|
||||
>`}
|
||||
`;
|
||||
}
|
||||
|
||||
getModelSizeString(addon) {
|
||||
return lazy.DownloadUtils.getTransferTotal(addon.totalSize ?? 0);
|
||||
}
|
||||
}
|
||||
|
||||
customElements.define(
|
||||
"mlmodel-card-header-additions",
|
||||
MLModelCardHeaderAdditions
|
||||
);
|
||||
@@ -8,7 +8,9 @@ const lazy = {};
|
||||
ChromeUtils.defineESModuleGetters(lazy, {
|
||||
AddonManager: "resource://gre/modules/AddonManager.sys.mjs",
|
||||
AddonManagerPrivate: "resource://gre/modules/AddonManager.sys.mjs",
|
||||
DownloadUtils: "resource://gre/modules/DownloadUtils.sys.mjs",
|
||||
computeSha256HashAsString:
|
||||
"resource://gre/modules/addons/crypto-utils.sys.mjs",
|
||||
ModelHub: "chrome://global/content/ml/ModelHub.sys.mjs",
|
||||
});
|
||||
|
||||
XPCOMUtils.defineLazyPreferenceGetter(
|
||||
@@ -16,56 +18,80 @@ XPCOMUtils.defineLazyPreferenceGetter(
|
||||
"MODELHUB_PROVIDER_ENABLED",
|
||||
"browser.ml.modelHubProvider",
|
||||
false,
|
||||
(_pref, _old, val) => ModelHubProvider[val ? "init" : "uninit"]()
|
||||
(_pref, _old, val) => ModelHubProvider[val ? "init" : "shutdown"]()
|
||||
);
|
||||
|
||||
const MODELHUB_ADDON_ID_SUFFIX = "-modelhub@mozilla.org";
|
||||
const MODELHUB_ADDON_ID_SUFFIX = "@modelhub.mozilla.org";
|
||||
const MODELHUB_ADDON_TYPE = "mlmodel";
|
||||
|
||||
class ModelHubAddonWrapper {
|
||||
constructor(id = MODELHUB_ADDON_ID_SUFFIX) {
|
||||
this.id = id;
|
||||
// TODO: use actual values (for now random 0 bytes to 100 GB)
|
||||
this.fileSize = lazy.DownloadUtils.getTransferTotal(
|
||||
Math.random() * Math.pow(10, 1 + Math.random() * 10)
|
||||
);
|
||||
export class ModelHubAddonWrapper {
|
||||
#provider;
|
||||
id;
|
||||
name;
|
||||
version;
|
||||
totalSize;
|
||||
lastUsed;
|
||||
updateDate;
|
||||
|
||||
constructor(params) {
|
||||
this.#provider = params.provider;
|
||||
this.id = params.id;
|
||||
this.name = params.name;
|
||||
this.version = params.version;
|
||||
this.totalSize = params.totalSize;
|
||||
this.lastUsed = params.lastUsed;
|
||||
this.updateDate = params.updateDate;
|
||||
}
|
||||
|
||||
async uninstall() {
|
||||
await this.#provider.modelHub.deleteModels({
|
||||
model: this.name,
|
||||
revision: this.version,
|
||||
});
|
||||
|
||||
await this.#provider.onUninstalled(this);
|
||||
}
|
||||
|
||||
get isActive() {
|
||||
return true;
|
||||
}
|
||||
|
||||
get isCompatible() {
|
||||
return true;
|
||||
}
|
||||
get name() {
|
||||
return "ModelHub";
|
||||
}
|
||||
|
||||
get permissions() {
|
||||
return lazy.AddonManager.PERM_CAN_UNINSTALL;
|
||||
}
|
||||
|
||||
get type() {
|
||||
return MODELHUB_ADDON_TYPE;
|
||||
}
|
||||
}
|
||||
|
||||
const ModelHubProvider = {
|
||||
export const ModelHubProvider = {
|
||||
cache: new Map(),
|
||||
modelHub: null,
|
||||
|
||||
get name() {
|
||||
return "ModelHubProvider";
|
||||
},
|
||||
|
||||
init() {
|
||||
// Activate lazy getter and initialize if necessary
|
||||
if (lazy.MODELHUB_PROVIDER_ENABLED && !this.initialized) {
|
||||
this.initialized = true;
|
||||
lazy.AddonManagerPrivate.registerProvider(this, [MODELHUB_ADDON_TYPE]);
|
||||
async getAddonsByTypes(types) {
|
||||
if (!lazy.MODELHUB_PROVIDER_ENABLED) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const match = types?.includes?.(MODELHUB_ADDON_TYPE);
|
||||
return match
|
||||
? await this.refreshAddonCache().then(() =>
|
||||
Array.from(this.cache.values())
|
||||
)
|
||||
: [];
|
||||
},
|
||||
|
||||
uninit() {
|
||||
if (this.initialized) {
|
||||
lazy.AddonManagerPrivate.unregisterProvider(this);
|
||||
this.initialized = false;
|
||||
}
|
||||
async getAddonByID(id) {
|
||||
return this.cache.get(id);
|
||||
},
|
||||
|
||||
observe(_subject, topic, _data) {
|
||||
@@ -77,16 +103,66 @@ const ModelHubProvider = {
|
||||
}
|
||||
},
|
||||
|
||||
async getAddonByID(id) {
|
||||
// TODO: should return a ModelHubAddonWrapper only if we have it
|
||||
return id == MODELHUB_ADDON_ID_SUFFIX ? new ModelHubAddonWrapper() : null;
|
||||
shutdown() {
|
||||
if (this.initialized) {
|
||||
lazy.AddonManagerPrivate.unregisterProvider(this);
|
||||
this.initialized = false;
|
||||
this.clearAddonCache();
|
||||
}
|
||||
},
|
||||
|
||||
async getAddonsByTypes(types) {
|
||||
// TODO: should return ModelHubAddonWrappers only if we have them
|
||||
return !types || types.includes(MODELHUB_ADDON_TYPE)
|
||||
? [new ModelHubAddonWrapper()]
|
||||
: [];
|
||||
async init() {
|
||||
if (!lazy.MODELHUB_PROVIDER_ENABLED || this.initialized) {
|
||||
return;
|
||||
}
|
||||
|
||||
this.initialized = true;
|
||||
lazy.AddonManagerPrivate.registerProvider(this, [MODELHUB_ADDON_TYPE]);
|
||||
this.modelHub = new lazy.ModelHub();
|
||||
await this.refreshAddonCache();
|
||||
},
|
||||
|
||||
async onUninstalled(addon) {
|
||||
if (!this.cache.has(addon.id)) {
|
||||
return;
|
||||
}
|
||||
this.cache.delete(addon.id);
|
||||
lazy.AddonManagerPrivate.callAddonListeners("onUninstalled", addon);
|
||||
},
|
||||
|
||||
async clearAddonCache() {
|
||||
this.cache.clear();
|
||||
},
|
||||
|
||||
getWrapperIdForModel(model) {
|
||||
return [
|
||||
lazy.computeSha256HashAsString(`${model.name}:${model.revision}`),
|
||||
MODELHUB_ADDON_ID_SUFFIX,
|
||||
].join("");
|
||||
},
|
||||
|
||||
async refreshAddonCache() {
|
||||
const models = await this.modelHub.listModels();
|
||||
|
||||
for (const model of models) {
|
||||
const { metadata } = await this.modelHub.listFiles({
|
||||
model: model.name,
|
||||
revision: model.revision,
|
||||
});
|
||||
|
||||
const id = this.getWrapperIdForModel(model);
|
||||
|
||||
const wrapper = new ModelHubAddonWrapper({
|
||||
provider: this,
|
||||
id,
|
||||
name: model.name,
|
||||
version: model.revision,
|
||||
totalSize: metadata.totalSize,
|
||||
lastUsed: new Date(metadata.lastUsed),
|
||||
updateDate: new Date(metadata.updateDate),
|
||||
});
|
||||
this.cache.set(wrapper.id, wrapper);
|
||||
}
|
||||
},
|
||||
};
|
||||
|
||||
|
||||
@@ -113,38 +113,150 @@ add_task(async function testModelHubProvider() {
|
||||
});
|
||||
|
||||
/**
|
||||
* Test model hub expanded details.
|
||||
* Test model hub card in the list view.
|
||||
*/
|
||||
add_task(async function testModelHubDetails() {
|
||||
const id = "mockmodel1@tests.mozilla.org";
|
||||
add_task(async function testModelHubCard() {
|
||||
const id1 = "mockmodel1-without-size@tests.mozilla.org";
|
||||
const id2 = "mockmodel2-with-size@tests.mozilla.org";
|
||||
|
||||
mockProvider.createAddons([
|
||||
{
|
||||
id,
|
||||
id: id1,
|
||||
name: "Model Mock 1",
|
||||
permissions: AddonManager.PERM_CAN_UNINSTALL,
|
||||
type: "mlmodel",
|
||||
totalSize: undefined,
|
||||
},
|
||||
{
|
||||
id: id2,
|
||||
name: "Model Mock 2",
|
||||
permissions: AddonManager.PERM_CAN_UNINSTALL,
|
||||
type: "mlmodel",
|
||||
totalSize: 5 * 1024 * 1024,
|
||||
},
|
||||
]);
|
||||
|
||||
let win = await loadInitialView("mlmodel");
|
||||
|
||||
let card = getAddonCard(win, id);
|
||||
// Card No Size
|
||||
let card1 = getAddonCard(win, id1);
|
||||
ok(card1, `Found addon card for model ${id1}`);
|
||||
verifyAddonCard(card1, "0 bytes");
|
||||
|
||||
// Card With Size
|
||||
let card2 = getAddonCard(win, id2);
|
||||
ok(card2, `Found addon card for model ${id2}`);
|
||||
verifyAddonCard(card2, "5.0 MB");
|
||||
|
||||
await closeView(win);
|
||||
|
||||
function verifyAddonCard(card, expectedTotalSize) {
|
||||
ok(!card.hasAttribute("expanded"), "The list card is not expanded");
|
||||
|
||||
let mlmodelTotalSizeBubble = card.querySelector(
|
||||
".mlmodel-total-size-bubble"
|
||||
);
|
||||
ok(mlmodelTotalSizeBubble, "Expect to see the mlmodel total size bubble");
|
||||
|
||||
is(
|
||||
mlmodelTotalSizeBubble?.textContent.trim(),
|
||||
expectedTotalSize,
|
||||
"Got the expected total size text"
|
||||
);
|
||||
|
||||
let mlmodelRemoveAddonButton = card.querySelector(
|
||||
".mlmodel-remove-addon-button"
|
||||
);
|
||||
|
||||
ok(
|
||||
!mlmodelRemoveAddonButton,
|
||||
"Expect to not see the mlmodel remove addon button"
|
||||
);
|
||||
|
||||
ok(
|
||||
BrowserTestUtils.isVisible(card.optionsButton),
|
||||
"Expect the card options button to be visible in the list view"
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
/**
|
||||
* Test model hub expanded details.
|
||||
*/
|
||||
add_task(async function testModelHubDetails() {
|
||||
const id1 = "mockmodel1-without-size@tests.mozilla.org";
|
||||
const id2 = "mockmodel2-with-size@tests.mozilla.org";
|
||||
|
||||
mockProvider.createAddons([
|
||||
{
|
||||
id: id1,
|
||||
name: "Model Mock 1",
|
||||
permissions: AddonManager.PERM_CAN_UNINSTALL,
|
||||
type: "mlmodel",
|
||||
totalSize: undefined,
|
||||
},
|
||||
{
|
||||
id: id2,
|
||||
name: "Model Mock 2",
|
||||
permissions: AddonManager.PERM_CAN_UNINSTALL,
|
||||
type: "mlmodel",
|
||||
totalSize: 5 * 1024 * 1024,
|
||||
},
|
||||
]);
|
||||
|
||||
await verifyAddonCardDetails(id1, "0 bytes");
|
||||
await verifyAddonCardDetails(id2, "5.0 MB");
|
||||
|
||||
async function verifyAddonCardDetails(id, expectedTotalSize) {
|
||||
let win = await loadInitialView("mlmodel");
|
||||
// Get the list view card DOM element for the given addon id.
|
||||
let card = getAddonCard(win, id);
|
||||
ok(card, `Found addon card for model ${id}`);
|
||||
ok(!card.hasAttribute("expanded"), "list view card is not expanded");
|
||||
|
||||
// Load the detail view.
|
||||
let loaded = waitForViewLoad(win);
|
||||
card.querySelector('[action="expand"]').click();
|
||||
await loaded;
|
||||
|
||||
// Get the detail view card DOM element for the given addon id.
|
||||
card = getAddonCard(win, id);
|
||||
ok(card.hasAttribute("expanded"), "The list card is expanded");
|
||||
ok(card.hasAttribute("expanded"), "detail view card is expanded");
|
||||
|
||||
let mlmodelTotalSizeBubble = card.querySelector(
|
||||
".mlmodel-total-size-bubble"
|
||||
);
|
||||
ok(
|
||||
!mlmodelTotalSizeBubble,
|
||||
"Expect to not see the mlmodel total size bubble"
|
||||
);
|
||||
|
||||
let mlmodelRemoveAddonButton = card.querySelector(
|
||||
".mlmodel-remove-addon-button"
|
||||
);
|
||||
ok(
|
||||
mlmodelRemoveAddonButton,
|
||||
"Expect to see the mlmodel remove addon button"
|
||||
);
|
||||
|
||||
ok(
|
||||
!card.querySelector(".addon-detail-mlmodel").hidden,
|
||||
"Expect to see the mlmodel details"
|
||||
);
|
||||
|
||||
ok(
|
||||
card.querySelector(".addon-detail-row-mlmodel-filesize"),
|
||||
"Expect to see the file size"
|
||||
BrowserTestUtils.isHidden(card.optionsButton),
|
||||
"Expect the card options button to be hidden in the detail view"
|
||||
);
|
||||
|
||||
let totalsizeEl = card.querySelector(".addon-detail-row-mlmodel-totalsize");
|
||||
ok(totalsizeEl, "Expect to see the total size");
|
||||
is(
|
||||
totalsizeEl?.querySelector("span")?.textContent.trim(),
|
||||
expectedTotalSize,
|
||||
"Got the expected total size text"
|
||||
);
|
||||
|
||||
await closeView(win);
|
||||
}
|
||||
});
|
||||
|
||||
@@ -2,7 +2,12 @@
|
||||
* http://creativecommons.org/publicdomain/zero/1.0/ */
|
||||
|
||||
"use strict";
|
||||
|
||||
const { ModelHubProvider } = ChromeUtils.importESModule(
|
||||
"resource://gre/modules/addons/ModelHubProvider.sys.mjs"
|
||||
);
|
||||
ChromeUtils.defineESModuleGetters(this, {
|
||||
sinon: "resource://testing-common/Sinon.sys.mjs",
|
||||
});
|
||||
createAppInfo("xpcshell@tests.mozilla.org", "XPCShell", "1");
|
||||
AddonTestUtils.init(this);
|
||||
|
||||
@@ -43,3 +48,198 @@ add_task(
|
||||
);
|
||||
}
|
||||
);
|
||||
|
||||
add_task(
|
||||
{
|
||||
pref_set: [[MODELHUBPROVIDER_PREF, true]],
|
||||
},
|
||||
async function test_modelhub_provider_addon_wrappers() {
|
||||
let sandbox = sinon.createSandbox();
|
||||
await ModelHubProvider.clearAddonCache();
|
||||
// Sanity checks.
|
||||
ok(
|
||||
AddonManager.hasProvider("ModelHubProvider"),
|
||||
"Expect ModelHubProvider to be registered"
|
||||
);
|
||||
Assert.deepEqual(
|
||||
await AddonManager.getAddonsByTypes(["mlmodel"]),
|
||||
[],
|
||||
"Expect getAddonsByTypes result to be initially empty"
|
||||
);
|
||||
Assert.ok(
|
||||
ModelHubProvider.modelHub,
|
||||
"Expect modelHub instance to be found"
|
||||
);
|
||||
|
||||
const mockModels = [
|
||||
{
|
||||
name: "MockName1",
|
||||
revision: "mockRevision1",
|
||||
engineIds: [],
|
||||
},
|
||||
{
|
||||
name: "MockName2",
|
||||
revision: "mockRevision2",
|
||||
engineIds: [],
|
||||
},
|
||||
];
|
||||
const mockListFilesResult = {
|
||||
metadata: {
|
||||
totalSize: 2048,
|
||||
lastUsed: 0,
|
||||
updateDate: 0,
|
||||
},
|
||||
};
|
||||
|
||||
const listModelsStub = sinon
|
||||
.stub(ModelHubProvider.modelHub, "listModels")
|
||||
.resolves(mockModels);
|
||||
|
||||
const listFilesStub = sinon
|
||||
.stub(ModelHubProvider.modelHub, "listFiles")
|
||||
.resolves(mockListFilesResult);
|
||||
|
||||
const modelWrappers = await AddonManager.getAddonsByTypes(["mlmodel"]);
|
||||
|
||||
// Check that the stubs were called the expected number of times.
|
||||
Assert.equal(
|
||||
listModelsStub.callCount,
|
||||
1,
|
||||
"listModels() getting all models only once"
|
||||
);
|
||||
Assert.equal(
|
||||
listFilesStub.callCount,
|
||||
mockModels.length,
|
||||
"listFiles() getting files and file metadata once for each model"
|
||||
);
|
||||
|
||||
// Verify that the listFiles was called with the expected arguments.
|
||||
for (let i = 0; i < mockModels.length; i++) {
|
||||
// ListFiles only has one argument, which is an config object
|
||||
const callArgs = listFilesStub.getCall(i).args[0];
|
||||
|
||||
Assert.ok(callArgs, `listFiles call ${i} received arguments`);
|
||||
// Compare the model name and revision arguments to the mockModels.
|
||||
Assert.equal(
|
||||
callArgs.model,
|
||||
mockModels[i].name,
|
||||
`Correct model name for call ${i}`
|
||||
);
|
||||
Assert.equal(
|
||||
callArgs.revision,
|
||||
mockModels[i].revision,
|
||||
`Correct revision for call ${i}`
|
||||
);
|
||||
}
|
||||
|
||||
Assert.equal(
|
||||
modelWrappers.length,
|
||||
mockModels.length,
|
||||
"Got the expected number of model AddonWrapper instances"
|
||||
);
|
||||
|
||||
for (const [idx, modelWrapper] of modelWrappers.entries()) {
|
||||
const { name, revision } = mockModels[idx];
|
||||
verifyModelAddonWrapper(modelWrapper, {
|
||||
name,
|
||||
version: revision,
|
||||
totalSize: mockListFilesResult.metadata.totalSize,
|
||||
});
|
||||
}
|
||||
|
||||
// Verify that the ModelHubProvider.getAddonsByTypes
|
||||
// doesn't return any entry if mlmodel isn't explicitly
|
||||
// requested.
|
||||
Assert.deepEqual(
|
||||
(await AddonManager.getAddonsByTypes()).filter(
|
||||
addon => addon.type === "mlmodel"
|
||||
),
|
||||
[],
|
||||
"Expect no mlmodel results with getAddonsByTypes()"
|
||||
);
|
||||
|
||||
Assert.deepEqual(
|
||||
(await AddonManager.getAddonsByTypes([])).filter(
|
||||
addon => addon.type === "mlmodel"
|
||||
),
|
||||
[],
|
||||
"Expect no mlmodel results with getAddonsByTypes([])"
|
||||
);
|
||||
|
||||
Assert.deepEqual(
|
||||
(await AddonManager.getAddonsByTypes(["extension"])).filter(
|
||||
addon => addon.type === "mlmodel"
|
||||
),
|
||||
[],
|
||||
"Expect no mlmodel result with getAddonsByTypes(['extension'])"
|
||||
);
|
||||
|
||||
Assert.equal(
|
||||
await AddonManager.getAddonByID(modelWrappers[0].id),
|
||||
modelWrappers[0],
|
||||
`Got the expected result from getAddonByID for ${modelWrappers[0].id}`
|
||||
);
|
||||
|
||||
// Selecting first model wrapper to test uninstall.
|
||||
const modelWrapper = modelWrappers[0];
|
||||
const uninstallPromise = AddonTestUtils.promiseAddonEvent("onUninstalled");
|
||||
await modelWrapper.uninstall();
|
||||
|
||||
const [uninstalled] = await uninstallPromise;
|
||||
equal(
|
||||
uninstalled,
|
||||
modelWrapper,
|
||||
"onUninstalled was called with that wrapper"
|
||||
);
|
||||
|
||||
// We expect getAddonByID for the removed model to not be found.
|
||||
Assert.equal(
|
||||
await AddonManager.getAddonByID(modelWrappers[0].id),
|
||||
null,
|
||||
`Got no model wrapper from getAddonByID for uninstalled ${modelWrappers[0].id}`
|
||||
);
|
||||
|
||||
// We expect getAddonByID for the non removed model to still be found.
|
||||
Assert.equal(
|
||||
await AddonManager.getAddonByID(modelWrappers[1].id),
|
||||
modelWrappers[1],
|
||||
`Got the expected result from getAddonByID for ${modelWrappers[1].id}`
|
||||
);
|
||||
|
||||
// Reset all sinon stubs.
|
||||
sandbox.restore();
|
||||
|
||||
function verifyModelAddonWrapper(modelWrapper, expected) {
|
||||
const { name, version } = expected;
|
||||
info(`Verify model addon wrapper for ${name}:${version}`);
|
||||
const expectedId = ModelHubProvider.getWrapperIdForModel({
|
||||
name,
|
||||
revision: version,
|
||||
});
|
||||
Assert.equal(modelWrapper.id, expectedId, "Got the expected id");
|
||||
Assert.equal(modelWrapper.type, "mlmodel", "Got the expected type");
|
||||
Assert.equal(
|
||||
modelWrapper.permissions,
|
||||
AddonManager.PERM_CAN_UNINSTALL,
|
||||
"Got the expected permissions"
|
||||
);
|
||||
Assert.equal(modelWrapper.name, name, "Got the expected name");
|
||||
Assert.equal(modelWrapper.version, version, "Got the expected version");
|
||||
Assert.equal(
|
||||
modelWrapper.totalSize,
|
||||
expected.totalSize,
|
||||
"Got the expected file size"
|
||||
);
|
||||
Assert.equal(
|
||||
modelWrapper.isActive,
|
||||
true,
|
||||
"Expect model AddonWrapper to be active"
|
||||
);
|
||||
Assert.equal(
|
||||
modelWrapper.isCompatible,
|
||||
true,
|
||||
"Expect model AddonWrapper to be compatible"
|
||||
);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
Reference in New Issue
Block a user