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

nodes/v3.py

@@ -187,5 +187,28 @@ class UltimateSDUpscaleDistributed(io.ComfyNode):
             master_url, enabled_worker_ids, worker_id, tile_indices, dynamic_threshold))
 
 
+class DistributedWorkerSink(io.ComfyNode):
+    """Internal output node that terminates pruned worker prompts.
+
+    Worker prompts drop everything downstream of the distributed node, which
+    leaves no output node for ComfyUI to execute towards. This sink provides
+    one without writing preview files (PreviewImage wrote every frame to the
+    worker's temp directory). Hidden from search/menus; inserted by the master.
+    """
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedWorkerSink', display_name='Distributed Worker Sink (internal)', category='utils',
+            inputs=[io.AnyType.Input('value')],
+            is_output_node=True, is_dev_only=True,
+        )
+
+    @classmethod
+    def execute(cls, value):
+        return io.NodeOutput()
+
+
 NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
-         ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed]
+         ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed,
+         DistributedWorkerSink]

tests/api/test_media_sync.py

@@ -245,3 +245,70 @@ class RewritePromptMediaInputsTests(unittest.TestCase):
 
 if __name__ == "__main__":
     unittest.main()
+
+
+# ---------------------------------------------------------------------------
+# fetch_worker_has_node
+# ---------------------------------------------------------------------------
+
+class _FakeResponse:
+    def __init__(self, status, payload):
+        self.status = status
+        self._payload = payload
+
+    async def json(self):
+        if isinstance(self._payload, Exception):
+            raise self._payload
+        return self._payload
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *exc):
+        return False
+
+
+class _FakeSession:
+    def __init__(self, response=None, error=None):
+        self.response = response
+        self.error = error
+        self.urls = []
+
+    def get(self, url, timeout=None):
+        self.urls.append(url)
+        if self.error:
+            raise self.error
+        return self.response
+
+
+class FetchWorkerHasNodeTests(unittest.TestCase):
+    def _probe(self, session):
+        import asyncio
+
+        async def _get_session():
+            return session
+
+        original = ms.get_client_session
+        ms.get_client_session = _get_session
+        try:
+            return asyncio.run(ms.fetch_worker_has_node({"id": "w1"}, "DistributedWorkerSink"))
+        finally:
+            ms.get_client_session = original
+
+    def test_registered_node_is_reported(self):
+        session = _FakeSession(_FakeResponse(200, {"DistributedWorkerSink": {"name": "DistributedWorkerSink"}}))
+        self.assertTrue(self._probe(session))
+        self.assertEqual(session.urls, ["http://localhost/object_info/DistributedWorkerSink"])
+
+    def test_unknown_node_returns_empty_object(self):
+        # ComfyUI answers 200 {} for classes that are not registered.
+        self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, {}))))
+
+    def test_http_error_counts_as_missing(self):
+        self.assertFalse(self._probe(_FakeSession(_FakeResponse(404, None))))
+
+    def test_bad_json_counts_as_missing(self):
+        self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, ValueError("bad json")))))
+
+    def test_connection_error_counts_as_missing(self):
+        self.assertFalse(self._probe(_FakeSession(error=OSError("unreachable"))))

tests/helpers/v3_runtime_check.py

@@ -43,6 +43,10 @@ ADDED_OPTIONAL_INPUTS = {
     },
 }
 
+# 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'}
+
 
 def normalized_input(value):
     kind = value[0]
@@ -222,7 +226,32 @@ async def check_execution(mapping, module, prompt_server, asset_manager):
             with Image.open(preview) as image:
                 assert image.size == (8, 8) and image.mode == 'RGB'
                 assert image.getextrema() == ((0, 0), (0, 0), (0, 0))
-    print('V3_EXECUTION_OK eight nodes; upscale GPU boundary mocked; preview artifact verified')
+
+    # A worker prompt terminated by the sink validates and runs through the
+    # real executor without writing anything to the temp directory.
+    sink_graph = {
+        '1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 4, 'color': 0}},
+        '2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
+                                                              'pass_through': True}},
+        '3': {'class_type': 'DistributedWorkerSink', 'inputs': {'value': ['2', 0]}},
+    }
+    valid = await execution.validate_prompt('v3-sink', sink_graph, None)
+    assert valid[0], valid
+    assert valid[2] == ['3'], valid[2]
+    collected = []
+    def record_run(self, **kwargs):
+        collected.append(kwargs['images'])
+        return (kwargs['images'][0], None)
+    with tempfile.TemporaryDirectory(prefix='v3-sink-', dir=scratch) as temp:
+        with patch.object(folder_paths, 'temp_directory', temp), patch.object(collector_runtime, 'run', record_run):
+            executor = execution.PromptExecutor(
+                prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
+            await asyncio.to_thread(executor.execute, sink_graph, 'v3-sink', {}, valid[2])
+            assert executor.success, executor.status_messages
+            assert len(collected) == 1 and collected[0][0].shape == (4, 8, 8, 3), collected
+            assert list(Path(temp).rglob('*')) == [], list(Path(temp).rglob('*'))
+            assert executor.history_result['outputs'].get('3', {}) == {}, executor.history_result
+    print('V3_EXECUTION_OK nine nodes; upscale GPU boundary mocked; preview artifact verified; sink writes nothing')
 
 
 async def main():
@@ -235,9 +264,15 @@ async def main():
     extension = await module.comfy_entrypoint()
     classes = await extension.get_node_list()
     mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
-    assert len(classes) == len(mapping) == len(BASELINE) == 8
-    assert set(mapping) == set(BASELINE)
+    assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES) == 9
+    assert set(mapping) == set(BASELINE) | ADDED_NODES
+    for node_id in ADDED_NODES:
+        assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id]
+    sink_info = mapping['DistributedWorkerSink'].GET_NODE_INFO_V1()
+    assert sink_info['output_node'] and sink_info['dev_only'], sink_info
     for node_id, cls in mapping.items():
+        if node_id in ADDED_NODES:
+            continue
         assert issubclass(cls, io.ComfyNode)
         assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls
         old = dict(BASELINE[node_id])

tests/test_prompt_transform.py

@@ -260,6 +260,40 @@ class PrunePromptForWorkerTests(unittest.TestCase):
         preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
         self.assertEqual(len(preview_nodes), 0)
 
+    def test_worker_sink_replaces_preview_image(self):
+        result = pt.prune_prompt_for_worker(_linear_prompt(), use_worker_sink=True)
+        sinks = [n for n in result.values() if n.get("class_type") == pt.WORKER_SINK_CLASS]
+        self.assertEqual(len(sinks), 1)
+        self.assertEqual(sinks[0]["inputs"], {"value": ["4", 0]})
+        self.assertFalse(any(n.get("class_type") in ("PreviewImage", "PreviewAudio") for n in result.values()))
+
+    def test_worker_sink_uses_audio_output_for_audio_only_collector(self):
+        result = pt.prune_prompt_for_worker(_audio_only_prompt(), use_worker_sink=True)
+        sinks = [n for n in result.values() if n.get("class_type") == pt.WORKER_SINK_CLASS]
+        self.assertEqual([s["inputs"] for s in sinks], [{"value": ["2", 1]}])
+        self.assertFalse(any(n.get("class_type") == "PreviewAudio" for n in result.values()))
+
+    def test_worker_sink_for_distributed_upscale(self):
+        prompt = {
+            "1": {"class_type": "LoadImage", "inputs": {}},
+            "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
+            "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
+        }
+        result = pt.prune_prompt_for_worker(prompt, use_worker_sink=True)
+        sinks = [n for n in result.values() if n.get("class_type") == pt.WORKER_SINK_CLASS]
+        self.assertEqual([s["inputs"] for s in sinks], [{"value": ["2", 0]}])
+        self.assertNotIn("3", result)
+
+    def test_worker_sink_ids_do_not_collide_with_original_prompt(self):
+        prompt = _linear_prompt()
+        result = pt.prune_prompt_for_worker(prompt, use_worker_sink=True)
+        sink_ids = [k for k, n in result.items() if n.get("class_type") == pt.WORKER_SINK_CLASS]
+        self.assertTrue(all(k not in prompt for k in sink_ids))
+
+    def test_no_worker_sink_when_no_downstream(self):
+        result = pt.prune_prompt_for_worker(_collector_only_prompt(), use_worker_sink=True)
+        self.assertFalse(any(n.get("class_type") == pt.WORKER_SINK_CLASS for n in result.values()))
+
     def test_unrelated_nodes_are_pruned(self):
         prompt = {
             "1": {"class_type": "DistributedCollector", "inputs": {}},