From 8be7b2206b75b896331cb2afaf1f670e4ebeec1a Mon Sep 17 00:00:00 2001 From: MadiatorLabs <236864052+MadiatorLabs@users.noreply.github.com> Date: Wed, 22 Apr 2026 15:52:10 +0200 Subject: [PATCH] v1.0.10 - Disable extensive models check on prompt que --- __init__.py | 155 +++++++++++++++++++++++++++++++++++- web/serverDownload.js | 180 +++++++++++++++++++++++++++++++++++++----- 2 files changed, 314 insertions(+), 21 deletions(-) diff --git a/__init__.py b/__init__.py index fdf1bc8..31fbb92 100644 --- a/__init__.py +++ b/__init__.py @@ -33,6 +33,7 @@ SOCK_READ_TIMEOUT = 120 HASH_READ_CHUNK_SIZE = 8 * 1024 * 1024 WS_KEEPALIVE_INTERVAL_SECONDS = 45 SETTINGS_FILENAME = "runpoddirect_settings.json" +PREQUEUE_CHECK_MAX_MODELS = 512 _ws_keepalive_task = None _ws_keepalive_enabled = True @@ -300,6 +301,61 @@ def _cleanup_partial_file(path): logging.warning(f"[RunpodDirect] Failed to remove partial file {path}: {e}") +def _sanitize_simple_filename(filename): + if not filename or not isinstance(filename, str): + raise ValueError("Missing filename") + + if "/" in filename or "\\" in filename or os.path.sep in filename: + raise ValueError("Invalid filename: must not contain path separators") + + if ".." in filename or filename.startswith("/") or filename.startswith("~"): + raise ValueError("Invalid filename: path traversal patterns detected") + + safe_filename = os.path.basename(filename) + if safe_filename != filename: + raise ValueError("Invalid filename: must be a simple filename without path components") + + return safe_filename + + +def _resolve_model_path(directory, safe_filename): + if directory not in folder_paths.folder_names_and_paths: + raise ValueError(f"Invalid directory: {directory}") + + output_dir = os.path.abspath(folder_paths.folder_names_and_paths[directory][0][0]) + model_path = os.path.abspath(os.path.join(output_dir, safe_filename)) + if not model_path.startswith(output_dir + os.sep): + raise ValueError("Security error: attempted directory escape") + + return output_dir, model_path + + +def _model_exists_in_directory(directory, safe_filename): + try: + _, model_path = _resolve_model_path(directory, safe_filename) + except Exception: + return False + return os.path.exists(model_path) + + +def _find_model_directory_anywhere(safe_filename, preferred_directories=None): + ordered_directories = [] + seen = set() + for directory in preferred_directories or []: + if directory in folder_paths.folder_names_and_paths and directory not in seen: + ordered_directories.append(directory) + seen.add(directory) + for directory in folder_paths.folder_names_and_paths.keys(): + if directory not in seen: + ordered_directories.append(directory) + seen.add(directory) + + for directory in ordered_directories: + if _model_exists_in_directory(directory, safe_filename): + return directory + return None + + async def _emit_download_progress(download_id, total_size, force=False): ctrl = download_control.get(download_id) if not ctrl: @@ -1153,6 +1209,103 @@ async def server_download_folder_paths(request): return web.json_response({"error": str(e)}, status=500) +@PromptServer.instance.routes.post("/server_download/check_missing_models") +async def check_missing_models(request): + """Batch-check workflow models without blocking queueing on many frontend round-trips.""" + try: + json_data = await request.json() + models = json_data.get("models") + verify_hashes = _parse_bool(json_data.get("verify_hashes"), default=False) + + if not isinstance(models, list): + return web.json_response( + {"error": "Missing required parameter: models"}, + status=400 + ) + + if len(models) > PREQUEUE_CHECK_MAX_MODELS: + return web.json_response( + {"error": f"Too many models requested. Max: {PREQUEUE_CHECK_MAX_MODELS}"}, + status=400 + ) + + missing = [] + unresolved = [] + + for raw_model in models: + if not isinstance(raw_model, dict): + continue + + report_model = dict(raw_model) + filename = raw_model.get("filename") or raw_model.get("name") + try: + safe_filename = _sanitize_simple_filename(filename) + except ValueError as e: + report_model["filename"] = str(filename or "") + report_model["reason"] = "invalid_filename" + report_model["error"] = str(e) + unresolved.append(report_model) + continue + + report_model["filename"] = safe_filename + + directory = raw_model.get("directory") + directory = str(directory) if isinstance(directory, str) and directory else None + if directory and directory not in folder_paths.folder_names_and_paths: + directory = None + + expected_hash = _normalize_expected_hash(raw_model.get("hash")) + expected_hash_type = raw_model.get("hash_type") + found_directory = None + + if directory: + if _model_exists_in_directory(directory, safe_filename): + found_directory = directory + else: + found_directory = _find_model_directory_anywhere( + safe_filename, + preferred_directories=[directory], + ) + if not found_directory: + report_model["directory"] = directory + report_model["reason"] = "missing" + missing.append(report_model) + continue + else: + found_directory = _find_model_directory_anywhere(safe_filename) + if not found_directory: + report_model["reason"] = "directory_unresolved" + unresolved.append(report_model) + continue + + if verify_hashes and expected_hash: + hash_type = _resolve_hash_type(expected_hash, expected_hash_type) + if not hash_type: + report_model["directory"] = found_directory + report_model["reason"] = "unsupported_hash_type" + unresolved.append(report_model) + continue + + _, model_path = _resolve_model_path(found_directory, safe_filename) + actual_hash = _compute_file_hash(model_path, hash_type) + if actual_hash.lower() != expected_hash.lower(): + report_model["directory"] = found_directory + report_model["corrupted"] = True + report_model["reason"] = "hash_mismatch" + missing.append(report_model) + + return web.json_response({ + "success": True, + "missing": missing, + "unresolved": unresolved, + "checked": len(models), + "verify_hashes": verify_hashes, + }) + except Exception as e: + logging.error(f"Error checking missing models: {e}") + return web.json_response({"error": str(e)}, status=500) + + @PromptServer.instance.routes.post("/server_download/verify_model_integrity") async def verify_model_integrity(request): """Verify whether a model file exists and, when hash is provided, whether it matches.""" @@ -1326,7 +1479,7 @@ async def serve_js_with_version(request): WEB_DIRECTORY = "./web" # Version for cache busting - increment this when you update the JS -__version__ = "1.0.8" +__version__ = "1.0.10" # Apply cgroup-aware RAM patch, then load persisted settings, then start keepalive. _patch_comfy_ram_detection_for_cgroups() diff --git a/web/serverDownload.js b/web/serverDownload.js index ee5f0b5..2bc7abb 100644 --- a/web/serverDownload.js +++ b/web/serverDownload.js @@ -31,14 +31,22 @@ function isAutoMissingCheckEnabled() { return getLocalFlag('runpoddirect_auto_check', true); } +function isPreQueueGuardEnabled() { + return getLocalFlag('runpoddirect_prequeue_guard', false); +} + +function isStrictPreQueueHashCheckEnabled() { + return getLocalFlag('runpoddirect_prequeue_hash_check', false); +} + function debugLog(...args) { if (!isVerboseLogsEnabled()) return; console.log(...args); } // ComfyUI RunpodDirect Extension -// Version: 1.0.8 -debugLog('[RunpodDirect] v1.0.8'); +// Version: 1.0.10 +debugLog('[RunpodDirect] v1.0.10'); // Track download states const downloadStates = new Map(); @@ -197,8 +205,11 @@ const modelExistsAnywhereCache = new Map(); const modelDiskExistsCache = new Map(); const modelSizeCache = new Map(); const modelIntegrityCache = new Map(); +const PREQUEUE_REPORT_TTL_MS = 10000; let queueGuardInstalled = false; let queueGuardBypassOnce = false; +let queueGuardOriginalPrompt = null; +let queueGuardWrappedPrompt = null; let autoMissingModalTimer = null; let autoMissingModalInFlight = false; let autoMissingModalSignature = ''; @@ -218,6 +229,7 @@ let runpodHubStylesInstalled = false; let managerModelIndexCache = null; let managerModelIndexPromise = null; let workflowGraphSnapshot = null; +let preQueueReportCache = null; const THEME = { // Status colors @@ -351,6 +363,7 @@ api.addEventListener("server_download_progress", ({ detail }) => { api.addEventListener("server_download_complete", ({ detail }) => { const { download_id, path, size } = detail; + invalidatePreQueueReportCache(); if (isDownloadingAll) { completedDownloads++; debugLog(`[RunpodDirect] Progress: ${completedDownloads}/${totalDownloads} completed`); @@ -369,6 +382,7 @@ api.addEventListener("server_download_complete", ({ detail }) => { api.addEventListener("server_download_error", ({ detail }) => { const { download_id, error } = detail; + invalidatePreQueueReportCache(); if (isDownloadingAll) { completedDownloads++; debugLog(`[RunpodDirect] Progress: ${completedDownloads}/${totalDownloads} completed (1 error)`); @@ -944,6 +958,29 @@ async function openRunpodHubPanel() { (enabled) => setLocalFlag('runpoddirect_auto_check', enabled), )); + settings.appendChild(makeCheckboxRow( + 'Pre-queue guard', + 'Intercept Run and block queueing when required models appear to be missing.', + () => isPreQueueGuardEnabled(), + (enabled) => { + setLocalFlag('runpoddirect_prequeue_guard', enabled); + invalidatePreQueueReportCache(); + syncPreQueueGuard(); + return true; + }, + )); + + settings.appendChild(makeCheckboxRow( + 'Strict pre-queue hash checks', + 'Re-hash matching files before queueing. Safer for corruption checks, slower on large models.', + () => isStrictPreQueueHashCheckEnabled(), + (enabled) => { + setLocalFlag('runpoddirect_prequeue_hash_check', enabled); + invalidatePreQueueReportCache(); + return true; + }, + )); + settings.appendChild(makeCheckboxRow( 'Connection keepalive', 'Keep websocket alive via silent server pings to reduce proxy reconnects.', @@ -1256,6 +1293,43 @@ function buildMissingReportSignature(report) { return rows.join('||'); } +function buildPreQueueCandidateSignature(candidates, verifyHashes = false) { + const rows = []; + for (const model of Array.isArray(candidates) ? candidates : []) { + const filename = sanitizeFilename(model?.filename).toLowerCase(); + const url = normalizeDownloadUrl(model?.url || ''); + const directory = String(model?.directory || '').toLowerCase(); + const hash = verifyHashes ? String(model?.hash || '').toLowerCase() : ''; + if (!filename) continue; + rows.push(`${filename}|${directory}|${url}|${hash}`); + } + rows.sort(); + return rows.join('||'); +} + +function clonePreQueueReport(report) { + return cloneWorkflowData(report); +} + +function getCachedPreQueueReport(signature) { + if (!preQueueReportCache) return null; + if (preQueueReportCache.signature !== signature) return null; + if ((Date.now() - preQueueReportCache.createdAt) > PREQUEUE_REPORT_TTL_MS) return null; + return clonePreQueueReport(preQueueReportCache.report); +} + +function setCachedPreQueueReport(signature, report) { + preQueueReportCache = { + signature, + createdAt: Date.now(), + report: clonePreQueueReport(report), + }; +} + +function invalidatePreQueueReportCache() { + preQueueReportCache = null; +} + function shouldWaitForNativeMissingDialog(reason = '') { const lower = String(reason || '').toLowerCase(); return lower.includes('afterconfiguregraph') @@ -2578,19 +2652,37 @@ async function collectPreQueueModelCandidates(folderPaths) { return await collectPreQueueModelCandidatesFromWorkflow(workflow, folderPaths); } -async function checkMissingModelsForWorkflow(workflow, piniaPaths = {}) { - const folderPaths = await getAvailableFolderPaths(piniaPaths); - // Force re-read from disk every check so manual deletes are detected. +async function checkMissingModelsViaBackend(candidates, folderPaths, options = {}) { + try { + const response = await api.fetchApi('/server_download/check_missing_models', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + models: candidates, + verify_hashes: !!options.verifyHashes, + }), + }); + if (!response.ok) return null; + const data = await response.json(); + return { + missing: Array.isArray(data?.missing) ? data.missing : [], + unresolved: Array.isArray(data?.unresolved) ? data.unresolved : [], + candidates, + folderPaths, + }; + } catch (_e) { + return null; + } +} + +async function checkMissingModelsForCandidatesFallback(candidates, folderPaths, options = {}) { + const verifyHashes = !!options.verifyHashes; + // Legacy path for older backends: keep behavior compatible if the batch route is unavailable. modelFolderCache.clear(); modelExistsAnywhereCache.clear(); modelDiskExistsCache.clear(); modelIntegrityCache.clear(); - const candidates = await collectPreQueueModelCandidatesFromWorkflow(workflow, folderPaths); - if (!candidates.length) { - return { missing: [], unresolved: [], candidates: [], folderPaths }; - } - const missing = []; const unresolved = []; @@ -2627,7 +2719,7 @@ async function checkMissingModelsForWorkflow(workflow, piniaPaths = {}) { missing.push(model); continue; } - if (exists === true && model.hash) { + if (verifyHashes && exists === true && model.hash) { const integrity = await verifyModelIntegrity( model.directory, model.filename, @@ -2643,11 +2735,36 @@ async function checkMissingModelsForWorkflow(workflow, piniaPaths = {}) { return { missing, unresolved, candidates, folderPaths }; } -async function checkMissingModelsBeforeQueue() { +async function checkMissingModelsForWorkflow(workflow, piniaPaths = {}, options = {}) { + const verifyHashes = !!options.verifyHashes; + const folderPaths = await getAvailableFolderPaths(piniaPaths); + const candidates = await collectPreQueueModelCandidatesFromWorkflow(workflow, folderPaths); + if (!candidates.length) { + return { missing: [], unresolved: [], candidates: [], folderPaths }; + } + + const signature = buildPreQueueCandidateSignature(candidates, verifyHashes); + if (!options.forceRefresh) { + const cachedReport = getCachedPreQueueReport(signature); + if (cachedReport) { + return cachedReport; + } + } + + let report = await checkMissingModelsViaBackend(candidates, folderPaths, { verifyHashes }); + if (!report) { + report = await checkMissingModelsForCandidatesFallback(candidates, folderPaths, { verifyHashes }); + } + + setCachedPreQueueReport(signature, report); + return clonePreQueueReport(report); +} + +async function checkMissingModelsBeforeQueue(options = {}) { const dialogState = getMissingModelsDialogState(); const piniaPaths = dialogState?.paths || {}; const workflow = getCurrentWorkflowData({ preferSnapshot: false, allowSerialize: true }); - const { missing, unresolved, candidates } = await checkMissingModelsForWorkflow(workflow, piniaPaths); + const { missing, unresolved, candidates } = await checkMissingModelsForWorkflow(workflow, piniaPaths, options); debugLog( `[RunpodDirect] Pre-queue check: ${candidates.length} candidates, ${missing.length} missing, ${unresolved.length} unresolved` @@ -2894,27 +3011,48 @@ function showPreQueueMissingModal(report, onQueueAnyway, options = {}) { document.body.appendChild(overlay); } -function installAlwaysCheckQueueGuard(retryCount = 0) { +function syncPreQueueGuard(retryCount = 0) { + if (!isPreQueueGuardEnabled()) { + if (queueGuardInstalled && app?.queuePrompt === queueGuardWrappedPrompt && typeof queueGuardOriginalPrompt === 'function') { + app.queuePrompt = queueGuardOriginalPrompt; + debugLog('[RunpodDirect] Pre-queue model check guard removed'); + } + queueGuardInstalled = false; + queueGuardWrappedPrompt = null; + queueGuardOriginalPrompt = null; + queueGuardBypassOnce = false; + closePreQueueMissingModal(); + return; + } + if (queueGuardInstalled) return; - const originalQueuePrompt = app?.queuePrompt; - if (typeof originalQueuePrompt !== 'function') { + const currentQueuePrompt = app?.queuePrompt; + if (typeof currentQueuePrompt !== 'function') { if (retryCount < 20) { - setTimeout(() => installAlwaysCheckQueueGuard(retryCount + 1), 500); + setTimeout(() => syncPreQueueGuard(retryCount + 1), 500); } else { debugLog('[RunpodDirect] queuePrompt hook not available, pre-queue guard not installed'); } return; } - app.queuePrompt = async function(number, batchCount = 1, queueNodeIds) { + const originalQueuePrompt = currentQueuePrompt; + queueGuardOriginalPrompt = originalQueuePrompt; + queueGuardWrappedPrompt = async function(number, batchCount = 1, queueNodeIds) { + if (!isPreQueueGuardEnabled()) { + return await originalQueuePrompt.call(this, number, batchCount, queueNodeIds); + } + if (queueGuardBypassOnce) { queueGuardBypassOnce = false; return await originalQueuePrompt.call(this, number, batchCount, queueNodeIds); } try { - const report = await checkMissingModelsBeforeQueue(); + const report = await checkMissingModelsBeforeQueue({ + verifyHashes: isStrictPreQueueHashCheckEnabled(), + }); if (report.missing.length > 0 || report.unresolved.length > 0) { debugLog( `[RunpodDirect] Blocking queue: ${report.missing.length} missing, ${report.unresolved.length} unresolved model(s)` @@ -2935,6 +3073,7 @@ function installAlwaysCheckQueueGuard(retryCount = 0) { return await originalQueuePrompt.call(this, number, batchCount, queueNodeIds); }; + app.queuePrompt = queueGuardWrappedPrompt; queueGuardInstalled = true; debugLog('[RunpodDirect] Pre-queue model check guard installed'); } @@ -4006,7 +4145,7 @@ app.registerExtension({ async setup() { debugLog("[RunpodDirect] Extension setup starting"); checkEnvHfToken(); - installAlwaysCheckQueueGuard(); + syncPreQueueGuard(); setupDialogObserver(); installAutoMissingCheckListeners(); installRunpodHubListeners(); @@ -4093,6 +4232,7 @@ api.addEventListener("server_download_resumed", ({ detail }) => { api.addEventListener("server_download_cancelled", ({ detail }) => { const { download_id } = detail; + invalidatePreQueueReportCache(); const prev = downloadStates.get(download_id) || {}; downloadStates.set(download_id, { ...prev, status: 'cancelled' }); window.dispatchEvent(new CustomEvent('serverDownloadUpdate', {