507 lines
25 KiB
Python
507 lines
25 KiB
Python
"""Master slot regressions, plus an opt-in real CPU ComfyUI subprocess probe."""
|
|
import asyncio
|
|
import importlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def graph():
|
|
return {
|
|
'source': {'class_type': 'Source', 'inputs': {}},
|
|
'start': {'class_type': 'DistributedImageListStart', 'inputs': {'images': ['source', 0]}},
|
|
'op': {'class_type': 'Operation', 'inputs': {'image': ['start', 0]}},
|
|
'end': {'class_type': 'DistributedImageListEnd', 'inputs': {
|
|
'images': ['op', 0], 'context': ['start', 3]}},
|
|
}
|
|
|
|
|
|
class Response:
|
|
status = 200
|
|
|
|
def __init__(self, data=None):
|
|
self.data = data
|
|
|
|
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):
|
|
pass
|
|
|
|
|
|
class MasterRuntimeTests(unittest.IsolatedAsyncioTestCase):
|
|
def setUp(self):
|
|
name = '_master_slot_tests'
|
|
self.modules = patch.dict(sys.modules)
|
|
self.modules.start()
|
|
self.addCleanup(self.modules.stop)
|
|
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
|
|
module = types.ModuleType(name + suffix)
|
|
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
|
sys.modules[name + suffix] = module
|
|
self.rt = importlib.import_module(name + '.image_list.runtime')
|
|
self.workers = []
|
|
self.preflight = AsyncMock(side_effect=lambda workers, **kw: {w['id']: True for w in workers})
|
|
self.dispatched = []
|
|
self.local_calls = []
|
|
self.tokens = []
|
|
self.closed = 0
|
|
self.local_error = None
|
|
self.local_delay = 0.005
|
|
self.interrupted = False
|
|
self.local_started = asyncio.Event()
|
|
self.local_finished = asyncio.Event()
|
|
self.deadline = 0
|
|
self.slot_instances = []
|
|
self.config = {'workers': self.workers, 'settings': {}}
|
|
fixture = self
|
|
|
|
class FakeSlot:
|
|
def __init__(self, model_callback=None):
|
|
fixture.slot_instances.append(self)
|
|
self.tasks = []
|
|
|
|
def submit(self, prompt, prompt_id, end_id, workflow_meta=None):
|
|
async def execute():
|
|
inputs = prompt[end_id]['inputs']
|
|
index = inputs['worker_item_index']
|
|
token = inputs['worker_result_token']
|
|
fixture.local_calls.append(index)
|
|
fixture.tokens.append(token)
|
|
fixture.local_started.set()
|
|
try:
|
|
await asyncio.sleep(fixture.local_delay)
|
|
if fixture.local_error:
|
|
raise fixture.local_error
|
|
item = fixture.rt.consume_worker_item(token)
|
|
fixture.rt.set_master_result(fixture.context, index, token, {
|
|
'image': item['image'], 'mask': item.get('mask'),
|
|
'text': item['values'].get('value_1')})
|
|
finally:
|
|
fixture.local_finished.set()
|
|
task = asyncio.create_task(execute())
|
|
self.tasks.append(task)
|
|
return task
|
|
|
|
def request_stop(self):
|
|
pass
|
|
|
|
async def close(self):
|
|
# Model a synchronous node that must finish before its thread can join.
|
|
await asyncio.gather(*self.tasks, return_exceptions=True)
|
|
fixture.closed += 1
|
|
|
|
dependencies = {
|
|
'.image_list.local_executor': types.SimpleNamespace(LocalExecutionSlot=FakeSlot),
|
|
'.api.orchestration.dispatch': types.SimpleNamespace(
|
|
preflight_image_list_workers=self.preflight,
|
|
dispatch_worker_prompt=AsyncMock(side_effect=self.dispatch)),
|
|
'.api.orchestration.prompt_transform': types.SimpleNamespace(
|
|
_resolve_direct_workers_for_image_list=lambda config, ids: self.workers),
|
|
'.api.orchestration.media_sync': types.SimpleNamespace(
|
|
sync_worker_media=AsyncMock(), fetch_worker_path_separator=AsyncMock(return_value=None)),
|
|
'.utils.config': types.SimpleNamespace(load_config=lambda: self.config,
|
|
get_worker_timeout_seconds=lambda: 60,
|
|
get_image_list_job_timeout_seconds=lambda: self.deadline),
|
|
'.utils.network': types.SimpleNamespace(
|
|
build_master_callback_url=lambda *a, **kw: 'http://master',
|
|
build_worker_url=lambda worker, path: 'http://remote' + path,
|
|
get_client_session=AsyncMock(return_value=self)),
|
|
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
|
|
}
|
|
for suffix, module in dependencies.items():
|
|
sys.modules[name + suffix] = module
|
|
self.network = dependencies['.utils.network']
|
|
sys.modules['aiohttp'] = types.SimpleNamespace(ClientTimeout=lambda **kw: kw, ClientError=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
|
|
|
|
def create(self, participates=True, count=5):
|
|
self.items = [{'image': torch.full((1, i + 1, 2, 3), float(i)),
|
|
'mask': torch.full((1, i + 1, 2), float(i)), 'values': {'value_1': str(i)},
|
|
'index': i}
|
|
for i in range(count)]
|
|
self.context = self.rt.create_master_job('start', 'end', graph(), self.items,
|
|
[w['id'] for w in self.workers], '', master_participates=participates)
|
|
self.state = self.rt.get_master_job(self.context)
|
|
return self.context
|
|
|
|
def post(self, url, *, data=None, **kwargs):
|
|
if url.endswith('/distributed/image-list/item'):
|
|
payload = json.loads(data)
|
|
self.rt.store_worker_item(payload['token'], payload['item'])
|
|
return Response({})
|
|
|
|
def delete(self, url, **kwargs):
|
|
self.rt.delete_worker_item(url.rsplit('/', 1)[-1])
|
|
return Response({})
|
|
|
|
def get(self, url, **kwargs):
|
|
pid = url.rsplit('/', 1)[-1]
|
|
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}})
|
|
|
|
async def dispatch(self, worker, prompt, meta, **kwargs):
|
|
inputs = prompt['end']['inputs']
|
|
index, token = inputs['worker_item_index'], inputs['worker_result_token']
|
|
self.dispatched.append(index)
|
|
# A remote callback is ready immediately; the local slot is slower.
|
|
item = self.rt.consume_worker_item(token)
|
|
self.rt.set_master_result(self.context, index, token, {
|
|
'image': item['image'], 'mask': item['mask'], 'text': item['values']['value_1']})
|
|
return f'remote-{index}'
|
|
|
|
async def test_local_only_preserves_order_and_uses_no_http(self):
|
|
result = await self.rt.run_image_list_job(self.create())
|
|
self.assertEqual(result[2], [str(i) for i in range(5)])
|
|
self.assertEqual(self.local_calls, list(range(5)))
|
|
self.assertEqual(len(self.slot_instances), 1)
|
|
self.assertEqual(len(set(self.tokens)), 5)
|
|
self.network.get_client_session.assert_not_awaited()
|
|
self.assertEqual(self.closed, 1)
|
|
self.assertFalse(self.rt._WORKER_ITEMS)
|
|
self.assertFalse(self.rt._MASTER_JOBS)
|
|
for i, (image, mask) in enumerate(zip(*result[:2])):
|
|
self.assertTrue(torch.equal(image, self.items[i]['image']))
|
|
self.assertTrue(torch.equal(mask, self.items[i]['mask']))
|
|
|
|
async def test_mixed_pool_claims_from_one_queue_and_relaxes_remote_idle(self):
|
|
self.workers.append({'id': 'master', 'type': 'local'}) # Not the in-process master slot.
|
|
self.local_delay = 0.04
|
|
result = await self.rt.run_image_list_job(self.create(count=9))
|
|
self.assertTrue(self.local_calls)
|
|
self.assertTrue(self.dispatched)
|
|
self.assertEqual(sorted(self.local_calls + self.dispatched), list(range(9)))
|
|
self.assertEqual(result[2], [str(i) for i in range(9)])
|
|
self.preflight.assert_awaited_once_with(self.workers, require_idle=False)
|
|
|
|
async def test_orchestrator_mode_never_instantiates_local_slot(self):
|
|
self.workers.append({'id': 'worker', 'type': 'remote'})
|
|
await self.rt.run_image_list_job(self.create(participates=False))
|
|
self.assertFalse(self.slot_instances)
|
|
self.preflight.assert_awaited_once_with(self.workers, require_idle=True)
|
|
with self.assertRaisesRegex(RuntimeError, 'worker|fallback'):
|
|
self.rt.create_master_job('start', 'end', graph(), self.items, [], '')
|
|
|
|
async def test_values_reconstructed_in_source_order_without_transporting_originals(self):
|
|
self.workers.append({'id': 'remote', 'type': 'remote'})
|
|
context = self.create(count=6)
|
|
tensor = torch.ones(4, 2, 2, 3, dtype=torch.float64)
|
|
for index, item in enumerate(self.items):
|
|
item['values'].update(value_2=tensor, value_3=True, value_4=index / 8)
|
|
self.local_delay = 0.1 # Remote results arrive ahead of master item zero.
|
|
with patch.object(self.rt, 'encode_result', side_effect=AssertionError('originals must not return')):
|
|
result = await self.rt.run_image_list_job(context, include_values=True)
|
|
self.assertFalse(self.rt._MASTER_JOBS)
|
|
self.assertEqual(len(result), 4)
|
|
for index, row in enumerate(result[3]):
|
|
self.assertEqual(row['value_1'], str(index))
|
|
self.assertIs(row['value_2'], tensor)
|
|
self.assertIs(row['value_3'], True)
|
|
self.assertEqual(row['value_4'], index / 8)
|
|
self.assertTrue(self.dispatched and self.local_calls)
|
|
|
|
async def test_overrides_merge_add_and_recheck_capability_before_every_item(self):
|
|
self.workers.append({'id': 'remote', 'type': 'remote'})
|
|
context = self.create(participates=False, count=3)
|
|
self.state.original_prompt['end']['inputs'].update(value_1=['start', 4], value_10=['start', 2])
|
|
original_set = self.rt.set_master_result
|
|
def override(context, index, token, result):
|
|
result['values'] = {'value_1': f'new-{index}', 'value_10': index}
|
|
return original_set(context, index, token, result)
|
|
with patch.object(self.rt, 'set_master_result', side_effect=override):
|
|
result = await self.rt.run_image_list_job(context, include_values=True)
|
|
self.assertEqual(result[3], [{'value_1': f'new-{index}', 'value_10': index} for index in range(3)])
|
|
self.assertEqual(self.preflight.await_count, 4)
|
|
self.assertTrue(all(call.kwargs['require_value_results'] is True
|
|
for call in self.preflight.await_args_list))
|
|
self.assertEqual(self.items[0]['values'], {'value_1': '0'})
|
|
|
|
async def test_override_capability_downgrade_stops_before_upload(self):
|
|
self.workers.append({'id': 'remote', 'type': 'remote'})
|
|
context = self.create(participates=False)
|
|
self.state.original_prompt['end']['inputs']['value_1'] = ['start', 4]
|
|
self.preflight.side_effect = [{'remote': True}, RuntimeError('value_results unsupported')]
|
|
with patch.object(self, 'post', wraps=self.post) as upload:
|
|
with self.assertRaisesRegex(RuntimeError, 'value_results'):
|
|
await self.rt.run_image_list_job(context, include_values=True)
|
|
upload.assert_not_called()
|
|
self.assertFalse(self.dispatched or self.rt._MASTER_JOBS)
|
|
|
|
async def test_missing_expected_override_fails_without_partial_output(self):
|
|
self.workers.append({'id': 'remote', 'type': 'remote'})
|
|
context = self.create(participates=False)
|
|
self.state.original_prompt['end']['inputs']['value_1'] = ['start', 4]
|
|
with self.assertRaisesRegex(ValueError, 'missing'):
|
|
await self.rt.run_image_list_job(context, include_values=True)
|
|
self.assertFalse(self.rt._MASTER_JOBS)
|
|
|
|
async def test_values_failure_and_cancellation_return_no_partial_lists(self):
|
|
self.local_error = RuntimeError('branch failed')
|
|
with self.assertRaisesRegex(RuntimeError, 'branch failed'):
|
|
await self.rt.run_image_list_job(self.create(), include_values=True)
|
|
self.local_error = None
|
|
self.local_started.clear()
|
|
self.local_delay = 0.1
|
|
task = asyncio.create_task(self.rt.run_image_list_job(self.create(), include_values=True))
|
|
await self.local_started.wait()
|
|
task.cancel()
|
|
with self.assertRaises(self.rt.ImageListCancelledError):
|
|
await task
|
|
self.assertFalse(self.rt._MASTER_JOBS)
|
|
|
|
async def test_local_failure_surfaces_and_reclaims_staged_token(self):
|
|
self.local_error = RuntimeError('local branch failed')
|
|
with self.assertRaisesRegex(RuntimeError, 'local branch failed'):
|
|
await self.rt.run_image_list_job(self.create())
|
|
self.assertEqual(self.closed, 1)
|
|
self.assertFalse(self.rt._WORKER_ITEMS)
|
|
self.assertFalse(self.rt._MASTER_JOBS)
|
|
|
|
async def test_deadline_joins_local_execution_before_removing_context(self):
|
|
self.deadline, self.local_delay = 0.03, 0.1
|
|
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
|
|
await self.rt.run_image_list_job(self.create())
|
|
self.assertTrue(self.local_finished.is_set())
|
|
self.assertEqual(self.closed, 1)
|
|
self.assertFalse(self.rt._WORKER_ITEMS)
|
|
self.assertFalse(self.rt._MASTER_JOBS)
|
|
|
|
|
|
async def test_cancelled_scheduler_joins_owned_execution_even_if_cancelled_again(self):
|
|
self.local_delay = 0.12
|
|
task = asyncio.create_task(self.rt.run_image_list_job(self.create()))
|
|
await self.local_started.wait()
|
|
task.cancel()
|
|
await asyncio.sleep(0.02)
|
|
task.cancel()
|
|
with self.assertRaises(self.rt.ImageListCancelledError):
|
|
await task
|
|
self.assertTrue(self.local_finished.is_set())
|
|
self.assertFalse(self.rt._WORKER_ITEMS)
|
|
self.assertFalse(self.rt._MASTER_JOBS)
|
|
|
|
|
|
class RealFrameworkTests(unittest.TestCase):
|
|
@unittest.skipUnless(os.environ.get('COMFYUI_SOURCE_ROOT'), 'set COMFYUI_SOURCE_ROOT for the real CPU framework probe')
|
|
def test_real_framework_local_slot(self):
|
|
result = subprocess.run([sys.executable, str(Path(__file__).resolve()), '--real-probe',
|
|
os.environ['COMFYUI_SOURCE_ROOT']], text=True, capture_output=True, timeout=60)
|
|
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
|
|
self.assertIn('MASTER_SLOT_REAL_OK', result.stdout)
|
|
|
|
|
|
def real_framework_probe(comfy_root):
|
|
import threading
|
|
import time
|
|
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 execution
|
|
import nodes
|
|
import comfy.memory_management as memory
|
|
import comfy.model_management as models
|
|
import comfy.utils as utils
|
|
from comfy_execution import progress
|
|
from comfy_execution.graph_utils import GraphBuilder
|
|
from comfy_api.v0_0_2 import io
|
|
|
|
name = '_real_master_slot'
|
|
for suffix in ('', '.image_list', '.utils', '.nodes'):
|
|
module = types.ModuleType(name + suffix)
|
|
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
|
sys.modules[name + suffix] = module
|
|
# No HTTP/service/extension bootstrap is needed by the local slot. Fail if
|
|
# either boundary tries to use a server-loop bridge instead of direct memory.
|
|
def forbidden(*args, **kwargs):
|
|
raise AssertionError('local boundary used the server loop or /prompt')
|
|
sys.modules[name + '.utils.async_helpers'] = types.SimpleNamespace(run_async_in_server_loop=forbidden)
|
|
config = {'settings': {}}
|
|
sys.modules[name + '.utils.config'] = types.SimpleNamespace(load_config=lambda: config)
|
|
rt = importlib.import_module(name + '.image_list.runtime')
|
|
local = importlib.import_module(name + '.image_list.local_executor')
|
|
boundary = importlib.import_module(name + '.nodes.image_list')
|
|
nodes.NODE_CLASS_MAPPINGS.update(DistributedImageListStart=boundary.DistributedImageListStart,
|
|
DistributedImageListEnd=boundary.DistributedImageListEnd)
|
|
loads, calls = [], []
|
|
entered, stopped = threading.Event(), threading.Event()
|
|
mode = ['normal']
|
|
model = object()
|
|
|
|
class Loader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {'required': {}}
|
|
RETURN_TYPES = ('MODEL',)
|
|
FUNCTION = 'run'
|
|
def run(self):
|
|
loads.append(True)
|
|
return (model,)
|
|
|
|
class Operation:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {'required': {'image': ('IMAGE',), 'model': ('MODEL',),
|
|
'mask': ('MASK',), 'text': ('STRING', {'forceInput': True})}}
|
|
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
|
|
FUNCTION = 'run'
|
|
async def run(self, image, model, mask, text):
|
|
calls.append(text)
|
|
entered.set()
|
|
if mode[0] == 'error':
|
|
raise ValueError('probe branch failure')
|
|
if mode[0] in ('cancel', 'interrupt'):
|
|
try:
|
|
while True:
|
|
utils.ProgressBar(1).update(1)
|
|
await asyncio.sleep(0.01)
|
|
finally:
|
|
await asyncio.sleep(0.02)
|
|
stopped.set()
|
|
if mode[0] == 'sync':
|
|
# No progress/interrupt checks: cleanup still must join, not
|
|
# release the parent just because an asyncio Future is cancelled.
|
|
time.sleep(0.08)
|
|
stopped.set()
|
|
await asyncio.sleep(0)
|
|
return image + 1, mask + 1, text + ':done'
|
|
|
|
class Expanded:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return Operation.INPUT_TYPES()
|
|
RETURN_TYPES = Operation.RETURN_TYPES
|
|
FUNCTION = 'run'
|
|
def run(self, image, model, mask, text):
|
|
builder = GraphBuilder()
|
|
node = builder.node('Operation', image=image, model=model, mask=mask, text=text)
|
|
return {'result': (node.out(0), node.out(1), node.out(2)), 'expand': builder.finalize()}
|
|
|
|
nodes.NODE_CLASS_MAPPINGS.update(Loader=Loader, Operation=Operation, Expanded=Expanded)
|
|
prompt = graph()
|
|
prompt['loader'] = {'class_type': 'Loader', 'inputs': {}}
|
|
prompt['op']['inputs'].update(model=['loader', 0], mask=['start', rt.MASK_OUT],
|
|
text=['start', rt.VALUE_OUT_BASE])
|
|
prompt['end']['inputs'].update(masks=['op', 1], texts=['op', 2])
|
|
items = [{'image': torch.full((1, i + 1, 2, 3), float(i)),
|
|
'mask': torch.full((1, i + 1, 2), float(i)), 'values': {'value_1': str(i)}, 'index': i}
|
|
for i in range(3)]
|
|
region = rt.validate_image_list_regions(prompt)[0]
|
|
parent_progress = progress.ProgressRegistry('outer', execution.DynamicPrompt({}))
|
|
parent_progress.nodes['outer'] = {'state': progress.NodeState.Running, 'value': 3, 'max': 7}
|
|
progress.global_progress_registry = parent_progress
|
|
parent_hook = lambda *a, **kw: None
|
|
utils.PROGRESS_BAR_HOOK = parent_hook
|
|
ram_callback = lambda *a, **kw: None
|
|
memory.set_ram_cache_release_state(ram_callback, 123)
|
|
GraphBuilder.set_default_prefix('outer-prefix', 7, 9)
|
|
|
|
def restored():
|
|
assert progress.global_progress_registry is parent_progress
|
|
assert parent_progress.nodes == {'outer': {'state': progress.NodeState.Running, 'value': 3, 'max': 7}}
|
|
assert utils.PROGRESS_BAR_HOOK is parent_hook
|
|
assert memory.extra_ram_release_callback is ram_callback and memory.RAM_CACHE_HEADROOM == 123
|
|
assert (GraphBuilder._default_prefix_root, GraphBuilder._default_prefix_call_index,
|
|
GraphBuilder._default_prefix_graph_index) == ('outer-prefix', 7, 9)
|
|
assert not any(t.name.startswith('image-list-master') for t in threading.enumerate())
|
|
|
|
async def probe():
|
|
# Verify actual V3 input parsing and list mapping, not only direct execute.
|
|
for supplied, delegated, expected in [(None, False, True), (None, True, False),
|
|
(False, False, False), (True, True, True)]:
|
|
config['settings']['master_delegate_only'] = delegated
|
|
inputs = {'image_list_enabled_worker_ids': '["remote"]'}
|
|
if supplied is not None:
|
|
inputs['image_list_master_participates'] = supplied
|
|
cls = boundary.DistributedImageListStart
|
|
data, _, v3 = execution.get_input_data(inputs, cls, 'start', dynprompt=execution.DynamicPrompt(prompt))
|
|
data['images'] = [torch.zeros((3, 2, 2, 3))]
|
|
result = await execution._async_map_node_over_list('outer', 'start', cls, data, cls.FUNCTION, v3_data=v3)
|
|
result = await execution.resolve_map_node_over_list_results(result)
|
|
context = result[0].result[rt.CONTEXT_OUT][0]
|
|
assert rt.get_master_job(context).master_participates is expected
|
|
rt.remove_master_job(context)
|
|
|
|
for behavior in ('normal', 'expand', 'error', 'cancel', 'interrupt', 'sync', 'already-interrupted'):
|
|
# Independent lifecycle cases start cold; retention is tested separately.
|
|
local._RETAINED_OUTPUTS = None
|
|
mode[0] = behavior
|
|
prompt['op']['class_type'] = 'Expanded' if behavior == 'expand' else 'Operation'
|
|
loads.clear()
|
|
calls.clear()
|
|
entered.clear()
|
|
stopped.clear()
|
|
tracked = []
|
|
slot = local.LocalExecutionSlot(model_callback=tracked.append)
|
|
context = rt.create_master_job('start', 'end', prompt, items, [], '', master_participates=True)
|
|
try:
|
|
for index in range(3 if behavior in ('normal', 'expand') else 1):
|
|
token = rt.create_worker_task_token()
|
|
rt.register_assignment(context, index, 'local', token)
|
|
rt._WORKER_ITEMS[token] = (time.monotonic() + 60, items[index])
|
|
worker_prompt = rt.build_worker_prompt(prompt, region, token, context, index, '')
|
|
for key in ('start', 'end'):
|
|
worker_prompt[key]['inputs']['image_list_local_execution'] = True
|
|
if behavior == 'already-interrupted':
|
|
models.interrupt_current_processing(True)
|
|
future = slot.submit(worker_prompt, f'local-{index}', 'end')
|
|
if behavior in ('cancel', 'interrupt', 'sync'):
|
|
assert await asyncio.to_thread(entered.wait, 5)
|
|
if behavior == 'interrupt':
|
|
models.interrupt_current_processing(True)
|
|
else:
|
|
slot.request_stop()
|
|
try:
|
|
await asyncio.wait_for(asyncio.shield(asyncio.wrap_future(future)), 5)
|
|
except rt.ImageListCancelledError:
|
|
assert behavior in ('cancel', 'interrupt', 'sync', 'already-interrupted'), behavior
|
|
except RuntimeError as exc:
|
|
assert behavior == 'error' and 'probe branch failure' in str(exc), str(exc)
|
|
else:
|
|
assert behavior in ('normal', 'expand'), behavior
|
|
result = rt.get_master_job(context).results[index]
|
|
assert torch.equal(result['image'], items[index]['image'] + 1)
|
|
assert torch.equal(result['mask'], items[index]['mask'] + 1)
|
|
assert result['text'] == f'{index}:done'
|
|
if behavior in ('normal', 'expand'):
|
|
assert len(loads) == 1 and calls == ['0', '1', '2'], (loads, calls)
|
|
assert any(model in outputs for outputs in tracked), tracked
|
|
if behavior in ('cancel', 'interrupt', 'sync'):
|
|
assert stopped.is_set(), behavior
|
|
if behavior == 'already-interrupted':
|
|
assert models.processing_interrupted() and not calls
|
|
finally:
|
|
await slot.close()
|
|
for token in rt.get_master_job(context).expected_tokens.values():
|
|
rt.delete_worker_item(token)
|
|
rt.remove_master_job(context)
|
|
models.interrupt_current_processing(False) # Only test setup/teardown.
|
|
restored()
|
|
assert not rt._MASTER_JOBS and not rt._WORKER_ITEMS
|
|
asyncio.run(probe())
|
|
print('MASTER_SLOT_REAL_OK V3 metadata; cache reuse; async/expanded graph; errors; cancellation; state restoration')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
if len(sys.argv) > 1 and sys.argv[1] == '--real-probe':
|
|
real_framework_probe(sys.argv[2])
|
|
else:
|
|
unittest.main()
|