nodes/robertvoy-ComfyUI-Distributed
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: preserve required output utilities in image-list execution

Browse Source
This commit is contained in:
Robert Wojciechowski
2026-10-13 23:26:25 +00:00
parent 41caf02481
commit c30036ae31
21 changed files with 985 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/helpers/image_list_network_check.py
+177 -57
View File
@@ -1,11 +1,12 @@
"""Model-free multi-process acceptance using real ComfyUI executors and loopback HTTP.
Only the transport test facade (/prompt, /queue, /history, /interrupt) is synthetic.
Isolated listeners use real core /prompt, /queue, /history and /interrupt handlers.
Start/End, graph extraction, scheduling, payload codecs and callbacks are production code.
No managed ComfyUI service, user workflow, model, or remote worker is touched.
"""
import asyncio
from collections import Counter
import copy
import importlib
import json
import os
@@ -36,6 +37,7 @@ import aiohttp
torch.set_num_threads(1)
SOURCE_CALLS, BRANCH_CALLS, CAPTURED, MODEL_CALLS = [], [], [], []
UTILITY_CALLS, ITEM_EXECUTIONS = [], []
ALLOW_MASTER = False
COUNT = 7
@@ -85,6 +87,51 @@ class TestCachedModel:
return (object(),)
def assert_item_execution(prompt, index):
"""Prove the operation belongs to an item, not merely a participating host."""
from comfy_execution.utils import get_executing_context
context = get_executing_context()
assert context is not None
assert prompt['2']['inputs'].get('worker_item_token'), 'per-item node ran in outer graph'
running, _ = server.PromptServer.instance.prompt_queue.get_current_queue_volatile()
assert len(running) == 1, running
entry = running[0]
if WORKER:
assert entry[1] == context.prompt_id and entry[2] == prompt
assert entry[4] == ['4'], entry[4]
else:
assert ALLOW_MASTER, 'per-item node ran on delegate-only master'
assert entry[1] != context.prompt_id, 'per-item node ran in outer executor'
assert set(entry[4]) == {'4', '5'}, entry[4]
return {'index': index, 'prompt_id': context.prompt_id,
'role': 'worker' if WORKER else 'master-item', 'queue_prompt_id': entry[1]}
class TestRequiredOutputUtility:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',), 'index': ('INT', {'forceInput': True})},
'hidden': {'prompt': 'PROMPT'}}
OUTPUT_NODE = True
RETURN_TYPES = ('IMAGE',)
FUNCTION = 'run'
def run(self, image, index, prompt):
assert_item_execution(prompt, index)
UTILITY_CALLS.append(index)
return (image,)
class TestDiscardedPreview:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',)}}
OUTPUT_NODE = True
RETURN_TYPES = ()
FUNCTION = 'run'
def run(self, image):
raise AssertionError('discarded per-item output preview executed')
class TestImageOperation:
@classmethod
def INPUT_TYPES(cls):
@@ -92,11 +139,13 @@ class TestImageOperation:
'text': ('STRING', {'forceInput': True}),
'strength': ('FLOAT', {'forceInput': True}),
'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True}),
'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)}}
'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)},
'hidden': {'prompt': 'PROMPT'}}
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
FUNCTION = 'run'
def run(self, image, reference_image, mask, text, strength, seed, index, flag, model):
def run(self, image, reference_image, mask, text, strength, seed, index, flag, model, prompt):
assert WORKER or ALLOW_MASTER, 'per-image branch executed in orchestrator-only mode'
ITEM_EXECUTIONS.append(assert_item_execution(prompt, index))
BRANCH_CALLS.append(index)
# Sleep in short chunks so the real Comfy interrupt signal can be tested.
until = time.monotonic() + DELAY
@@ -120,6 +169,7 @@ class TestCapture:
RETURN_TYPES = ()
FUNCTION = 'run'
def run(self, images, masks, texts):
assert not WORKER, 'post-End capture leaked into a worker prompt'
CAPTURED.append((images, masks, texts))
return ()
@@ -131,7 +181,9 @@ async def setup():
module = sys.modules[str(ROOT).replace('.', '_x_')]
nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource,
TestImageOperation=TestImageOperation, TestCapture=TestCapture,
TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag)
TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag,
TestRequiredOutputUtility=TestRequiredOutputUtility,
TestDiscardedPreview=TestDiscardedPreview)
# Match the CORS-enabled direct-worker launch contract for canvas probes.
app = web.Application(middlewares=[server.create_cors_middleware('*')])
# Register the actual feature routes, not test replacements.
@@ -147,61 +199,56 @@ async def listen(app):
return runner, site._server.sockets[0].getsockname()[1]
def finish_queue_item(ps, item_id, executor):
ps.prompt_queue.task_done(item_id, executor.history_result,
status=execution.PromptQueue.ExecutionStatus(
status_str='success' if executor.success else 'error',
completed=executor.success, messages=executor.status_messages),
process_item=lambda item: item[:5] + item[6:])
async def worker_main():
ps, assets, module, app = await setup()
pending, running, history = [], {}, {}
wake = asyncio.Event()
interrupted_ids = []
interrupted_ids, submissions = [], []
core_handlers = {(route.method, route.path): route.handler for route in ps.routes}
async def prompt(request):
data = await request.json()
prompt_id = str(uuid.uuid4())
graph = data['prompt']
valid = await execution.validate_prompt(prompt_id, graph, None)
if not valid[0]:
return web.json_response({'error': valid[1], 'node_errors': valid[3]}, status=400)
pending.append((prompt_id, graph, valid[2]))
wake.set()
return web.json_response({'prompt_id': prompt_id})
async def queue(request):
if request.method == 'POST':
deleted = (await request.json()).get('delete', [])
pending[:] = [entry for entry in pending if entry[0] not in deleted]
return web.json_response({})
return web.json_response({'queue_running': [[0, key] for key in running],
'queue_pending': [[0, entry[0]] for entry in pending]})
async def get_history(request):
pid = request.match_info['prompt_id']
return web.json_response({pid: history[pid]} if pid in history else {})
# Use core's actual validation/enqueue/response schema, including
# partial_execution_targets; only observe it and wake our CPU consumer.
response = await core_handlers['POST', '/prompt'](request)
if response.status == 200:
accepted = json.loads(response.body)
_, pending = ps.prompt_queue.get_current_queue_volatile()
entry = next(item for item in pending if item[1] == accepted['prompt_id'])
submissions.append({'prompt_id': entry[1], 'requested': data.get('partial_execution_targets'),
'queued': entry[4], 'node_ids': sorted(entry[2]),
'end_ids': [key for key, node in entry[2].items()
if node['class_type'] == 'DistributedImageListEnd']})
wake.set()
return response
async def interrupt(request):
pid = (await request.json())['prompt_id']
interrupted_ids.append(pid)
if pid in running:
comfy.model_management.interrupt_current_processing(True)
return web.json_response({})
interrupted_ids.append((await request.json())['prompt_id'])
return await core_handlers['POST', '/interrupt'](request)
async def stats(request):
return web.json_response({'calls': BRANCH_CALLS, 'running': list(running),
'pending': [e[0] for e in pending], 'interrupted': interrupted_ids,
'history': history})
running, pending = ps.prompt_queue.get_current_queue_volatile()
return web.json_response({'calls': BRANCH_CALLS, 'utility_calls': UTILITY_CALLS,
'item_executions': ITEM_EXECUTIONS, 'submissions': submissions,
'running': [e[1] for e in running], 'pending': [e[1] for e in pending],
'interrupted': interrupted_ids, 'history': ps.prompt_queue.get_history()})
async def consume():
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
while True:
await wake.wait()
while pending:
pid, graph, outputs = pending.pop(0)
running[pid] = True
await asyncio.to_thread(executor.execute, graph, pid, {}, outputs)
history[pid] = {'status': {'completed': True,
'status_str': 'success' if executor.success else 'error',
'messages': str(executor.status_messages)}}
running.pop(pid, None)
while (queued := ps.prompt_queue.get(timeout=0)) is not None:
entry, item_id = queued
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
finish_queue_item(ps, item_id, executor)
wake.clear()
async def probe(_request):
return web.json_response({'exec_info': {'queue_remaining': len(pending) + len(running)}})
app.router.add_get('/prompt', probe)
app.add_routes([route for route in ps.routes if
(route.path in ('/queue', '/history/{prompt_id}') or
(route.path == '/prompt' and route.method == 'GET'))])
app.router.add_post('/prompt', prompt)
app.router.add_get('/queue', queue)
app.router.add_post('/queue', queue)
app.router.add_get('/history/{prompt_id}', get_history)
app.router.add_post('/interrupt', interrupt)
app.router.add_get('/test-stats', stats)
runner, port = await listen(app)
@@ -228,12 +275,14 @@ def graph(callback, worker_ids):
'3': {'class_type': 'TestImageOperation', 'inputs': {
key: ['2', index] for key, index in [('image', 0), ('mask', 1), ('index', 2),
('reference_image', 4), ('text', 5), ('strength', 6), ('seed', 7), ('flag', 8)]}
| {'model': ['6', 0]}},
| {'image': ['8', 0], 'model': ['6', 0]}},
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 3]}},
'5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2]}},
'6': {'class_type': 'TestCachedModel', 'inputs': {}},
'7': {'class_type': 'TestSharedFlag', 'inputs': {}},
'8': {'class_type': 'TestRequiredOutputUtility', 'inputs': {'image': ['2', 0], 'index': ['2', 2]}},
'9': {'class_type': 'TestDiscardedPreview', 'inputs': {'image': ['8', 0]}},
}
@@ -241,6 +290,8 @@ async def coordinator():
global ALLOW_MASTER
ps, assets, module, app = await setup()
runtime = importlib.import_module(module.__name__ + '.image_list.runtime')
execution_plan = importlib.import_module(module.__name__ + '.image_list.execution_plan')
lifecycle = importlib.import_module(module.__name__ + '.image_list.lifecycle')
config_mod = importlib.import_module(module.__name__ + '.utils.config')
orchestration = importlib.import_module(module.__name__ + '.api.queue_orchestration')
runner, port = await listen(app)
@@ -294,16 +345,49 @@ async def coordinator():
prompt = graph(f'http://127.0.0.1:{port}', [w['id'] for w in workers])
valid = await execution.validate_prompt('network-acceptance', prompt, None)
assert valid[0], valid
assert {'4', '5', '8', '9'} == set(valid[2]), valid
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
# Exercise production queue preparation; only master enqueue is redirected
# to this isolated executor, not the managed server's live queue.
prompt['2']['inputs'].pop('image_list_enabled_worker_ids')
prompt['2']['inputs'].pop('image_list_master_callback')
saved_prompt = copy.deepcopy(prompt)
outer_executions = []
async def enqueue(prepared, workflow, client_id, **kwargs):
checked = await execution.validate_prompt('network-acceptance', prepared, None)
targets = kwargs['partial_execution_targets']
planned, planned_targets = execution_plan.prepare_image_list_execution(prepared)
assert planned == prepared and targets == planned_targets
if '9' in saved_prompt and 'TestImageOperation' in {
node['class_type'] for node in prepared.values()}:
assert set(prepared) == set(saved_prompt) - {'9'}, prepared
assert prepared['8'] == saved_prompt['8'], 'required utility changed'
assert prepared['5'] == saved_prompt['5'], 'post-End output changed'
assert set(targets) == {'4', '5'}, targets
assert nodes.NODE_CLASS_MAPPINGS['TestRequiredOutputUtility'].OUTPUT_NODE is True
assert nodes.NODE_CLASS_MAPPINGS['TestDiscardedPreview'].OUTPUT_NODE is True
pid = 'network-acceptance' if not outer_executions else str(uuid.uuid4())
checked = await execution.validate_prompt(pid, prepared, targets)
assert checked[0], checked
await asyncio.to_thread(executor.execute, prepared, 'network-acceptance', {}, checked[2])
return {'prompt_id': 'network-acceptance', 'number': 1, 'node_errors': {}}
assert not checked[3] and Counter(checked[2]) == Counter(targets), checked
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
# Exactly the core queue tuple/put/get/task_done lifecycle. Start
# must prove its roots from this real running record, not a mock.
number = len(outer_executions)
ps.prompt_queue.put((number, pid, prepared, {'client_id': client_id}, checked[2], {}))
entry, item_id = ps.prompt_queue.get(timeout=0)
assert lifecycle.has_safe_execution_targets(ps, pid, prepared)
before_captures = len(CAPTURED)
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
finish_queue_item(ps, item_id, executor)
assert ps.prompt_queue.get_current_queue_volatile() == ([], [])
history = ps.prompt_queue.get_history(pid)[pid]
assert Counter(history['prompt'][4]) == Counter(targets)
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
outer_executions.append({'prompt_id': pid, 'targets': targets,
'captures': len(CAPTURED) - before_captures})
if executor.success:
assert len(CAPTURED) == before_captures + 1, 'post-End Capture must run exactly once'
return {'prompt_id': pid, 'number': number, 'node_errors': checked[3]}
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
@@ -312,11 +396,12 @@ async def coordinator():
prompt, None, 'test', [w['id'] for w in workers]), 45)
assert queued[0] == 'network-acceptance' and queued[2] == 2, queued
assert not ps.distributed_pending_jobs, 'image-list created unused Collector queues'
assert prompt == saved_prompt, 'queue preparation changed the saved/source graph'
if not executor.success:
for err in stderr_files:
err.flush(); err.seek(0); print(err.read(), file=sys.stderr)
assert executor.success, executor.status_messages
assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS
assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS and not UTILITY_CALLS
assert len(CAPTURED) == 1
images, masks, texts = CAPTURED[0]
assert len(images) == len(masks) == len(texts) == COUNT
@@ -332,6 +417,14 @@ async def coordinator():
counts = [len(row['calls']) for row in stats]
assert counts[0] > counts[1] > 0, counts
assert Counter(i for row in stats for i in row['calls']) == Counter(range(COUNT))
for row in stats:
assert Counter(row['utility_calls']) == Counter(row['calls']), row
assert len(row['submissions']) == len(row['calls'])
assert len({event['prompt_id'] for event in row['item_executions']}) == len(row['calls'])
for submission in row['submissions']:
assert submission['requested'] == submission['queued'] == ['4'], submission
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
assert '5' not in submission['node_ids']
assert not runtime._MASTER_JOBS
# Three sequential regions on one worker, including optional outputs
# left unconnected. Region N+1 must not race the prior worker End.
@@ -346,12 +439,14 @@ async def coordinator():
previous = end
chained['8'] = {'class_type': 'TestCapture', 'inputs': {
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2]}}
saved_chained = copy.deepcopy(chained)
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
chained, None, 'test', ['slow']), 45)
assert chained == saved_chained
assert executor.success, executor.status_messages
assert len(CAPTURED) == 2 and len(CAPTURED[1][0]) == 2
assert all(torch.equal(image, torch.zeros((1, 8, 8, 3))) for image in CAPTURED[1][0])
@@ -366,29 +461,43 @@ async def coordinator():
for mode, selected in [('mixed', [w['id'] for w in workers]), ('master-only', [])]:
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
before_sources, before_local, before_models = len(SOURCE_CALLS), len(BRANCH_CALLS), len(MODEL_CALLS)
before_utilities, before_events = len(UTILITY_CALLS), len(ITEM_EXECUTIONS)
async with aiohttp.ClientSession() as session:
before_worker_calls = []
before_worker_stats = []
for worker in workers:
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
before_worker_calls.append(len((await response.json())['calls']))
before_worker_stats.append(await response.json())
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
prompt, None, 'test', selected), 45)
assert prompt == saved_prompt, (mode, 'saved/source graph changed')
assert executor.success, (mode, executor.status_messages)
assert queued[2] == len(selected)
assert len(SOURCE_CALLS) == before_sources + 1, 'source was reevaluated per local item'
local_calls = BRANCH_CALLS[before_local:]
assert Counter(UTILITY_CALLS[before_utilities:]) == Counter(local_calls), mode
assert len({event['prompt_id'] for event in ITEM_EXECUTIONS[before_events:]}) == len(local_calls)
assert local_calls, f'{mode}: master did not process any items'
assert len(MODEL_CALLS) == before_models + 1, 'local model dependency did not stay cached'
remote_calls = []
async with aiohttp.ClientSession() as session:
for worker, offset in zip(workers, before_worker_calls):
for worker, before_stats in zip(workers, before_worker_stats):
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
current = await response.json()
remote_calls.extend(current['calls'][offset:])
calls = current['calls'][len(before_stats['calls']):]
utilities = current['utility_calls'][len(before_stats['utility_calls']):]
submissions = current['submissions'][len(before_stats['submissions']):]
events = current['item_executions'][len(before_stats['item_executions']):]
assert Counter(utilities) == Counter(calls), (mode, current)
assert len(submissions) == len(calls) == len({event['prompt_id'] for event in events})
for submission in submissions:
assert submission['requested'] == submission['queued'] == ['4'], submission
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
assert '5' not in submission['node_ids']
remote_calls.extend(calls)
assert not current['running'] and not current['pending'], current
if mode == 'mixed':
assert remote_calls, 'mixed execution failed to use direct workers'
@@ -404,8 +513,19 @@ async def coordinator():
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_calls)}
print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts))
assert len(outer_executions) == 4 and all(row['captures'] == 1 for row in outer_executions)
print('IMAGE_LIST_TARGETS_OK ' + json.dumps({
'modes': ['delegate-only', 'mixed', 'master-only'],
'outer_targets': outer_executions[0]['targets'], 'worker_targets': ['4'],
'required_utility_once_per_item': True, 'preview_discarded': True,
'saved_prompt_unchanged': prompt == saved_prompt, 'capture_once_per_outer_prompt': True}))
ALLOW_MASTER = False
config['settings']['master_delegate_only'] = True
# Preserve the original direct-runtime cancellation/failure fixture:
# it has no outer queue and must not claim target-control evidence.
direct_prompt = copy.deepcopy(prompt)
del direct_prompt['8'], direct_prompt['9']
direct_prompt['3']['inputs']['image'] = ['2', 0]
# Real Comfy user-interrupt flag while a real worker executor is active.
slow = workers[1]
async with aiohttp.ClientSession() as session:
@@ -420,7 +540,7 @@ async def coordinator():
await asyncio.sleep(0.02)
raise AssertionError('worker state transition timed out')
before = await until(lambda state: not state['running'] and not state['pending'])
context = runtime.create_master_job('2', '4', prompt, source_rows(), ['slow'],
context = runtime.create_master_job('2', '4', direct_prompt, source_rows(), ['slow'],
f'http://127.0.0.1:{port}')
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
@@ -444,7 +564,7 @@ async def coordinator():
# A worker exception before End must fail, not hang or return a partial list.
bad_rows = source_rows()
bad_rows[0]['seed'] = -1
context = runtime.create_master_job('2', '4', prompt, bad_rows, ['fast'],
context = runtime.create_master_job('2', '4', direct_prompt, bad_rows, ['fast'],
f'http://127.0.0.1:{port}')
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):