"""Offline image-list queue/scheduler integration; no provider or GPU calls.""" import asyncio import copy import json import os from pathlib import Path import subprocess import sys import types import unittest from unittest.mock import AsyncMock, patch import torch import test_image_list_master_runtime as direct_fixture from test_serverless_integration import PKG, load_modules def pool(**updates): return dict({"id": "pool", "name": "Cloud", "type": "serverless", "enabled": True, "provider": "runpod", "endpoint_id": "endpoint", "jobs_per_run": 16, "timeout_seconds": 30, "media_profile": "comfy-media-v1"}, **updates) class MixedRuntimeTests(unittest.IsolatedAsyncioTestCase): def setUp(self): # Reuse only fixture setup, not inherited tests. The transport/slot doubles # execute real runtime staging, assignment and result ordering code. self.fixture = direct_fixture.MasterRuntimeTests() self.fixture.setUp() # This reused TestCase has no asyncio runner: invoke its patch cleanup # directly, not doCleanups(), which would silently fail and leak aiohttp. self.addCleanup(self.fixture.modules.stop) self.f = f = self.fixture self.rt = f.rt self.configured_pool = pool() f.config["workers"] = [self.configured_pool] self.cloud_calls = [] self.completed = [] self.closed = 0 self.hold = False self.missing = False self.dispatch_error = False self.poll_entered = asyncio.Event() self.dispatch_entered = asyncio.Event() self.hold_dispatch = False self.slot_count = 2 fixture = self class Controller: def __init__(self, state, event): self.state, self.event = state, event self.slots = [("serverless", "pool", index) for index in range(fixture.slot_count)] self.tasks = [] self.closed = False async def dispatch(self, slot, index): assert not self.closed fixture.cloud_calls.append((slot, index)) fixture.dispatch_entered.set() if fixture.dispatch_error: raise RuntimeError("serverless dispatch failed") async def execute(): if fixture.hold: await asyncio.Event().wait() # Make a later source item complete before the first local one. await asyncio.sleep(0.001) fixture.completed.append(index) if fixture.missing: raise RuntimeError("serverless completed without an image result") item = state.items[index] return {"image": item["image"], "mask": item["mask"], "text": item["values"]["value_1"]} state = self.state task = asyncio.create_task(execute()) self.tasks.append(task) task.add_done_callback(lambda _: self.event.set()) # Own the task before any await that might be cancelled. if fixture.hold_dispatch: await asyncio.Event().wait() return {"index": index, "task": task} def poll(self, inflight): fixture.poll_entered.set() for slot, handle in inflight.items(): if slot in self.slots and handle["task"].done(): return slot, handle["index"], handle["task"].result() return None async def close(self): self.closed = True for task in self.tasks: if not task.done(): task.cancel() await asyncio.gather(*self.tasks, return_exceptions=True) fixture.closed += 1 async def create(state, region, config, event): self.assertEqual(region["end_id"], "end") self.assertEqual(config["workers"][-1]["id"], "pool") self.controller = Controller(state, event) return self.controller self.create_controller = AsyncMock(side_effect=create) sys.modules['_master_slot_tests.image_list.serverless'] = types.SimpleNamespace( create_serverless_slots=self.create_controller) self.direct_dispatch = sys.modules['_master_slot_tests.api.orchestration.dispatch'].dispatch_worker_prompt def create(self, *, master=True, direct=True, count=9): f = self.f if direct: f.workers.append({"id": "peer", "type": "local"}) f.config["workers"] = f.workers + [self.configured_pool] f.local_delay = 0.04 context = f.create(participates=master, count=count) if (master or direct) else None if context is None: # Fixture.create assumes only direct peers; cloud-only uses a pool ID. f.items = self.rt.normalize_image_list_inputs(torch.zeros((count, 2, 2, 3)), mask=torch.zeros((count, 2, 2)), values={"value_1": [str(i) for i in range(count)]}) context = self.rt.create_master_job('start', 'end', direct_fixture.graph(), f.items, ['pool'], '') f.context = context f.state = self.rt.get_master_job(context) else: f.state.worker_ids.append('pool') return context async def test_mixed_roles_share_one_queue_and_preserve_pairing(self): context = self.create() original = copy.deepcopy(self.f.state.original_prompt) result = await self.rt.run_image_list_job(context) indices = self.f.local_calls + self.f.dispatched + [i for _, i in self.cloud_calls] self.assertEqual(sorted(indices), list(range(9))) self.assertTrue(self.f.local_calls and self.f.dispatched and self.cloud_calls) self.assertNotEqual(self.completed[0], 0) self.assertEqual(result[2], [str(i) for i in range(9)]) for index, (image, mask) in enumerate(zip(*result[:2])): self.assertTrue(torch.equal(image, self.f.items[index]['image'])) self.assertTrue(torch.equal(mask, self.f.items[index]['mask'])) self.assertEqual(self.f.state.original_prompt, original) self.f.preflight.assert_awaited_once_with(self.f.workers, require_idle=False) for call in self.direct_dispatch.await_args_list: self.assertEqual(call.kwargs['partial_execution_targets'], ['end']) self.assertEqual(self.closed, 1) self.assertFalse(self.rt._MASTER_JOBS) async def test_cloud_only_one_job_per_item_ignores_images_per_run(self): for jobs_per_run in (1, 16): with self.subTest(jobs_per_run=jobs_per_run): self.configured_pool['jobs_per_run'] = jobs_per_run self.cloud_calls.clear() result = await self.rt.run_image_list_job(self.create(master=False, direct=False, count=7)) self.assertEqual([i for _, i in self.cloud_calls], list(range(7))) self.assertEqual(result[2], [str(i) for i in range(7)]) self.assertEqual(len({slot for slot, _ in self.cloud_calls}), self.slot_count) self.f.preflight.assert_not_awaited() self.f.network.get_client_session.assert_not_awaited() self.assertFalse(self.f.slot_instances) async def test_missing_cloud_result_fails_without_refill(self): self.missing = True with self.assertRaisesRegex(RuntimeError, 'without an image result'): await self.rt.run_image_list_job(self.create(master=False, direct=False)) self.assertEqual(len(self.cloud_calls), self.slot_count) self.assertEqual(self.closed, 1) self.assertFalse(self.rt._MASTER_JOBS) async def test_controller_closed_when_direct_preflight_fails(self): self.f.preflight.side_effect = RuntimeError('direct preflight failed') with self.assertRaisesRegex(RuntimeError, 'direct preflight failed'): await self.rt.run_image_list_job(self.create()) self.assertEqual(self.closed, 1) self.assertFalse(self.cloud_calls) async def test_controller_closed_when_local_setup_fails(self): slot_module = sys.modules['_master_slot_tests.image_list.local_executor'] with patch.object(slot_module, 'LocalExecutionSlot', side_effect=RuntimeError('local setup failed')): with self.assertRaisesRegex(RuntimeError, 'local setup failed'): await self.rt.run_image_list_job(self.create()) self.assertEqual(self.closed, 1) self.assertFalse(self.cloud_calls) async def test_controller_closed_on_dispatch_failure(self): self.dispatch_error = True with self.assertRaisesRegex(RuntimeError, 'serverless dispatch failed'): await self.rt.run_image_list_job(self.create(master=False, direct=False)) self.assertEqual(len(self.cloud_calls), 1) self.assertEqual(self.closed, 1) async def test_cancelled_dispatch_without_returned_handle_is_owned(self): self.hold = self.hold_dispatch = True with patch.object(self.rt, 'SUBMISSION_CANCEL_GRACE_SECONDS', 0.01): task = asyncio.create_task(self.rt.run_image_list_job(self.create(master=False, direct=False))) await asyncio.wait_for(self.dispatch_entered.wait(), 1) task.cancel() with self.assertRaises(self.rt.ImageListCancelledError): await asyncio.wait_for(task, 1) self.assertEqual(len(self.cloud_calls), 1) self.assertEqual(self.closed, 1) self.assertTrue(all(t.done() for t in self.controller.tasks)) async def interrupt_probe(self, models): self.hold = True task = asyncio.create_task(self.rt.run_image_list_job(self.create(master=False, direct=False))) try: await asyncio.wait_for(self.poll_entered.wait(), 1) assigned = list(self.cloud_calls) models.interrupt_current_processing(True) with self.assertRaises(self.rt.ImageListCancelledError): await asyncio.wait_for(task, 1) self.assertEqual(self.cloud_calls, assigned) self.assertEqual(self.closed, 1) self.assertTrue(all(t.done() for t in self.controller.tasks)) self.assertFalse(self.rt._MASTER_JOBS) finally: models.interrupt_current_processing(False) if not task.done(): task.cancel() await asyncio.gather(task, return_exceptions=True) class QueueSelectionTests(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.q, modules, self.ps = load_modules() self.modules_patch = patch.dict(sys.modules, modules) self.modules_patch.start() self.addCleanup(self.modules_patch.stop) self.pool = pool() self.peer = {'id': 'peer', 'name': 'Peer', 'type': 'local', 'enabled': True} self.config = {'workers': [self.pool], 'settings': {'master_delegate_only': True}} self.q.load_config = lambda: self.config self.q.select_active_workers = AsyncMock(side_effect=lambda workers, ws, delegate, **kw: (workers, delegate)) self.preflight = AsyncMock(side_effect=lambda workers, **kw: {w['id']: True for w in workers}) sys.modules[PKG + '.api.orchestration.dispatch'].preflight_image_list_workers = self.preflight self.q.queue_prompt_payload = AsyncMock(return_value={'prompt_id': 'outer', 'number': 1}) self.prompt = direct_fixture.graph() async def queue(self, **kwargs): return await self.q.orchestrate_distributed_execution(self.prompt, None, 'client', **kwargs) async def test_delegate_pool_selected_once_no_direct_probe_or_source_mutation(self): original = copy.deepcopy(self.prompt) result = await self.queue() selected = self.q.select_active_workers.await_args.args[0] self.assertEqual([w['id'] for w in selected], ['pool']) prepared = self.q.queue_prompt_payload.await_args.args[0] inputs = prepared['start']['inputs'] self.assertEqual(json.loads(inputs['image_list_enabled_worker_ids']), ['pool']) self.assertIs(inputs['image_list_master_participates'], False) self.assertEqual(self.q.queue_prompt_payload.await_args.kwargs['partial_execution_targets'], ['end']) self.assertEqual(result[2], 1) self.assertEqual(self.prompt, original) self.preflight.assert_not_awaited() self.q.dispatch_worker_prompt.assert_not_awaited() self.assertEqual(self.ps.distributed_pending_jobs, {}) async def test_mixed_selection_preflights_only_direct_peers(self): self.config['workers'].insert(0, self.peer) result = await self.queue(delegate_master=False) self.assertEqual([w['id'] for w in self.preflight.await_args.args[0]], ['peer']) self.assertFalse(self.preflight.await_args.kwargs['require_idle']) inputs = self.q.queue_prompt_payload.await_args.args[0]['start']['inputs'] self.assertEqual(json.loads(inputs['image_list_enabled_worker_ids']), ['peer', 'pool']) self.assertEqual(result[2], 2) async def test_legacy_pool_rejected_before_probe_or_queue(self): self.pool['media_profile'] = 'legacy-image-v1' with self.assertRaisesRegex((ValueError, RuntimeError), 'comfy-media-v1'): await self.queue() self.q.select_active_workers.assert_not_awaited() self.q.queue_prompt_payload.assert_not_awaited() async def test_unknown_selection_rejected_before_queue(self): with self.assertRaisesRegex((ValueError, RuntimeError), 'Unknown|unknown'): await self.queue(enabled_worker_ids=['missing']) self.q.select_active_workers.assert_not_awaited() self.q.queue_prompt_payload.assert_not_awaited() async def test_disabled_pool_can_be_explicitly_selected(self): self.pool['enabled'] = False await self.queue(enabled_worker_ids=['pool']) self.assertEqual([w['id'] for w in self.q.select_active_workers.await_args.args[0]], ['pool']) async def test_explicit_empty_selection_does_not_import_cloud_adapter(self): sys.modules[PKG + '.image_list.serverless'] = None result = await self.queue(enabled_worker_ids=[], delegate_master=False) self.assertEqual(result[2], 0) self.preflight.assert_not_awaited() async def test_cloud_auth_or_capacity_failure_never_queues_master(self): for error in ('authentication failed', 'capacity exhausted'): with self.subTest(error=error): self.q.select_active_workers.side_effect = ValueError(error) with self.assertRaisesRegex(ValueError, error): await self.queue() self.q.queue_prompt_payload.assert_not_awaited() async def test_real_selection_uses_manager_health_not_direct_http(self): from test_dispatch_selection import _load_dispatch_module dispatch = _load_dispatch_module() health = AsyncMock(return_value={'ok': True, 'queueRemaining': 3}) sys.modules['dist_dispatch_testpkg.api.orchestration.serverless_runtime'] = types.SimpleNamespace( get_manager=lambda: types.SimpleNamespace(health=health)) self.q.select_active_workers = dispatch.select_active_workers with patch.object(dispatch, 'get_client_session', AsyncMock(side_effect=AssertionError('no HTTP'))), \ patch.object(dispatch, 'probe_worker', AsyncMock(side_effect=AssertionError('no direct probe'))): await self.queue() health.assert_awaited_once() self.assertEqual(health.await_args.args[0]['id'], 'pool') self.preflight.assert_not_awaited() self.q.queue_prompt_payload.reset_mock() health.side_effect = ValueError('provider authentication failed') with self.assertRaisesRegex(ValueError, 'authentication failed'): await self.queue() self.q.queue_prompt_payload.assert_not_awaited() async def test_unreachable_pool_never_enables_master_fallback(self): self.q.select_active_workers.side_effect = None self.q.select_active_workers.return_value = ([], False) with self.assertRaisesRegex(RuntimeError, 'orchestrator|Orchestrator'): await self.queue() self.q.queue_prompt_payload.assert_not_awaited() class RealInterruptTests(unittest.TestCase): @unittest.skipUnless(os.environ.get('COMFYUI_SOURCE_ROOT'), 'set COMFYUI_SOURCE_ROOT for CPU interrupt probe') def test_processing_interrupted_signal_during_cloud_result_wait(self): result = subprocess.run([sys.executable, str(Path(__file__).resolve()), '--interrupt-probe', os.environ['COMFYUI_SOURCE_ROOT']], text=True, capture_output=True, timeout=45) self.assertEqual(result.returncode, 0, result.stdout + result.stderr) self.assertIn('SERVERLESS_INTERRUPT_OK', result.stdout) def real_interrupt_probe(comfy_root): sys.path.insert(0, str(Path(comfy_root).resolve())) import comfy.cli_args comfy.cli_args.args.cpu = True comfy.cli_args.args.disable_assets = True import comfy.model_management as models real_comfy = sys.modules['comfy'] case = MixedRuntimeTests() case.setUp() try: # Keep actual thread-safe ComfyUI interruption state, not a task.cancel # surrogate or a fake processing_interrupted callback. sys.modules['comfy'] = real_comfy sys.modules['comfy.model_management'] = models asyncio.run(case.interrupt_probe(models)) finally: case.f.modules.stop() print('SERVERLESS_INTERRUPT_OK no refill, owned close, actual ComfyUI flag') if __name__ == '__main__': if len(sys.argv) > 1 and sys.argv[1] == '--interrupt-probe': real_interrupt_probe(sys.argv[2]) else: unittest.main()