40678a51c4
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>
325 lines
18 KiB
Python
325 lines
18 KiB
Python
"""Exercise the actual loader, schemas and executor in a fresh CPU process.
|
|
|
|
No HTTP listener, workers, model downloads or live installation changes.
|
|
The original contracts were captured from 32ac027 using the same core.
|
|
"""
|
|
import asyncio
|
|
import inspect
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
from unittest.mock import patch
|
|
|
|
from PIL import Image
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
COMFY_ROOT = Path(sys.argv[1]).resolve()
|
|
sys.path.insert(0, str(COMFY_ROOT))
|
|
os.environ['COMFYUI_IS_WORKER'] = '1'
|
|
import comfy.cli_args
|
|
comfy.cli_args.args.cpu = True
|
|
comfy.cli_args.args.disable_assets = True
|
|
from app.assets.manager import default_asset_manager
|
|
from comfy_api.v0_0_2 import io
|
|
import execution
|
|
import nodes
|
|
import server
|
|
import torch
|
|
|
|
BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text())
|
|
|
|
# Inputs intentionally added after the V1 baseline was captured. They must be
|
|
# optional, appended after the baseline inputs, and keep legacy behaviour as
|
|
# their default so saved workflows are unaffected.
|
|
ADDED_OPTIONAL_INPUTS = {
|
|
'DistributedCollector': {
|
|
'transfer_format': [
|
|
['legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1'],
|
|
{'default': 'legacy_png'},
|
|
],
|
|
'transfer_quality': ['INT', {'default': 90, 'min': 1, 'max': 100, 'step': 1}],
|
|
},
|
|
}
|
|
|
|
# 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]
|
|
opts = dict(value[1]) if len(value) > 1 else {}
|
|
if kind == 'STRING':
|
|
# V3 serializes the same default single-line widget explicitly.
|
|
opts.setdefault('multiline', False)
|
|
if kind == 'COMBO':
|
|
kind = opts['options']
|
|
opts = {key: val for key, val in opts.items() if key != 'options'}
|
|
if isinstance(kind, list):
|
|
# Single-selection is the V1 dropdown default as well.
|
|
opts.setdefault('multiselect', False)
|
|
return [kind, opts]
|
|
|
|
|
|
async def map_node(cls, values, extra=None):
|
|
"""Use the actual input parser and sanitized V3 executor clones."""
|
|
prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe', extra_data=extra or {})
|
|
assert not missing, missing
|
|
# Use a worker thread for synchronous nodes: collector/upscale bridge back
|
|
# to PromptServer's running loop, just as ComfyUI's prompt worker does.
|
|
def run():
|
|
return asyncio.run(execution._async_map_node_over_list(
|
|
'v3-acceptance', 'probe', cls, prepared, cls.FUNCTION, v3_data=hidden))
|
|
return await asyncio.to_thread(run)
|
|
|
|
|
|
def check_saved_workflows(mapping):
|
|
checked = 0
|
|
occurrences = 0
|
|
for path in sorted((ROOT / 'workflows').glob('*.json')):
|
|
workflow = json.loads(path.read_text())
|
|
local_nodes = {node['id']: node for node in workflow['nodes']}
|
|
for saved in workflow['nodes']:
|
|
if saved['type'] not in mapping:
|
|
continue
|
|
occurrences += 1
|
|
info = mapping[saved['type']].GET_NODE_INFO_V1()
|
|
inputs = {**info['input'].get('required', {}), **info['input'].get('optional', {})}
|
|
for socket in saved.get('inputs', []):
|
|
assert socket['name'] in inputs, (path.name, saved['id'], socket)
|
|
assert socket['type'] == inputs[socket['name']][0], (path.name, socket)
|
|
for index, output in enumerate(saved.get('outputs', [])):
|
|
assert output['type'] == info['output'][index], (path.name, index)
|
|
for link in workflow['links']:
|
|
if link[1] == saved['id']:
|
|
assert link[2] < len(info['output']), (path.name, link)
|
|
target = local_nodes[link[3]]['inputs'][link[4]]
|
|
assert target['type'] == info['output'][link[2]], (path.name, link)
|
|
checked += 1
|
|
assert checked == 5 and occurrences == 10, (checked, occurrences)
|
|
print('SAVED_WORKFLOW_CONTRACTS_OK', checked, occurrences)
|
|
|
|
|
|
async def check_execution(mapping, module, prompt_server, asset_manager):
|
|
result = await map_node(mapping['DistributedSeed'],
|
|
{'seed': 123, 'is_worker': True, 'worker_id': 'worker_2'})
|
|
assert result[0].result == (126,)
|
|
values = json.dumps({'_type': 'INT', '3': '17'})
|
|
result = await map_node(mapping['DistributedValue'],
|
|
{'default_value': '4', 'worker_values': values,
|
|
'is_worker': True, 'worker_id': 'worker_2'})
|
|
assert result[0].result == (17,)
|
|
metadata = {'workflow': {'nodes': [{'id': 'probe', 'widgets_values': []}]}}
|
|
result = await map_node(mapping['DistributedModelName'], {'text': 'model.ckpt'},
|
|
{'extra_pnginfo': metadata})
|
|
assert result[0].result == ('model.ckpt',)
|
|
assert result[0].ui == {'text': ['model.ckpt']}
|
|
assert metadata['workflow']['nodes'][0]['widgets_values'] == [['model.ckpt']]
|
|
result = await map_node(mapping['DistributedEmptyImage'],
|
|
{'height': 8, 'width': 8, 'channels': 3})
|
|
empty = result[0].result[0]
|
|
assert empty.shape == (0, 8, 8, 3) and empty.numel() == 0
|
|
images = torch.arange(10 * 8 * 8 * 3, dtype=torch.float32).reshape(10, 8, 8, 3)
|
|
for count in (1, 3, 10):
|
|
result = await map_node(mapping['ImageBatchDivider'], {'images': images, 'divide_by': count})
|
|
assert len(result[0].result) == 10
|
|
assert torch.equal(torch.cat(result[0].result[:count]), images)
|
|
audio = {'waveform': torch.arange(33, dtype=torch.float32).reshape(1, 1, 33), 'sample_rate': 24000}
|
|
result = await map_node(mapping['AudioBatchDivider'], {'audio': audio, 'divide_by': 3})
|
|
assert len(result[0].result) == 10
|
|
assert torch.equal(torch.cat([item['waveform'] for item in result[0].result[:3]], dim=-1), audio['waveform'])
|
|
# INPUT_IS_LIST must preserve all images/audio but unwrap transport scalars.
|
|
collector = mapping['DistributedCollector']
|
|
prepared, missing, hidden = execution.get_input_data(
|
|
{'load_balance': False, 'multi_job_id': '', 'pass_through': True}, collector, 'collector-list')
|
|
assert not missing
|
|
prepared['images'] = [images[:2], images[2:]]
|
|
prepared['audio'] = [audio, audio]
|
|
result = await execution._async_map_node_over_list(
|
|
'v3-acceptance', 'collector-list', collector, prepared, collector.FUNCTION, v3_data=hidden)
|
|
assert torch.equal(result[0].result[0], images)
|
|
assert result[0].result[1]['waveform'].shape[-1] == 66
|
|
# Exercise real async aggregation without an HTTP endpoint or worker.
|
|
queue = asyncio.Queue()
|
|
await queue.put({'worker_id': 'worker_2', 'tensor': images[2:3], 'image_index': 0, 'is_last': True})
|
|
prompt_server.distributed_pending_jobs['v3-aggregate'] = queue
|
|
result = await map_node(collector,
|
|
{'images': images[:2], 'load_balance': False,
|
|
'multi_job_id': 'v3-aggregate', 'enabled_worker_ids': '["worker_2"]'})
|
|
assert torch.equal(result[0].result[0], images[:3])
|
|
assert 'v3-aggregate' not in prompt_server.distributed_pending_jobs
|
|
|
|
# The registered node must forward the transfer options to the runtime
|
|
# helper, and default to the legacy transfer when they are not set.
|
|
collector_runtime = sys.modules[module.__name__ + '.nodes.collector'].DistributedCollectorNode
|
|
forwarded = []
|
|
def fake_run(self, **kwargs):
|
|
forwarded.append(kwargs)
|
|
return (kwargs['images'], kwargs['audio'])
|
|
with patch.object(collector_runtime, 'run', fake_run):
|
|
await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True,
|
|
'transfer_format': 'av1', 'transfer_quality': 70})
|
|
await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True})
|
|
# Set values arrive list-wrapped (is_input_list); unset ones use the
|
|
# execute() defaults. The runtime helper unwraps both forms.
|
|
assert [(kw['transfer_format'], kw['transfer_quality']) for kw in forwarded] == [
|
|
(['av1'], [70]), ('legacy_png', 90)], forwarded
|
|
|
|
# Prove V3's per-execution class clones do not leak helper instance state;
|
|
# replace only the GPU/model boundary, not the input parser or executor.
|
|
upscale_runtime = sys.modules[module.__name__ + '.nodes.distributed_upscale'].UltimateSDUpscaleDistributed
|
|
seen = []
|
|
def fake_upscale(self, *args):
|
|
assert not hasattr(self, 'acceptance_marker')
|
|
self.acceptance_marker = True
|
|
seen.append(args)
|
|
return (args[0],)
|
|
inputs = {'upscaled_image': images, 'model': object(), 'positive': [], 'negative': [],
|
|
'vae': object(), 'seed': 1, 'steps': 1, 'cfg': 1.0,
|
|
'sampler_name': 'euler', 'scheduler': 'normal', 'denoise': 0.5,
|
|
'tile_width': 64, 'tile_height': 64, 'padding': 0, 'mask_blur': 0,
|
|
'force_uniform_tiles': True, 'tiled_decode': False,
|
|
'multi_job_id': 'tile-job', 'is_worker': True, 'master_url': 'http://master.invalid',
|
|
'enabled_worker_ids': '["worker_2"]', 'worker_id': 'worker_2',
|
|
'tile_indices': '[2,3]', 'dynamic_threshold': 9}
|
|
with patch.object(upscale_runtime, 'run', fake_upscale):
|
|
for _ in range(2):
|
|
result = await map_node(mapping['UltimateSDUpscaleDistributed'], inputs)
|
|
assert result[0].result[0] is images
|
|
assert len(seen) == 2 and seen[0][-7:] == tuple(inputs[name] for name in (
|
|
'multi_job_id', 'is_worker', 'master_url', 'enabled_worker_ids', 'worker_id',
|
|
'tile_indices', 'dynamic_threshold')), seen[0]
|
|
import math
|
|
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id='tile-job'))
|
|
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id=''))
|
|
|
|
graph = {
|
|
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 10, 'color': 0}},
|
|
'2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
|
|
'pass_through': True}},
|
|
'3': {'class_type': 'ImageBatchDivider', 'inputs': {'images': ['2', 0], 'divide_by': 10}},
|
|
'4': {'class_type': 'PreviewImage', 'inputs': {'images': ['3', 9]}},
|
|
}
|
|
valid = await execution.validate_prompt('v3-graph', graph, None)
|
|
assert valid[0], valid
|
|
invalid = {'1': {'class_type': 'UltimateSDUpscaleDistributed',
|
|
'inputs': {**inputs, 'sampler_name': 'INVALID_ENUM', 'scheduler': 'INVALID_ENUM'}},
|
|
'2': {'class_type': 'PreviewImage', 'inputs': {'images': ['1', 0]}}}
|
|
rejected = await execution.validate_prompt('v3-invalid', invalid, None)
|
|
assert not rejected[0]
|
|
errors = [error for entry in rejected[3].values() for error in entry['errors']]
|
|
enum_errors = [error['extra_info']['input_name'] for error in errors if error['type'] == 'value_not_in_list']
|
|
assert {'sampler_name', 'scheduler'} <= set(enum_errors), errors
|
|
import folder_paths
|
|
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
|
|
with tempfile.TemporaryDirectory(prefix='v3-preview-', dir=scratch) as temp:
|
|
with patch.object(folder_paths, 'temp_directory', temp):
|
|
executor = execution.PromptExecutor(
|
|
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
|
|
await asyncio.to_thread(executor.execute, graph, 'v3-graph', {}, valid[2])
|
|
assert executor.success, executor.status_messages
|
|
history = executor.history_result
|
|
record = history['outputs']['4']['images'][0]
|
|
preview = Path(temp) / record.get('subfolder', '') / record['filename']
|
|
with Image.open(preview) as image:
|
|
assert image.size == (8, 8) and image.mode == 'RGB'
|
|
assert image.getextrema() == ((0, 0), (0, 0), (0, 0))
|
|
|
|
# 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():
|
|
asset_manager = default_asset_manager()
|
|
prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager)
|
|
assert not (ROOT / 'distributed.py').exists(), 'obsolete root bootstrap remains'
|
|
assert await nodes.load_custom_node(str(ROOT)), 'ComfyUI loader rejected the pack'
|
|
module = sys.modules[str(ROOT).replace('.', '_x_')]
|
|
assert not hasattr(module, 'NODE_CLASS_MAPPINGS'), 'V1 map shadows V3 entrypoint'
|
|
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) + 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])
|
|
if node_id in ('ImageBatchDivider', 'AudioBatchDivider'):
|
|
# V1's ByPassTypeTuple advertises '*' when indexed, while its
|
|
# underlying tuple and existing frontend declare IMAGE/AUDIO.
|
|
# Native V3 declares all ten existing typed sockets explicitly.
|
|
assert old['output'] == ['*'] * 10
|
|
old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10
|
|
new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
|
|
for group in ('required', 'optional'):
|
|
old_inputs = old['input'].get(group, {})
|
|
new_inputs = dict(new['input'].get(group, {}))
|
|
if group == 'optional':
|
|
for name, expected in ADDED_OPTIONAL_INPUTS.get(node_id, {}).items():
|
|
assert name in new_inputs, (node_id, name, 'added input missing')
|
|
actual = normalized_input(new_inputs.pop(name))
|
|
kind, opts = normalized_input(expected)
|
|
assert actual[0] == kind, (node_id, name, actual)
|
|
for key, value in opts.items():
|
|
assert actual[1].get(key) == value, (node_id, name, key, actual)
|
|
assert list(old_inputs) == list(new_inputs), (node_id, group, 'input order')
|
|
for name, original in old_inputs.items():
|
|
assert normalized_input(original) == normalized_input(new_inputs[name]), (node_id, name, original, new_inputs[name])
|
|
for key in ('output', 'output_name', 'output_is_list', 'is_input_list', 'output_node', 'category', 'display_name'):
|
|
assert old[key] == new[key], (node_id, key, old[key], new[key])
|
|
# Standard context lives in cls.hidden; orchestrator metadata remains
|
|
# accepted by its original kwarg name, without creating new widgets.
|
|
signature = inspect.signature(cls.execute)
|
|
for name, field in old['input'].get('hidden', {}).items():
|
|
if isinstance(field, list):
|
|
assert cls.GET_SCHEMA().accept_all_inputs, node_id
|
|
assert name in signature.parameters, (node_id, name)
|
|
assert signature.parameters[name].default == field[1]['default'], (node_id, name)
|
|
else:
|
|
assert name in new['input']['hidden'], (node_id, name)
|
|
expected_hidden = []
|
|
if node_id == 'DistributedModelName':
|
|
expected_hidden.extend(['unique_id', 'extra_pnginfo'])
|
|
if old['output_node']:
|
|
expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden)
|
|
assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden'))
|
|
print('SCHEMA_PARITY_OK', len(mapping))
|
|
check_saved_workflows(mapping)
|
|
await check_execution(mapping, module, prompt_server, asset_manager)
|
|
print('V3_ACCEPTANCE_OK')
|
|
|
|
|
|
asyncio.run(main())
|