Skip to content

Commit

Permalink
Add installation validation (#616)
Browse files Browse the repository at this point in the history 
* Allow process kill signals to be checked

Add basic validation for venv, UV, and python

Until UI is in place to perform auto-fix / ignore issues, validation should only log issues.

Add basic validation for venv, UV, and python

Instantiate venv object earlier: allows validation

[Refactor] Simplify code: pass ComfyInstallation

Ensure users can still start app if validation fails

Add install validation

Validate ComfyUI core venv reqs on startup
Add terminal IPC
Fix ipc return value passed

Add VC redist & handlers

Add uv clear and reset

Add current working directory for uv command.

Helps it find the correct .venv to run pip in.

nit - Remove completed TODO [skip ci]

nit

nit

Remove redundant code

Fix basePath check ignored when upgrading config

Load maintenance page on first validation error

Allow uninstall / reinstall if startup fails

nit - ESLint

nit

nit

[Refactor] Remove redundant code

* Inject telemetry in ComfyInstallation constructor

---------

Co-authored-by: huchenlei <[email protected]>
  • Loading branch information
@webfiltered @huchenlei
webfiltered and huchenlei authored Jan 20, 2025
1 parent 7458104 commit 95bc799
Show file tree
Hide file tree
Showing 6 changed files with 376 additions and 108 deletions.
152 changes: 107 additions & 45 deletions src/install/installationManager.ts
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { app, ipcMain } from 'electron';
import { app, dialog, ipcMain } from 'electron';
import log from 'electron-log/main';

import { IPC_CHANNELS } from '../constants';
Expand All @@ -21,32 +21,99 @@ export class InstallationManager {
* Ensures that ComfyUI is installed and ready to run.
*
* First checks for an existing installation and validates it. If missing or invalid, a fresh install is started.
* Will not resolve until the installation is valid.
* @returns A valid {@link ComfyInstallation} object.
*/
async ensureInstalled(): Promise<ComfyInstallation> {
const installation = ComfyInstallation.fromConfig();
log.verbose(`Install state: ${installation?.state ?? 'not installed'}`);
log.info(`Install state: ${installation?.state ?? 'not installed'}`);

// Fresh install
if (!installation) return await this.freshInstall();

// Validate installation
const state = await installation.validate();
log.verbose(`Validated install state: ${state}`);
if (state !== 'installed') await this.resumeInstallation(installation);
try {
// Send updates to renderer
this.#setupIpc(installation);

// Resolve issues and re-run validation
if (installation.issues.size > 0) {
await this.resolveIssues(installation);
await installation.validate();
// Validate installation
const state = await installation.validate();
if (state !== 'installed') await this.resumeInstallation(installation);

// Resolve issues and re-run validation
if (installation.hasIssues) {
while (!(await this.resolveIssues(installation))) {
// Re-run validation
log.verbose('Re-validating installation.');
}
}

// Return validated installation
return installation;
} finally {
delete installation.onUpdate;
this.#removeIpcHandlers();
}
}

/** Removes all handlers created by {@link #setupIpc} */
#removeIpcHandlers() {
ipcMain.removeHandler(IPC_CHANNELS.GET_VALIDATION_STATE);
ipcMain.removeHandler(IPC_CHANNELS.VALIDATE_INSTALLATION);
ipcMain.removeHandler(IPC_CHANNELS.UV_INSTALL_REQUIREMENTS);
ipcMain.removeHandler(IPC_CHANNELS.UV_CLEAR_CACHE);
ipcMain.removeHandler(IPC_CHANNELS.UV_RESET_VENV);
}

// TODO: Confirm this is no longer possible after resolveIssues and remove.
if (!installation.basePath) throw new Error('Base path was invalid after installation validation.');
if (installation.issues.size > 0) throw new Error('Installation issues remain after validation.');
/** Set to `true` the first time an error is found during validation. @todo Move to app state singleton once impl. */
#onMaintenancePage = false;

/** Creates IPC handlers for the installation instance. */
#setupIpc(installation: ComfyInstallation) {
this.#onMaintenancePage = false;
installation.onUpdate = (data) => {
this.appWindow.send(IPC_CHANNELS.VALIDATION_UPDATE, data);

// Load maintenance page the first time any error is found.
if (!this.#onMaintenancePage && Object.values(data).includes('error')) {
this.#onMaintenancePage = true;

log.info('Validation error - loading maintenance page.');
this.appWindow.loadRenderer('maintenance').catch((error) => {
log.error('Error loading maintenance page.', error);
const message = `An error was detected with your installation, and the maintenance page could not be loaded to resolve it. The app will close now. Please reinstall if this issue persists.\n\nError message:\n\n${error}`;
dialog.showErrorBox('Critical Error', message);
app.quit();
});
}
};
const sendLogIpc = (data: string) => this.appWindow.send(IPC_CHANNELS.LOG_MESSAGE, data);

// Return validated installation
return installation;
ipcMain.handle(IPC_CHANNELS.GET_VALIDATION_STATE, () => {
installation.onUpdate?.(installation.validation);
return installation.validation;
});
ipcMain.handle(IPC_CHANNELS.VALIDATE_INSTALLATION, async () => await installation.validate());
ipcMain.handle(IPC_CHANNELS.UV_INSTALL_REQUIREMENTS, () =>
installation.virtualEnvironment.reinstallRequirements(sendLogIpc)
);
ipcMain.handle(IPC_CHANNELS.UV_CLEAR_CACHE, async () => await installation.virtualEnvironment.clearUvCache());
ipcMain.handle(IPC_CHANNELS.UV_RESET_VENV, async (): Promise<boolean> => {
const venv = installation.virtualEnvironment;
const deleted = await venv.removeVenvDirectory();
if (!deleted) return false;

const created = await venv.createVenv(sendLogIpc);
if (!created) return false;

return await venv.upgradePip({ onStdout: sendLogIpc, onStderr: sendLogIpc });
});

// Replace the reinstall IPC handler.
ipcMain.removeHandler(IPC_CHANNELS.REINSTALL);
ipcMain.handle(IPC_CHANNELS.REINSTALL, async () => {
log.info('Reinstalling...');
await InstallationManager.reinstall(installation);
});
}

/**
Expand Down Expand Up @@ -116,7 +183,7 @@ export class InstallationManager {
useDesktopConfig().set('migrateCustomNodesFrom', installWizard.migrationSource);
}

const installation = new ComfyInstallation('installed', installWizard.basePath, device);
const installation = new ComfyInstallation('installed', installWizard.basePath, this.telemetry, device);
installation.setState('installed');
return installation;
}
Expand All @@ -135,40 +202,35 @@ export class InstallationManager {
return filePaths[0];
}

/** Notify user that the provided base apth is not valid. */
async #showInvalidBasePathMessage() {
await this.appWindow.showMessageBox({
title: 'Invalid base path',
message:
'ComfyUI needs a valid directory set as its base path. Inside, models, custom nodes, etc will be stored.\n\nClick OK, then selected a new base path.',
type: 'error',
});
}

/**
* Resolves any issues found during installation validation.
* @param installation The installation to resolve issues for
* @throws If the base path is invalid or cannot be saved
*/
async resolveIssues(installation: ComfyInstallation) {
const issues = [...installation.issues];
for (const issue of issues) {
switch (issue) {
// TODO: Other issues (uv mising, venv etc)
case 'invalidBasePath': {
// TODO: Add IPC listeners and proper UI for this
await this.#showInvalidBasePathMessage();

const path = await this.showBasePathPicker();
if (!path) return;

const success = await installation.updateBasePath(path);
if (!success) throw new Error('No base path selected or failed to save in config.');

installation.issues.delete('invalidBasePath');
break;
}
}
}
log.verbose('Resolving issues - awaiting user response:', installation.validation);

// Await user close window request, validate if any errors remain
const isValid = await new Promise<boolean>((resolve) => {
ipcMain.handleOnce(IPC_CHANNELS.COMPLETE_VALIDATION, async (): Promise<boolean> => {
log.verbose('Attempting to close validation window');
// Check if issues have been resolved externally
if (!installation.isValid) await installation.validate();

// Resolve main thread & renderer
const { isValid } = installation;
resolve(isValid);
return isValid;
});
});

log.verbose('Resolution complete:', installation.validation);
return isValid;
}

static async reinstall(installation: ComfyInstallation): Promise<void> {
await installation.uninstall();
app.relaunch();
app.quit();
}
}
17 changes: 5 additions & 12 deletions src/main-process/comfyDesktopApp.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@ import * as Sentry from '@sentry/electron/main';
import todesktop from '@todesktop/runtime';
import { Notification, type TitleBarOverlayOptions, app, dialog, ipcMain } from 'electron';
import log from 'electron-log/main';
import { rm } from 'node:fs/promises';
import path from 'node:path';
import { graphics } from 'systeminformation';

import { ComfyServerConfig } from '../config/comfyServerConfig';
import { ComfySettings } from '../config/comfySettings';
import { IPC_CHANNELS, ProgressStatus, ServerArgs } from '../constants';
import { InstallationManager } from '../install/installationManager';
import { DownloadManager } from '../models/DownloadManager';
import { type ElectronContextMenuOptions } from '../preload';
import { CmCli } from '../services/cmCli';
Expand Down Expand Up @@ -122,9 +122,12 @@ export class ComfyDesktopApp implements HasTelemetry {
ipcMain.handle(IPC_CHANNELS.IS_FIRST_TIME_SETUP, () => {
return !ComfyServerConfig.exists();
});

// Replace the reinstall IPC handler.
ipcMain.removeHandler(IPC_CHANNELS.REINSTALL);
ipcMain.handle(IPC_CHANNELS.REINSTALL, async () => {
log.info('Reinstalling...');
await this.reinstall();
await InstallationManager.reinstall(this.installation);
});
// Restart core
ipcMain.handle(IPC_CHANNELS.RESTART_CORE, async (): Promise<boolean> => {
Expand Down Expand Up @@ -203,16 +206,6 @@ export class ComfyDesktopApp implements HasTelemetry {
}
}

async uninstall(): Promise<void> {
await rm(ComfyServerConfig.configPath);
await useDesktopConfig().permanentlyDeleteConfigFile();
}

async reinstall(): Promise<void> {
await this.uninstall();
this.restart();
}

restart({ customMessage, delay }: { customMessage?: string; delay?: number } = {}): void {
function relaunchApplication(delay?: number) {
if (delay) {
Expand Down
Loading

0 comments on commit 95bc799

Please sign in to comment.