From 5eecce2667c1dd793b0f0c9e74bc747a8927fa6a Mon Sep 17 00:00:00 2001 From: robertvoy Date: Thu, 3 Jul 2025 08:13:28 +1000 Subject: [PATCH] Clean up workflow sent to workers --- __init__.py | 23 +++ gpu_config.json | 34 ++++ multigpu.py | 438 ++++++++++++++++++++++++++++++++++++++++++ web/gpupanel.js | 493 ++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 988 insertions(+) create mode 100644 __init__.py create mode 100644 gpu_config.json create mode 100644 multigpu.py create mode 100644 web/gpupanel.js diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..bc655ec --- /dev/null +++ b/__init__.py @@ -0,0 +1,23 @@ +import os +import sys + +# Add the directory to Python path +sys.path.append(os.path.dirname(__file__)) + +# Import everything needed from the main module +from .multigpu import ( + NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS, + ensure_config_exists, + CONFIG_FILE +) + +WEB_DIRECTORY = "./web" + +# Use the centralized config management from multigpu.py +ensure_config_exists() + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + +print("[MultiGPU] Loaded Multi-GPU nodes.") +print(f"[MultiGPU] Config file: {CONFIG_FILE}") \ No newline at end of file diff --git a/gpu_config.json b/gpu_config.json new file mode 100644 index 0000000..ef1de58 --- /dev/null +++ b/gpu_config.json @@ -0,0 +1,34 @@ +{ + "master": { + "port": 8080, + "cuda_device": 0, + "extra_args": "--mmap-torch-files --listen" + }, + "workers": [ + { + "id": 1, + "name": "GPU 1", + "cuda_device": 1, + "port": 8180, + "enabled": true, + "extra_args": "--mmap-torch-files --listen" + }, + { + "id": 2, + "name": "GPU 2", + "cuda_device": 2, + "port": 8280, + "enabled": true, + "extra_args": "--mmap-torch-files --listen" + }, + { + "id": 3, + "name": "GPU 3", + "cuda_device": 3, + "port": 8380, + "enabled": false, + "extra_args": "--mmap-torch-files --listen --reserve-vram 2" + } + ], + "settings": {} +} \ No newline at end of file diff --git a/multigpu.py b/multigpu.py new file mode 100644 index 0000000..c4e6f7b --- /dev/null +++ b/multigpu.py @@ -0,0 +1,438 @@ +import torch +import numpy as np +from PIL import Image +import folder_paths +import os +import json +import asyncio +import aiohttp +from aiohttp import web +import io +import server +from concurrent.futures import Future + +# --- Config Management --- +CONFIG_FILE = os.path.join(os.path.dirname(__file__), "gpu_config.json") + +def get_default_config(): + """Returns the default configuration dictionary. Single source of truth.""" + return { + "master": {"port": 8188, "cuda_device": 0}, + "workers": [], + "settings": {} + } + +def load_config(): + """Loads the config, falling back to defaults if the file is missing or invalid.""" + if os.path.exists(CONFIG_FILE): + try: + with open(CONFIG_FILE, 'r') as f: + return json.load(f) + except Exception as e: + print(f"[MultiGPU Config] Error loading config, using defaults: {e}") + return get_default_config() + +def save_config(config): + """Saves the configuration to file.""" + try: + with open(CONFIG_FILE, 'w') as f: + json.dump(config, f, indent=2) + return True + except Exception as e: + print(f"[MultiGPU Config] Error saving config: {e}") + return False + +def ensure_config_exists(): + """Creates default config file if it doesn't exist. Used by __init__.py""" + if not os.path.exists(CONFIG_FILE): + default_config = get_default_config() + if save_config(default_config): + print("[MultiGPU] Created default config file") + else: + print("[MultiGPU] Could not create default config file") + +# --- Helper Functions --- +async def handle_api_error(request, error, status=500): + """Standardized error response handler""" + print(f"[MultiGPU] API Error: {error}") + return web.json_response({"status": "error", "message": str(error)}, status=status) + +# --- API Endpoints --- +@server.PromptServer.instance.routes.get("/multigpu/config") +async def get_config_endpoint(request): + config = load_config() + return web.json_response(config) + +@server.PromptServer.instance.routes.post("/multigpu/config/update_worker") +async def update_worker_endpoint(request): + try: + data = await request.json() + worker_id = data.get("worker_id") + enabled = data.get("enabled", False) + + if worker_id is None: + return await handle_api_error(request, "Missing worker_id", 400) + + config = load_config() + worker_found = False + + for worker in config.get("workers", []): + if worker["id"] == worker_id: + worker["enabled"] = enabled + worker_found = True + break + + if not worker_found: + return await handle_api_error(request, f"Worker {worker_id} not found", 404) + + if save_config(config): + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Failed to save config") + except Exception as e: + return await handle_api_error(request, e, 400) + +@server.PromptServer.instance.routes.post("/multigpu/config/update_setting") +async def update_setting_endpoint(request): + """Updates a specific key in the settings object.""" + try: + data = await request.json() + key = data.get("key") + value = data.get("value") + + if not key or value is None: + return await handle_api_error(request, "Missing 'key' or 'value' in request", 400) + + config = load_config() + if 'settings' not in config: + config['settings'] = {} + + config['settings'][key] = value + + if save_config(config): + return web.json_response({"status": "success", "message": f"Setting '{key}' updated."}) + else: + return await handle_api_error(request, "Failed to save config") + except Exception as e: + return await handle_api_error(request, e, 400) + + +@server.PromptServer.instance.routes.post("/multigpu/prepare_job") +async def prepare_job_endpoint(request): + try: + data = await request.json() + multi_job_id = data.get('multi_job_id') + if not multi_job_id: + return await handle_api_error(request, "Missing multi_job_id", 400) + + async with PENDING_JOBS_LOCK: + if multi_job_id not in PENDING_JOBS: + PENDING_JOBS[multi_job_id] = asyncio.Queue() + + print(f"[MultiGPU] Prepared queue for job {multi_job_id}") + return web.json_response({"status": "success"}) + except Exception as e: + return await handle_api_error(request, e) + +@server.PromptServer.instance.routes.post("/multigpu/clear_memory") +async def clear_memory_endpoint(request): + print("[MultiGPU] Received request to clear VRAM.") + try: + # Use ComfyUI's prompt server queue system like the /free endpoint does + if hasattr(server.PromptServer.instance, 'prompt_queue'): + server.PromptServer.instance.prompt_queue.set_flag("unload_models", True) + server.PromptServer.instance.prompt_queue.set_flag("free_memory", True) + print("[MultiGPU] Set queue flags for memory clearing.") + + # Wait a bit for the queue to process + await asyncio.sleep(0.5) + + # Also do direct cleanup as backup, but with error handling + import gc + import comfy.model_management as mm + + try: + mm.unload_all_models() + except AttributeError as e: + print(f"[MultiGPU] Warning during model unload: {e}") + + try: + mm.soft_empty_cache() + except Exception as e: + print(f"[MultiGPU] Warning during cache clear: {e}") + + for _ in range(3): + gc.collect() + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + print("[MultiGPU] VRAM cleared successfully.") + return web.json_response({"status": "success", "message": "GPU memory cleared."}) + except Exception as e: + # Even if there's an error, try to do basic cleanup + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + print(f"[MultiGPU] Partial VRAM clear completed with warning: {e}") + return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"}) + + +# --- Global State & Callback Endpoint --- +PENDING_JOBS = {} +PENDING_JOBS_LOCK = asyncio.Lock() + +@server.PromptServer.instance.routes.post("/multigpu/job_complete") +async def job_complete_endpoint(request): + try: + data = await request.post() + multi_job_id = data.get('multi_job_id') + image_file = data.get('image') + worker_id = data.get('worker_id') + image_index = data.get('image_index') + is_last = data.get('is_last', 'False').lower() == 'true' + + if not all([multi_job_id, image_file]): + return await handle_api_error(request, "Missing job_id or image data", 400) + + # Process image + img_data = image_file.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + async with PENDING_JOBS_LOCK: + if multi_job_id in PENDING_JOBS: + await PENDING_JOBS[multi_job_id].put({ + 'tensor': tensor, + 'worker_id': worker_id, + 'image_index': int(image_index) if image_index else 0, + 'is_last': is_last + }) + print(f"[MultiGPU] Received result for job {multi_job_id} from worker {worker_id} (last: {is_last})") + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Job not found or already complete", 404) + except Exception as e: + return await handle_api_error(request, e) + +# --- Collector Node --- +class MultiGPUCollectorNode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { "images": ("IMAGE",) }, + "hidden": { + "multi_job_id": ("STRING", {"default": ""}), + "is_worker": ("BOOLEAN", {"default": False}), + "master_url": ("STRING", {"default": ""}), + "enabled_worker_ids": ("STRING", {"default": "[]"}), + "worker_batch_size": ("INT", {"default": 1, "min": 1, "max": 1024}), + "worker_id": ("STRING", {"default": ""}), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "run" + CATEGORY = "image" + + def run(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id=""): + if not multi_job_id: + return (images,) + + main_loop = server.PromptServer.instance.loop + future = Future() + + async def run_and_set_future(): + try: + result = await self.execute(images, multi_job_id, is_worker, master_url, enabled_worker_ids, worker_batch_size, worker_id) + future.set_result(result) + except Exception as e: + future.set_exception(e) + + asyncio.run_coroutine_threadsafe(run_and_set_future(), main_loop) + return future.result() + + async def send_image_to_master(self, image_tensor, multi_job_id, master_url, image_index, worker_id, is_last=False): + """Helper method to send a single image to the master""" + img_np = (image_tensor.cpu().numpy() * 255.).astype(np.uint8) + img = Image.fromarray(img_np) + byte_io = io.BytesIO() + img.save(byte_io, format='PNG') + byte_io.seek(0) + + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('image_index', str(image_index)) + data.add_field('is_last', str(is_last)) + data.add_field('image', byte_io, filename=f'image_{image_index}.png', content_type='image/png') + + try: + async with aiohttp.ClientSession() as session: + async with session.post(f"{master_url}/multigpu/job_complete", data=data) as response: + response.raise_for_status() + except Exception as e: + print(f"[MultiGPU Worker] Failed to send image {image_index+1} to master: {e}") + + async def execute(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id=""): + if is_worker: + # Worker mode: send images to master + print(f"[MultiGPU Worker] Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master") + + # Send all images, marking the last one + for i in range(images.shape[0]): + is_last = (i == images.shape[0] - 1) + await self.send_image_to_master(images[i], multi_job_id, master_url, i, worker_id, is_last) + + return (images,) + else: + # Master mode: collect images from workers + enabled_workers = json.loads(enabled_worker_ids) + num_workers = len(enabled_workers) + if num_workers == 0: + return (images,) + + images_on_cpu = images.cpu() + print(f"[MultiGPU Master] Job {multi_job_id}: Collecting images from {num_workers} workers...") + + # Store master images first + master_batch_size = images.shape[0] + + # Initialize storage for collected images + worker_images = {} # Dict to store images by worker_id and index + + async with PENDING_JOBS_LOCK: + if multi_job_id not in PENDING_JOBS: + PENDING_JOBS[multi_job_id] = asyncio.Queue() + q = PENDING_JOBS[multi_job_id] + + # Collect images until all workers report they're done + collected_count = 0 + workers_done = set() + + # Use a reasonable timeout for the first image + timeout = 10.0 + + while len(workers_done) < num_workers: + try: + result = await asyncio.wait_for(q.get(), timeout=timeout) + worker_id = result['worker_id'] + image_index = result['image_index'] + tensor = result['tensor'] + is_last = result.get('is_last', False) + + if worker_id not in worker_images: + worker_images[worker_id] = {} + worker_images[worker_id][image_index] = tensor + + collected_count += 1 + + # Once we start receiving images, use shorter timeout + timeout = 10.0 + + if is_last: + workers_done.add(worker_id) + print(f"[MultiGPU Master] Worker {worker_id} done. Collected {len(worker_images[worker_id])} images") + else: + print(f"[MultiGPU Master] Collected image {image_index + 1} from worker {worker_id}") + + except asyncio.TimeoutError: + missing_workers = set(enabled_workers) - workers_done + print(f"[MultiGPU Master] Timeout. Still waiting for workers: {list(missing_workers)}") + break + + total_collected = sum(len(imgs) for imgs in worker_images.values()) + print(f"[MultiGPU Master] Collection complete. Received {total_collected} images from {len(workers_done)} workers") + + # Clean up job queue + async with PENDING_JOBS_LOCK: + if multi_job_id in PENDING_JOBS: + del PENDING_JOBS[multi_job_id] + + # Reorder images according to seed distribution pattern + # Pattern: master img 1, master img 2, worker 1 img 1, worker 1 img 2, worker 2 img 1, worker 2 img 2, etc. + ordered_tensors = [] + + # Add master images first + for i in range(master_batch_size): + ordered_tensors.append(images_on_cpu[i:i+1]) + + # Add worker images in order + for worker_id in enabled_workers: + # Convert worker_id to string since it comes as string from form data + worker_id_str = str(worker_id) + if worker_id_str in worker_images: + # Sort by image index for each worker + for idx in sorted(worker_images[worker_id_str].keys()): + ordered_tensors.append(worker_images[worker_id_str][idx]) + + combined = torch.cat(ordered_tensors, dim=0) + print(f"[MultiGPU Master] Job {multi_job_id} complete. Combined {combined.shape[0]} images total.") + return (combined,) + +# --- Distributor Node --- +class MultiGPUDistributor: + """ + Distributes seed values across multiple GPUs. + On master: passes through the original seed. + On workers: adds offset based on worker ID. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "seed": ("INT", { + "default": 0, + "min": 0, + "max": 1125899906842624, + "forceInput": False # Widget by default, can be converted to input + }), + }, + "hidden": { + "is_worker": ("BOOLEAN", {"default": False}), + "worker_id": ("STRING", {"default": ""}), + }, + } + + RETURN_TYPES = ("INT",) + RETURN_NAMES = ("seed",) + FUNCTION = "distribute" + CATEGORY = "utils" + + def distribute(self, seed, is_worker=False, worker_id=""): + if not is_worker: + # Master node: pass through original values + print(f"[MultiGPU Distributor] Master: seed={seed}") + return (seed,) + else: + # Worker node: apply offset based on worker index + # Find worker index from enabled_worker_ids + try: + # Worker IDs are passed as "worker_0", "worker_1", etc. + if worker_id.startswith("worker_"): + worker_index = int(worker_id.split("_")[1]) + else: + # Fallback: try to parse as direct index + worker_index = int(worker_id) + + offset = worker_index + 1 + new_seed = seed + offset + print(f"[MultiGPU Distributor] Worker {worker_index}: seed={seed} → {new_seed}") + return (new_seed,) + except (ValueError, IndexError) as e: + print(f"[MultiGPU Distributor] Error parsing worker_id '{worker_id}': {e}") + # Fallback: return original seed + return (seed,) + +NODE_CLASS_MAPPINGS = { + "MultiGPUCollector": MultiGPUCollectorNode, + "MultiGPUDistributor": MultiGPUDistributor +} +NODE_DISPLAY_NAME_MAPPINGS = { + "MultiGPUCollector": "Multi-GPU Collector", + "MultiGPUDistributor": "Multi-GPU Seed" +} \ No newline at end of file diff --git a/web/gpupanel.js b/web/gpupanel.js new file mode 100644 index 0000000..01765c2 --- /dev/null +++ b/web/gpupanel.js @@ -0,0 +1,493 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +class MultiGPUExtension { + // Button styling constants + static BUTTON_STYLES = { + base: "width: 100%; padding: 8px; font-size: 12px; color: white; border: none; border-radius: 4px; cursor: pointer; transition: background-color 0.2s;", + clearMemory: "background-color: #555;", + success: "background-color: #3ca03c;", + error: "background-color: #c04c4c;" + }; + + constructor() { + this.config = null; + this.originalQueuePrompt = api.queuePrompt.bind(api); + + this.loadConfig().then(() => { + this.registerSidebarTab(); + this.setupInterceptor(); + }); + } + + // --- State & Config Management (Single Source of Truth) --- + + get enabledWorkers() { + return this.config?.workers?.filter(w => w.enabled) || []; + } + + get isEnabled() { + return this.enabledWorkers.length > 0; + } + + async loadConfig() { + try { + const response = await fetch(`${window.location.origin}/multigpu/config`); + if (response.ok) { + this.config = await response.json(); + console.log("[MultiGPU] Loaded config:", this.config); + } else { + console.error("[MultiGPU] Failed to load config"); + this.config = { workers: [], settings: {} }; + } + } catch (error) { + console.error("[MultiGPU] Error loading config:", error); + this.config = { workers: [], settings: {} }; + } + } + + async updateWorkerEnabled(workerId, enabled) { + const worker = this.config.workers.find(w => w.id === workerId); + if (worker) { + worker.enabled = enabled; + } + + try { + await fetch(`${window.location.origin}/multigpu/config/update_worker`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ worker_id: workerId, enabled: enabled }) + }); + } catch (error) { + console.error("[MultiGPU] Error updating worker:", error); + } + } + + async _updateSetting(key, value) { + try { + await fetch(`${window.location.origin}/multigpu/config/update_setting`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ key, value }) + }); + } catch (error) { + console.error(`[MultiGPU] Error updating setting '${key}':`, error); + } + } + + // --- UI Rendering --- + + registerSidebarTab() { + app.extensionManager.registerSidebarTab({ + id: "multiGPU", + icon: "pi pi-server", + title: "Multi-GPU", + tooltip: "Multi-GPU Control Panel", + type: "custom", + render: (el) => this.renderSidebarContent(el) + }); + } + + _createButton(text, onClick, style) { + const button = document.createElement("button"); + button.textContent = text; + button.style.cssText = MultiGPUExtension.BUTTON_STYLES.base + style; + button.onclick = onClick; + return button; + } + + async _handleClearMemory(button) { + const originalText = button.textContent; + button.textContent = "Clearing..."; + button.disabled = true; + + try { + const urlsToClear = this.enabledWorkers.map(w => ({ + name: w.name, + url: `http://${w.host || window.location.hostname}:${w.port}` + })); + + if (urlsToClear.length === 0) { + button.textContent = "No Workers"; + button.style.backgroundColor = MultiGPUExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); + setTimeout(() => { + button.textContent = originalText; + button.style.cssText = MultiGPUExtension.BUTTON_STYLES.base + MultiGPUExtension.BUTTON_STYLES.clearMemory; + }, 3000); + return; + } + + const promises = urlsToClear.map(target => + fetch(`${target.url}/multigpu/clear_memory`, { + method: 'POST', + mode: 'cors' + }) + .then(response => ({ ok: response.ok, name: target.name })) + .catch(() => ({ ok: false, name: target.name })) + ); + + const results = await Promise.all(promises); + const failures = results.filter(r => !r.ok); + + if (failures.length === 0) { + button.textContent = "Success!"; + button.style.backgroundColor = MultiGPUExtension.BUTTON_STYLES.success.split(':')[1].trim().replace(';', ''); + } else { + button.textContent = "Error! See Console"; + button.style.backgroundColor = MultiGPUExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); + console.error("[MultiGPU] Failed to clear memory on:", failures.map(f => f.name).join(", ")); + } + + setTimeout(() => { + button.textContent = originalText; + button.style.cssText = MultiGPUExtension.BUTTON_STYLES.base + MultiGPUExtension.BUTTON_STYLES.clearMemory; + }, 3000); + } finally { + button.disabled = false; + } + } + + + renderSidebarContent(el) { + el.innerHTML = ''; + const container = document.createElement("div"); + container.style.cssText = "padding: 15px; display: flex; flex-direction: column; height: 100%;"; + + const title = document.createElement("h3"); + title.textContent = "Multi-GPU Control"; + title.style.cssText = "margin: 0 0 15px 0; padding-bottom: 10px; border-bottom: 1px solid #444;"; + container.appendChild(title); + + const gpuSection = document.createElement("div"); + gpuSection.style.cssText = "flex: 1; overflow-y: auto; margin-bottom: 15px;"; + const gpuTitle = document.createElement("h4"); + gpuTitle.textContent = "Available Workers"; + gpuTitle.style.cssText = "margin: 0 0 10px 0; font-size: 14px;"; + gpuSection.appendChild(gpuTitle); + + const gpuList = document.createElement("div"); + (this.config?.workers || []).forEach(worker => { + const gpuDiv = document.createElement("div"); + gpuDiv.style.cssText = "margin-bottom: 8px; padding: 8px; background: #2a2a2a; border-radius: 4px;"; + const checkbox = document.createElement("input"); + checkbox.type = "checkbox"; + checkbox.id = `gpu-${worker.id}`; + checkbox.checked = worker.enabled; + checkbox.onchange = async (e) => { + await this.updateWorkerEnabled(worker.id, e.target.checked); + this.updateSummary(); + }; + const label = document.createElement("label"); + label.htmlFor = `gpu-${worker.id}`; + label.style.cssText = "cursor: pointer; display: flex; align-items: center; gap: 8px;"; + const gpuInfo = document.createElement("span"); + const hostInfo = worker.host ? ` • ${worker.host}` : ''; + const cudaInfo = worker.cuda_device !== undefined ? `CUDA ${worker.cuda_device} • ` : ''; + gpuInfo.innerHTML = `${worker.name}
${cudaInfo}Port ${worker.port}${hostInfo}`; + label.appendChild(checkbox); + label.appendChild(gpuInfo); + gpuDiv.appendChild(label); + gpuList.appendChild(gpuDiv); + }); + gpuSection.appendChild(gpuList); + container.appendChild(gpuSection); + + const actionsSection = document.createElement("div"); + actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;"; + + const clearMemButton = this._createButton("Clear Worker VRAM", (e) => this._handleClearMemory(e.target), MultiGPUExtension.BUTTON_STYLES.clearMemory); + clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)"; + actionsSection.appendChild(clearMemButton); + + container.appendChild(actionsSection); + + const summarySection = document.createElement("div"); + summarySection.style.cssText = "border-top: 1px solid #444; padding-top: 10px;"; + const summary = document.createElement("div"); + summary.id = "multigpu-summary"; + summary.style.cssText = "font-size: 11px; color: #888;"; + summarySection.appendChild(summary); + container.appendChild(summarySection); + el.appendChild(container); + this.updateSummary(); + } + + updateSummary() { + const summaryEl = document.getElementById('multigpu-summary'); + if (summaryEl) { + const totalGPUs = this.enabledWorkers.length + 1; + if (this.isEnabled) { + summaryEl.textContent = `If Collector node is present, total generation = (${totalGPUs} GPUs × Batch Size)`; + } else { + summaryEl.textContent = "Only the master GPU will be used."; + } + } + } + + // --- Helper Methods --- + + // --- Core Logic & Execution --- + + setupInterceptor() { + api.queuePrompt = async (number, prompt) => { + if (this.isEnabled && this.findNodesByClass(prompt.output, "MultiGPUCollector").length > 0) { + return await this.executeParallelMultiGPU(prompt); + } + return this.originalQueuePrompt(number, prompt); + }; + } + + async executeParallelMultiGPU(promptWrapper) { + try { + const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution + const enabledWorkers = this.enabledWorkers; + + // Find all collector nodes in the workflow + const collectorNodes = this.findNodesByClass(promptWrapper.output, "MultiGPUCollector"); + + // Map original node IDs to truly unique job IDs for this specific run + const job_id_map = new Map(collectorNodes.map(node => [node.id, `${executionPrefix}_${node.id}`])); + + // Prepare a separate job queue on the backend for each unique job ID + const preparePromises = Array.from(job_id_map.values()).map(uniqueId => this._prepareMultiGpuJob(uniqueId)); + await Promise.all(preparePromises); + + const jobs = []; + const participants = ['master', ...enabledWorkers.map(w => w.id)]; + + for (const participantId of participants) { + const jobApiPrompt = this._prepareApiPromptForParticipant( + promptWrapper.output, participantId, + { + enabled_worker_ids: enabledWorkers.map(w => w.id), + workflow: promptWrapper.workflow, + job_id_map: job_id_map // Pass the map of unique IDs + } + ); + + if (participantId === 'master') { + jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } }); + } else { + const worker = this.config.workers.find(w => w.id === participantId); + if (worker) jobs.push({ type: 'worker', worker, prompt: jobApiPrompt, workflow: promptWrapper.workflow }); + } + } + + const result = await this._executeJobs(jobs); + return result; + } catch (error) { + console.error("[MultiGPU] Parallel execution failed:", error); + alert(`[MultiGPU] Parallel execution failed: ${error.message}`); + throw error; + } + } + + + async _executeJobs(jobs) { + let masterPromptId = null; + const promises = jobs.map(job => { + if (job.type === 'master') { + return this.originalQueuePrompt(0, job.promptWrapper).then(result => { + masterPromptId = result; + return result; + }); + } else { + return this._dispatchToWorker(job.worker, job.prompt, job.workflow); + } + }); + await Promise.all(promises); + return masterPromptId || { "prompt_id": "multi-gpu-job-dispatched" }; + } + + // --- Helper Methods --- + + findNodesByClass(apiPrompt, className) { + return Object.entries(apiPrompt) + .filter(([, nodeData]) => nodeData.class_type === className) + .map(([nodeId, nodeData]) => ({ id: nodeId, data: nodeData })); + } + + /** + * Find all nodes connected to collector nodes (both upstream and downstream) + * @param {Object} apiPrompt - The workflow API prompt + * @param {Array} collectorIds - Array of collector node IDs + * @returns {Set} Set of node IDs connected to any collector + */ + findCollectorConnectedNodes(apiPrompt, collectorIds) { + const connected = new Set(collectorIds); // Include all collectors + const toProcess = [...collectorIds]; + + // First, build a map of which nodes output to which nodes + const outputMap = new Map(); // sourceNode -> Set of target nodes + for (const [nodeId, node] of Object.entries(apiPrompt)) { + if (node.inputs) { + for (const [inputName, inputValue] of Object.entries(node.inputs)) { + if (Array.isArray(inputValue) && inputValue.length === 2) { + const sourceNodeId = String(inputValue[0]); + if (!outputMap.has(sourceNodeId)) { + outputMap.set(sourceNodeId, new Set()); + } + outputMap.get(sourceNodeId).add(nodeId); + } + } + } + } + + // Traverse both upstream and downstream + while (toProcess.length > 0) { + const nodeId = toProcess.pop(); + const node = apiPrompt[nodeId]; + + // Traverse upstream (inputs) + if (node && node.inputs) { + for (const [inputName, inputValue] of Object.entries(node.inputs)) { + if (Array.isArray(inputValue) && inputValue.length === 2) { + const sourceNodeId = String(inputValue[0]); + if (!connected.has(sourceNodeId)) { + connected.add(sourceNodeId); + toProcess.push(sourceNodeId); + } + } + } + } + + // Traverse downstream (outputs) + if (outputMap.has(nodeId)) { + for (const targetNodeId of outputMap.get(nodeId)) { + if (!connected.has(targetNodeId)) { + connected.add(targetNodeId); + toProcess.push(targetNodeId); + } + } + } + } + + return connected; + } + + /** + * Prune workflow to only include nodes connected to collectors + * @param {Object} apiPrompt - The full workflow API prompt + * @param {Array} collectorNodes - Array of collector nodes (optional, will find if not provided) + * @returns {Object} Pruned API prompt with only required nodes + */ + pruneWorkflowForWorker(apiPrompt, collectorNodes = null) { + // Find all collector nodes if not provided + if (!collectorNodes) { + collectorNodes = this.findNodesByClass(apiPrompt, "MultiGPUCollector"); + } + + if (collectorNodes.length === 0) { + // No collectors, return full workflow + return apiPrompt; + } + + // Get all nodes connected to collectors (both upstream and downstream) + const collectorIds = collectorNodes.map(node => node.id); + const connectedNodes = this.findCollectorConnectedNodes(apiPrompt, collectorIds); + + console.log(`[MultiGPU] Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes`); + + // Create pruned prompt with only required nodes + const prunedPrompt = {}; + for (const nodeId of connectedNodes) { + prunedPrompt[nodeId] = JSON.parse(JSON.stringify(apiPrompt[nodeId])); + } + + return prunedPrompt; + } + + + _prepareApiPromptForParticipant(baseApiPrompt, participantId, options = {}) { + let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt)); + const isMaster = participantId === 'master'; + + // Find collector nodes once (before pruning) + const collectorNodes = this.findNodesByClass(jobApiPrompt, "MultiGPUCollector"); + + // For workers, prune the workflow to only include collector dependencies + if (!isMaster && collectorNodes.length > 0) { + jobApiPrompt = this.pruneWorkflowForWorker(jobApiPrompt, collectorNodes); + } + + // Handle MultiGPU distributor nodes + const distributorNodes = this.findNodesByClass(jobApiPrompt, "MultiGPUDistributor"); + if (distributorNodes.length > 0) { + console.log(`[MultiGPU] Found ${distributorNodes.length} distributor node(s)`); + } + + for (const distributor of distributorNodes) { + const { inputs } = jobApiPrompt[distributor.id]; + inputs.is_worker = !isMaster; + if (!isMaster) { + const workerIndex = options.enabled_worker_ids.indexOf(participantId); + inputs.worker_id = `worker_${workerIndex}`; + console.log(`[MultiGPU] Set distributor ${distributor.id} for worker ${workerIndex}`); + } + } + + // Handle MultiGPU collector nodes (already found above) + for (const collector of collectorNodes) { + const { inputs } = jobApiPrompt[collector.id]; + + // Get the unique job ID from the map created for this execution + const uniqueJobId = options.job_id_map ? options.job_id_map.get(collector.id) : collector.id; + + // Use the truly unique ID for this execution + inputs.multi_job_id = uniqueJobId; + inputs.is_worker = !isMaster; + if (isMaster) { + inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); + } else { + inputs.master_url = window.location.origin; + // Also make the worker_job_id unique to prevent potential caching issues + inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`; + inputs.worker_id = participantId; + } + } + + return jobApiPrompt; + } + + async _prepareMultiGpuJob(multi_job_id) { + try { + await fetch(`${window.location.origin}/multigpu/prepare_job`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ multi_job_id }) + }); + } catch (error) { + console.error("[MultiGPU] Error preparing job:", error); + throw error; + } + } + + async _dispatchToWorker(worker, prompt, workflow) { + const workerUrl = `http://${worker.host || window.location.hostname}:${worker.port}`; + + const promptToSend = { + prompt, + extra_data: { extra_pnginfo: { workflow } }, + client_id: api.clientId + }; + + + try { + await fetch(`${workerUrl}/prompt`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + mode: 'cors', + body: JSON.stringify(promptToSend) + }); + } catch (e) { + console.error(`[MultiGPU] Failed to connect to worker ${worker.name} at ${workerUrl}`, e); + } + } +} + +app.registerExtension({ + name: "MultiGPU.Panel", + async setup() { + new MultiGPUExtension(); + } +}); \ No newline at end of file