From 3c5e361d1da4d305481a72184947bbdd81a1d0eb Mon Sep 17 00:00:00 2001
From: speedyrulz <137450672+speedyrulz@users.noreply.github.com>
Date: Tue, 20 Oct 2026 20:32:44 -0400
Subject: [PATCH 1/3] fix: audit follow-ups for the next-free-GPU dispatcher

Lifecycle:
- Queue clears and deletes leave a run whose master prompt is currently
  executing alone, as ComfyUI leaves the running prompt alone; previously an
  executing orchestrator-only head run could be cancelled out from under its
  waiting Collector.
- The sweep now also drops worker-assigned runs whose queued delegate prompt
  left the queue without the Collector finishing, releasing their result
  queues instead of leaking them and polling forever.
- Completed runs release their result queues, and orchestration no longer
  pre-creates result queues when load balance picks the master alone.

Transport:
- A websocket timeout during connect, before anything was sent, is a plain
  failure the dispatcher may retry elsewhere; only timeouts after the prompt
  was sent remain AmbiguousDispatchError.

Efficiency and cleanup:
- One WorkerSnapshot per worker per cycle answers every run's "is my job
  there" question (queue first, history only when needed).
- Sweep and cancellation snapshot queued/running prompt ids once per pass.
- probe_worker_queues reuses dispatch._probe_worker_queue; the dead
  aggregate idle/busy probe helper is gone.

Rebased on "never replay load-balanced runs on missing history": an
unresolved submission stays on hold until delivery or positive evidence,
so there is no rejection path to fix and no bounded hold.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 api/orchestration/dispatch.py            |   6 +
 api/orchestration/load_balance_queue.py  | 233 ++++++++++++++---------
 api/queue_orchestration.py               |   6 +-
 tests/test_dispatch_selection.py         |  13 ++
 tests/test_load_balance_orchestration.py |  19 +-
 tests/test_load_balance_queue.py         | 154 +++++++++++++--
 6 files changed, 321 insertions(+), 110 deletions(-)

diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py
index fc542a1..7b1a439 100644
--- a/api/orchestration/dispatch.py
+++ b/api/orchestration/dispatch.py
@@ -98,6 +98,12 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
     try:
         async with asyncio.timeout_at(deadline):
             ws = await session.ws_connect(ws_url, heartbeat=20)
+    except asyncio.TimeoutError as exc:
+        # Nothing was sent: a definite failure the caller may retry elsewhere.
+        raise RuntimeError("Worker websocket connect timed out before dispatch; nothing was sent") from exc
+
+    try:
+        async with asyncio.timeout_at(deadline):
             await ws.send_json(ws_payload)
             while True:
                 msg = await ws.receive()
diff --git a/api/orchestration/load_balance_queue.py b/api/orchestration/load_balance_queue.py
index a8e7a32..e2f3133 100644
--- a/api/orchestration/load_balance_queue.py
+++ b/api/orchestration/load_balance_queue.py
@@ -14,9 +14,9 @@ A delegate prompt blocks ComfyUI's single executor while it waits for its
 worker, which would stall every full workflow queued behind it on a
 participating master. So on a participating master the delegate prompt is
 queued only once the worker's final result has arrived (or the worker is
-seen idle or unreachable without delivering, so the Collector can report
-it); an orchestrator-only master has nothing else to run and queues it at
-once.
+seen to have finished or lost the run without delivering, so the Collector
+can report it); an orchestrator-only master has nothing else to run and
+queues it at once.
 
 The claim is atomic through ComfyUI's queue mutex: a run can be removed from
 the queue by the dispatcher or popped by the master executor, never both.
@@ -37,7 +37,8 @@ Two lifecycle rules keep a run from executing twice or vanishing:
   run itself. Work that was not yet submitted is dropped along with its
   result queues; work a worker already accepted stays tracked until its
   result arrives or the worker is seen done, and is then discarded instead
-  of reappearing.
+  of reappearing. A run whose master prompt is currently executing is left
+  alone by clears and deletes, exactly as ComfyUI leaves the running prompt.
 """
 import asyncio
 import threading
@@ -47,8 +48,8 @@ from collections import OrderedDict
 import server
 
 from ...utils.logging import debug_log, log
-from ...utils.network import build_worker_url, probe_worker
-from .dispatch import _extract_queue_remaining
+from ...utils.network import build_worker_url
+from .dispatch import _probe_worker_queue
 from .dispatch import _select_idle_round_robin
 
 try:
@@ -74,7 +75,6 @@ REQUEUE_RETRY_MIN_SECONDS = 1.0
 REQUEUE_RETRY_MAX_SECONDS = 30.0
 # Consecutive failed probes before a dispatched worker counts as gone.
 WORKER_UNREACHABLE_PROBES = 3
-
 # How far back the worker's history is searched for an unresolved run.
 HISTORY_LOOKBACK_ITEMS = 200
 # Job ids whose result queues were discarded; a late worker must not recreate them.
@@ -90,19 +90,9 @@ def _worker_key(worker):
 async def probe_worker_queues(workers, probe_concurrency=8, probe_timeout=3.0):
     """Return ``[{"worker", "queue_remaining"}]`` for every worker that answered."""
     semaphore = asyncio.Semaphore(max(1, int(probe_concurrency or 1)))
-
-    async def _probe(worker):
-        async with semaphore:
-            if worker.get("type") == "serverless":
-                from .serverless_runtime import get_manager
-                status = await get_manager().health(worker)
-                return {"worker": worker, "queue_remaining": status["queueRemaining"]}
-            payload = await probe_worker(build_worker_url(worker), timeout=probe_timeout)
-            if payload is None:
-                return None
-            return {"worker": worker, "queue_remaining": _extract_queue_remaining(payload)}
-
-    statuses = await asyncio.gather(*[_probe(worker) for worker in workers])
+    statuses = await asyncio.gather(
+        *[_probe_worker_queue(worker, semaphore, probe_timeout) for worker in workers]
+    )
     return [status for status in statuses if status is not None]
 
 
@@ -118,51 +108,87 @@ def _graph_has_job(graph, job_ids):
     return False
 
 
-async def fetch_worker_job_state(worker, job_ids, timeout=5.0):
-    """Ask the worker itself whether a specific run is there.
+class WorkerSnapshot:
+    """What one worker says about the runs it holds, fetched at most once.
 
-    Returns ``"running"`` when a prompt carrying one of ``job_ids`` is in the
-    worker's running or pending queue, ``"finished"`` when it is in the
-    worker's recent history, ``"absent"`` when neither snapshot contains it,
-    and ``None`` when the worker could not be asked. Absence is not proof that
-    a submission was rejected: history can be truncated, cleared or lost on
-    restart. Direct workers only; serverless pools never reach this path.
+    The queue (running and pending prompts) is fetched on first use; the
+    history, which is larger, only when a run is not found in the queue.
+    One snapshot serves every run on that worker within a dispatcher cycle.
+    Direct workers only; serverless pools never reach this path.
     """
-    import aiohttp
-    from ...utils.network import get_client_session
 
-    job_ids = {str(job_id) for job_id in job_ids}
-    base_url = build_worker_url(worker)
-    session = await get_client_session()
-    try:
-        async with session.get(f"{base_url}/queue", timeout=aiohttp.ClientTimeout(total=timeout)) as response:
-            if response.status != 200:
-                return None
-            queue_info = await response.json()
-        if not isinstance(queue_info, dict):
-            return None
-        for key in ("queue_running", "queue_pending"):
-            for item in queue_info.get(key) or []:
-                if isinstance(item, (list, tuple)) and len(item) > 2 and _graph_has_job(item[2], job_ids):
-                    return "running"
+    def __init__(self, worker, timeout=5.0):
+        self.worker = worker
+        self.timeout = timeout
+        self._queue_graphs = None
+        self._history_graphs = None
+        self.failed = False
+
+    async def _fetch(self, endpoint, params=None):
+        import aiohttp
+        from ...utils.network import get_client_session
+
+        session = await get_client_session()
         async with session.get(
-            f"{base_url}/history",
-            params={"max_items": str(HISTORY_LOOKBACK_ITEMS)},
-            timeout=aiohttp.ClientTimeout(total=timeout),
+            f"{build_worker_url(self.worker)}{endpoint}",
+            params=params,
+            timeout=aiohttp.ClientTimeout(total=self.timeout),
         ) as response:
             if response.status != 200:
-                return None
-            history = await response.json()
-        if not isinstance(history, dict):
+                raise RuntimeError(f"{endpoint} returned HTTP {response.status}")
+            payload = await response.json()
+        if not isinstance(payload, dict):
+            raise RuntimeError(f"{endpoint} returned a non-object body")
+        return payload
+
+    async def _queue(self):
+        if self._queue_graphs is None:
+            queue_info = await self._fetch("/queue")
+            graphs = []
+            for key in ("queue_running", "queue_pending"):
+                for item in queue_info.get(key) or []:
+                    if isinstance(item, (list, tuple)) and len(item) > 2:
+                        graphs.append(item[2])
+            self._queue_graphs = graphs
+        return self._queue_graphs
+
+    async def _history(self):
+        if self._history_graphs is None:
+            history = await self._fetch("/history", params={"max_items": str(HISTORY_LOOKBACK_ITEMS)})
+            graphs = []
+            for entry in history.values():
+                item = entry.get("prompt") if isinstance(entry, dict) else None
+                if isinstance(item, (list, tuple)) and len(item) > 2:
+                    graphs.append(item[2])
+            self._history_graphs = graphs
+        return self._history_graphs
+
+    async def job_state(self, job_ids):
+        """``"running"`` (in the worker's queue), ``"finished"`` (in its recent
+        history), ``"absent"`` (in neither), or ``None`` when the worker could
+        not be asked.
+
+        Absence is not proof that a submission was rejected: history can be
+        truncated, cleared or lost on restart.
+        """
+        if self.failed:
             return None
-        for entry in history.values():
-            item = entry.get("prompt") if isinstance(entry, dict) else None
-            if isinstance(item, (list, tuple)) and len(item) > 2 and _graph_has_job(item[2], job_ids):
+        job_ids = {str(job_id) for job_id in job_ids}
+        try:
+            if any(_graph_has_job(graph, job_ids) for graph in await self._queue()):
+                return "running"
+            if any(_graph_has_job(graph, job_ids) for graph in await self._history()):
                 return "finished"
-        return "absent"
-    except Exception as exc:
-        debug_log(f"Load balance: could not check worker {_worker_key(worker)} for a run: {exc}")
-        return None
+            return "absent"
+        except Exception as exc:
+            self.failed = True
+            debug_log(f"Load balance: could not check worker {_worker_key(self.worker)} for a run: {exc}")
+            return None
+
+
+async def fetch_worker_job_state(worker, job_ids, timeout=5.0):
+    """Ask the worker itself whether a specific run is there (see ``WorkerSnapshot``)."""
+    return await WorkerSnapshot(worker, timeout).job_state(job_ids)
 
 
 async def select_idle_worker(workers, trace_execution_id=None, probe_concurrency=8, probe_timeout=3.0):
@@ -260,6 +286,11 @@ class PendingRun:
         """Out of the master queue for a hand-off: cancellation must be recorded here."""
         return self.claiming or self.on_worker
 
+    @property
+    def collector_unqueued(self):
+        """Out of the master queue on purpose: being handed off, held unresolved, or awaiting its delegate."""
+        return self.claiming or self.unresolved is not None or self.requeue_pending
+
     def accepts(self, worker):
         key = _worker_key(worker)
         return any(_worker_key(candidate) == key for candidate in self.candidates)
@@ -275,6 +306,7 @@ class LoadBalanceQueue:
         self._by_job = {}  # job_id -> PendingRun
         self._recent_dispatch = {}  # worker id -> monotonic time
         self._cancelled_jobs = OrderedDict()
+        self._snapshots = {}  # worker id -> WorkerSnapshot, valid for one cycle
         self._wake = asyncio.Event()
         self._cycle_lock = asyncio.Lock()
         self._task = None
@@ -377,13 +409,17 @@ class LoadBalanceQueue:
             if self._by_job.get(job_id) is record:
                 self._by_job.pop(job_id, None)
 
+    def _release_result_queues(self, record):
+        pending_jobs = self._pending_jobs_getter()
+        if isinstance(pending_jobs, dict):
+            for job_id in record.job_ids:
+                pending_jobs.pop(job_id, None)
+
     def _discard(self, record):
         """Drop the run and the result queues nobody will consume any more."""
         self._remove(record)
-        pending_jobs = self._pending_jobs_getter()
+        self._release_result_queues(record)
         for job_id in record.job_ids:
-            if isinstance(pending_jobs, dict):
-                pending_jobs.pop(job_id, None)
             self._cancelled_jobs[job_id] = True
             self._cancelled_jobs.move_to_end(job_id)
         while len(self._cancelled_jobs) > CANCELLED_JOB_MEMORY:
@@ -430,10 +466,15 @@ class LoadBalanceQueue:
         self._hooks_installed = True
 
     def cancel_matching(self, predicate):
-        """Cancel in-flight runs whose queue item ``predicate`` would have deleted."""
+        """Cancel in-flight runs whose queue item ``predicate`` would have deleted.
+
+        A run whose master prompt is currently executing is left alone, as
+        ComfyUI's clear and delete leave the running prompt alone.
+        """
         cancelled = []
+        running = self._running_ids()
         for record in list(self._records.values()):
-            if record.cancelled or not record.in_flight:
+            if record.cancelled or not record.in_flight or record.prompt_id in running:
                 continue
             item = record.unresolved_item or (0, record.prompt_id, {}, {}, [])
             try:
@@ -503,6 +544,8 @@ class LoadBalanceQueue:
         On failure or interruption the run is cancelled: a hand-off still in
         flight is stopped before submission, and work a worker already
         accepted is discarded when it ends, so no result queue is left behind.
+        On success the run is forgotten once every job finished, including the
+        result queues a master-run never needed.
         """
         if failed:
             if record.prompt_id in self._records:
@@ -511,20 +554,27 @@ class LoadBalanceQueue:
         record.finished_jobs.add(str(job_id))
         if record.finished_jobs >= record.job_ids:
             self._remove(record)
+            self._release_result_queues(record)
 
     # ----- queue inspection ---------------------------------------------
     def _prompt_queue(self):
         return self._prompt_queue_getter()
 
-    def _is_running(self, prompt_id):
+    def _running_ids(self):
         prompt_queue = self._prompt_queue()
         with prompt_queue.mutex:
-            return any(item[1] == prompt_id for item in prompt_queue.currently_running.values())
+            return {item[1] for item in prompt_queue.currently_running.values()}
+
+    def _queued_ids(self):
+        prompt_queue = self._prompt_queue()
+        with prompt_queue.mutex:
+            return {item[1] for item in prompt_queue.queue}
+
+    def _is_running(self, prompt_id):
+        return prompt_id in self._running_ids()
 
     def _is_queued(self, prompt_id):
-        prompt_queue = self._prompt_queue()
-        with prompt_queue.mutex:
-            return any(item[1] == prompt_id for item in prompt_queue.queue)
+        return prompt_id in self._queued_ids()
 
     def _take_from_queue(self, prompt_id):
         """Remove the prompt from the master queue; return the item or None."""
@@ -579,13 +629,20 @@ class LoadBalanceQueue:
             self._task = None
 
     def sweep(self):
-        """Drop runs that left the master queue without being collected."""
+        """Drop runs that left the master queue without being collected.
+
+        Runs that are out of the queue on purpose (being handed off, held
+        unresolved, or waiting for their delegate prompt to be queued) and
+        cancelled runs are left to the dispatcher's own bookkeeping.
+        """
+        queued = self._queued_ids()
+        running = self._running_ids()
         for record in list(self._records.values()):
-            if not record.queued or record.in_flight or record.cancelled:
+            if not record.queued or record.collector_unqueued or record.cancelled:
                 continue
-            if self._is_queued(record.prompt_id):
+            if record.prompt_id in queued:
                 continue
-            if self._is_running(record.prompt_id):
+            if record.prompt_id in running:
                 if record.assigned is None and record.master_may_run:
                     record._assign(MASTER)
                     debug_log(f"Load balance: master is running run {record.prompt_id}")
@@ -605,7 +662,11 @@ class LoadBalanceQueue:
     async def run_cycle(self, probe_concurrency=8, probe_timeout=3.0):
         """One dispatcher pass: sweep, resolve, retry, then hand runs to idle workers."""
         async with self._cycle_lock:
-            return await self._run_cycle(probe_concurrency, probe_timeout)
+            self._snapshots = {}
+            try:
+                return await self._run_cycle(probe_concurrency, probe_timeout)
+            finally:
+                self._snapshots = {}
 
     async def _run_cycle(self, probe_concurrency, probe_timeout):
         self.sweep()
@@ -662,14 +723,14 @@ class LoadBalanceQueue:
         record.delegate_needed = False
         return True
 
-    async def _probe_one(self, record, worker):
-        """Return 'unreachable', 'idle' or 'busy' for the worker running ``record``."""
-        statuses = await probe_worker_queues([worker])
-        if not statuses:
-            record.missed_probes += 1
-            return "unreachable"
-        record.missed_probes = 0
-        return "idle" if statuses[0]["queue_remaining"] == 0 else "busy"
+    async def _job_state(self, record, worker):
+        """Ask ``worker`` about ``record``'s jobs, sharing one snapshot per worker per cycle."""
+        key = _worker_key(worker)
+        snapshot = self._snapshots.get(key)
+        if snapshot is None:
+            snapshot = WorkerSnapshot(worker)
+            self._snapshots[key] = snapshot
+        return await snapshot.job_state(record.job_ids)
 
     async def _worker_done_without_result(self, record):
         """True when the worker no longer has this run and never delivered.
@@ -680,7 +741,7 @@ class LoadBalanceQueue:
         """
         if time.monotonic() - record.dispatched_at < RECENT_DISPATCH_TTL_SECONDS:
             return False
-        state = await fetch_worker_job_state(record.assigned, record.job_ids)
+        state = await self._job_state(record, record.assigned)
         if state is None:
             record.missed_probes += 1
             return record.missed_probes >= WORKER_UNREACHABLE_PROBES
@@ -695,7 +756,9 @@ class LoadBalanceQueue:
         an idle worker may have run the job and lost the delivery, a busy one
         may be on unrelated work. Missing history is not rejection evidence,
         even on repeated checks: accepted work may have aged out or been cleared.
+        A worker that cannot be asked keeps the run on hold as well.
         """
+        now = time.monotonic()
         for record in list(self._records.values()):
             worker = record.unresolved
             if worker is None:
@@ -703,9 +766,9 @@ class LoadBalanceQueue:
             if record.delivered:
                 self._accept_unresolved(record, "delivered its result")
                 continue
-            if time.monotonic() - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS:
+            if now - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS:
                 continue
-            state = await fetch_worker_job_state(worker, record.job_ids)
+            state = await self._job_state(record, worker)
             if state == "running":
                 self._accept_unresolved(record, "has it queued or running")
             elif state == "finished":
@@ -764,8 +827,8 @@ class LoadBalanceQueue:
                 if not await self._worker_done_without_result(record):
                     continue
                 log(
-                    f"[Distributed] Load balance: worker {record.assigned_id} is idle or unreachable without "
-                    f"delivering run {record.prompt_id[:8]}; queueing its collector so the run can report."
+                    f"[Distributed] Load balance: worker {record.assigned_id} no longer has run "
+                    f"{record.prompt_id[:8]} and never delivered it; queueing its collector so the run can report."
                 )
                 record.wait_for_result = False
             if record.requeue_retry_at > now:
diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py
index b9d47eb..428712e 100644
--- a/api/queue_orchestration.py
+++ b/api/queue_orchestration.py
@@ -566,6 +566,9 @@ async def orchestrate_distributed_execution(
     # collect-only prompt once the worker delivers, not now, so the master
     # executor stays free for the next full workflow.
     hold_delegate_for_worker = None
+    # Load balance picked the master alone: no worker will ever deliver, so
+    # no result queues are needed for this run.
+    master_runs_alone = False
     if load_balance_requested and not image_list_mode:
         candidate_workers = list(active_workers)
         if not delegate_master:
@@ -617,6 +620,7 @@ async def orchestrate_distributed_execution(
             # Master selected as least busy; run master workload only.
             active_workers = []
             delegate_master = False
+            master_runs_alone = True
             trace_debug(
                 execution_trace_id,
                 "Load-balance selected master for execution (workers skipped).",
@@ -711,7 +715,7 @@ async def orchestrate_distributed_execution(
         manager.ensure_capacity(len(serverless_jobs))
 
     # Do not leave collector queues behind when offline compilation/media validation fails.
-    if not image_list_mode:
+    if not image_list_mode and not master_runs_alone:
         for job_id in job_id_map.values():
             await _ensure_distributed_queue(job_id)
 
diff --git a/tests/test_dispatch_selection.py b/tests/test_dispatch_selection.py
index 9bb62eb..5658b92 100644
--- a/tests/test_dispatch_selection.py
+++ b/tests/test_dispatch_selection.py
@@ -196,6 +196,19 @@ class ExplicitTargetTransportTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(calls[0]['partial_execution_targets'], ['9'])
         self.assertNotIn('partial_execution_targets', calls[1])
 
+    async def test_websocket_connect_timeout_is_a_definite_failure(self):
+        """A timeout before anything was sent must not read as ambiguous acceptance."""
+        never = asyncio.Event()
+
+        async def hang(*_args, **_kwargs):
+            await never.wait()
+
+        session = types.SimpleNamespace(ws_connect=hang)
+        with patch.object(dispatch, "get_client_session", AsyncMock(return_value=session)):
+            with self.assertRaisesRegex(RuntimeError, "nothing was sent") as caught:
+                await dispatch._dispatch_via_websocket("http://worker", {"prompt": {}}, "client", timeout=0.05)
+        self.assertNotIsInstance(caught.exception, dispatch.AmbiguousDispatchError)
+
     async def test_websocket_sends_targets_all_the_way_to_wire(self):
         import json
         sent = []
diff --git a/tests/test_load_balance_orchestration.py b/tests/test_load_balance_orchestration.py
index 6f1f91d..73902b3 100644
--- a/tests/test_load_balance_orchestration.py
+++ b/tests/test_load_balance_orchestration.py
@@ -188,6 +188,14 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(self.lb_queue.records[0].unresolved["id"], "w1")
         self.q.queue_prompt_payload.assert_not_awaited()
 
+    async def test_master_selected_alone_creates_no_result_queues(self):
+        self.select_idle_worker.return_value = ({"id": "master", "name": "Master", "type": "local"}, 2)
+        _, _, worker_count, _ = await self._run()
+        self.assertEqual(worker_count, 0)
+        self.q.queue_prompt_payload.assert_awaited_once()
+        self.assertEqual(self.ps.distributed_pending_jobs, {})
+        self.assertEqual(self.lb_queue.records, [])
+
     async def test_orchestrator_only_master_queues_delegate_immediately(self):
         self.settings["master_delegate_only"] = True
         self.select_idle_worker.return_value = (self.worker, 1)
@@ -223,18 +231,17 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
         network.build_worker_url = lambda worker, *a: f"http://{worker['id']}"
         network.probe_worker = AsyncMock(return_value=None)
         dispatch = sys.modules[PKG + ".api.orchestration.dispatch"]
-        dispatch._extract_queue_remaining = lambda payload: int(payload["exec_info"]["queue_remaining"])
+
+        async def always_busy_probe(worker, semaphore, probe_timeout):
+            return {"worker": worker, "queue_remaining": 1}
+
+        dispatch._probe_worker_queue = always_busy_probe
         dispatch._select_idle_round_robin = lambda statuses: statuses[0]
         spec = importlib.util.spec_from_file_location(
             PKG + ".api.orchestration.load_balance_queue", ROOT / "api" / "orchestration" / "load_balance_queue.py")
         lb = importlib.util.module_from_spec(spec)
         spec.loader.exec_module(lb)
         lb.select_idle_worker = AsyncMock(return_value=(None, 1))
-
-        async def always_busy(url, timeout=3.0):
-            return {"exec_info": {"queue_remaining": 1}}
-
-        lb.probe_worker = always_busy
         sys.modules[PKG + ".api.orchestration.load_balance_queue"] = lb
         queue = lb.get_load_balance_queue(self.ps)
         queue.poll_interval = 0.01
diff --git a/tests/test_load_balance_queue.py b/tests/test_load_balance_queue.py
index e90f161..5feb762 100644
--- a/tests/test_load_balance_queue.py
+++ b/tests/test_load_balance_queue.py
@@ -12,6 +12,8 @@ from unittest.mock import AsyncMock, patch
 
 ROOT = Path(__file__).resolve().parents[1]
 PKG = "dist_lb_queue_testpkg"
+# The running test case; the dispatch stub's probe reads its ``remaining`` table.
+_ACTIVE = {"case": None}
 
 
 def _module(name, **attrs):
@@ -84,9 +86,15 @@ def _load_module(prompt_queue):
     sys.modules[PKG + ".utils.network"] = _module(
         "network", build_worker_url=lambda worker, *a: f"http://{worker['id']}", probe_worker=AsyncMock(return_value=None)
     )
+    async def _probe_worker_queue(worker, semaphore, probe_timeout):
+        async with semaphore:
+            case = _ACTIVE["case"]
+            value = case.remaining.get(worker["id"]) if case is not None else None
+            return None if value is None else {"worker": worker, "queue_remaining": value}
+
     sys.modules[PKG + ".api.orchestration.dispatch"] = _module(
         "dispatch",
-        _extract_queue_remaining=lambda payload: int(payload.get("exec_info", {}).get("queue_remaining", 0)),
+        _probe_worker_queue=_probe_worker_queue,
         _select_idle_round_robin=lambda statuses: statuses[0],
     )
     spec = importlib.util.spec_from_file_location(
@@ -111,35 +119,37 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.dispatched = []
         self.prepared = []
         self.requeued = []
-
-        async def fake_probe(url, timeout=3.0):
-            worker_id = url.rsplit("/", 1)[-1]
-            value = self.remaining.get(worker_id)
-            return None if value is None else {"exec_info": {"queue_remaining": value}}
-
-        self.lb.probe_worker = fake_probe
+        _ACTIVE["case"] = self
         # What a worker answers when asked about a specific run (its job ids):
         # "running" | "finished" | "absent" | None (could not be asked).
         self.job_state = {}
         self.job_state_queries = []
+        self.snapshots_taken = []
+        self.real_snapshot_class = self.lb.WorkerSnapshot
+        case = self
 
-        async def fake_job_state(worker, job_ids, timeout=5.0):
-            self.job_state_queries.append((worker["id"], set(job_ids)))
-            return self.job_state.get(worker["id"], "absent")
+        class FakeSnapshot:
+            def __init__(self, worker, timeout=5.0):
+                self.worker = worker
+                case.snapshots_taken.append(worker["id"])
 
-        self.fetch_worker_job_state = self.lb.fetch_worker_job_state
-        self.lb.fetch_worker_job_state = fake_job_state
+            async def job_state(self, job_ids):
+                case.job_state_queries.append((self.worker["id"], set(job_ids)))
+                return case.job_state.get(self.worker["id"], "absent")
+
+        self.lb.WorkerSnapshot = FakeSnapshot
 
     def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None, requeue_errors=0, queued=True,
                    prepare_error=None, prepare_gate=None, ambiguous=False, requeue_gate=None):
         """Register a run. ``prepare_gate`` (an Event) delays preparation until set;
-        ``ambiguous`` makes the send raise AmbiguousDispatchError after the worker
-        has in fact accepted the job (its acknowledgement was lost)."""
+        ``ambiguous`` makes the first send raise AmbiguousDispatchError after the
+        worker has in fact accepted the job (its acknowledgement was lost)."""
         if queued:
             number = len(self.prompt_queue.queue) + len(self.prompt_queue.currently_running)
             self.prompt_queue.put((number, prompt_id, {"full": prompt_id}, {}, [], {}))
         self.pending_jobs[f"job-{prompt_id}"] = asyncio.Queue()
         remaining_requeue_errors = [requeue_errors]
+        ambiguous_sends_left = [1 if ambiguous else 0]
 
         async def prepare_for_worker(worker):
             if prepare_gate is not None:
@@ -154,7 +164,8 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
             if dispatch_error is not None:
                 raise dispatch_error
             self.dispatched.append((prompt_id, worker["id"]))
-            if ambiguous:
+            if ambiguous_sends_left[0] > 0:
+                ambiguous_sends_left[0] -= 1
                 raise self.lb.AmbiguousDispatchError("acknowledgement lost")
 
         async def requeue_as_delegate(worker):
@@ -193,7 +204,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(record.unresolved["id"], "w1")
         self.assertIsNone(record.assigned)
         self.assertEqual(self.prompt_queue.prompt_ids(), [])  # not restored: the master cannot pick it up
-        record.unresolved_since = 0  # past the grace period
+        record.unresolved_since = time.monotonic() - self.lb.RECENT_DISPATCH_TTL_SECONDS - 1  # past the grace period
         self.queue._recent_dispatch.clear()
         return record
 
@@ -275,7 +286,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
                 raise AssertionError(f"Unexpected request: {url}")
 
         network = sys.modules[PKG + ".utils.network"]
-        self.lb.fetch_worker_job_state = self.fetch_worker_job_state
+        self.lb.WorkerSnapshot = self.real_snapshot_class
         with patch.object(network, "get_client_session", AsyncMock(return_value=Session()), create=True):
             for _ in range(5):
                 self.queue._recent_dispatch.clear()
@@ -523,6 +534,113 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.assertNotIn("job-p1", self.pending_jobs)
         self.assertEqual(self.prompt_queue.prompt_ids(), [])
 
+    # ----- audit regressions ---------------------------------------------------
+    async def test_clear_queue_leaves_an_executing_head_run_alone(self):
+        """ComfyUI's clear removes pending prompts only; the running one, whose Collector
+        is waiting for a worker, keeps its hand-off."""
+        gate = asyncio.Event()
+        record = self._queue_run("p1", master_may_run=False, prepare_gate=gate)
+        self.prompt_queue.start_next()
+        waiter = asyncio.ensure_future(self.queue.wait_for_assignment(record))
+        await asyncio.sleep(0)
+        self.remaining["w1"] = 0
+        cycle = asyncio.ensure_future(self.queue.run_cycle())
+        for _ in range(200):
+            if record.claiming:
+                break
+            await asyncio.sleep(0.005)
+        self.assertTrue(record.claiming)
+        self.prompt_queue.wipe_queue()
+        self.prompt_queue.delete_queue_item(lambda item: item[1] == "p1")
+        self.assertFalse(record.cancelled)
+        gate.set()
+        await cycle
+        self.assertEqual(self.dispatched, [("p1", "w1")])
+        assignment = await asyncio.wait_for(waiter, timeout=2)
+        self.assertEqual(assignment["id"], "w1")
+        self.assertIs(self.queue.lookup_job("job-p1"), record)
+
+    async def test_sweep_drops_assigned_run_whose_delegate_left_the_queue(self):
+        record = self._queue_run("p1", master_may_run=False)
+        self.remaining["w1"] = 0
+        await self.queue.run_cycle()
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])  # delegate queued
+        self.assertEqual(record.assigned_id, "w1")
+        # The delegate prompt is gone without its Collector ever finishing
+        # (an execution error before the Collector node, for example).
+        self.prompt_queue.queue = []
+        self.queue.sweep()
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+        self.assertNotIn("job-p1", self.pending_jobs)
+        self.assertEqual(len(self.queue), 0)
+
+    async def test_sweep_keeps_runs_whose_delegate_is_executing_or_still_held(self):
+        running = self._queue_run("p1", master_may_run=False)
+        held = self._queue_run("p2")
+        self.remaining = {"w1": 0, "w2": 0}
+        await self.queue.run_cycle()
+        self.assertEqual(sorted(self.dispatched), [("p1", "w1"), ("p2", "w2")])
+        self.prompt_queue.start_next()  # p1's delegate executes
+        self.queue.sweep()
+        self.assertIs(self.queue.lookup_job("job-p1"), running)
+        self.assertIs(self.queue.lookup_job("job-p2"), held)  # delegate held until delivery
+
+    async def test_master_run_completion_releases_its_result_queues(self):
+        record = self._queue_run("p1")
+        self.prompt_queue.start_next()
+        self.assertEqual(await self.queue.wait_for_assignment(record), self.lb.MASTER)
+        self.assertIn("job-p1", self.pending_jobs)
+        self.queue.finish_job(record, "job-p1")
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+        self.assertNotIn("job-p1", self.pending_jobs)
+        self.assertFalse(self.queue.is_cancelled_job("job-p1"))  # a normal completion, not a cancel
+
+    async def test_one_worker_snapshot_serves_every_run_in_a_cycle(self):
+        first = self._queue_run("p1")
+        second = self._queue_run("p2")
+        self.remaining["w1"] = 0
+        await self.queue.run_cycle()
+        self.queue._recent_dispatch.clear()
+        await self.queue.run_cycle()
+        self.assertEqual(sorted(self.dispatched), [("p1", "w1"), ("p2", "w1")])
+        first.dispatched_at = second.dispatched_at = 0
+        self.job_state["w1"] = "running"
+        self.snapshots_taken.clear()
+        self.job_state_queries.clear()
+        await self.queue.run_cycle()
+        self.assertEqual(self.snapshots_taken, ["w1"])  # one fetch, two questions
+        self.assertEqual(sorted(query[1] for query in self.job_state_queries), [{"job-p1"}, {"job-p2"}])
+        self.assertEqual(self.requeued, [])
+
+    async def test_worker_snapshot_reads_queue_then_history_lazily(self):
+        calls = []
+
+        class Snapshot(self.real_snapshot_class):
+            async def _fetch(self, endpoint, params=None):
+                calls.append(endpoint)
+                if endpoint == "/queue":
+                    return {"queue_running": [[1, "x", {"n": {"inputs": {"multi_job_id": "job-a"}}}]],
+                            "queue_pending": []}
+                return {"h1": {"prompt": [0, "y", {"n": {"inputs": {"multi_job_id": "job-b"}}}]}}
+
+        snapshot = Snapshot({"id": "w1"})
+        self.assertEqual(await snapshot.job_state({"job-a"}), "running")
+        self.assertEqual(calls, ["/queue"])
+        self.assertEqual(await snapshot.job_state({"job-b"}), "finished")
+        self.assertEqual(calls, ["/queue", "/history"])
+        self.assertEqual(await snapshot.job_state({"job-c"}), "absent")
+        self.assertEqual(calls, ["/queue", "/history"])  # both cached
+
+        class Broken(self.real_snapshot_class):
+            async def _fetch(self, endpoint, params=None):
+                calls.append(endpoint)
+                raise RuntimeError("connection refused")
+
+        broken = Broken({"id": "w2"})
+        self.assertIsNone(await broken.job_state({"job-a"}))
+        self.assertIsNone(await broken.job_state({"job-a"}))
+        self.assertEqual(calls.count("/queue"), 2)  # a failed worker is not asked again this cycle
+
     # ----- interrupted assignment on an orchestrator-only master ---------------
     async def test_interrupt_during_assignment_cleans_result_queue_and_stops_dispatch(self):
         gate = asyncio.Event()

From 18ff75500194021fe534ca7dd82bae5fae1bee21 Mon Sep 17 00:00:00 2001
From: speedyrulz <137450672+speedyrulz@users.noreply.github.com>
Date: Sat, 31 Oct 2026 16:22:45 -0400
Subject: [PATCH 2/3] feat: per-GPU preview outputs on the Distributed
 Collector

Append master and worker_1..worker_8 IMAGE outputs after images/audio so
each participant's own batch can be previewed. worker_N follows the Nth
configured worker so a slot keeps its GPU when others are toggled; slots
with no images return ExecutionBlocker(None) so attached nodes are
skipped silently. The frontend labels worker_N with the worker's name.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 nodes/collector.py                      | 64 ++++++++++++++++++++-
 nodes/v3.py                             | 19 +++++--
 tests/helpers/v3_runtime_check.py       | 13 ++++-
 tests/test_collector_image_transport.py | 60 +++++++++++++++++++
 tests/test_v3_collector_transfer.py     | 38 ++++++++++++-
 web/collectorOutputs.js                 | 76 +++++++++++++++++++++++++
 web/tests/collectorOutputs.test.js      | 45 +++++++++++++++
 7 files changed, 308 insertions(+), 7 deletions(-)
 create mode 100644 web/collectorOutputs.js
 create mode 100644 web/tests/collectorOutputs.test.js

diff --git a/nodes/collector.py b/nodes/collector.py
index 7aef1e8..7e982cb 100644
--- a/nodes/collector.py
+++ b/nodes/collector.py
@@ -21,6 +21,11 @@ from ..utils import transfer_codec
 
 prompt_server = _server.PromptServer.instance
 
+# Per-GPU preview outputs follow the combined images/audio outputs. worker_N is
+# the Nth worker in the config, so a slot keeps its GPU when others are toggled.
+MAX_WORKER_OUTPUTS = 8
+PER_GPU_OUTPUT_NAMES = ("master",) + tuple(f"worker_{i + 1}" for i in range(MAX_WORKER_OUTPUTS))
+
 
 def _encode_all_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes):
     """Materialize every transfer chunk (runs in an executor)."""
@@ -113,7 +118,12 @@ class DistributedCollectorNode:
     RETURN_NAMES = ("images", "audio")
     FUNCTION = "run"
     CATEGORY = "image"
-    
+
+    def __init__(self):
+        # Image batch each participant produced in this run, keyed by
+        # PER_GPU_OUTPUT_NAMES slot. Runtime instances are per-execution.
+        self.participant_images = {}
+
     @staticmethod
     def _unwrap_list_input(value):
         """Unwrap scalar inputs when ComfyUI passes them via INPUT_IS_LIST."""
@@ -232,6 +242,9 @@ class DistributedCollectorNode:
         if not multi_job_id or pass_through:
             if pass_through:
                 debug_log("Collector: pass-through mode enabled, returning images unchanged")
+            elif images is not None:
+                # Not distributed: everything was produced on this GPU.
+                self.participant_images = {"master": images}
             return (images, audio if audio is not None else empty_audio)
 
         # Use async helper to run in server loop
@@ -703,6 +716,49 @@ class DistributedCollectorNode:
             return self._combine_image_tensors([fallback_images], ["master fallback"])
         return None
 
+    @staticmethod
+    def _worker_output_slots(participant_ids):
+        """Map worker ids to worker_N preview slots.
+
+        Configured workers keep their config position; other participants
+        (e.g. serverless) take the free slots after the configured workers.
+        """
+        try:
+            config_ids = [str(worker.get("id")) for worker in load_config().get("workers", [])]
+        except Exception as e:
+            debug_log(f"Collector - Could not read worker order for preview outputs: {e}")
+            config_ids = []
+        slots = {}
+        next_free = len(config_ids)
+        for worker_id in participant_ids:
+            if worker_id in config_ids:
+                index = config_ids.index(worker_id)
+            else:
+                index, next_free = next_free, next_free + 1
+            if index < MAX_WORKER_OUTPUTS:
+                slots[worker_id] = PER_GPU_OUTPUT_NAMES[index + 1]
+            else:
+                debug_log(f"Collector - No preview output for worker {worker_id}; only in the combined images")
+        return slots
+
+    def _split_participant_images(self, worker_images, worker_order, master_images):
+        """Return each participant's own image batch keyed by preview slot."""
+        participants = {}
+        if master_images is not None:
+            participants["master"] = master_images
+        ordered = [str(w) for w in (worker_order or [])]
+        ordered += sorted(str(w) for w in worker_images if str(w) not in ordered)
+        images_by_id = {str(w): frames for w, frames in worker_images.items()}
+        for worker_id, slot in self._worker_output_slots(ordered).items():
+            frames = images_by_id.get(worker_id)
+            if not frames:
+                continue
+            tensors = [ensure_contiguous(frames[idx].cpu()) for idx in sorted(frames)]
+            participants[slot] = self._combine_image_tensors(
+                tensors, [f"worker {worker_id!r} image {idx}" for idx in sorted(frames)]
+            )
+        return participants
+
     async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
         if is_worker:
             # Worker mode: send images and audio to master in a single batch
@@ -781,6 +837,8 @@ class DistributedCollectorNode:
             exact_requested = False
         num_workers = len(expected_workers)
         if num_workers == 0:
+            if images is not None:
+                self.participant_images = {"master": images}
             return (images, audio if audio is not None else self.EMPTY_AUDIO)
 
         # Create the queue before any expensive local work to avoid job_complete race.
@@ -1021,6 +1079,10 @@ class DistributedCollectorNode:
                 else:
                     debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
                               f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
+                    self.participant_images = self._split_participant_images(
+                        worker_images, enabled_workers,
+                        None if delegate_mode or master_batch_size == 0 else images_on_cpu,
+                    )
                 return (combined, combined_audio)
             except Exception as e:
                 if serverless_results or exact_requested:
diff --git a/nodes/v3.py b/nodes/v3.py
index 05bbd97..39b05ba 100644
--- a/nodes/v3.py
+++ b/nodes/v3.py
@@ -5,9 +5,10 @@ This avoids sharing mutable instance state through V3's sanitized class clones.
 """
 import comfy.samplers
 from comfy_api.v0_0_2 import io
+from comfy_execution.graph_utils import ExecutionBlocker
 
 from . import utilities as _utilities
-from .collector import DistributedCollectorNode as _CollectorRuntime
+from .collector import DistributedCollectorNode as _CollectorRuntime, PER_GPU_OUTPUT_NAMES
 from ..utils import transfer_codec as _transfer_codec
 from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
 from .image_list import DistributedImageListEnd, DistributedImageListStart, DistributedImageListGetValue
@@ -130,7 +131,13 @@ class DistributedCollector(io.ComfyNode):
                     io.Int.Input('transfer_quality', default=_transfer_codec.DEFAULT_TRANSFER_QUALITY,
                                  min=1, max=100, step=1, optional=True,
                                  tooltip=_transfer_codec.TRANSFER_QUALITY_TOOLTIP)],
-            outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
+            # Per-GPU outputs are appended so saved links to images/audio keep their slots.
+            outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')] + [
+                io.Image.Output(id=name, display_name=name,
+                                tooltip='Images produced by this GPU only, for previewing. '
+                                        'worker_N is the Nth worker in the Distributed panel. '
+                                        'Nodes attached to a GPU that did not take part are skipped.')
+                for name in PER_GPU_OUTPUT_NAMES],
             is_input_list=True, accept_all_inputs=True,
         )
 
@@ -140,12 +147,16 @@ class DistributedCollector(io.ComfyNode):
                 worker_id='', pass_through=False, delegate_only=False,
                 transfer_format=_transfer_codec.DEFAULT_TRANSFER_FORMAT,
                 transfer_quality=_transfer_codec.DEFAULT_TRANSFER_QUALITY):
-        return io.NodeOutput(*_CollectorRuntime().run(
+        runtime = _CollectorRuntime()
+        result = runtime.run(
             images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
             is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
             worker_batch_size=worker_batch_size, worker_id=worker_id,
             pass_through=pass_through, delegate_only=delegate_only,
-            transfer_format=transfer_format, transfer_quality=transfer_quality))
+            transfer_format=transfer_format, transfer_quality=transfer_quality)
+        # A silent blocker skips previews wired to a GPU that produced nothing.
+        per_gpu = [runtime.participant_images.get(name, ExecutionBlocker(None)) for name in PER_GPU_OUTPUT_NAMES]
+        return io.NodeOutput(*result, *per_gpu)
 
 
 class UltimateSDUpscaleDistributed(io.ComfyNode):
diff --git a/tests/helpers/v3_runtime_check.py b/tests/helpers/v3_runtime_check.py
index a75cc52..e2dc2f9 100644
--- a/tests/helpers/v3_runtime_check.py
+++ b/tests/helpers/v3_runtime_check.py
@@ -43,6 +43,11 @@ ADDED_OPTIONAL_INPUTS = {
     },
 }
 
+# Outputs appended after the V1 baseline outputs, so saved links keep their slots.
+ADDED_OUTPUTS = {
+    'DistributedCollector': ['master'] + [f'worker_{index + 1}' for index in range(8)],
+}
+
 # Nodes added after the V1 baseline; not part of saved-workflow contracts.
 # DistributedWorkerSink is internal (dev-only) and only inserted by the master.
 ADDED_NODES = {'DistributedWorkerSink', 'DistributedImageListStart', 'DistributedImageListEnd',
@@ -374,7 +379,8 @@ async def main():
         assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id]
     assert 'DistributedMediaCollector' not in mapping
     collector_info = mapping['DistributedCollector'].GET_NODE_INFO_V1()
-    assert collector_info['output'] == ['IMAGE', 'AUDIO'], collector_info
+    assert collector_info['output'] == ['IMAGE', 'AUDIO'] + ['IMAGE'] * 9, collector_info
+    assert collector_info['output_name'][:2] == ['images', 'audio'], collector_info
     assert not any(collector_info['output_is_list']), collector_info
     assert list(collector_info['input']['optional']) == [
         'images', 'audio', 'transfer_format', 'transfer_quality']
@@ -393,6 +399,11 @@ async def main():
             assert old['output'] == ['*'] * 10
             old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10
         new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
+        added_outputs = ADDED_OUTPUTS.get(node_id, [])
+        if added_outputs:
+            assert new['output_name'][len(old['output']):] == added_outputs, (node_id, new['output_name'])
+            for key in ('output', 'output_name', 'output_is_list'):
+                new[key] = new[key][:len(old['output'])]
         for group in ('required', 'optional'):
             old_inputs = old['input'].get(group, {})
             new_inputs = dict(new['input'].get(group, {}))
diff --git a/tests/test_collector_image_transport.py b/tests/test_collector_image_transport.py
index 6994167..9ee24fb 100644
--- a/tests/test_collector_image_transport.py
+++ b/tests/test_collector_image_transport.py
@@ -223,3 +223,63 @@ def test_empty_image_results_raise_instead_of_reaching_save_image(media_modules)
     with pytest.raises(ValueError, match="empty image batch"):
         asyncio.run(_master_collect(module, torch.empty(0, 64, 64, 3),
                                     [torch.empty(0, 8, 12, 3)], delegate=True))
+
+
+def test_master_collection_records_each_participant_in_its_config_slot(media_modules):
+    module, _ = media_modules
+    module.load_config = lambda: {"workers": [{"id": "worker-off"}, {"id": "worker-a"}]}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    worker_frames = [torch.full((1, 8, 12, 3), value) for value in (0.25, 0.5)]
+    # Out-of-order delivery must still preview in image order.
+    for index in (1, 0):
+        queue.put_nowait({"worker_id": "worker-a", "image_index": index,
+                          "tensor": worker_frames[index], "is_last": index == 0})
+    queue.put_nowait({"worker_id": "serverless:pool", "image_index": 0,
+                      "tensor": torch.full((1, 8, 12, 3), 0.75), "is_last": True})
+    module.prompt_server.distributed_pending_jobs = {"image-test": queue}
+    master = torch.zeros(1, 8, 12, 3)
+    collector = module.DistributedCollectorNode()
+
+    combined, _ = asyncio.run(collector.execute(
+        master, None, multi_job_id="image-test", enabled_worker_ids='["worker-a", "serverless:pool"]'))
+
+    previews = collector.participant_images
+    assert combined.shape[0] == 4
+    # worker-a is the second configured worker; the unconfigured serverless
+    # participant takes the first slot after the configured workers.
+    assert set(previews) == {"master", "worker_2", "worker_3"}
+    assert torch.equal(previews["master"], master)
+    assert previews["worker_2"][:, 0, 0, 0].tolist() == [0.25, 0.5]
+    assert previews["worker_3"][:, 0, 0, 0].tolist() == [0.75]
+
+
+def test_delegate_master_has_no_master_preview(media_modules):
+    module, _ = media_modules
+    module.load_config = lambda: {"workers": [{"id": "worker-a"}]}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    worker = torch.full((1, 8, 12, 3), 0.5)
+    queue.put_nowait({"worker_id": "worker-a", "image_index": 0, "tensor": worker, "is_last": True})
+    module.prompt_server.distributed_pending_jobs = {"image-test": queue}
+    collector = module.DistributedCollectorNode()
+
+    asyncio.run(collector.execute(
+        torch.empty(0, 64, 64, 3), None, multi_job_id="image-test",
+        enabled_worker_ids='["worker-a"]', delegate_only=True))
+
+    assert list(collector.participant_images) == ["worker_1"]
+    assert torch.equal(collector.participant_images["worker_1"], worker)
+
+
+def test_undistributed_run_previews_everything_as_master(media_modules):
+    module, _ = media_modules
+    images = torch.zeros(2, 8, 12, 3)
+    collector = module.DistributedCollectorNode()
+    collector.run(images=[images])
+    assert list(collector.participant_images) == ["master"]
+    assert torch.equal(collector.participant_images["master"], images)
+
+    pass_through = module.DistributedCollectorNode()
+    pass_through.run(images=[images], multi_job_id=["job"], pass_through=[True])
+    assert pass_through.participant_images == {}
diff --git a/tests/test_v3_collector_transfer.py b/tests/test_v3_collector_transfer.py
index cd66fd8..352893c 100644
--- a/tests/test_v3_collector_transfer.py
+++ b/tests/test_v3_collector_transfer.py
@@ -13,6 +13,7 @@ import pytest
 
 ROOT = Path(__file__).resolve().parents[1]
 PKG = "dist_v3_transfer_testpkg"
+PER_GPU_OUTPUT_NAMES = ("master", "worker_1", "worker_2")
 
 
 class _Field:
@@ -51,10 +52,20 @@ def v3(monkeypatch):
     calls = []
 
     class CollectorRuntime:
+        participant_images_for_run = {}
+
+        def __init__(self):
+            self.participant_images = {}
+
         def run(self, **kwargs):
             calls.append(kwargs)
+            self.participant_images = dict(self.participant_images_for_run)
             return (kwargs["images"], kwargs["audio"])
 
+    class ExecutionBlocker:
+        def __init__(self, message):
+            self.message = message
+
     def module(name, **attrs):
         mod = types.ModuleType(name)
         mod.__dict__.update(attrs)
@@ -68,13 +79,16 @@ def v3(monkeypatch):
     module(f"{PKG}.utils", __path__=[])
     module(f"{PKG}.nodes.utilities")
     module(f"{PKG}.nodes.distributed_upscale", UltimateSDUpscaleDistributed=object)
-    module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime)
+    module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime,
+           PER_GPU_OUTPUT_NAMES=PER_GPU_OUTPUT_NAMES)
     module(f"{PKG}.nodes.image_list", DistributedImageListStart=object, DistributedImageListEnd=object,
            DistributedImageListGetValue=object)
     module("comfy", __path__=[])
     module("comfy.samplers", KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["normal"]))
     module("comfy_api", __path__=[])
     module("comfy_api.v0_0_2", io=_fake_io())
+    module("comfy_execution", __path__=[])
+    module("comfy_execution.graph_utils", ExecutionBlocker=ExecutionBlocker)
 
     codec_spec = importlib.util.spec_from_file_location(f"{PKG}.utils.transfer_codec", ROOT / "utils" / "transfer_codec.py")
     codec = importlib.util.module_from_spec(codec_spec)
@@ -85,6 +99,7 @@ def v3(monkeypatch):
     mod = importlib.util.module_from_spec(spec)
     monkeypatch.setitem(sys.modules, f"{PKG}.nodes.v3", mod)
     spec.loader.exec_module(mod)
+    mod._TestRuntime = CollectorRuntime
     return mod, codec, calls
 
 
@@ -127,3 +142,24 @@ def test_registered_collector_defaults_to_legacy_transfer(v3):
 
     assert calls[-1]["transfer_format"] == "legacy_png"
     assert calls[-1]["transfer_quality"] == 90
+
+
+def test_registered_collector_appends_per_gpu_outputs_after_combined_outputs(v3):
+    mod, _, _ = v3
+    schema = mod.DistributedCollector.define_schema()
+
+    assert [field.kind for field in schema.outputs] == ["Image", "Audio", "Image", "Image", "Image"]
+    assert [field.kwargs["display_name"] for field in schema.outputs] == [
+        "images", "audio", "master", "worker_1", "worker_2"]
+
+
+def test_registered_collector_blocks_per_gpu_outputs_without_images(v3):
+    mod, _, _ = v3
+    mod._TestRuntime.participant_images_for_run = {"master": "m", "worker_2": "w2"}
+
+    _, values, _ = mod.DistributedCollector.execute(images=["img"], audio=["aud"])
+
+    assert values[:2] == (["img"], ["aud"])
+    assert values[2] == "m"
+    assert type(values[3]).__name__ == "ExecutionBlocker" and values[3].message is None
+    assert values[4] == "w2"
diff --git a/web/collectorOutputs.js b/web/collectorOutputs.js
new file mode 100644
index 0000000..8973a68
--- /dev/null
+++ b/web/collectorOutputs.js
@@ -0,0 +1,76 @@
+import { app } from "/scripts/app.js";
+import { ENDPOINTS, NODE_CLASSES } from "./constants.js";
+
+// worker_N previews the Nth configured worker (see nodes/collector.py).
+const WORKER_OUTPUT_PREFIX = "worker_";
+const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
+
+const trackedNodes = new Set();
+let workers = [];
+let listenerAttached = false;
+
+async function fetchWorkers() {
+    try {
+        const resp = await fetch(ENDPOINTS.CONFIG);
+        if (!resp.ok) return [];
+        const config = await resp.json();
+        return Array.isArray(config.workers) ? config.workers : [];
+    } catch {
+        return [];
+    }
+}
+
+// Only the displayed label changes; slot names stay fixed so links survive.
+function labelWorkerOutputs(node, configWorkers) {
+    for (const output of node.outputs || []) {
+        if (!output?.name?.startsWith(WORKER_OUTPUT_PREFIX)) continue;
+        const index = parseInt(output.name.slice(WORKER_OUTPUT_PREFIX.length), 10) - 1;
+        const worker = configWorkers[index];
+        output.label = worker?.name ? `${output.name} (${worker.name})` : output.name;
+    }
+    node.setDirtyCanvas?.(true, true);
+}
+
+function relabelAll() {
+    for (const node of trackedNodes) labelWorkerOutputs(node, workers);
+}
+
+function attachListener() {
+    if (listenerAttached || typeof window === "undefined") return;
+    window.addEventListener(WORKERS_CHANGED_EVENT, (event) => {
+        if (Array.isArray(event?.detail?.workers)) {
+            workers = event.detail.workers;
+            relabelAll();
+        }
+    });
+    listenerAttached = true;
+}
+
+app.registerExtension({
+    name: "Distributed.CollectorOutputs",
+    async nodeCreated(node) {
+        if (node.comfyClass !== NODE_CLASSES.DISTRIBUTED_COLLECTOR) return;
+        try {
+            attachListener();
+            trackedNodes.add(node);
+
+            const originalConfigure = node.configure;
+            node.configure = function (data) {
+                const result = originalConfigure ? originalConfigure.call(this, data) : undefined;
+                labelWorkerOutputs(this, workers);
+                return result;
+            };
+
+            const originalOnRemoved = node.onRemoved;
+            node.onRemoved = function () {
+                trackedNodes.delete(this);
+                return originalOnRemoved ? originalOnRemoved.call(this) : undefined;
+            };
+
+            if (workers.length === 0) workers = await fetchWorkers();
+            labelWorkerOutputs(node, workers);
+        } catch (error) {
+            console.error("Error in DistributedCollector outputs extension:", error);
+        }
+    },
+});
diff --git a/web/tests/collectorOutputs.test.js b/web/tests/collectorOutputs.test.js
new file mode 100644
index 0000000..d6434f2
--- /dev/null
+++ b/web/tests/collectorOutputs.test.js
@@ -0,0 +1,45 @@
+import { readFileSync } from "node:fs";
+import vm from "node:vm";
+import { describe, expect, it } from "vitest";
+
+const source = readFileSync(new URL("../collectorOutputs.js", import.meta.url), "utf8")
+    .replace(/^import .*;\r?\n/gm, "");
+
+function setup(workers) {
+    let extension;
+    const listeners = new Map();
+    vm.runInNewContext(source, {
+        app: { registerExtension: (ext) => { extension = ext; } },
+        ENDPOINTS: { CONFIG: "/distributed/config" },
+        NODE_CLASSES: { DISTRIBUTED_COLLECTOR: "DistributedCollector" },
+        fetch: async () => ({ ok: true, json: async () => ({ workers }) }),
+        window: { addEventListener: (name, fn) => listeners.set(name, fn) },
+        parseInt,
+        Array,
+        console,
+    });
+    const node = {
+        comfyClass: "DistributedCollector",
+        outputs: ["images", "audio", "master", "worker_1", "worker_2", "worker_3"].map((name) => ({ name })),
+    };
+    return { extension, listeners, node };
+}
+
+describe("DistributedCollector per-GPU outputs", () => {
+    it("labels worker_N with the Nth configured worker and keeps slot names", async () => {
+        const { extension, node } = setup([{ id: "a", name: "GPU A", enabled: false }, { id: "b", name: "GPU B" }]);
+        await extension.nodeCreated(node);
+        expect(node.outputs.map((o) => o.name)).toEqual(
+            ["images", "audio", "master", "worker_1", "worker_2", "worker_3"]);
+        expect(node.outputs.slice(3).map((o) => o.label)).toEqual(
+            ["worker_1 (GPU A)", "worker_2 (GPU B)", "worker_3"]);
+        expect(node.outputs[2].label).toBeUndefined();
+    });
+
+    it("relabels when the worker list changes", async () => {
+        const { extension, listeners, node } = setup([{ id: "a", name: "GPU A" }]);
+        await extension.nodeCreated(node);
+        listeners.get("distributed:workers-changed")({ detail: { workers: [{ id: "c", name: "Renamed" }] } });
+        expect(node.outputs[3].label).toBe("worker_1 (Renamed)");
+    });
+});

From 339719ece701db64aacb9f20cca4efbc369d7fa9 Mon Sep 17 00:00:00 2001
From: speedyrulz <137450672+speedyrulz@users.noreply.github.com>
Date: Sat, 31 Oct 2026 16:42:40 -0400
Subject: [PATCH 3/3] fix: harden per-GPU preview outputs after review

- Typed serverless pools validate the unpruned workflow, so a preview on
  any new slot failed the whole run; both serverless compilers now accept
  every Collector output slot (COLLECTOR_OUTPUT_COUNT).
- Serverless pools are config workers: their serverless:<pool>:<job>
  participants now share the pool's worker_N slot instead of overflowing.
- A worker returning only empty frames no longer raises in the preview
  split (the combined batch already skipped them).
- Worker previews are views into the combined batch, not a second copy,
  and a frame-count guard skips them rather than misattribute frames.
- Frontend: one shared config fetch per load; unnamed slots drop their
  saved label instead of storing a redundant one.
- Tooltip notes that the per-GPU outputs are unavailable downstream of
  Ultimate SD Upscale Distributed.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 api/orchestration/serverless_media_prompt.py |   3 +-
 api/orchestration/serverless_prompt.py       |   6 +-
 nodes/collector.py                           | 126 ++++++++++++-------
 nodes/v3.py                                  |   4 +-
 tests/test_collector_image_transport.py      |  65 ++++++++++
 tests/test_serverless_media_prompt.py        |  18 +++
 web/collectorOutputs.js                      |  22 +++-
 web/tests/collectorOutputs.test.js           |  23 +++-
 8 files changed, 217 insertions(+), 50 deletions(-)

diff --git a/api/orchestration/serverless_media_prompt.py b/api/orchestration/serverless_media_prompt.py
index a86446e..de7ee32 100644
--- a/api/orchestration/serverless_media_prompt.py
+++ b/api/orchestration/serverless_media_prompt.py
@@ -16,6 +16,7 @@ from graphlib import CycleError, TopologicalSorter
 from ...serverless_bridge.contract import SOCKET_KINDS, VALUE_SOCKETS, identifier, validate_bindings
 from ...serverless_bridge.image_transfer import DEFAULT_FORMAT, validate_options
 from .serverless_prompt import (
+    COLLECTOR_OUTPUT_COUNT,
     _PRIMITIVE_NODES,
     _REMOVED_SCALAR_NODES,
     _is_link,
@@ -99,7 +100,7 @@ def _validate_prompt(prepared_prompt, *, require_collectors=True):
             if source_class in _REMOVED_SCALAR_NODES or source_class in _PRIMITIVE_NODES:
                 if slot != 0:
                     raise ValueError(f"invalid output slot {slot} for scalar node {source_id!r}")
-            elif source_class == _COLLECTOR and slot >= len(COLLECTOR_SOCKETS):
+            elif source_class == _COLLECTOR and slot >= COLLECTOR_OUTPUT_COUNT:
                 raise ValueError(f"invalid output slot {slot} for collector {source_id!r}")
             parents.add(source_id)
         dependencies[node_id] = parents
diff --git a/api/orchestration/serverless_prompt.py b/api/orchestration/serverless_prompt.py
index 1168a56..32be508 100644
--- a/api/orchestration/serverless_prompt.py
+++ b/api/orchestration/serverless_prompt.py
@@ -29,8 +29,12 @@ _PRIMITIVE_NODES = {
     "PrimitiveInt": 0,
     "PrimitiveString": "",
 }
+# DistributedCollector outputs: images, audio, then the per-GPU previews
+# (PER_GPU_OUTPUT_NAMES in nodes/collector.py). Previews are valid link
+# sources in a workflow even though a pool only ever returns images and audio.
+COLLECTOR_OUTPUT_COUNT = 2 + 9
 _REMOVED_NODE_OUTPUTS = {
-    "DistributedCollector": {0, 1},
+    "DistributedCollector": set(range(COLLECTOR_OUTPUT_COUNT)),
     "DistributedSeed": {0},
     "DistributedValue": {0},
     "DistributedModelName": {0},
diff --git a/nodes/collector.py b/nodes/collector.py
index 7e982cb..89001c7 100644
--- a/nodes/collector.py
+++ b/nodes/collector.py
@@ -169,6 +169,9 @@ class DistributedCollectorNode:
                 # RGB is opaque; never discard alpha from another participant.
                 tensor = torch.cat((tensor, torch.ones_like(tensor[..., :1])), dim=-1)
             normalized.append(ensure_contiguous(tensor))
+        if len(normalized) == 1:
+            # torch.cat would copy a single batch for nothing.
+            return normalized[0]
         return ensure_contiguous(torch.cat(normalized, dim=0))
 
     def _normalize_images_input(self, images):
@@ -668,6 +671,28 @@ class DistributedCollectorNode:
             worker_audio[worker_id] = self._normalize_audio_input(media["audio"])
         return stored
 
+    @staticmethod
+    def _ordered_worker_frames(worker_images, worker_order):
+        """Yield (worker_id, image_index, tensor) in combined-batch order.
+
+        Expected workers come in ``worker_order``; any unexpected worker ids
+        follow in sorted order.
+        """
+        seen = set()
+        for worker_id_str in [str(worker_id) for worker_id in (worker_order or [])]:
+            seen.add(worker_id_str)
+            frames = worker_images.get(worker_id_str)
+            if not frames:
+                continue
+            for idx in sorted(frames):
+                yield worker_id_str, idx, frames[idx]
+        for worker_id_str in sorted(worker_images):
+            if worker_id_str in seen:
+                continue
+            frames = worker_images[worker_id_str]
+            for idx in sorted(frames):
+                yield worker_id_str, idx, frames[idx]
+
     def _reorder_and_combine_tensors(
         self,
         worker_images: dict,
@@ -685,23 +710,9 @@ class DistributedCollectorNode:
                 ordered_tensors.append(images_on_cpu[i:i+1])
                 sources.append(f"master image {i}")
 
-        ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
-        seen = set()
-        for worker_id_str in ordered_worker_ids:
-            seen.add(worker_id_str)
-            if worker_id_str not in worker_images:
-                continue
-            for idx in sorted(worker_images[worker_id_str].keys()):
-                ordered_tensors.append(worker_images[worker_id_str][idx])
-                sources.append(f"worker {worker_id_str!r} image {idx}")
-
-        # Append any unexpected worker ids deterministically.
-        for worker_id_str in sorted(worker_images.keys()):
-            if worker_id_str in seen:
-                continue
-            for idx in sorted(worker_images[worker_id_str].keys()):
-                ordered_tensors.append(worker_images[worker_id_str][idx])
-                sources.append(f"worker {worker_id_str!r} image {idx}")
+        for worker_id_str, idx, tensor in self._ordered_worker_frames(worker_images, worker_order):
+            ordered_tensors.append(tensor)
+            sources.append(f"worker {worker_id_str!r} image {idx}")
 
         cpu_tensors = []
         for t in ordered_tensors:
@@ -717,11 +728,22 @@ class DistributedCollectorNode:
         return None
 
     @staticmethod
-    def _worker_output_slots(participant_ids):
-        """Map worker ids to worker_N preview slots.
+    def _config_worker_id(participant_id):
+        """Config entry a participant belongs to.
 
-        Configured workers keep their config position; other participants
-        (e.g. serverless) take the free slots after the configured workers.
+        Serverless pools run as ``serverless:<pool id>:<job>`` participants
+        (see queue_orchestration) and all preview on their pool's slot.
+        """
+        if participant_id.startswith("serverless:"):
+            return participant_id[len("serverless:"):].rsplit(":", 1)[0]
+        return participant_id
+
+    @staticmethod
+    def _worker_output_slots(participant_ids):
+        """Map participant ids to worker_N preview slots (many-to-one).
+
+        Configured workers keep their config position; participants outside
+        the config take the free slots after the configured workers.
         """
         try:
             config_ids = [str(worker.get("id")) for worker in load_config().get("workers", [])]
@@ -729,34 +751,53 @@ class DistributedCollectorNode:
             debug_log(f"Collector - Could not read worker order for preview outputs: {e}")
             config_ids = []
         slots = {}
-        next_free = len(config_ids)
-        for worker_id in participant_ids:
-            if worker_id in config_ids:
-                index = config_ids.index(worker_id)
+        overflow = {}
+        for participant_id in participant_ids:
+            config_id = DistributedCollectorNode._config_worker_id(participant_id)
+            if config_id in config_ids:
+                index = config_ids.index(config_id)
             else:
-                index, next_free = next_free, next_free + 1
+                index = overflow.setdefault(config_id, len(config_ids) + len(overflow))
             if index < MAX_WORKER_OUTPUTS:
-                slots[worker_id] = PER_GPU_OUTPUT_NAMES[index + 1]
+                slots[participant_id] = PER_GPU_OUTPUT_NAMES[index + 1]
             else:
-                debug_log(f"Collector - No preview output for worker {worker_id}; only in the combined images")
+                debug_log(f"Collector - No preview output for {participant_id}; only in the combined images")
         return slots
 
-    def _split_participant_images(self, worker_images, worker_order, master_images):
-        """Return each participant's own image batch keyed by preview slot."""
+    def _split_participant_images(self, worker_images, worker_order, master_images, combined, master_frames):
+        """Return each participant's own images keyed by preview slot.
+
+        ``combined`` holds the master's ``master_frames`` frames and then the
+        worker frames in ``_ordered_worker_frames`` order, so worker previews
+        are views into it rather than a second copy of every frame.
+        """
         participants = {}
         if master_images is not None:
             participants["master"] = master_images
-        ordered = [str(w) for w in (worker_order or [])]
-        ordered += sorted(str(w) for w in worker_images if str(w) not in ordered)
-        images_by_id = {str(w): frames for w, frames in worker_images.items()}
-        for worker_id, slot in self._worker_output_slots(ordered).items():
-            frames = images_by_id.get(worker_id)
-            if not frames:
+        frames = list(self._ordered_worker_frames(worker_images, worker_order))
+        slots = self._worker_output_slots(list(dict.fromkeys(worker_id for worker_id, _, _ in frames)))
+        segments = {}  # slot -> [start, end) frame runs within combined
+        offset = master_frames
+        for worker_id, _idx, tensor in frames:
+            count = tensor.shape[0]
+            # Empty placeholders were dropped from the combined batch too.
+            if count == 0:
                 continue
-            tensors = [ensure_contiguous(frames[idx].cpu()) for idx in sorted(frames)]
-            participants[slot] = self._combine_image_tensors(
-                tensors, [f"worker {worker_id!r} image {idx}" for idx in sorted(frames)]
-            )
+            slot = slots.get(worker_id)
+            if slot is not None:
+                runs = segments.setdefault(slot, [])
+                if runs and runs[-1][1] == offset:
+                    runs[-1][1] = offset + count
+                else:
+                    runs.append([offset, offset + count])
+            offset += count
+        if offset != combined.shape[0]:
+            # Never attribute frames to the wrong GPU; the combined output is unaffected.
+            log(f"Master - Per-GPU previews skipped: expected {offset} combined frames, got {combined.shape[0]}")
+            return participants
+        for slot, runs in segments.items():
+            views = [combined[start:end] for start, end in runs]
+            participants[slot] = views[0] if len(views) == 1 else torch.cat(views, dim=0)
         return participants
 
     async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
@@ -1079,9 +1120,10 @@ class DistributedCollectorNode:
                 else:
                     debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
                               f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
+                    master_frames = 0 if delegate_mode or images_on_cpu is None else master_batch_size
                     self.participant_images = self._split_participant_images(
-                        worker_images, enabled_workers,
-                        None if delegate_mode or master_batch_size == 0 else images_on_cpu,
+                        worker_images, enabled_workers, images_on_cpu if master_frames else None,
+                        combined, master_frames,
                     )
                 return (combined, combined_audio)
             except Exception as e:
diff --git a/nodes/v3.py b/nodes/v3.py
index 39b05ba..a71c2e7 100644
--- a/nodes/v3.py
+++ b/nodes/v3.py
@@ -136,7 +136,9 @@ class DistributedCollector(io.ComfyNode):
                 io.Image.Output(id=name, display_name=name,
                                 tooltip='Images produced by this GPU only, for previewing. '
                                         'worker_N is the Nth worker in the Distributed panel. '
-                                        'Nodes attached to a GPU that did not take part are skipped.')
+                                        'Nodes attached to a GPU that did not take part are skipped. '
+                                        'After Ultimate SD Upscale Distributed every GPU contributes '
+                                        'tiles, so only the combined images output is available.')
                 for name in PER_GPU_OUTPUT_NAMES],
             is_input_list=True, accept_all_inputs=True,
         )
diff --git a/tests/test_collector_image_transport.py b/tests/test_collector_image_transport.py
index 9ee24fb..b37967b 100644
--- a/tests/test_collector_image_transport.py
+++ b/tests/test_collector_image_transport.py
@@ -283,3 +283,68 @@ def test_undistributed_run_previews_everything_as_master(media_modules):
     pass_through = module.DistributedCollectorNode()
     pass_through.run(images=[images], multi_job_id=["job"], pass_through=[True])
     assert pass_through.participant_images == {}
+
+
+def test_empty_worker_batch_keeps_master_preview_and_combined_result(media_modules):
+    module, _ = media_modules
+    module.load_config = lambda: {"workers": [{"id": "worker-a"}]}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    queue.put_nowait({"worker_id": "worker-a", "image_index": 0,
+                      "tensor": torch.empty(0, 8, 12, 3), "is_last": True})
+    module.prompt_server.distributed_pending_jobs = {"image-test": queue}
+    master = torch.zeros(1, 8, 12, 3)
+    collector = module.DistributedCollectorNode()
+
+    combined, _ = asyncio.run(collector.execute(
+        master, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]'))
+
+    assert combined.shape[0] == 1
+    assert list(collector.participant_images) == ["master"]
+
+
+def test_serverless_pool_jobs_share_their_pools_preview_slot(media_modules):
+    module, _ = media_modules
+    module.load_config = lambda: {"workers": [{"id": "gpu-a"}, {"id": "pool", "type": "serverless"}]}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    for job, value in ((2, 0.75), (1, 0.5)):
+        queue.put_nowait({"worker_id": f"serverless:pool:{job}", "image_index": 0,
+                          "tensor": torch.full((1, 8, 12, 3), value), "is_last": True})
+    module.prompt_server.distributed_pending_jobs = {"image-test": queue}
+    collector = module.DistributedCollectorNode()
+
+    asyncio.run(collector.execute(
+        torch.zeros(1, 8, 12, 3), None, multi_job_id="image-test",
+        enabled_worker_ids='["serverless:pool:1", "serverless:pool:2"]'))
+
+    previews = collector.participant_images
+    assert set(previews) == {"master", "worker_2"}
+    # Pool jobs appear in their enabled order, not delivery order.
+    assert previews["worker_2"][:, 0, 0, 0].tolist() == [0.5, 0.75]
+
+
+def test_serverless_validators_accept_every_collector_output(media_modules):
+    module, _ = media_modules
+    legacy = _load_file("serverless_prompt_slot_check", ROOT / "api/orchestration/serverless_prompt.py")
+    assert legacy.COLLECTOR_OUTPUT_COUNT == 2 + len(module.PER_GPU_OUTPUT_NAMES)
+    assert legacy._REMOVED_NODE_OUTPUTS["DistributedCollector"] == set(range(legacy.COLLECTOR_OUTPUT_COUNT))
+
+
+def test_worker_previews_are_views_of_the_combined_batch(media_modules):
+    module, _ = media_modules
+    module.load_config = lambda: {"workers": [{"id": "worker-a"}]}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    for index in range(2):
+        queue.put_nowait({"worker_id": "worker-a", "image_index": index,
+                          "tensor": torch.full((1, 8, 12, 3), 0.5), "is_last": index == 1})
+    module.prompt_server.distributed_pending_jobs = {"image-test": queue}
+    collector = module.DistributedCollectorNode()
+
+    combined, _ = asyncio.run(collector.execute(
+        torch.zeros(1, 8, 12, 3), None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]'))
+
+    preview = collector.participant_images["worker_1"]
+    assert preview.shape == (2, 8, 12, 3) and preview.is_contiguous()
+    assert preview.data_ptr() == combined[1:].data_ptr()
diff --git a/tests/test_serverless_media_prompt.py b/tests/test_serverless_media_prompt.py
index f17be44..b0a0dc8 100644
--- a/tests/test_serverless_media_prompt.py
+++ b/tests/test_serverless_media_prompt.py
@@ -424,3 +424,21 @@ assert graph[bindings[0]["node_id"]]["class_type"] == "MediaExportV1"
 '''
     result = subprocess.run([sys.executable, "-c", script, str(ROOT)], text=True, capture_output=True)
     assert result.returncode == 0, result.stderr
+
+
+def test_links_from_collector_preview_outputs_are_pruned_not_rejected():
+    # Per-GPU preview sockets follow images/audio; a pool never returns them,
+    # but a workflow previewing them must still compile for typed pools.
+    prompt = _prompt()
+    prompt["preview:1"] = _node("PreviewImage", images=["collector:1", 2])
+    prompt["preview:2"] = _node("PreviewImage", images=["collector:1", legacy.COLLECTOR_OUTPUT_COUNT - 1])
+    graph, bindings = _run(prompt)
+    assert [binding["socket"] for binding in bindings] == ["images"]
+    assert not any(node["class_type"] == "PreviewImage" for node in graph.values())
+
+
+def test_rejects_collector_output_slots_beyond_the_schema():
+    prompt = _prompt()
+    prompt["preview:1"] = _node("PreviewImage", images=["collector:1", legacy.COLLECTOR_OUTPUT_COUNT])
+    with pytest.raises(ValueError, match="invalid output slot"):
+        _run(prompt)
diff --git a/web/collectorOutputs.js b/web/collectorOutputs.js
index 8973a68..81b0a63 100644
--- a/web/collectorOutputs.js
+++ b/web/collectorOutputs.js
@@ -7,6 +7,7 @@ const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
 
 const trackedNodes = new Set();
 let workers = [];
+let workersPromise = null;
 let listenerAttached = false;
 
 async function fetchWorkers() {
@@ -20,13 +21,30 @@ async function fetchWorkers() {
     }
 }
 
+// Collectors created together (a workflow load) share one config fetch.
+function loadWorkers() {
+    if (!workersPromise) {
+        workersPromise = fetchWorkers().then((list) => {
+            workers = list;
+            if (list.length === 0) workersPromise = null;
+            return list;
+        });
+    }
+    return workersPromise;
+}
+
 // Only the displayed label changes; slot names stay fixed so links survive.
+// Unnamed slots drop the label so saved workflows carry no redundant copies.
 function labelWorkerOutputs(node, configWorkers) {
     for (const output of node.outputs || []) {
         if (!output?.name?.startsWith(WORKER_OUTPUT_PREFIX)) continue;
         const index = parseInt(output.name.slice(WORKER_OUTPUT_PREFIX.length), 10) - 1;
         const worker = configWorkers[index];
-        output.label = worker?.name ? `${output.name} (${worker.name})` : output.name;
+        if (worker?.name) {
+            output.label = `${output.name} (${worker.name})`;
+        } else {
+            delete output.label;
+        }
     }
     node.setDirtyCanvas?.(true, true);
 }
@@ -67,7 +85,7 @@ app.registerExtension({
                 return originalOnRemoved ? originalOnRemoved.call(this) : undefined;
             };
 
-            if (workers.length === 0) workers = await fetchWorkers();
+            if (workers.length === 0) await loadWorkers();
             labelWorkerOutputs(node, workers);
         } catch (error) {
             console.error("Error in DistributedCollector outputs extension:", error);
diff --git a/web/tests/collectorOutputs.test.js b/web/tests/collectorOutputs.test.js
index d6434f2..bdee2fe 100644
--- a/web/tests/collectorOutputs.test.js
+++ b/web/tests/collectorOutputs.test.js
@@ -5,14 +5,15 @@ import { describe, expect, it } from "vitest";
 const source = readFileSync(new URL("../collectorOutputs.js", import.meta.url), "utf8")
     .replace(/^import .*;\r?\n/gm, "");
 
-function setup(workers) {
+function setup(workers, onFetch = () => {}) {
     let extension;
     const listeners = new Map();
     vm.runInNewContext(source, {
         app: { registerExtension: (ext) => { extension = ext; } },
         ENDPOINTS: { CONFIG: "/distributed/config" },
         NODE_CLASSES: { DISTRIBUTED_COLLECTOR: "DistributedCollector" },
-        fetch: async () => ({ ok: true, json: async () => ({ workers }) }),
+        fetch: async () => { onFetch(); return { ok: true, json: async () => ({ workers }) }; },
+        Promise,
         window: { addEventListener: (name, fn) => listeners.set(name, fn) },
         parseInt,
         Array,
@@ -32,10 +33,26 @@ describe("DistributedCollector per-GPU outputs", () => {
         expect(node.outputs.map((o) => o.name)).toEqual(
             ["images", "audio", "master", "worker_1", "worker_2", "worker_3"]);
         expect(node.outputs.slice(3).map((o) => o.label)).toEqual(
-            ["worker_1 (GPU A)", "worker_2 (GPU B)", "worker_3"]);
+            ["worker_1 (GPU A)", "worker_2 (GPU B)", undefined]);
         expect(node.outputs[2].label).toBeUndefined();
     });
 
+    it("shares one config fetch across collectors created together", async () => {
+        let fetches = 0;
+        const { extension, node } = setup([{ id: "a", name: "GPU A" }], () => { fetches += 1; });
+        const second = { ...node, outputs: node.outputs.map((o) => ({ ...o })) };
+        await Promise.all([extension.nodeCreated(node), extension.nodeCreated(second)]);
+        expect(fetches).toBe(1);
+        expect(second.outputs[3].label).toBe("worker_1 (GPU A)");
+    });
+
+    it("clears a stale saved label when the slot has no configured worker", async () => {
+        const { extension, node } = setup([]);
+        node.outputs[3].label = "worker_1 (Old PC)";
+        await extension.nodeCreated(node);
+        expect("label" in node.outputs[3]).toBe(false);
+    });
+
     it("relabels when the worker list changes", async () => {
         const { extension, listeners, node } = setup([{ id: "a", name: "GPU A" }]);
         await extension.nodeCreated(node);