359 lines
18 KiB
Python
359 lines
18 KiB
Python
"""Exercise production per-item orchestration through controlled provider seams."""
|
|
import asyncio
|
|
import importlib.util
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
|
|
import torch
|
|
|
|
from test_image_list import graph_one_region
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
class Response:
|
|
def __init__(self, data=None, status=200):
|
|
self.data, self.status = data, status
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *_args):
|
|
return None
|
|
|
|
async def json(self):
|
|
return self.data
|
|
|
|
def raise_for_status(self):
|
|
if self.status >= 400:
|
|
raise RuntimeError(f'HTTP {self.status}')
|
|
|
|
|
|
class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
name = '_image_list_boundaries'
|
|
self.modules = {}
|
|
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
|
|
module = types.ModuleType(name + suffix)
|
|
module.__path__ = []
|
|
self.modules[name + suffix] = module
|
|
self.patch_modules = patch.dict(sys.modules, self.modules)
|
|
self.patch_modules.start()
|
|
self.addCleanup(self.patch_modules.stop)
|
|
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
|
|
self.rt = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = self.rt
|
|
spec.loader.exec_module(self.rt)
|
|
self.worker = {'id': 'worker', 'host': 'worker', 'type': 'remote', 'port': 8188}
|
|
self.config = {'workers': [self.worker], 'settings': {}}
|
|
self.prompts, self.uploads, self.deletions, self.cancel_requests = [], [], [], []
|
|
self.reject = False
|
|
self.interrupted = False
|
|
self.completed = True
|
|
self.queue = {'queue_running': [], 'queue_pending': []}
|
|
self.callback = Mock(side_effect=lambda worker, **_kwargs:
|
|
'http://127.0.0.1:8188' if worker['type'] == 'local' else 'http://public-master')
|
|
self.sync = AsyncMock(side_effect=self.sync_media)
|
|
self.dispatch = AsyncMock(side_effect=self.dispatch_prompt)
|
|
self.logs = []
|
|
dependencies = {
|
|
'.api.orchestration.dispatch': types.SimpleNamespace(
|
|
preflight_image_list_workers=AsyncMock(return_value={'worker': True}),
|
|
dispatch_worker_prompt=self.dispatch),
|
|
'.api.orchestration.prompt_transform': types.SimpleNamespace(
|
|
_resolve_direct_workers_for_image_list=lambda config, ids: config['workers']),
|
|
'.api.orchestration.media_sync': types.SimpleNamespace(
|
|
sync_worker_media=self.sync, fetch_worker_path_separator=AsyncMock(return_value='/'),
|
|
convert_paths_for_platform=lambda prompt, separator: prompt),
|
|
'.utils.config': types.SimpleNamespace(load_config=lambda: self.config,
|
|
get_worker_timeout_seconds=lambda: 60,
|
|
get_image_list_job_timeout_seconds=lambda config=None: 0),
|
|
'.utils.network': types.SimpleNamespace(build_master_url=lambda **kwargs: 'http://public-master',
|
|
build_master_callback_url=self.callback,
|
|
build_worker_url=lambda worker, path='': f"http://{worker['id']}{path}",
|
|
get_client_session=AsyncMock(return_value=self)),
|
|
'.utils.logging': types.SimpleNamespace(log=self.logs.append),
|
|
}
|
|
for suffix, module in dependencies.items():
|
|
sys.modules[name + suffix] = module
|
|
self.config_module = dependencies['.utils.config']
|
|
# Other legacy test modules replace aiohttp during collection. This
|
|
# controlled provider fixture needs only its explicit timeout/error seam.
|
|
sys.modules['aiohttp'] = types.SimpleNamespace(
|
|
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError, ClientConnectionError=ConnectionError)
|
|
comfy = types.ModuleType('comfy')
|
|
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted)
|
|
sys.modules['comfy'] = comfy
|
|
sys.modules['comfy.model_management'] = comfy.model_management
|
|
|
|
async def sync_media(self, worker, prompt):
|
|
prompt['7']['inputs']['image'] = 'shared/ref.png'
|
|
|
|
async def dispatch_prompt(self, worker, prompt, meta, **kwargs):
|
|
self.assertEqual(kwargs.get('partial_execution_targets'), ['5'])
|
|
if self.reject:
|
|
raise RuntimeError('prompt rejected')
|
|
self.prompts.append((json.loads(json.dumps(prompt)), meta))
|
|
inputs = prompt['5']['inputs']
|
|
token = inputs['worker_result_token']
|
|
item = self.rt.consume_worker_item(token)
|
|
self.rt.set_master_result(self.context, inputs['worker_item_index'], token, {'image': item['image']})
|
|
return f'prompt-{len(self.prompts)}'
|
|
|
|
def post(self, url, *, json=None, data=None, **kwargs):
|
|
if data is not None:
|
|
json = __import__('json').loads(data)
|
|
if url.endswith('/distributed/image-list/item'):
|
|
self.uploads.append(json['token'])
|
|
self.rt.store_worker_item(json['token'], json['item'])
|
|
else:
|
|
self.cancel_requests.append((url, json))
|
|
return Response({})
|
|
|
|
def delete(self, url, **kwargs):
|
|
token = url.rsplit('/', 1)[-1]
|
|
self.deletions.append(token)
|
|
self.rt.delete_worker_item(token)
|
|
return Response({})
|
|
|
|
def get(self, url, **kwargs):
|
|
if url.endswith('/queue'):
|
|
return Response(self.queue)
|
|
pid = url.rsplit('/', 1)[-1]
|
|
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}}
|
|
if self.completed else {})
|
|
|
|
def create(self, count=2):
|
|
prompt = graph_one_region()
|
|
prompt['7'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared/ref.png [output]'}}
|
|
prompt['4']['inputs']['reference'] = ['7', 0]
|
|
self.metadata = {'nodes': [{'id': 42}], 'extra': {'label': 'generic fixture'}}
|
|
self.context = self.rt.create_master_job('3', '5', prompt,
|
|
[{'image': torch.zeros((1, 2, 2, 3)), 'index': i} for i in range(count)],
|
|
['worker'], 'http://public-master', self.metadata)
|
|
return self.context
|
|
|
|
async def test_runtime_preflight_rejects_old_target_capability_before_upload(self):
|
|
from test_dispatch_selection import _load_dispatch_module, _FakeSession, _valid_capability
|
|
real_dispatch = _load_dispatch_module()
|
|
capability = _valid_capability()
|
|
capability.pop('execution_targets')
|
|
session = _FakeSession({('worker', 'capabilities'): (200, capability)})
|
|
module = sys.modules['_image_list_boundaries.api.orchestration.dispatch']
|
|
module.preflight_image_list_workers = real_dispatch.preflight_image_list_workers
|
|
with patch.object(real_dispatch, 'get_client_session', AsyncMock(return_value=session)), patch.object(
|
|
real_dispatch, 'build_worker_url', side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
|
):
|
|
with self.assertRaisesRegex(RuntimeError, 'execution_targets.*upgrade'):
|
|
await self.rt.run_image_list_job(self.create(1))
|
|
self.assertFalse(self.uploads)
|
|
self.dispatch.assert_not_awaited()
|
|
self.assertFalse(self.rt._MASTER_JOBS)
|
|
|
|
async def test_each_assignment_reuses_media_paths_with_fresh_tokens_and_metadata(self):
|
|
context = self.create()
|
|
await self.rt.run_image_list_job(context)
|
|
self.assertEqual(self.sync.await_count, 1)
|
|
self.assertEqual([p['7']['inputs']['image'] for p, _ in self.prompts],
|
|
['shared/ref.png', 'shared/ref.png'])
|
|
tokens = [p['5']['inputs']['worker_result_token'] for p, _ in self.prompts]
|
|
self.assertEqual(len(set(tokens)), 2)
|
|
self.assertEqual([p['5']['inputs']['worker_item_index'] for p, _ in self.prompts], [0, 1])
|
|
self.assertEqual([meta for _, meta in self.prompts], [self.metadata, self.metadata])
|
|
|
|
async def test_annotated_media_roots_survive_windows_prep_and_repeated_items(self):
|
|
import tempfile
|
|
|
|
context = self.create(2)
|
|
state = self.rt.get_master_job(context)
|
|
state.original_prompt['8'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared\\ref.png [temp]'}}
|
|
state.original_prompt['4']['inputs']['second_reference'] = ['8', 0]
|
|
name = '_image_list_boundaries.api.orchestration.media_sync'
|
|
spec = importlib.util.spec_from_file_location(name, ROOT / 'api/orchestration/media_sync.py')
|
|
assert spec is not None and spec.loader is not None
|
|
media = importlib.util.module_from_spec(spec)
|
|
sys.modules['_image_list_boundaries.utils.logging'].debug_log = self.logs.append
|
|
sys.modules['_image_list_boundaries.utils.trace_logger'] = types.SimpleNamespace(
|
|
trace_debug=lambda *args: None, trace_info=lambda *args: None)
|
|
spec.loader.exec_module(media)
|
|
media.fetch_worker_path_separator = AsyncMock(return_value='\\')
|
|
converted = Mock(wraps=media.convert_paths_for_platform)
|
|
media.convert_paths_for_platform = converted
|
|
uploads = []
|
|
|
|
async def upload(worker, filename, content, digest, mime, **kwargs):
|
|
uploads.append((filename, content))
|
|
return True, f'worker-input/{len(uploads)}.png'
|
|
|
|
media._upload_media_to_worker = AsyncMock(side_effect=upload)
|
|
media.sync_worker_media = AsyncMock(wraps=media.sync_worker_media)
|
|
with tempfile.TemporaryDirectory(prefix='image-list-media-') as directory:
|
|
root = Path(directory)
|
|
for kind in ('input', 'output', 'temp'):
|
|
path = root / kind / 'shared/ref.png'
|
|
path.parent.mkdir(parents=True)
|
|
path.write_bytes(kind.encode())
|
|
|
|
def resolve_annotated(filename):
|
|
for kind in ('output', 'temp', 'input'):
|
|
suffix = f' [{kind}]'
|
|
if filename.endswith(suffix):
|
|
return str(root / kind / filename[:-len(suffix)])
|
|
return str(root / 'input' / filename)
|
|
|
|
with patch.dict(sys.modules, {name: media, 'folder_paths': types.SimpleNamespace(
|
|
get_annotated_filepath=resolve_annotated)}):
|
|
await self.rt.run_image_list_job(context)
|
|
self.assertEqual(uploads, [('shared/ref.png [output]', b'output'),
|
|
('shared/ref.png [temp]', b'temp')])
|
|
media.sync_worker_media.assert_awaited_once()
|
|
media.fetch_worker_path_separator.assert_awaited_once()
|
|
converted.assert_called_once()
|
|
self.assertEqual(converted.call_args.args[1], '\\')
|
|
for prompt, _ in self.prompts:
|
|
self.assertEqual(prompt['7']['inputs']['image'], 'worker-input/1.png')
|
|
self.assertEqual(prompt['8']['inputs']['image'], 'worker-input/2.png')
|
|
self.assertEqual(state.original_prompt['7']['inputs']['image'], 'shared/ref.png [output]')
|
|
|
|
async def test_local_worker_callback_does_not_use_public_master_address(self):
|
|
self.worker['type'] = 'local'
|
|
await self.rt.run_image_list_job(self.create(1))
|
|
self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://127.0.0.1:8188')
|
|
self.callback.assert_called()
|
|
|
|
async def test_remote_worker_uses_topology_callback_helper(self):
|
|
await self.rt.run_image_list_job(self.create(1))
|
|
self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://public-master')
|
|
self.callback.assert_called()
|
|
|
|
def test_metadata_is_forwarded_into_master_start_prompt(self):
|
|
from test_prompt_transform import pt
|
|
prompt = graph_one_region()
|
|
metadata = {'nodes': [{'id': 1}]}
|
|
index = pt.PromptIndex(prompt)
|
|
transformed = pt.apply_participant_overrides(
|
|
prompt, 'master', ['worker'], {'3': 'run'}, 'http://public-master', False,
|
|
index, workflow_meta=metadata)
|
|
self.assertEqual(transformed['3']['inputs']['image_list_workflow_meta'], metadata)
|
|
|
|
async def test_aggregate_input_limit_fails_before_upload(self):
|
|
self.rt.MAX_REQUEST_BYTES = 64
|
|
with self.assertRaisesRegex(ValueError, 'complete.*request'):
|
|
await self.rt.run_image_list_job(self.create(1))
|
|
self.assertFalse(self.uploads)
|
|
self.assertFalse(self.rt._WORKER_ITEMS)
|
|
|
|
def test_tensor_and_combined_request_limits_preserve_dtype(self):
|
|
item = {'image': torch.zeros((1, 4, 4, 3)), 'index': 0}
|
|
encoded = self.rt.encode_item(item)
|
|
payload = {'context_id': 'example', 'token': 't' * 48, 'item': encoded}
|
|
single = self.rt.encode_request_payload(payload)
|
|
self.assertEqual(json.loads(single)['item']['image']['dtype'], 'float32')
|
|
self.rt.MAX_REQUEST_BYTES = len(single)
|
|
self.assertEqual(self.rt.encode_request_payload(payload), single)
|
|
item['values'] = {'value_1': item['image']}
|
|
item['mask'] = torch.zeros((1, 4, 4))
|
|
payload['item'] = self.rt.encode_item(item)
|
|
with self.assertRaisesRegex(ValueError, 'complete.*request'):
|
|
self.rt.encode_request_payload(payload)
|
|
# Size is checked before allocating/copying a 4K tensor onto the CPU.
|
|
with self.assertRaisesRegex(ValueError, 'transport limit'):
|
|
self.rt.encode_value(torch.empty((1, 2160, 3840, 3), device='meta'), 'IMAGE')
|
|
|
|
async def test_rejected_submission_reclaims_uploaded_token_without_prompt_id(self):
|
|
self.reject = True
|
|
with self.assertRaisesRegex(RuntimeError, 'prompt rejected'):
|
|
await self.rt.run_image_list_job(self.create(1))
|
|
self.assertEqual(self.deletions, self.uploads)
|
|
self.assertFalse(self.rt._WORKER_ITEMS)
|
|
self.assertFalse(self.rt._MASTER_JOBS)
|
|
self.assertFalse(self.cancel_requests)
|
|
|
|
async def test_user_interrupt_during_pending_submission_is_bounded(self):
|
|
entered = asyncio.Event()
|
|
async def never_ack(*args, **kwargs):
|
|
entered.set()
|
|
await asyncio.Event().wait()
|
|
self.dispatch.side_effect = never_ack
|
|
self.rt.SUBMISSION_CANCEL_GRACE_SECONDS = 0.02
|
|
task = asyncio.create_task(self.rt.run_image_list_job(self.create(2)))
|
|
await asyncio.wait_for(entered.wait(), 1)
|
|
self.interrupted = True
|
|
with self.assertRaises(self.rt.ImageListCancelledError):
|
|
await asyncio.wait_for(task, 0.8)
|
|
self.assertEqual(len(self.uploads), 1)
|
|
self.assertEqual(self.deletions, self.uploads)
|
|
self.assertFalse(self.rt._MASTER_JOBS)
|
|
self.assertTrue(any('ambiguous' in line for line in self.logs), self.logs)
|
|
|
|
async def test_owned_running_prompt_survives_cold_load_and_old_hour_cap(self):
|
|
self.completed = False
|
|
context = self.create(2)
|
|
clock = [0.0]
|
|
calls = [0]
|
|
original_get = self.get
|
|
def get(url, **kwargs):
|
|
if '/history/' in url:
|
|
calls[0] += 1
|
|
clock[0] += 3700
|
|
self.completed = calls[0] >= 3
|
|
if url.endswith('/queue'):
|
|
return Response({'queue_running': [[0, 'prompt-1']] if calls[0] == 2 else [],
|
|
'queue_pending': [[0, 'prompt-1']] if calls[0] == 1 else []})
|
|
return original_get(url, **kwargs)
|
|
self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0])
|
|
with patch.object(self, 'get', side_effect=get):
|
|
result = await self.rt.run_image_list_job(context)
|
|
self.assertEqual(len(result[0]), 2)
|
|
self.assertGreaterEqual(clock[0], 7200)
|
|
self.assertFalse(self.cancel_requests)
|
|
|
|
async def test_unrelated_busy_prompt_does_not_refresh_owned_liveness(self):
|
|
self.completed = False
|
|
clock = [0.0]
|
|
original_get = self.get
|
|
def get(url, **kwargs):
|
|
if '/history/' in url:
|
|
clock[0] += 70
|
|
if url.endswith('/queue'):
|
|
return Response({'queue_running': [[0, 'not-our-prompt']], 'queue_pending': []})
|
|
return original_get(url, **kwargs)
|
|
self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0])
|
|
with patch.object(self, 'get', side_effect=get):
|
|
with self.assertRaisesRegex(TimeoutError, 'liveness'):
|
|
await self.rt.run_image_list_job(self.create(1))
|
|
self.assertEqual(self.cancel_requests, [
|
|
('http://worker/queue', {'delete': ['prompt-1']}),
|
|
('http://worker/interrupt', {'prompt_id': 'prompt-1'})])
|
|
|
|
async def test_failed_queue_cleanup_still_interrupts_and_reclaims_staging(self):
|
|
self.completed = False
|
|
self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03
|
|
original_post = self.post
|
|
def post(url, **kwargs):
|
|
result = original_post(url, **kwargs)
|
|
return Response(status=500) if url.endswith('/queue') else result
|
|
with patch.object(self, 'post', side_effect=post):
|
|
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
|
|
await self.rt.run_image_list_job(self.create(2))
|
|
self.assertEqual(self.cancel_requests, [
|
|
('http://worker/queue', {'delete': ['prompt-1']}),
|
|
('http://worker/interrupt', {'prompt_id': 'prompt-1'})])
|
|
self.assertEqual(self.deletions, self.uploads)
|
|
self.assertTrue(any('cleanup unconfirmed' in line for line in self.logs), self.logs)
|
|
self.assertEqual(self.dispatch.await_count, 1)
|
|
|
|
async def test_configured_total_deadline_cleans_owned_work(self):
|
|
self.completed = False
|
|
self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03
|
|
self.queue = {'queue_running': [[0, 'prompt-1']], 'queue_pending': []}
|
|
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
|
|
await asyncio.wait_for(self.rt.run_image_list_job(self.create(2)), 1)
|
|
self.assertEqual(len(self.uploads), 1)
|
|
self.assertEqual(self.deletions, self.uploads)
|
|
self.assertFalse(self.rt._MASTER_JOBS)
|