From 40678a51c450c95243e902e1a230a64b01fdb1ce Mon Sep 17 00:00:00 2001 From: ymiyad <74828771+ymiyad@users.noreply.github.com> Date: Wed, 14 Oct 2026 05:27:19 +0900 Subject: [PATCH] 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 Co-authored-by: Claude Opus 5.5 --- api/orchestration/media_sync.py | 24 ++++++++++ api/orchestration/prompt_transform.py | 23 +++++++-- api/queue_orchestration.py | 27 ++++++++--- nodes/v3.py | 25 +++++++++- tests/api/test_media_sync.py | 67 +++++++++++++++++++++++++++ tests/helpers/v3_runtime_check.py | 41 ++++++++++++++-- tests/test_prompt_transform.py | 34 ++++++++++++++ 7 files changed, 228 insertions(+), 13 deletions(-) diff --git a/api/orchestration/media_sync.py b/api/orchestration/media_sync.py index 84e51fe..0064859 100644 --- a/api/orchestration/media_sync.py +++ b/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() diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index ef25de8..498c18d 100644 --- a/api/orchestration/prompt_transform.py +++ b/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]}, diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index ab18e7f..4bfe02f 100644 --- a/api/queue_orchestration.py +++ b/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"], diff --git a/nodes/v3.py b/nodes/v3.py index 5441a0d..97976ea 100644 --- a/nodes/v3.py +++ b/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] diff --git a/tests/api/test_media_sync.py b/tests/api/test_media_sync.py index 8c40efc..ddae1f7 100644 --- a/tests/api/test_media_sync.py +++ b/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")))) diff --git a/tests/helpers/v3_runtime_check.py b/tests/helpers/v3_runtime_check.py index 309599b..b8c1ada 100644 --- a/tests/helpers/v3_runtime_check.py +++ b/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]) diff --git a/tests/test_prompt_transform.py b/tests/test_prompt_transform.py index 8e64c65..d6b7a1e 100644 --- a/tests/test_prompt_transform.py +++ b/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": {}},