Skip to content

fix: init models catalog synchronously #3017

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix: init models catalog synchronously #3017

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 11 additions & 11 deletions packages/backend/src/managers/modelsManager.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,7 @@ test('getModelsInfo should get models in local directory', async () => {
modelHandlerRegistry,
);
modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir));
manager.init();
await manager.init();
await manager.loadLocalModels();
expect(manager.getModelsInfo()).toEqual([
{
Expand Down Expand Up @@ -277,7 +277,7 @@ test('getModelsInfo should return an empty array if the models folder does not e
modelHandlerRegistry,
);
modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir));
manager.init();
await manager.init();
await manager.getLocalModelsFromDisk();
expect(manager.getModelsInfo()).toEqual([]);
if (process.platform === 'win32') {
Expand Down Expand Up @@ -318,7 +318,7 @@ test('getLocalModelsFromDisk should return undefined Date and size when stat fai
modelHandlerRegistry,
);
modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir));
manager.init();
await manager.init();
await manager.loadLocalModels();
expect(manager.getModelsInfo()).toEqual([
{
Expand Down Expand Up @@ -377,7 +377,7 @@ test('getLocalModelsFromDisk should skip folders containing tmp files', async ()
modelHandlerRegistry,
);
modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir));
manager.init();
await manager.init();
await manager.loadLocalModels();
expect(manager.getModelsInfo()).toEqual([
{
Expand Down Expand Up @@ -417,7 +417,7 @@ test('loadLocalModels should post a message with the message on disk and on cata
modelHandlerRegistry,
);
modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir));
manager.init();
await manager.init();
await manager.loadLocalModels();
expect(rpcExtensionMock.fire).toHaveBeenNthCalledWith(2, MSG_NEW_MODELS_STATE, [
{
Expand Down Expand Up @@ -464,7 +464,7 @@ test('deleteModel deletes the model folder', async () => {
modelHandlerRegistry,
);
modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir));
manager.init();
await manager.init();
await manager.loadLocalModels();
await manager.deleteModel('model-id-1');
// check that the model's folder is removed from disk
Expand Down Expand Up @@ -525,7 +525,7 @@ describe('deleting models', () => {
modelHandlerRegistry,
);
modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir));
manager.init();
await manager.init();
await manager.loadLocalModels();
await manager.deleteModel('model-id-1');
// check that the model's folder is removed from disk
Expand Down Expand Up @@ -899,7 +899,7 @@ describe('getModelMetadata', () => {
modelHandlerRegistry,
);

manager.init();
await manager.init();

const fakeMetadata: Record<string, string> = {
hello: 'world',
Expand Down Expand Up @@ -939,7 +939,7 @@ describe('getModelMetadata', () => {
modelHandlerRegistry,
);

manager.init();
await manager.init();

const fakeMetadata: Record<string, string> = {
hello: 'world',
Expand Down Expand Up @@ -995,7 +995,7 @@ describe('uploadModelToPodmanMachine', () => {
modelHandlerRegistry,
);

manager.init();
await manager.init();
const result = await manager.uploadModelToPodmanMachine(connectionMock, modelMock);
expect(result).toBe('uploader-result');
expect(performMock).toHaveBeenCalledWith(modelMock.id);
Expand Down Expand Up @@ -1028,7 +1028,7 @@ describe('uploadModelToPodmanMachine', () => {
modelHandlerRegistry,
);

manager.init();
await manager.init();
await manager.uploadModelToPodmanMachine(connectionMock, modelMock);
expect(Uploader).not.toHaveBeenCalled();
});
Expand Down
8 changes: 5 additions & 3 deletions packages/backend/src/managers/modelsManager.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,17 +63,19 @@ export class ModelsManager implements Disposable {
this.modelHandlerRegistry.getAll().forEach(handler => handler.onUpdate(this.loadLocalModels));
}

init(): void {
async init(): Promise<void> {
const disposable = this.catalogManager.onUpdate(() => {
this.loadLocalModels().catch((err: unknown) => {
console.error(`Something went wrong when loading local models`, err);
});
});
this.#disposables.push(disposable);

this.loadLocalModels().catch((err: unknown) => {
try {
await this.loadLocalModels();
} catch (err: unknown) {
console.error('Something went wrong while trying to load local models', err);
});
}
}

dispose(): void {
Expand Down
2 changes: 1 addition & 1 deletion packages/backend/src/studio.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,7 +251,7 @@ export class Studio {
const hfModelHandler = new HuggingFaceModelHandler(this.#modelsManager);
this.#extensionContext.subscriptions.push(hfModelHandler);
this.#extensionContext.subscriptions.push(modelHandlerRegistry.register(hfModelHandler));
this.#modelsManager.init();
await this.#modelsManager.init();
this.#extensionContext.subscriptions.push(this.#modelsManager);

/**
Expand Down
Loading