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>

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