Merge pull request #50 from IT-BillDeng/feature/disable-master

Feature/disable master

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",
 }

docs/worker-setup-guides.md

@@ -6,6 +6,14 @@
 
 <img width="600" src="https://github.com/user-attachments/assets/609c42aa-8a1c-4a3f-939e-f3552fa1d54f" />
 
+### 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

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

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;
     

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";
             }
         }
     }

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 `<strong id="master-name-display">${data?.name || extension.config?.master?.name || "Master"}</strong><br><small style="color: ${UI_COLORS.MUTED_TEXT};"><span id="master-cuda-info">${cudaInfo}Port ${port}</span></small>`;
+            const participationEnabled = extension.isMasterParticipationEnabled();
+            const fallbackActive = extension.isMasterFallbackActive();
+            let delegateBadge = '';
+            if (!participationEnabled) {
+                delegateBadge = fallbackActive
+                    ? `<br><small style="color: #6bd06b;">Fallback active • Master executing</small>`
+                    : `<br><small style="color: ${UI_COLORS.SECONDARY_TEXT};">Orchestrator-only mode</small>`;
+            }
+            return `<strong id="master-name-display">${data?.name || extension.config?.master?.name || "Master"}</strong><br><small style="color: ${UI_COLORS.MUTED_TEXT};"><span id="master-cuda-info">${cudaInfo}Port ${port}</span></small>${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";

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<string>} collectorIds
+ * @returns {Set<string>}
+ */
+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