diff --git a/TESTING.md b/TESTING.md index 7238ebb2..b831dd31 100644 --- a/TESTING.md +++ b/TESTING.md @@ -313,6 +313,26 @@ from the tests' package unload (`cleanup_package`), and it removes its own ComfyUI-Manager restart through `os.execv`, a closed Windows console, or ComfyUI Desktop stopping its server. SQLite's WAL journal covers those. +Jobs running at exit (#376): before any `atexit` callback, Python joins every +executor thread (`concurrent.futures` registers that join with +`threading._register_atexit`). The package registers its own callback there +(`utils/exit_signal.install`, after `concurrent.futures`, so it runs first): +it sets the exit signal and kills the running encoder processes +(`shutdown_encoders`, the ffmpeg that never sees Ctrl+C on Windows). The call in +flight then stops at its next safe point: `Job.cancelled` is True, so every +check a job already makes between files stops it, `map_parallel` starts no +further item (`ExitRequested`), the scan walk stops between folders, LoRA +import staging stops before its next sidecar (also for the legacy scan) and a +tagging model download stops before its next chunk where huggingface_hub takes +`tqdm_class` (1.x removes its temp file; 0.x `hf_hub_download` does not take it +and finishes the file). Xet transfers ignore that checkpoint, so the hook also +calls hub's own `abort_xet_session` when huggingface_hub is loaded. The +trigger-word cache load (small JSON sidecars) has no checkpoint. No check sits +inside a library write. That hook is +CPython-internal; without it (or when it is refused) `install` returns False and +exit waits for the call in flight, as before. `tests/test_exit_cancel.py` runs a +child process with three such calls running and requires its exit within 15 s. + Stop contract: `PromptTracker.shutdown(timeout)`, `ImageGenerationHandler.shutdown(timeout)` (restartable), `ImageGenerationHandler.close(timeout)` (final: `schedule()` refuses and a @@ -4171,6 +4191,19 @@ legacy page workflows), not by this check. | Log files list and download | admin | `#downloadLogsBtn`, `#logFilesList` | `loadLogFiles()`, `downloadLogs()`, `downloadLogFile()` | each listed file downloads by its returned name and a link has no Download (workspace-system-diagnostics.spec.mjs) | | Clear/truncate logs | admin | `#clearLogsBtn` | `clearLogs()` | Clear all logs states both counts, confirms first and reads everything again (workspace-system-diagnostics.spec.mjs) | | Set log level (incl. CRITICAL), console logging | admin | `#updateLogConfigBtn`, `#setLogLevel`, `#consoleLogging` | `loadLogConfig()`, `updateLogConfig()` | runtime level (including CRITICAL) and console logging save and say they apply until the server restarts (workspace-system-diagnostics.spec.mjs) | +| Image scan directories add/remove | admin | `#scanPathsList`, `#addScanPathBtn` | `renderScanPaths()`, `addScanPath()`, `removeScanPath()`, `_collectScanPaths()` | Add source validates on the server, then marks the host draft dirty; Save persists (workspace-system-sources.spec.mjs) | +| Currently monitoring list + Refresh | admin | `#refreshMonitoringStatus`, `#monitoringStatus` | `updateMonitoringStatus()` | watching only with a live observer; restart pending otherwise (workspace-system-sources.spec.mjs) | +| Search results auto-hide `result_timeout` | admin | `#resultTimeout` | — | host fields round-trip through /settings with revision (workspace-system.spec.mjs) | +| Web UI display mode | admin | `#webuiDisplayMode` | — | host fields round-trip through /settings with revision (workspace-system.spec.mjs) | +| Worker threads | admin | `#workerThreads`, `#workerThreadsValue` | `renderWorkerThreads()` | worker range max is cpu_count (workspace-system-sources.spec.mjs) | +| LoraManager badge, enable, path, trigger words, CivitAI key, clear key | admin | `#loraDetectionBadge`, `#loraEnabled`, `#loraSettings`, `#loraManagerPath`, `#loraTriggerWords`, `#civitaiApiKey`, `#clearCivitaiKey` | `detectLoraManager()`, `_bindLoraEvents()`, `saveLoraSettings()` | LoRA rows show detected / enabled / cache-loaded separately (workspace-system-integrations.spec.mjs) | +| LoRA data import with progress | admin | `#loraImportModal`, `#loraImportBtn`, `#loraImportStatus`, `#loraImportPercent`, `#loraImportBar`, `#loraImportProcessed`, `#loraImportImported` | `runLoraImport()` | Review import shows replacement, blocked items and network effects with the backup box; Start opens the job (D107) (workspace-system-integrations.spec.mjs) | +| Maintenance cleanup_duplicates, vacuum, cleanup_orphaned_images, prune_orphaned_prompts | admin | `#maintenanceModal`, `#cleanup_duplicates`, `#vacuum`, `#cleanup_orphaned_images`, `#prune_orphaned_prompts`, `#runMaintenanceBtn`, `#maintenanceResults` | `showMaintenanceModal()`, `closeMaintenance()`, `runMaintenance()`, `renderMaintenanceResults()` | Apply reviewed operation closes, toasts, opens the job in Activity (workspace-system-maintenance.spec.mjs) | +| Maintenance check_hash_duplicates, check_consistency | admin | `#check_hash_duplicates`, `#check_consistency` | — | Run health checks opens Health report and updates metric 4 (workspace-system.spec.mjs) | +| Maintenance statistics (database info) | admin | `#statistics` | — | meta-grid shows real database and source status, the pinned Viewer.js version and the host renderers (workspace-system-diagnostics.spec.mjs) | +| Maintenance select all / clear all / run | admin | `#selectAllMaintenanceBtn`, `#clearAllMaintenanceBtn` | `selectAllMaintenance()`, `clearAllMaintenance()` | plan dialog shows the server's scope, protected, backup policy and media lines with the backup box (D100); Cancel writes nothing (workspace-system-maintenance.spec.mjs) | +| Database backup | admin | — | `backupDatabase()` | Back up downloads through the api client and revokes the object URL (D103) (workspace-system-backups.spec.mjs) | +| Restore from file | admin | `#restoreModal`, `#restoreFileInput`, `#confirmRestore`, `#cancelRestore` | `showRestoreModal()`, `confirmRestore()` | Review restore stages the chosen file and shows the real summary with the backup box (D105) (workspace-system-backups.spec.mjs) | ### Browser cache ownership (#222) diff --git a/__init__.py b/__init__.py index a398acfd..1365a2d9 100644 --- a/__init__.py +++ b/__init__.py @@ -355,6 +355,44 @@ async def _stop_background_work_on_cleanup(app): _stop_background_work() +def _kill_encoders_at_exit(): + """Run by the exit hook before Python joins the executor threads (#376): + a JPEG XL or ffmpeg encoder (no console on Windows, so Ctrl+C never + reaches it) is killed instead of being waited for.""" + from .py.api.thumbnail_encode import shutdown_encoders + + shutdown_encoders() + + +def _abort_downloads_at_exit(): + """Run by the exit hook (#376): huggingface_hub's Xet transfers ignore + the chunk checkpoint, so they are aborted the way hub aborts on Ctrl+C.""" + import sys + + if "huggingface_hub" not in sys.modules: + return # no download ever ran: import nothing at exit + from .py.autotag import abort_hub_downloads + + abort_hub_downloads() + + +try: + # Before the executor join: running jobs stop at their next file, item or + # chunk (#376). The teardown below still runs from atexit afterwards. + from .utils import exit_signal as _exit_signal + + _exit_signal.install(_kill_encoders_at_exit) + _exit_signal.install(_abort_downloads_at_exit) +except Exception as e: + try: + from .utils.logging_config import get_logger + + get_logger("prompt_manager.init").error( + f"Failed to register the exit signal: {e}" + ) + except Exception: + pass + try: atexit.register(_stop_background_work) from .py.config import server_instance as _teardown_server diff --git a/py/api/admin.py b/py/api/admin.py index 047c6b8a..a96d07d2 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -34,6 +34,11 @@ try: except ImportError: from utils.parallel import map_parallel +try: + from ...utils import exit_signal +except ImportError: + from utils import exit_signal + try: from ...database.pair_ingestion import blank_positive, ingest_resolved_pair from ...database.run_attribution import ( @@ -118,6 +123,7 @@ def _collect_media_files(output_dirs, extensions): if not os.path.isdir(root): continue for current, dirnames, filenames in os.walk(root, followlinks=False): + exit_signal.checkpoint() # #376: stop between directories at exit depth = 0 if current == root else len(Path(current).relative_to(root).parts) dirnames[:] = sorted( d diff --git a/py/api/jobs.py b/py/api/jobs.py index 39b74b36..de92fba8 100644 --- a/py/api/jobs.py +++ b/py/api/jobs.py @@ -47,6 +47,11 @@ import time import uuid from pathlib import PurePosixPath, PureWindowsPath +try: + from ...utils import exit_signal +except ImportError: # tests import the package top-level + from utils import exit_signal + logger = logging.getLogger(__name__) # Start and finish order: wall-clock times can tie (coarse Windows clocks). _sequence = itertools.count() @@ -195,11 +200,21 @@ class Job: self.counters = dict.fromkeys(groups, 0) self.outcomes = {group: [] for group in groups} self.log = [] - self.cancelled = False + self._cancel_requested = False self.last_event = None self.task = None self._subscribers = set() + @property + def cancelled(self): + """Cancel was requested, or the interpreter is exiting (#376): every + check a job already makes between files also stops it at exit.""" + return self._cancel_requested or exit_signal.requested() + + @cancelled.setter + def cancelled(self, value): + self._cancel_requested = bool(value) + @property def record_changing(self): return self.admission == "records" diff --git a/py/api/thumbnail_auto.py b/py/api/thumbnail_auto.py index 0c9095cf..3f425dd9 100644 --- a/py/api/thumbnail_auto.py +++ b/py/api/thumbnail_auto.py @@ -92,6 +92,11 @@ try: except ImportError: # tests import the package top-level from utils.parallel import map_parallel, note_refusal, refused_at_shutdown +try: + from ...utils.exit_signal import ExitRequested +except ImportError: # tests import the package top-level + from utils.exit_signal import ExitRequested + # The placeholder frame a video without a preview is answered with (JPEG # bytes, never written to disk); written previews follow the settings. PLACEHOLDER_SIZE = (thumbnail_settings.DEFAULT_EDGE, thumbnail_settings.DEFAULT_EDGE) @@ -460,6 +465,8 @@ class AutoThumbnailer: def _run(self): try: self._drain() + except ExitRequested: + pass # #376: the interpreter is exiting; queued files are dropped except Exception: self._logger.warning("The automatic preview worker stopped", exc_info=True) finally: diff --git a/py/autotag.py b/py/autotag.py index 8938ed5f..d002f626 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -7,6 +7,7 @@ Refactored from standalone_tagger.py for integration with PromptManager API. import errno import gc +import inspect import os import sys from pathlib import Path @@ -32,6 +33,54 @@ try: except ImportError: from py.tagging_defaults import DEFAULT_TAG_TEMPLATE +try: + from ..utils import exit_signal +except ImportError: + from utils import exit_signal + + +def abort_hub_downloads() -> None: + """Abort huggingface_hub's Xet transfers (run by the exit hook, #376). + + Xet downloads report progress from native code, which ignores an + exception raised in ``tqdm.update``; huggingface_hub aborts them itself + on Ctrl+C with ``abort_xet_session``. Nothing happens when + huggingface_hub was never imported or has no Xet support. + """ + if "huggingface_hub" not in sys.modules: + return + try: + from huggingface_hub.utils import _xet + except ImportError: + return + abort = getattr(_xet, "abort_xet_session", None) + if callable(abort): + abort() + + +def _stop_at_exit(download) -> Dict[str, Any]: + """Keyword arguments that stop a huggingface_hub download at its next + chunk once the interpreter is exiting (#376), or none when this + huggingface_hub does not take ``tqdm_class``. + + huggingface_hub 1.x removes its partial temp file on any failure; older + versions keep their resumable ``.incomplete`` file. The model's own path + is only written by the final rename either way. + """ + try: + if "tqdm_class" not in inspect.signature(download).parameters: + return {} + from huggingface_hub.utils import tqdm as hub_tqdm + except (ImportError, TypeError, ValueError): + return {} + + class StopAtExit(hub_tqdm): + def update(self, n=1): + exit_signal.checkpoint() + return super().update(n) + + return {"tqdm_class": StopAtExit} + def _autotag_config(): """The saved defaults; imported late, as config needs ComfyUI's server.""" @@ -397,6 +446,7 @@ class AutoTagService: filename=filename, local_dir=str(model_dir), local_dir_use_symlinks=False, + **_stop_at_exit(hf_hub_download), ) if progress_callback: @@ -541,6 +591,13 @@ class AutoTagService: else: return self._download_hf_model(progress_callback) except Exception as e: + if exit_signal.requested(): + # ExitRequested between chunks, or the Xet session aborted + # by the exit hook (abort_hub_downloads). + self.logger.info( + f"Model download stopped: ComfyUI is exiting ({model_type})" + ) + return False self.logger.error(f"Download failed: {e}") if progress_callback: progress_callback(f"Error: {str(e)}", 0) @@ -570,6 +627,7 @@ class AutoTagService: filename=config["filename"], local_dir=str(gguf_dir), local_dir_use_symlinks=False, + **_stop_at_exit(hf_hub_download), ) self.logger.info("GGUF model downloaded") @@ -587,6 +645,7 @@ class AutoTagService: filename=config["mmproj_filename"], local_dir=str(gguf_dir), local_dir_use_symlinks=False, + **_stop_at_exit(hf_hub_download), ) self.logger.info("mmproj downloaded") @@ -613,6 +672,7 @@ class AutoTagService: repo_id=config["repo"], local_dir=str(model_dir), local_dir_use_symlinks=False, + **_stop_at_exit(snapshot_download), ) self.logger.info("HF model downloaded") diff --git a/py/lora_import.py b/py/lora_import.py index 23f053cf..c906eb8d 100644 --- a/py/lora_import.py +++ b/py/lora_import.py @@ -44,8 +44,10 @@ from .lora_manager_layout import ( ) try: + from ..utils import exit_signal from ..utils.hashing import generate_prompt_hash except ImportError: # pragma: no cover - flat test imports + from utils import exit_signal from utils.hashing import generate_prompt_hash LORA_CATEGORY = "lora-manager" @@ -189,14 +191,15 @@ def stage( downloads its CivitAI example images. A sidecar that cannot be read, or whose record fails to build, is skipped by its name, as are the centralized sidecars discovery could not parse. ``cancelled()`` is asked - before each sidecar (raises ImportCancelled); ``on_record(index, + before each sidecar, and so is the exit signal (#376: the legacy scan's + ``cancelled`` does not read it); either raises ImportCancelled. ``on_record(index, record_or_None, staged)`` follows each one. """ sidecars = discovery.sidecars staged = StagedImport(total=len(sidecars)) staged.skipped.extend(discovery.unreadable) for index, sidecar in enumerate(sidecars): - if cancelled(): + if cancelled() or exit_signal.requested(): raise ImportCancelled() record = None path = sidecar.metadata_path diff --git a/tests/browser/legacy-system-links.spec.mjs b/tests/browser/legacy-system-links.spec.mjs index 50b1c420..c1d98672 100644 --- a/tests/browser/legacy-system-links.spec.mjs +++ b/tests/browser/legacy-system-links.spec.mjs @@ -18,6 +18,10 @@ const LABELS = Object.freeze({ // One row per System link of the design's §2 map: the page, the control's label, how to find it, the section it opens // and what a test must arrange before the control shows. const LINKS = Object.freeze([ + { page: 'admin', label: 'Settings', locate: page => page.locator('#settingsBtn'), section: 'overview' }, + { page: 'admin', label: 'Maintenance', locate: page => page.locator('#maintenanceBtn'), section: 'maintenance' }, + { page: 'admin', label: 'Backup DB', locate: page => page.locator('#backupBtn'), section: 'backups' }, + { page: 'admin', label: 'Restore DB', locate: page => page.locator('#restoreBtn'), section: 'backups' }, { page: 'admin', label: 'Diagnostics', locate: page => page.locator('#diagnosticsBtn'), section: 'diagnostics' }, { page: 'admin', label: 'View Logs', locate: page => page.locator('#logsBtn'), section: 'diagnostics' }, ]); @@ -81,7 +85,8 @@ test('a System link sends no request from the legacy page', async ({ page, conte // The page's own start-up reads (list, film strips, settings) finish first; then the click must add none. await page.waitForLoadState('networkidle'); const sent = []; - const record = item => sent.push(`${item.method()} ${item.url()}`); + // Lazy thumbnails may still load whenever the layout changes; they are not the click's. + const record = item => { if (item.resourceType() !== 'image') sent.push(`${item.method()} ${item.url()}`); }; page.on('request', record); const opened = context.waitForEvent('page'); await control.click(); @@ -161,3 +166,31 @@ test('kept admin flows still work', async ({ page, context }) => { await tab.close(); } }); + +test('Page options keeps the page-only settings admin', async ({ page, request }) => { + const before = (await (await request.get('/prompt_manager/settings')).json()).settings; + await page.goto(PAGES.admin); + await expect(page.locator('#resultsList .prompt-text').first()).toBeVisible(); + await page.locator('#pageOptionsBtn').click(); + const dialog = page.locator('#settingsModal'); + await expect(dialog).toBeVisible(); + await expect(dialog.locator('h3')).toHaveText('Page options'); + // Exactly the page-only setting (G5): the list's Infinite Scrolling default. + expect(await dialog.locator('input, select, textarea').evaluateAll(fields => fields.map(field => field.id))) + .toEqual(['infiniteScrollDefault']); + await expect(page.locator('#infiniteScrollDefault')).not.toBeChecked(); + await page.locator('#infiniteScrollDefault').check(); + const saved = page.waitForRequest(item => item.method() === 'POST' && item.url().endsWith('/prompt_manager/settings')); + await page.locator('#saveSettings').click(); + expect((await saved).postDataJSON()).toEqual({ infinite_scroll: true }); + await expect(dialog).toBeHidden(); + // The list header follows the saved default at once; nothing else changed. + await expect(page.locator('#infiniteScrollSession')).toBeChecked(); + const after = (await (await request.get('/prompt_manager/settings')).json()).settings; + expect(after).toEqual({ ...before, infinite_scroll: true }); + await page.reload(); + // The dialog shows the saved default once the page has read it (the header follows it first). + await expect(page.locator('#infiniteScrollSession')).toBeChecked(); + await page.locator('#pageOptionsBtn').click(); + await expect(page.locator('#infiniteScrollDefault')).toBeChecked(); +}); diff --git a/tests/browser/workflows.spec.mjs b/tests/browser/workflows.spec.mjs index e822d998..e68e0695 100644 --- a/tests/browser/workflows.spec.mjs +++ b/tests/browser/workflows.spec.mjs @@ -246,37 +246,6 @@ async function promptSnapshot(request) { })).sort((a, b) => a.id - b.id); } -async function downloadBackup(page) { - const downloading = page.waitForEvent('download'); - await page.locator('#backupBtn').click(); - const download = await downloading; - expect(download.suggestedFilename()).toMatch(/^prompts_backup_.*\.db$/); - expect(await download.failure()).toBeNull(); - const chunks = []; - for await (const chunk of await download.createReadStream()) chunks.push(chunk); - const buffer = Buffer.concat(chunks); - expect(buffer.subarray(0, 16).toString()).toBe('SQLite format 3\0'); - return { name: download.suggestedFilename(), mimeType: 'application/octet-stream', buffer }; -} - -async function chooseRestore(page, file) { - await page.locator('#restoreBtn').click(); - await expect(page.locator('#restoreModal')).toBeVisible(); - await page.locator('#restoreFileInput').setInputFiles(file); -} - -// Confirms the chosen restore and waits for the server's answer, not for a timer: a restore is disk-bound and can take -// longer than an expect timeout on a slow runner (#359). The modal closes only after that answer. -async function confirmRestoreAndWait(page) { - const restored = page.waitForResponse(response => response.url().endsWith('/prompt_manager/restore')); - await page.locator('#confirmRestore').click(); - const response = await restored; - const body = await response.text(); - expect(response.status(), body).toBe(200); - expect(JSON.parse(body)).toMatchObject({ success: true, prompt_count: 2, backup_created: expect.any(String) }); - await expect(page.locator('#restoreModal')).toBeHidden(); -} - function confirmNextDialog(page, accept, message) { page.once('dialog', async dialog => { expect(dialog.type()).toBe('confirm'); @@ -292,88 +261,6 @@ async function deletePromptButton(page, text) { return page.locator(`[data-action="delete"][data-prompt-id="${id}"]`); } -test('downloaded backup restores deleted prompts, tags, and image links through the actual upload form', async ({ page, request }) => { - await page.goto('/prompt_manager/admin'); - await expect(page.locator('.prompt-text')).toHaveCount(2); - const before = await promptSnapshot(request); - expect(before.map(prompt => prompt.tags)).toEqual([['cats'], ['dogs']]); - expect(before.every(prompt => prompt.images.length === 1)).toBe(true); - const backup = await downloadBackup(page); - const button = await deletePromptButton(page, 'Cat portrait'); - confirmNextDialog(page, true, 'delete this prompt'); - await button.click(); - await expect(page.locator('.prompt-text')).toHaveCount(1); - await expect(await request.post('/prompt_manager/save', { data: { text: 'Added after backup', tags: ['temporary'] } })).toBeOK(); - expect(await promptSnapshot(request)).not.toEqual(before); - await chooseRestore(page, backup); - confirmNextDialog(page, true, backup.name); - await confirmRestoreAndWait(page); - await expect(page.locator('.prompt-text')).toHaveCount(2); - await page.reload(); - await expect(page.locator('.prompt-text').filter({ hasText: 'Cat portrait' })).toBeVisible(); - await expect(page.locator('.prompt-text').filter({ hasText: 'Added after backup' })).toHaveCount(0); - expect(await promptSnapshot(request)).toEqual(before); - await page.goto('/prompt_manager/admin#/tags'); - await page.locator('[data-tag-name="cats"]').click(); - await expect(page.locator('#tagPromptsGrid')).toContainText('Cat portrait'); - await page.locator('#tagPromptsGrid img').first().click(); - await expect(page.locator('.viewer-canvas img')).toBeVisible(); - await expect.poll(() => page.locator('.viewer-canvas img').evaluate(image => image.complete && image.naturalWidth > 0)).toBe(true); -}); - -test('restore validation and cancellation preserve live data and allow a valid retry', async ({ page, request }) => { - await page.goto('/prompt_manager/admin'); - const before = await promptSnapshot(request); - const backup = await downloadBackup(page); - let uploads = 0; - page.on('request', request => { if (request.url().endsWith('/prompt_manager/restore')) uploads++; }); - await page.locator('#restoreBtn').click(); - await page.locator('#confirmRestore').click(); - await expect(page.getByText('Please select a database file to restore', { exact: true })).toBeVisible(); - await page.locator('#restoreFileInput').setInputFiles({ ...backup, name: 'backup.txt' }); - await page.locator('#confirmRestore').click(); - await expect(page.getByText('Please select a valid .db file', { exact: true })).toBeVisible(); - expect(uploads).toBe(0); - await page.locator('#restoreFileInput').setInputFiles(backup); - confirmNextDialog(page, false, 'replace your current database'); - await page.locator('#confirmRestore').click(); - await expect(page.locator('#restoreModal')).toBeVisible(); - expect(uploads).toBe(0); - await page.locator('#restoreFileInput').setInputFiles({ name: 'corrupt.db', mimeType: 'application/octet-stream', buffer: Buffer.from('not a SQLite database') }); - confirmNextDialog(page, true, 'corrupt.db'); - const rejected = page.waitForResponse(response => response.url().endsWith('/prompt_manager/restore')); - await page.locator('#confirmRestore').click(); - expect((await rejected).status()).toBe(400); - await expect(page.getByText(/Restore failed:/)).toBeVisible(); - await expect(page.locator('#restoreModal')).toBeVisible(); - await expect(page.locator('#confirmRestore')).toBeEnabled(); - expect(await promptSnapshot(request)).toEqual(before); - await page.locator('#restoreFileInput').setInputFiles(backup); - confirmNextDialog(page, true, backup.name); - await confirmRestoreAndWait(page); - expect(await promptSnapshot(request)).toEqual(before); -}); - -test('backup and restore recover from interrupted network requests without losing data', async ({ page, request }) => { - await page.goto('/prompt_manager/admin'); - const before = await promptSnapshot(request); - await page.route('**/prompt_manager/backup', route => route.abort('failed'), { times: 1 }); - await page.locator('#backupBtn').click(); - await expect(page.getByText(/Backup failed:/)).toBeVisible(); - const backup = await downloadBackup(page); - await chooseRestore(page, backup); - await page.route('**/prompt_manager/restore', route => route.abort('failed'), { times: 1 }); - confirmNextDialog(page, true, backup.name); - await page.locator('#confirmRestore').click(); - await expect(page.getByText(/Restore failed:/)).toBeVisible(); - await expect(page.locator('#confirmRestore')).toBeEnabled(); - await expect(page.locator('#restoreModal')).toBeVisible(); - expect(await promptSnapshot(request)).toEqual(before); - confirmNextDialog(page, true, backup.name); - await confirmRestoreAndWait(page); - expect(await promptSnapshot(request)).toEqual(before); -}); - test('prompt deletion requires confirmation and can retry after a network failure', async ({ page, request }) => { await page.goto('/prompt_manager/admin'); const before = await promptSnapshot(request); diff --git a/tests/browser/workspace-system-backups.spec.mjs b/tests/browser/workspace-system-backups.spec.mjs index 6251d0f1..edc369d1 100644 --- a/tests/browser/workspace-system-backups.spec.mjs +++ b/tests/browser/workspace-system-backups.spec.mjs @@ -395,14 +395,15 @@ test('a failed settings read leaves the policy disabled with the reason', async await expect(row.locator('input')).toBeDisabled(); }); -test('the legacy admin Backup and Restore keep working', async ({ page, request }) => { +test('the legacy Backup and Restore routes keep working and the admin page links to System', async ({ page, request }) => { const bytes = await backupBytes(request); const restored = await request.post('/prompt_manager/restore', { multipart: { database_file: { name: 'prompts_backup.db', mimeType: 'application/octet-stream', buffer: bytes } } }); await expect(restored).toBeOK(); expect((await restored.json()).success).toBe(true); + // #330: the admin page's Backup DB and Restore DB open System › Backup & export. await page.goto('/prompt_manager/admin'); - await expect(page.locator('#restoreFileInput')).toHaveCount(1); + for (const id of ['#backupBtn', '#restoreBtn']) await expect(page.locator(id)).toHaveAttribute('href', 'workspace#/system/backups'); }); /** Start a fake running job (a metadata scan another tab started), which an exclusive restore must wait for. */ diff --git a/tests/js/legacy-refusals.test.js b/tests/js/legacy-refusals.test.js index c06daedc..f43f912b 100644 --- a/tests/js/legacy-refusals.test.js +++ b/tests/js/legacy-refusals.test.js @@ -146,28 +146,12 @@ test('gallery cache clear says so when browser storage is denied', async () => { 'Thumbnails cleared; original images untouched. This browser refused access to its storage, so nothing there was cleared.']]); }); -test('admin LoRA import shows a refusal frame and closes its modal', async () => { - const { page, dom, timers, notifications } = load('admin.js', 'PromptAdmin', - async () => sseResponse([{ type: 'error', status: BUSY, progress: 0, message: BUSY }])); - page.saveLoraSettings = async () => true; - page.loadStatistics = () => assert.fail('a refused import must not reload statistics'); - - await page.runLoraImport(); - - assert.equal(dom.nodes.loraImportStatus.textContent, `Error: ${BUSY}`); - assert.deepEqual(notifications.filter(([type]) => type === 'error'), [['error', `LoRA import failed: ${BUSY}`]]); - assert.equal(timers.length, 1); - timers[0](); - assert.deepEqual(notifications.at(-1), ['hide', 'loraImportModal']); -}); - test('admin settings save shows why config.json refused it', async () => { const message = 'config.json cannot be read, so nothing was saved. Fix or remove the file, then save again.'; const { page, notifications } = load('admin.js', 'PromptAdmin', async () => jsonResponse(409, { success: false, persisted: false, code: 'config-unreadable', error: message })); page.settings = {}; - page._collectScanPaths = () => []; - page.saveLoraSettings = async () => assert.fail('nothing else is saved after a refusal'); + page.setInfiniteScroll = () => assert.fail('nothing changes after a refusal'); await page.saveSettings(); @@ -178,27 +162,30 @@ test('admin settings save keeps its plain message when the reply has no reason', const { page, notifications } = load('admin.js', 'PromptAdmin', async () => ({ ok: false, status: 500, json: async () => { throw new SyntaxError('not json'); } })); page.settings = {}; - page._collectScanPaths = () => []; await page.saveSettings(); assert.deepEqual(notifications, [['error', 'Failed to save settings']]); }); -test('admin maintenance names the job that refused it', async () => { - const { page, dom, notifications } = load('admin.js', 'PromptAdmin', - async () => jsonResponse(409, { success: false, error: BUSY }), { - escapeHtml: text => String(text), - document: { readyState: 'loading', addEventListener() {}, getElementById: id => dom.getElementById(id), - querySelectorAll: selector => (selector === '.maintenance-option:checked' ? [{ id: 'vacuum' }] : []) }, - }); +// #330 S13-9: the admin page's reduced Page options save sends only the list's scroll default, so a stale legacy form +// can no longer overwrite scan folders or other settings edited in System. +test('admin Page options save posts only infinite_scroll', async () => { + const sent = []; + const { page, dom, notifications } = load('admin.js', 'PromptAdmin', async (url, init) => { + sent.push([url, init.method, JSON.parse(init.body)]); + return jsonResponse(200, { success: true, restart_required: true }); + }); + page.settings = { infiniteScroll: false }; + const scrolled = []; + page.setInfiniteScroll = (on, options) => scrolled.push([on, { ...options }]); + dom.getElementById('infiniteScrollDefault').checked = true; - await page.runMaintenance(); + await page.saveSettings(); - const results = dom.nodes.maintenanceResults.innerHTML; - assert.match(results, /Maintenance Failed/); - assert.match(results, new RegExp(BUSY)); - assert.doesNotMatch(results, /HTTP 409/); - assert.deepEqual(notifications, [['error', '❌ Maintenance failed']]); - assert.equal(dom.nodes.runMaintenanceBtn.disabled, false); + assert.deepEqual(sent, [['/prompt_manager/settings', 'POST', { infinite_scroll: true }]]); + assert.deepEqual(scrolled, [[true, { session: false }]]); + assert.equal(page.settings.infiniteScroll, true); + assert.deepEqual(notifications, [['success', 'Page options saved'], ['hide', 'settingsModal']]); }); + diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py index e9b451cf..32da612d 100644 --- a/tests/test_api_backup_restore.py +++ b/tests/test_api_backup_restore.py @@ -264,6 +264,41 @@ class TestRestoreRoute(BackupRestoreTestCase): texts = sorted(p["text"] for p in self.api.db.search_prompts()) self.assertEqual(texts, ["restored a", "restored b"]) + async def test_restore_round_trip_keeps_tags_and_image_links(self): + """A downloaded backup brings back every prompt with its tags and images. + + The legacy admin page's Restore form left for System (#330); this keeps + its browser test's data check at the route. + """ + db = self.api.db + originals = {} + for text, tags, image in ( + ("round trip cat", ["cats", "pets"], "cat.png"), + ("round trip dog", ["dogs"], "dog.png"), + ): + prompt_id = db.save_prompt( + text=text, tags=tags, prompt_hash=generate_prompt_hash(text) + ) + path = self._write(image, b"\x89PNG\r\n\x1a\n") + self.assertTrue(db.link_image_to_prompt(str(prompt_id), path)) + originals[text] = (prompt_id, sorted(tags), [path]) + + resp = await self.client.get("/prompt_manager/backup") + self.assertEqual(resp.status, 200) + backup = self._write("downloaded.db", await resp.read()) + self.assertTrue(db.delete_prompt(originals["round trip dog"][0])) + + resp = await self._restore(backup) + self.assertEqual(resp.status, 200, await resp.text()) + + restored = {p["text"]: p for p in self.api.db.search_prompts()} + self.assertEqual(sorted(restored), sorted(originals)) + for text, (prompt_id, tags, paths) in originals.items(): + self.assertEqual(restored[text]["id"], prompt_id) + self.assertEqual(sorted(restored[text]["tags"]), tags) + images = self.api.db.get_prompt_images(str(prompt_id)) + self.assertEqual([image["image_path"] for image in images], paths) + async def test_restore_leaves_no_temp_files_behind(self): upload = self._make_db_file(["restored"]) created = [] diff --git a/tests/test_exit_cancel.py b/tests/test_exit_cancel.py new file mode 100644 index 00000000..e3fb5ba5 --- /dev/null +++ b/tests/test_exit_cancel.py @@ -0,0 +1,463 @@ +"""Work running at Ctrl+C stops at its next safe point (#376). + +ComfyUI catches Ctrl+C in ``main.py`` and returns without closing its event +loop. Python then joins every executor thread (``concurrent.futures``'s +``_python_exit``, a ``threading._register_atexit`` callback) before any +``atexit`` callback runs, so a call still running on the default executor +used to keep the process alive until it returned. The package registers its +own ``threading._register_atexit`` hook, which runs before that join: it sets +the exit signal that ``Job.cancelled`` and ``map_parallel`` read, and kills +the encoder processes. + +The child-process test starts three long calls on the default executor (a +registry job checking ``job.cancelled`` between files, a ``map_parallel`` +batch and an encoder process that ignores the exit, as ffmpeg does on +Windows), returns from its main module and must exit far sooner than any of +them would end. +""" + +import json +import os +import subprocess +import sys +import tempfile +import textwrap +import threading +import unittest +from unittest import mock + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, ROOT) + +try: + from utils import exit_signal # noqa: E402 +except ImportError: # before #376: only the child-process test can run + exit_signal = None + +# The parent never waits longer than this for the child (the hang watchdog +# allows 120 s per test). Each fake job alone would run 300 s or more. +CHILD_CEILING_SECONDS = 60 +# From the child's return to its last atexit callback: the stopped calls end +# at their next file, item or chunk (the encoder is killed). +EXIT_BOUND_SECONDS = 15 +JOB_SECONDS = 600 +# The fake encoder outlives the ceiling, yet a failing run leaks it only +# briefly (the parent cannot reach the child's own child to kill it). +ENCODER_SECONDS = 90 +REPORT_PREFIX = "PM_EXIT_REPORT " + +CHILD = textwrap.dedent(""" + import asyncio + import atexit + import contextlib + import importlib + import io + import json + import os + import sys + import threading + import time + + root, tmp = sys.argv[1], sys.argv[2] + job_seconds, encoder_seconds = float(sys.argv[3]), int(sys.argv[4]) + sys.path.insert(0, root) + state = {"stopped": {}} + + def report(): + returned = state.get("returned") + print("PM_EXIT_REPORT " + json.dumps({ + "exit_seconds": None if returned is None else time.monotonic() - returned, + "stopped": state["stopped"], + })) + sys.stdout.flush() + + # Registered first, so it runs after every other atexit callback, and + # every atexit callback runs after the executor threads were joined. + atexit.register(report) + + from tests.test_init_package import _import_package, fake_server_module + + output = os.path.join(tmp, "output") + with contextlib.redirect_stdout(io.StringIO()): + package, _ = _import_package( + "pm_exit_child", + os.path.join(tmp, "prompts.db"), + output, + fake_server_module(queue_task_done=True), + ) + parallel = importlib.import_module("pm_exit_child.utils.parallel") + encode = importlib.import_module("pm_exit_child.py.api.thumbnail_encode") + registry = package.api._jobs + entered = {name: threading.Event() for name in ("job", "batch", "encoder")} + + def job_work(job): + # A job's thread work: one "file" at a time, checking its cancel + # between files the way the thumbnail and import jobs do. + entered["job"].set() + deadline = time.monotonic() + job_seconds + while time.monotonic() < deadline: + if job.cancelled: + state["stopped"]["job"] = True + return + time.sleep(0.05) + + def batch_item(index): + entered["batch"].set() + time.sleep(0.5) + return index + + def batch_work(): + try: + parallel.map_parallel(batch_item, range(int(job_seconds)), 2) + except Exception as error: + state["stopped"]["batch"] = type(error).__name__ + + def encoder_work(): + # Stands for ffmpeg on Windows, which runs without a console and + # never sees Ctrl+C: only a kill ends it before its timeout. + cmd = [sys.executable, "-c", "import time; time.sleep(%d)" % encoder_seconds] + try: + encode.run_encoder( + cmd, + timeout=encoder_seconds, + track=lambda _process: entered["encoder"].set(), + ) + except Exception as error: + state["stopped"]["encoder"] = type(error).__name__ + + async def begin(): + loop = asyncio.get_running_loop() + + async def runner(job): + await loop.run_in_executor(None, job_work, job) + + registry.start("exit-test", "Exit test", admission="read", runner=runner) + loop.run_in_executor(None, batch_work) + loop.run_in_executor(None, encoder_work) + for name, event in entered.items(): + if not await loop.run_in_executor(None, event.wait, 30): + sys.exit("not started: " + name) + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(begin()) + state["returned"] = time.monotonic() + # Return with the three calls running and the loop open, as ComfyUI does + # after Ctrl+C. + """) + + +class TestExitStopsRunningWork(unittest.TestCase): + def test_running_executor_work_stops_at_exit(self): + with tempfile.TemporaryDirectory() as tmp: + os.makedirs(os.path.join(tmp, "output")) + script = os.path.join(tmp, "child.py") + with open(script, "w", encoding="utf-8") as handle: + handle.write(CHILD) + env = dict(os.environ) + env.pop("PROMPT_MANAGER_CONFIG_PATH", None) + env["PROMPT_MANAGER_DATA_DIR"] = tmp + env["PYTHONIOENCODING"] = "utf-8" + try: + completed = subprocess.run( + [ + sys.executable, + script, + ROOT, + tmp, + str(JOB_SECONDS), + str(ENCODER_SECONDS), + ], + cwd=ROOT, + env=env, + capture_output=True, + timeout=CHILD_CEILING_SECONDS, + encoding="utf-8", + errors="replace", + ) + except subprocess.TimeoutExpired as hung: + self.fail( + f"the child did not exit within {CHILD_CEILING_SECONDS} s\n" + f"stdout:\n{hung.stdout}\nstderr:\n{hung.stderr}" + ) + detail = f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + self.assertEqual(completed.returncode, 0, detail) + reports = [ + line[len(REPORT_PREFIX) :] + for line in completed.stdout.splitlines() + if line.startswith(REPORT_PREFIX) + ] + self.assertEqual(len(reports), 1, detail) + report = json.loads(reports[0]) + self.assertLess(report["exit_seconds"], EXIT_BOUND_SECONDS, detail) + self.assertEqual( + report["stopped"], + {"job": True, "batch": "ExitRequested", "encoder": "EncoderClosed"}, + detail, + ) + for text in ("Traceback", "Background work not stopped in time"): + self.assertNotIn(text, completed.stderr, detail) + + +class _Isolated(unittest.TestCase): + """Each test starts with the signal clear and no hook installed.""" + + def setUp(self): + exit_signal.reset() + self.addCleanup(exit_signal.reset) + + +class TestInstall(_Isolated): + def test_the_hook_is_registered_once_before_the_executor_join(self): + calls = [] + with mock.patch.object( + exit_signal.threading, + "_register_atexit", + side_effect=lambda func: calls.append(func), + create=True, + ): + self.assertTrue(exit_signal.install()) + self.assertTrue(exit_signal.install()) + self.assertEqual(len(calls), 1) + self.assertFalse(exit_signal.requested()) + calls[0]() + self.assertTrue(exit_signal.requested()) + + def test_a_missing_internal_hook_degrades_to_the_old_exit(self): + with mock.patch.object(exit_signal, "threading") as fake: + del fake._register_atexit + self.assertFalse(exit_signal.install()) + self.assertFalse(exit_signal.requested()) + + def test_a_refused_registration_degrades_to_the_old_exit(self): + refuse = RuntimeError("can't register atexit after shutdown") + with mock.patch.object( + exit_signal.threading, "_register_atexit", side_effect=refuse, create=True + ): + self.assertFalse(exit_signal.install()) + + def test_a_failing_exit_action_never_escapes_the_hook(self): + # An exception out of a threading atexit callback would abort + # threading._shutdown and skip the executor join itself. + calls = [] + + def boom(): + raise KeyboardInterrupt # a second Ctrl+C + + with mock.patch.object( + exit_signal.threading, + "_register_atexit", + side_effect=lambda func: calls.append(func), + create=True, + ): + exit_signal.install(boom) + with self.assertLogs(exit_signal.logger, "ERROR"): + calls[0]() + self.assertTrue(exit_signal.requested()) + + +class TestCheckpoints(_Isolated): + def test_checked_runs_until_the_signal_then_raises(self): + seen = [] + record = exit_signal.checked(seen.append) + record(1) + exit_signal.request() + with self.assertRaises(exit_signal.ExitRequested): + record(2) + self.assertEqual(seen, [1]) + + def test_map_parallel_starts_no_item_after_the_signal(self): + from utils.parallel import map_parallel + + started = [] + gate = threading.Event() + + def item(index): + started.append(index) + if index == 0: + exit_signal.request() + gate.set() + return index + + with self.assertRaises(exit_signal.ExitRequested): + map_parallel(item, range(50), 1) + self.assertTrue(gate.is_set()) + self.assertEqual(started, [0]) + + def test_job_cancelled_follows_the_signal_and_its_own_cancel(self): + from py.api.jobs import Job + + job = Job("scan", "Scan", admission="records") + self.assertFalse(job.cancelled) + job.cancelled = True + self.assertTrue(job.cancelled) + other = Job("scan", "Scan", admission="records", cancellable=False) + exit_signal.request() + self.assertTrue(other.cancelled) + + def test_the_scan_walk_stops_between_directories(self): + from py.api.admin import _collect_media_files + + with tempfile.TemporaryDirectory() as tmp: + for name in ("a", "b"): + os.makedirs(os.path.join(tmp, name)) + with open(os.path.join(tmp, name, "x.png"), "wb") as image: + image.write(b"x") + self.assertEqual(len(_collect_media_files([tmp], [".png"])), 2) + exit_signal.request() + with self.assertRaises(exit_signal.ExitRequested): + _collect_media_files([tmp], [".png"]) + + +class TestLoraImportStopsAtExit(_Isolated): + def test_staging_stops_before_the_next_sidecar_whatever_the_caller_checks(self): + import types + + from py import lora_import + + discovery = types.SimpleNamespace( + sidecars=[object(), object()], unreadable=[], layout=None + ) + exit_signal.request() + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(lora_import.ImportCancelled): + lora_import.stage( + discovery, + tmp, + media_extensions=frozenset(), + cancelled=lambda: False, # the legacy scan's own check + ) + + +class _FakeTqdm: + """huggingface_hub.utils.tqdm as far as a download uses it.""" + + def __init__(self, *args, **kwargs): + self.n = 0 + + def update(self, n=1): + self.n += n + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def _fake_hub(chunks_before_exit): + """A ``huggingface_hub`` whose download writes chunks to a temp file and, + like huggingface_hub 1.x, removes it on any failure before the rename.""" + import types + + def hf_hub_download( + repo_id, filename, local_dir=None, local_dir_use_symlinks=None, tqdm_class=None + ): + final = os.path.join(local_dir, filename) + partial = final + ".abc.incomplete" + bar = (tqdm_class or _FakeTqdm)(total=10) + try: + with open(partial, "wb") as handle: + for index in range(10): + if index == chunks_before_exit: + exit_signal.request() + bar.update(1) + handle.write(b"x") + os.replace(partial, final) + finally: + if os.path.exists(partial): + os.unlink(partial) + return final + + hub = types.ModuleType("huggingface_hub") + hub.hf_hub_download = hf_hub_download + hub.snapshot_download = lambda **kwargs: None + utils = types.ModuleType("huggingface_hub.utils") + utils.tqdm = _FakeTqdm + hub.utils = utils + return {"huggingface_hub": hub, "huggingface_hub.utils": utils} + + +class TestModelDownloadStopsAtExit(_Isolated): + def test_a_download_stops_between_chunks_and_leaves_no_file(self): + from py.autotag import AutoTagService + + with tempfile.TemporaryDirectory() as tmp: + service = AutoTagService(models_dir=tmp) + with mock.patch.dict(sys.modules, _fake_hub(chunks_before_exit=3)): + with self.assertLogs(service.logger, "INFO") as logs: + self.assertFalse(service.download_model("wd14-vit")) + folder = os.path.join(tmp, "WD14-ViT") + self.assertEqual(os.listdir(folder), []) + self.assertTrue(any("exiting" in line for line in logs.output), logs.output) + self.assertFalse(any("Download failed" in line for line in logs.output)) + + def test_a_download_without_the_signal_finishes(self): + from py.autotag import AutoTagService + + with tempfile.TemporaryDirectory() as tmp: + service = AutoTagService(models_dir=tmp) + with mock.patch.dict(sys.modules, _fake_hub(chunks_before_exit=None)): + self.assertTrue(service.download_model("wd14-vit")) + folder = os.path.join(tmp, "WD14-ViT") + self.assertEqual( + sorted(os.listdir(folder)), ["model.onnx", "selected_tags.csv"] + ) + + def test_xet_transfers_are_aborted_the_way_hub_aborts_on_ctrl_c(self): + import types + + from py.autotag import abort_hub_downloads + + aborted = [] + xet = types.ModuleType("huggingface_hub.utils._xet") + xet.abort_xet_session = lambda: aborted.append(True) + hub = _fake_hub(chunks_before_exit=None) + hub["huggingface_hub.utils._xet"] = xet + hub["huggingface_hub.utils"]._xet = xet + with mock.patch.dict(sys.modules, hub): + abort_hub_downloads() + self.assertEqual(aborted, [True]) + with mock.patch.dict(sys.modules): + sys.modules.pop("huggingface_hub", None) + abort_hub_downloads() # never imported: nothing to abort + self.assertEqual(aborted, [True]) + + def test_an_aborted_xet_download_is_logged_as_stopped(self): + from py.autotag import AutoTagService + + hub = _fake_hub(chunks_before_exit=None) + + def aborted(*args, **kwargs): + exit_signal.request() # the hook ran, then hub raised + raise RuntimeError("xet session aborted") + + hub["huggingface_hub"].hf_hub_download = aborted + with tempfile.TemporaryDirectory() as tmp: + service = AutoTagService(models_dir=tmp) + with mock.patch.dict(sys.modules, hub): + with self.assertLogs(service.logger, "INFO") as logs: + self.assertFalse(service.download_model("wd14-vit")) + self.assertTrue(any("exiting" in line for line in logs.output), logs.output) + self.assertFalse(any("Download failed" in line for line in logs.output)) + + def test_an_old_hub_without_tqdm_class_downloads_as_before(self): + from py.autotag import AutoTagService + + hub = _fake_hub(chunks_before_exit=None) + full = hub["huggingface_hub"].hf_hub_download + + def old_download( + repo_id, filename, local_dir=None, local_dir_use_symlinks=None + ): + return full(repo_id, filename, local_dir, local_dir_use_symlinks) + + hub["huggingface_hub"].hf_hub_download = old_download + with tempfile.TemporaryDirectory() as tmp: + service = AutoTagService(models_dir=tmp) + with mock.patch.dict(sys.modules, hub): + self.assertTrue(service.download_model("wd14-vit")) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/exit_signal.py b/utils/exit_signal.py new file mode 100644 index 00000000..fa7d82b6 --- /dev/null +++ b/utils/exit_signal.py @@ -0,0 +1,110 @@ +"""The process exit signal: running work stops at its next safe point (#376). + +ComfyUI catches Ctrl+C in ``main.py`` and returns without closing its event +loop. The interpreter then runs ``threading._shutdown``, which first calls the +``threading._register_atexit`` callbacks, newest first. One of them, +``concurrent.futures``'s ``_python_exit``, joins every executor thread, and +only after that do the ``atexit`` callbacks run (the package teardown). A call +still running on the default executor (a scan batch, a model download, a +thumbnail batch) used to keep the process alive until it returned. + +``install()`` registers a callback in that same list after ``concurrent.futures`` +did, so it runs before the join. It sets the signal that ``requested()`` reads +(``Job.cancelled`` and ``map_parallel`` read it) and runs ``on_exit``. + +``threading._register_atexit`` is CPython-internal. It has the same contract +on 3.12, 3.13 and 3.14. Without it, or when it refuses the registration, exit +behaves as before: the call in flight runs to its end. +""" + +import logging +import threading + +logger = logging.getLogger("prompt_manager.exit") + +_requested = threading.Event() +_install_lock = threading.Lock() +_installed = False +_on_exit = [] + + +class ExitRequested(Exception): + """The interpreter is exiting: work not yet started is not started.""" + + +def requested(): + """True once the interpreter has started exiting.""" + return _requested.is_set() + + +def request(): + """Set the signal (the exit hook; tests).""" + _requested.set() + + +def checkpoint(): + """Raise ExitRequested once the interpreter is exiting.""" + if _requested.is_set(): + raise ExitRequested("ComfyUI is exiting") + + +def checked(func): + """``func`` behind a checkpoint: a call that starts after the signal + raises ExitRequested instead of running.""" + + def run(*args, **kwargs): + checkpoint() + return func(*args, **kwargs) + + return run + + +def _hook(): + """The threading atexit callback. It never raises: an exception here + would end ``threading._shutdown`` before it shuts the executors.""" + _requested.set() + for action in tuple(_on_exit): + try: + action() + except (Exception, KeyboardInterrupt): + logger.error("Stopping work at exit failed", exc_info=True) + + +def install(on_exit=None): + """Register the exit hook once per process; True when it is registered. + + ``on_exit`` (optional) runs inside the hook, after the signal is set. + False (logged at DEBUG) when this Python has no ``_register_atexit`` or + refuses it; exit then waits for running work as before. + """ + global _installed + with _install_lock: + if on_exit is not None and on_exit not in _on_exit: + _on_exit.append(on_exit) + if _installed: + return True + # Imported first so its pool join is registered before our hook: + # the callbacks run newest first, so ours runs before that join. + import concurrent.futures.thread # noqa: F401 + + register = getattr(threading, "_register_atexit", None) + if not callable(register): + logger.debug("threading._register_atexit is missing; exit waits for jobs") + return False + try: + register(_hook) + except RuntimeError as error: # already shutting down + logger.debug(f"Exit hook not registered: {error}") + return False + _installed = True + return True + + +def reset(): + """Clear the signal and forget the registration (tests only; a hook + already registered stays in ``threading`` and sets the signal at exit).""" + global _installed + with _install_lock: + _requested.clear() + _installed = False + _on_exit.clear() diff --git a/utils/parallel.py b/utils/parallel.py index 65dcca44..b790794b 100644 --- a/utils/parallel.py +++ b/utils/parallel.py @@ -21,6 +21,8 @@ import logging import threading from concurrent.futures import ThreadPoolExecutor +from .exit_signal import checked + SHUTDOWN_REFUSAL = "interpreter shutdown" _logger = logging.getLogger("prompt_manager.parallel") _noted = set() @@ -49,8 +51,11 @@ def map_parallel(func, items, workers): ``workers`` of 1 (or fewer) runs on the calling thread with no pool. An exception in ``func`` propagates to the caller, so callers that must survive bad input wrap ``func`` themselves. A pool refused at interpreter - exit runs the items it did not start on the calling thread (#375). + exit runs the items it did not start on the calling thread (#375). An item + that has not started when the interpreter starts exiting raises + ExitRequested instead (#376). """ + func = checked(func) items = list(items) workers = max(1, int(workers or 1)) if workers == 1 or len(items) <= 1: diff --git a/web/admin.html b/web/admin.html index 2f064a3c..4d0e4789 100644 --- a/web/admin.html +++ b/web/admin.html @@ -164,27 +164,31 @@ class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors"> Export - + Diagnostics - - - + Loading... @@ -369,61 +374,13 @@