diff --git a/distributed.py b/distributed.py index 1c2cf57..cdc699d 100644 --- a/distributed.py +++ b/distributed.py @@ -23,7 +23,15 @@ from comfy.utils import ProgressBar # Import shared utilities from .utils.logging import debug_log, log -from .utils.config import CONFIG_FILE, get_default_config, load_config, save_config, ensure_config_exists, get_worker_timeout_seconds +from .utils.config import ( + CONFIG_FILE, + get_default_config, + load_config, + save_config, + ensure_config_exists, + get_worker_timeout_seconds, + is_master_delegate_only, +) from .utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous from .utils.process import is_process_alive, terminate_process, get_python_executable from .utils.network import handle_api_error, get_server_port, get_server_loop, get_client_session, cleanup_client_session @@ -1813,6 +1821,7 @@ class DistributedCollectorNode: "worker_batch_size": ("INT", {"default": 1, "min": 1, "max": 1024}), "worker_id": ("STRING", {"default": ""}), "pass_through": ("BOOLEAN", {"default": False}), + "delegate_only": ("BOOLEAN", {"default": False}), }, } @@ -1820,7 +1829,7 @@ class DistributedCollectorNode: 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="", pass_through=False): + def run(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False): if not multi_job_id or pass_through: if pass_through: print(f"[Distributed Collector] Pass-through mode enabled, returning images unchanged") @@ -1828,7 +1837,7 @@ class DistributedCollectorNode: # Use async helper to run in server loop result = run_async_in_server_loop( - self.execute(images, multi_job_id, is_worker, master_url, enabled_worker_ids, worker_batch_size, worker_id) + self.execute(images, multi_job_id, is_worker, master_url, enabled_worker_ids, worker_batch_size, worker_id, delegate_only) ) return result @@ -1876,25 +1885,31 @@ class DistributedCollectorNode: raise # Re-raise to handle at caller level - async def execute(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id=""): + async def execute(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False): if is_worker: # Worker mode: send images to master in a single batch debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master") await self.send_batch_to_master(images, multi_job_id, master_url, worker_id) return (images,) else: + delegate_mode = delegate_only or is_master_delegate_only() # 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() - master_batch_size = images.shape[0] - debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...") - - # Ensure master images are contiguous - images_on_cpu = ensure_contiguous(images_on_cpu) + if delegate_mode: + master_batch_size = 0 + images_on_cpu = None + debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers") + else: + images_on_cpu = images.cpu() + master_batch_size = images.shape[0] + debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...") + + # Ensure master images are contiguous + images_on_cpu = ensure_contiguous(images_on_cpu) # Initialize storage for collected images @@ -2101,9 +2116,10 @@ class DistributedCollectorNode: # 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 master images first (if any) + if not delegate_mode and images_on_cpu is not None: + for i in range(master_batch_size): + ordered_tensors.append(images_on_cpu[i:i+1]) # Add worker images in order # The worker IDs in worker_images are already strings (e.g., "1", "2") @@ -2123,7 +2139,13 @@ class DistributedCollectorNode: cpu_tensors.append(t) try: - combined = torch.cat(cpu_tensors, dim=0) + if cpu_tensors: + combined = torch.cat(cpu_tensors, dim=0) + else: + # No tensors collected (likely delegate mode with no worker output) + combined = ensure_contiguous(images) if images is not None else images + if combined is None: + raise ValueError("No image data collected from master or workers") # Ensure the combined tensor is contiguous and properly formatted combined = ensure_contiguous(combined) debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total (master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})") @@ -2262,13 +2284,40 @@ class ImageBatchDivider: return tuple(outputs) +class DistributedEmptyImage: + """Produces an empty IMAGE batch used when the master delegates all work.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "height": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}), + "width": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}), + "channels": ("INT", {"default": 3, "min": 1, "max": 4, "step": 1}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "create" + CATEGORY = "image" + + def create(self, height, width, channels): + import torch + + shape = (0, height, width, channels) + tensor = torch.zeros(shape, dtype=torch.float32) + return (tensor,) + + NODE_CLASS_MAPPINGS = { "DistributedCollector": DistributedCollectorNode, "DistributedSeed": DistributedSeed, - "ImageBatchDivider": ImageBatchDivider + "ImageBatchDivider": ImageBatchDivider, + "DistributedEmptyImage": DistributedEmptyImage, } NODE_DISPLAY_NAME_MAPPINGS = { "DistributedCollector": "Distributed Collector", "DistributedSeed": "Distributed Seed", - "ImageBatchDivider": "Image Batch Divider" + "ImageBatchDivider": "Image Batch Divider", + "DistributedEmptyImage": "Distributed Empty Image", } diff --git a/docs/worker-setup-guides.md b/docs/worker-setup-guides.md index 20e35ef..1d1dba6 100644 --- a/docs/worker-setup-guides.md +++ b/docs/worker-setup-guides.md @@ -6,6 +6,14 @@ +### Master participation modes + +The master can either contribute GPU work or stay in **orchestrator-only** mode: + +- **Participating**: Master renders alongside workers, useful when you want every available GPU. +- **Orchestrator-only**: Master sends jobs to selected workers but skips local rendering. Enable this by opening the Distributed panel and unchecking the master toggle. The master card will display *“Master disabled: running as orchestrator only.”* +- **Fallback**: If orchestrator-only is enabled but no workers remain selected, the master automatically re-enables execution to guarantee the workflow still runs. The UI shows a green *“Master fallback execution active”* badge so you know work is executing locally again. + ### Types of Workers - **Local workers**: Additional GPUs on the same machine as the master diff --git a/utils/config.py b/utils/config.py index 2569663..8c9aa49 100644 --- a/utils/config.py +++ b/utils/config.py @@ -18,7 +18,8 @@ def get_default_config(): "settings": { "debug": False, "auto_launch_workers": False, - "stop_workers_on_master_exit": True + "stop_workers_on_master_exit": True, + "master_delegate_only": False } } @@ -69,3 +70,12 @@ def get_worker_timeout_seconds(default: int = HEARTBEAT_TIMEOUT) -> int: return max(1, val) except Exception: return max(1, int(default)) + + +def is_master_delegate_only() -> bool: + """Returns True when master should skip local workload and act as orchestrator only.""" + try: + cfg = load_config() + return bool(cfg.get('settings', {}).get('master_delegate_only', False)) + except Exception: + return False diff --git a/web/executionUtils.js b/web/executionUtils.js index 9faea12..c53a016 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -1,5 +1,5 @@ import { api } from "../../scripts/api.js"; -import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo } from './workerUtils.js'; +import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo, findCollectorDownstreamNodes } from './workerUtils.js'; import { TIMEOUTS } from './constants.js'; /** @@ -131,6 +131,21 @@ export async function executeParallelDistributed(extension, promptWrapper) { const collectorNodes = findNodesByClass(promptWrapper.output, "DistributedCollector"); const upscaleNodes = findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed"); const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; + + const wantsDelegate = Boolean(extension.config?.settings?.master_delegate_only); + let masterDelegateActive = false; + if (wantsDelegate) { + if (activeWorkers.length === 0) { + extension.log("Master delegate-only mode disabled: no active workers detected. Falling back to master execution.", "debug"); + } else if (upscaleNodes.length > 0) { + extension.log("Master delegate-only mode is not yet supported for UltimateSDUpscaleDistributed nodes. Falling back to normal execution.", "warn"); + } else if (collectorNodes.length === 0) { + extension.log("Master delegate-only mode requested but no DistributedCollector nodes found. Falling back to normal execution.", "debug"); + } else { + masterDelegateActive = true; + extension.log("Master delegate-only mode enabled: master will orchestrate without running upstream nodes.", "debug"); + } + } // Map original node IDs to truly unique job IDs for this specific run const job_id_map = new Map(allDistributedNodes.map(node => [node.id, `${executionPrefix}_${node.id}`])); @@ -147,7 +162,8 @@ export async function executeParallelDistributed(extension, promptWrapper) { const options = { enabled_worker_ids: activeWorkers.map(w => w.id), workflow: promptWrapper.workflow, - job_id_map: job_id_map // Pass the map of unique IDs + job_id_map: job_id_map, // Pass the map of unique IDs + delegate_master: masterDelegateActive && participantId === 'master' }; const jobApiPrompt = await prepareApiPromptForParticipant( @@ -187,11 +203,24 @@ export async function executeParallelDistributed(extension, promptWrapper) { export async function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) { let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt)); const isMaster = participantId === 'master'; + const delegateMaster = Boolean(options.delegate_master); // Find all distributed nodes once (before pruning) - const collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector"); + let collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector"); const upscaleNodes = findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed"); - const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; + let allDistributedNodes = [...collectorNodes, ...upscaleNodes]; + + if (isMaster && delegateMaster) { + if (upscaleNodes.length > 0) { + extension.log("Delegate-only master mode does not support UltimateSDUpscaleDistributed nodes yet. Using full prompt.", "warn"); + } else if (collectorNodes.length === 0) { + extension.log("Delegate-only master mode requested but no collectors found in master prompt. Using full prompt.", "debug"); + } else { + jobApiPrompt = prepareMasterDelegatePrompt(extension, jobApiPrompt, collectorNodes); + collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector"); + allDistributedNodes = [...collectorNodes, ...upscaleNodes]; + } + } // For workers, handle platform-specific path conversion if (!isMaster) { @@ -282,6 +311,9 @@ export async function prepareApiPromptForParticipant(extension, baseApiPrompt, p inputs.is_worker = !isMaster; if (isMaster) { inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); + if (delegateMaster) { + inputs.delegate_only = true; + } } else { inputs.master_url = extension.getMasterUrl(); // Also make the worker_job_id unique to prevent potential caching issues @@ -323,6 +355,75 @@ export async function prepareDistributedJob(extension, multi_job_id) { } } +function createNumericIdGenerator(promptObj) { + let maxId = 0; + for (const key of Object.keys(promptObj)) { + const numeric = parseInt(key, 10); + if (!Number.isNaN(numeric)) { + maxId = Math.max(maxId, numeric); + } + } + return () => { + maxId += 1; + return String(maxId); + }; +} + +function prepareMasterDelegatePrompt(extension, apiPrompt, collectorNodes) { + const collectorIds = collectorNodes.map((node) => node.id); + const nodesToKeep = findCollectorDownstreamNodes(apiPrompt, collectorIds); + + // Ensure collectors themselves are present + collectorIds.forEach((id) => nodesToKeep.add(id)); + + const prunedPrompt = {}; + nodesToKeep.forEach((nodeId) => { + const nodeData = apiPrompt[nodeId]; + if (nodeData) { + prunedPrompt[nodeId] = JSON.parse(JSON.stringify(nodeData)); + } + }); + + const prunedIds = new Set(Object.keys(prunedPrompt)); + + // Remove dangling references to trimmed nodes + for (const [nodeId, node] of Object.entries(prunedPrompt)) { + if (!node.inputs) continue; + for (const [inputName, inputValue] of Object.entries(node.inputs)) { + if (Array.isArray(inputValue) && inputValue.length === 2) { + const sourceId = String(inputValue[0]); + if (!prunedIds.has(sourceId)) { + delete node.inputs[inputName]; + extension.log(`Removed upstream reference '${inputName}' from node ${nodeId} for delegate-only master prompt.`, "debug"); + } + } + } + } + + const nextId = createNumericIdGenerator(prunedPrompt); + collectorIds.forEach((collectorId) => { + const collectorEntry = prunedPrompt[collectorId]; + if (!collectorEntry) return; + + const placeholderId = nextId(); + prunedPrompt[placeholderId] = { + class_type: "DistributedEmptyImage", + inputs: { + height: 64, + width: 64, + channels: 3 + } + }; + prunedIds.add(placeholderId); + + collectorEntry.inputs = collectorEntry.inputs || {}; + collectorEntry.inputs.images = [placeholderId, 0]; + }); + + extension.log(`Prepared delegate-only master prompt with ${Object.keys(prunedPrompt).length} nodes (kept ${nodesToKeep.size}).`, "debug"); + return prunedPrompt; +} + export async function executeJobs(extension, jobs) { let masterPromptId = null; diff --git a/web/main.js b/web/main.js index 1f0b007..f2c89ae 100644 --- a/web/main.js +++ b/web/main.js @@ -90,6 +90,34 @@ class DistributedExtension { return this.enabledWorkers.length > 0; } + isMasterParticipationEnabled() { + return !Boolean(this.config?.settings?.master_delegate_only); + } + + isMasterFallbackActive() { + return Boolean(this.config?.settings?.master_delegate_only) && this.enabledWorkers.length === 0; + } + + isMasterParticipating() { + return this.isMasterParticipationEnabled() || this.isMasterFallbackActive(); + } + + async updateMasterParticipation(enabled) { + if (!this.config?.settings) { + this.config.settings = {}; + } + const delegateOnly = !enabled; + if (this.config.settings.master_delegate_only === delegateOnly) { + return; + } + + await this._updateSetting('master_delegate_only', delegateOnly); + + if (this.panelElement) { + renderSidebarContent(this, this.panelElement); + } + } + async loadConfig() { try { this.config = await this.api.getConfig(); @@ -143,6 +171,10 @@ class DistributedExtension { } catch (error) { this.log("Error updating worker: " + error.message, "error"); } + + if (this.panelElement) { + renderSidebarContent(this, this.panelElement); + } } async _updateSetting(key, value) { @@ -300,11 +332,19 @@ class DistributedExtension { // Update master status dot const statusDot = document.getElementById('master-status'); if (statusDot) { - if (isProcessing) { - statusDot.style.backgroundColor = "#f0ad4e"; + if (!this.isMasterParticipating()) { + if (isProcessing) { + statusDot.style.backgroundColor = STATUS_COLORS.PROCESSING_YELLOW; + statusDot.title = `Orchestrating (${queueRemaining} in queue)`; + } else { + statusDot.style.backgroundColor = STATUS_COLORS.DISABLED_GRAY; + statusDot.title = "Master orchestrator only"; + } + } else if (isProcessing) { + statusDot.style.backgroundColor = STATUS_COLORS.PROCESSING_YELLOW; statusDot.title = `Processing (${queueRemaining} in queue)`; } else { - statusDot.style.backgroundColor = "#4CAF50"; + statusDot.style.backgroundColor = STATUS_COLORS.ONLINE_GREEN; statusDot.title = "Online"; } } @@ -313,8 +353,8 @@ class DistributedExtension { // Master is always online (we're running on it), so keep it green const statusDot = document.getElementById('master-status'); if (statusDot) { - statusDot.style.backgroundColor = "#4CAF50"; - statusDot.title = "Online"; + statusDot.style.backgroundColor = this.isMasterParticipating() ? STATUS_COLORS.ONLINE_GREEN : STATUS_COLORS.DISABLED_GRAY; + statusDot.title = this.isMasterParticipating() ? "Online" : "Master orchestrator only"; } } } diff --git a/web/ui.js b/web/ui.js index 46fbe73..04d5a8d 100644 --- a/web/ui.js +++ b/web/ui.js @@ -3,28 +3,32 @@ import { BUTTON_STYLES, UI_STYLES, STATUS_COLORS, UI_COLORS, TIMEOUTS } from './ const cardConfigs = { master: { checkbox: { - enabled: false, - checked: true, - disabled: true, - opacity: 0.6, - title: "Master node is always enabled" + enabled: true, + masterToggle: true, + title: "Toggle master participation in workloads" }, statusDot: { - color: STATUS_COLORS.ONLINE_GREEN, - title: 'Online', id: 'master-status', + initialColor: (_, extension) => extension.isMasterParticipating() ? STATUS_COLORS.ONLINE_GREEN : STATUS_COLORS.DISABLED_GRAY, + initialTitle: (_, extension) => extension.isMasterParticipating() ? 'Master participating' : 'Master orchestrator only', dynamic: true }, infoText: (data, extension) => { const cudaDevice = extension.config?.master?.cuda_device ?? extension.masterCudaDevice; const cudaInfo = cudaDevice !== undefined ? `CUDA ${cudaDevice} • ` : ''; const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); - return `${data?.name || extension.config?.master?.name || "Master"}
${cudaInfo}Port ${port}`; + const participationEnabled = extension.isMasterParticipationEnabled(); + const fallbackActive = extension.isMasterFallbackActive(); + let delegateBadge = ''; + if (!participationEnabled) { + delegateBadge = fallbackActive + ? `
Fallback active • Master executing` + : `
Orchestrator-only mode`; + } + return `${data?.name || extension.config?.master?.name || "Master"}
${cudaInfo}Port ${port}${delegateBadge}`; }, controls: { - type: 'info', - text: 'Master', - style: "background-color: #333; color: #999;" + type: 'master' }, settings: { formType: 'master', @@ -868,7 +872,36 @@ export class DistributedUI { if (config?.opacity) checkbox.style.opacity = config.opacity; if (config?.title) column.title = config.title; - if (config?.enabled && !config?.disabled && data?.id) { + const isMasterToggle = config?.masterToggle && typeof extension.isMasterParticipating === 'function'; + if (isMasterToggle) { + const participationEnabled = extension.isMasterParticipationEnabled(); + const fallbackActive = extension.isMasterFallbackActive(); + const buildTitle = (enabled, fallback) => { + if (enabled) { + return "Master participating • Click to switch to orchestrator-only"; + } + if (fallback) { + return "No workers selected • Master fallback execution active"; + } + return "Master orchestrator-only • Click to re-enable participation"; + }; + + checkbox.checked = participationEnabled; + checkbox.style.pointerEvents = "none"; + column.style.cursor = "pointer"; + column.title = buildTitle(participationEnabled, fallbackActive); + column.onclick = async (event) => { + if (event) { + event.stopPropagation(); + event.preventDefault(); + } + const nextState = !extension.isMasterParticipationEnabled(); + const nextFallback = !nextState && extension.enabledWorkers.length === 0; + checkbox.checked = nextState; + column.title = buildTitle(nextState, nextFallback); + await extension.updateMasterParticipation(nextState); + }; + } else if (config?.enabled && !config?.disabled && data?.id) { checkbox.style.pointerEvents = "none"; column.style.cursor = "pointer"; column.onclick = async () => { @@ -890,10 +923,10 @@ export class DistributedUI { let id = config.id; if (typeof config.initialColor === 'function') { - color = config.initialColor(data); + color = config.initialColor(data, extension); } if (typeof config.initialTitle === 'function') { - title = config.initialTitle(data); + title = config.initialTitle(data, extension); } if (typeof config.id === 'function') { id = config.id(data); @@ -940,7 +973,29 @@ export class DistributedUI { const controlsWrapper = document.createElement("div"); controlsWrapper.style.cssText = this.styles.controlsWrapper; - if (config.dynamic && data) { + if (config.type === 'master') { + const participationEnabled = extension.isMasterParticipationEnabled(); + const fallbackActive = extension.isMasterFallbackActive(); + let message; + const badge = document.createElement("div"); + badge.style.cssText = this.styles.infoBox; + if (fallbackActive) { + message = "No workers selected. Master fallback execution active."; + badge.textContent = message; + badge.style.backgroundColor = "#243024"; + badge.style.color = "#6bd06b"; + badge.style.border = "1px solid #335533"; + } else if (!participationEnabled) { + message = "Master disabled: running as orchestrator only."; + badge.textContent = message; + badge.style.backgroundColor = "#3a3a3a"; + badge.style.color = "#ffcc66"; + } else { + message = "Master participating in workflows."; + badge.textContent = message; + } + controlsWrapper.appendChild(badge); + } else if (config.dynamic && data) { if (isRemote) { const isCloud = data.type === 'cloud'; const workerTypeText = isCloud ? "Cloud worker" : "Remote worker"; diff --git a/web/workerUtils.js b/web/workerUtils.js index 194f523..d2f8831 100644 --- a/web/workerUtils.js +++ b/web/workerUtils.js @@ -155,6 +155,47 @@ export function findCollectorUpstreamNodes(apiPrompt, collectorIds) { return connected; } +/** + * Find all downstream nodes from the provided collector IDs. + * Used to keep only post-collection nodes when master runs in orchestrator mode. + * @param {Object} apiPrompt + * @param {Array} collectorIds + * @returns {Set} + */ +export function findCollectorDownstreamNodes(apiPrompt, collectorIds) { + const adjacency = new Map(); + + for (const [nodeId, node] of Object.entries(apiPrompt)) { + if (!node.inputs) continue; + for (const inputValue of Object.values(node.inputs)) { + if (Array.isArray(inputValue) && inputValue.length === 2) { + const sourceId = String(inputValue[0]); + if (!adjacency.has(sourceId)) { + adjacency.set(sourceId, new Set()); + } + adjacency.get(sourceId).add(String(nodeId)); + } + } + } + + const connected = new Set(collectorIds); + const queue = [...collectorIds]; + + while (queue.length > 0) { + const current = queue.shift(); + const dependents = adjacency.get(current); + if (!dependents) continue; + dependents.forEach((dependentId) => { + if (!connected.has(dependentId)) { + connected.add(dependentId); + queue.push(dependentId); + } + }); + } + + return connected; +} + /** * Prune workflow to only include nodes connected to distributed nodes * @param {Object} apiPrompt - The full workflow API prompt