feat(worker): terminate pruned worker prompts with a hidden no-op sink (#99)

Worker prompts drop everything downstream of DistributedCollector /
UltimateSDUpscaleDistributed, so the master re-added a PreviewImage
(or PreviewAudio) to give ComfyUI an output node. That preview wrote
every result frame as a PNG into the worker's temp directory: for a
362-frame 1344x768 video, ~2 GB per run (RAM when temp is on tmpfs)
plus CPU time spent PNG-encoding frames nobody looks at.

- Add DistributedWorkerSink: an output node that discards its input.
  It is is_dev_only, so it stays out of node search and menus.
- The master probes each worker with GET /object_info/DistributedWorkerSink
  while preparing its prompt and uses the sink when it is registered.
  Workers on older versions (or a failed probe) keep the PreviewImage /
  PreviewAudio terminator, so mixed versions keep working.

Verified with a real master plus one updated and one old worker:
the updated worker ran EmptyImage -> Collector -> Sink and left its
temp directory empty; the old worker still received PreviewImage.


Claude-Session: https://claude.ai/code/session_0166t8T3CVaxZkJbP2RtSpo6

Co-authored-by: ymiyad <cab-13-might@icloud.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

api/orchestration/media_sync.py

@@ -143,6 +143,30 @@ async def fetch_worker_path_separator(worker, trace_execution_id=None):
         return None
 
 
+async def fetch_worker_has_node(worker, node_class, trace_execution_id=None):
+    """Best-effort check that a worker has ``node_class`` registered.
+
+    Uses ComfyUI's ``GET /object_info/{node_class}``, which returns ``{}`` for
+    unknown classes. Any failure counts as "not available" so callers fall
+    back to nodes every ComfyUI install has.
+    """
+    url = build_worker_url(worker, f"/object_info/{node_class}")
+    session = await get_client_session()
+    try:
+        async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
+            if resp.status != 200:
+                return False
+            payload = await resp.json()
+            return isinstance(payload, dict) and node_class in payload
+    except Exception as exc:
+        message = f"Failed to probe node '{node_class}' on worker {worker.get('id')}: {exc}"
+        if trace_execution_id:
+            trace_debug(trace_execution_id, message)
+        else:
+            debug_log(f"[Distributed] {message}")
+        return False
+
+
 async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type, trace_execution_id=None):
     """Upload one media file to worker iff missing or hash-mismatched."""
     session = await get_client_session()

api/orchestration/prompt_transform.py

@@ -328,8 +328,19 @@ def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
     return connected
 
 
-def prune_prompt_for_worker(prompt_obj):
-    """Prune worker prompt to distributed nodes and their upstream dependencies."""
+WORKER_SINK_CLASS = "DistributedWorkerSink"
+
+
+def prune_prompt_for_worker(prompt_obj, use_worker_sink=False):
+    """Prune worker prompt to distributed nodes and their upstream dependencies.
+
+    Removing downstream nodes can leave the worker prompt without an output
+    node. One is re-added after each such distributed node: the hidden
+    ``DistributedWorkerSink`` when ``use_worker_sink`` is set (the worker must
+    have it registered), otherwise PreviewImage/PreviewAudio for workers
+    running older versions. The preview nodes write every frame to the
+    worker's temp directory, which the sink avoids.
+    """
     collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
     upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
     distributed_ids = collector_ids + upscale_ids
@@ -364,7 +375,13 @@ def prune_prompt_for_worker(prompt_obj):
                 and len(inputs["audio"]) == 2
             )
 
-            if image_connected:
+            if use_worker_sink and (image_connected or audio_connected):
+                pruned_prompt[next_id()] = {
+                    "inputs": {"value": [dist_id, 0 if image_connected else 1]},
+                    "class_type": WORKER_SINK_CLASS,
+                    "_meta": {"title": "Distributed Worker Sink (auto-added)"},
+                }
+            elif image_connected:
                 preview_id = next_id()
                 pruned_prompt[preview_id] = {
                     "inputs": {"images": [dist_id, 0]},

api/queue_orchestration.py

@@ -21,7 +21,12 @@ from .orchestration.dispatch import (
     select_active_workers,
     select_least_busy_worker,
 )
-from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media
+from .orchestration.media_sync import (
+    convert_paths_for_platform,
+    fetch_worker_has_node,
+    fetch_worker_path_separator,
+    sync_worker_media,
+)
 from .orchestration.prompt_transform import (
     PromptIndex,
     apply_participant_overrides,
@@ -29,6 +34,7 @@ from .orchestration.prompt_transform import (
     generate_job_id_map,
     prepare_delegate_master_prompt,
     prune_prompt_for_worker,
+    WORKER_SINK_CLASS,
 )
 
 
@@ -162,12 +168,21 @@ async def _prepare_worker_payload(
 
         worker_type = str(worker.get("type") or "local").strip().lower()
         is_remote_like = bool(worker.get("host")) and worker_type != "local"
-        if is_remote_like:
-            path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id)
-            if path_separator:
-                worker_prompt = convert_paths_for_platform(worker_prompt, path_separator)
+        # Older workers lack the sink node; they keep the PreviewImage terminator.
+        has_sink_task = asyncio.ensure_future(
+            fetch_worker_has_node(worker, WORKER_SINK_CLASS, trace_execution_id=trace_execution_id)
+        )
+        try:
+            if is_remote_like:
+                path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id)
+                if path_separator:
+                    worker_prompt = convert_paths_for_platform(worker_prompt, path_separator)
+            use_worker_sink = await has_sink_task
+        finally:
+            if not has_sink_task.done():
+                has_sink_task.cancel()
 
-        worker_prompt = prune_prompt_for_worker(worker_prompt)
+        worker_prompt = prune_prompt_for_worker(worker_prompt, use_worker_sink=use_worker_sink)
         worker_prompt = apply_participant_overrides(
             worker_prompt,
             worker["id"],