132 lines
7.1 KiB
Python
132 lines
7.1 KiB
Python
"""Image-list participation uses the existing master orchestrator setting."""
|
|
import json
|
|
import sys
|
|
import unittest
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from test_serverless_integration import PKG, load_modules
|
|
|
|
|
|
class MasterSelectionTests(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.worker = {"id": "worker", "name": "Worker", "type": "local", "enabled": True}
|
|
self.config = {"workers": [self.worker], "settings": {}}
|
|
self.q.load_config = lambda: self.config
|
|
# Model the legacy helper's fallback: image-list must not inherit it.
|
|
self.q.select_active_workers = AsyncMock(side_effect=lambda workers, ws, delegate, **kw:
|
|
(workers, delegate if workers else False))
|
|
self.preflight = AsyncMock(return_value={"worker": True})
|
|
sys.modules[PKG + ".api.orchestration.dispatch"].preflight_image_list_workers = self.preflight
|
|
self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1})
|
|
self.prompt = {
|
|
"1": {"class_type": "EmptyImage", "inputs": {}},
|
|
"2": {"class_type": "DistributedImageListStart", "inputs": {"images": ["1", 0]}},
|
|
"3": {"class_type": "DistributedImageListEnd", "inputs": {"images": ["2", 0], "context": ["2", 3]}},
|
|
}
|
|
|
|
async def queue(self, **kwargs):
|
|
return await self.q.orchestrate_distributed_execution(self.prompt, None, "client", **kwargs)
|
|
|
|
def prepared_inputs(self):
|
|
return self.q.queue_prompt_payload.call_args.args[0]["2"]["inputs"]
|
|
|
|
async def test_prepared_master_prunes_side_outputs_without_mutating_input(self):
|
|
import copy
|
|
self.prompt['op'] = {'class_type': 'Flagged', 'inputs': {'image': ['2', 0]}}
|
|
self.prompt['3']['inputs']['images'] = ['op', 0]
|
|
self.prompt['preview'] = {'class_type': 'Flagged', 'inputs': {'image': ['op', 0]}}
|
|
self.prompt['save'] = {'class_type': 'Flagged', 'inputs': {'image': ['3', 0]}}
|
|
original = copy.deepcopy(self.prompt)
|
|
registry = type('Registry', (), {'NODE_CLASS_MAPPINGS': {
|
|
'Flagged': type('Flagged', (), {'OUTPUT_NODE': True})}})
|
|
with patch.dict(sys.modules, {'nodes': registry}):
|
|
await self.queue()
|
|
queued = self.q.queue_prompt_payload.call_args
|
|
self.assertEqual(queued.kwargs['partial_execution_targets'], ['3', 'save'])
|
|
self.assertIn('op', queued.args[0])
|
|
self.assertNotIn('preview', queued.args[0])
|
|
self.assertEqual(self.prompt, original)
|
|
|
|
async def test_default_includes_master_and_direct_workers(self):
|
|
result = await self.queue()
|
|
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
|
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
|
|
self.assertEqual(result[2], 1) # Public worker_count still counts remote participants only.
|
|
self.preflight.assert_awaited_once_with(
|
|
[{"id": "worker", "name": "Worker", "type": "local", "host": None, "port": 8188}],
|
|
require_idle=False)
|
|
self.q.dispatch_worker_prompt.assert_not_awaited()
|
|
self.assertEqual(self.ps.distributed_pending_jobs, {})
|
|
|
|
async def test_orchestrator_setting_excludes_master_processing(self):
|
|
self.config["settings"]["master_delegate_only"] = True
|
|
await self.queue()
|
|
self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
|
|
self.assertTrue(self.preflight.call_args.kwargs["require_idle"])
|
|
# Upstream source stays on master even in orchestrator-only mode.
|
|
self.assertIn("1", self.q.queue_prompt_payload.call_args.args[0])
|
|
|
|
async def test_explicit_override_wins_over_config_in_both_directions(self):
|
|
for configured, requested in [(True, False), (False, True)]:
|
|
with self.subTest(configured=configured, requested=requested):
|
|
self.config["settings"]["master_delegate_only"] = configured
|
|
await self.queue(delegate_master=requested)
|
|
self.assertIs(self.prepared_inputs()["image_list_master_participates"], not requested)
|
|
|
|
async def test_master_only_with_explicit_empty_worker_selection(self):
|
|
result = await self.queue(enabled_worker_ids=[])
|
|
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
|
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
|
|
self.assertEqual(result[2], 0)
|
|
self.preflight.assert_not_awaited()
|
|
|
|
async def test_master_only_when_no_workers_are_configured(self):
|
|
self.config["workers"] = []
|
|
await self.queue()
|
|
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
|
|
|
async def test_unavailable_workers_do_not_prevent_enabled_master(self):
|
|
self.q.select_active_workers.return_value = ([], False)
|
|
self.q.select_active_workers.side_effect = None
|
|
self.preflight.return_value = {}
|
|
await self.queue()
|
|
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
|
|
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
|
|
|
async def test_busy_workers_do_not_prevent_enabled_master(self):
|
|
self.preflight.return_value = {}
|
|
await self.queue()
|
|
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
|
|
|
|
async def test_orchestrator_without_workers_fails_before_queue(self):
|
|
self.config["settings"]["master_delegate_only"] = True
|
|
with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
|
|
await self.queue(enabled_worker_ids=[])
|
|
self.q.queue_prompt_payload.assert_not_awaited()
|
|
|
|
async def test_orchestrator_never_inherits_legacy_offline_master_fallback(self):
|
|
self.config["settings"]["master_delegate_only"] = True
|
|
self.q.select_active_workers.side_effect = None
|
|
self.q.select_active_workers.return_value = ([], False)
|
|
self.preflight.return_value = {}
|
|
with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
|
|
await self.queue()
|
|
self.q.queue_prompt_payload.assert_not_awaited()
|
|
|
|
async def test_serverless_only_uses_master_without_provider_dispatch(self):
|
|
self.config["workers"] = [{"id": "pool", "type": "serverless", "enabled": True}]
|
|
result = await self.queue()
|
|
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
|
self.assertEqual(result[2], 0)
|
|
self.q.dispatch_worker_prompt.assert_not_awaited()
|
|
|
|
async def test_incompatible_worker_is_not_silently_hidden_by_master(self):
|
|
self.preflight.side_effect = RuntimeError("upgrade worker")
|
|
with self.assertRaisesRegex(RuntimeError, "upgrade worker"):
|
|
await self.queue()
|
|
self.q.queue_prompt_payload.assert_not_awaited()
|