From bbd8936c6522024e129cc9c1d93e7364f47b906c Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sun, 26 Apr 2026 22:51:01 +0000 Subject: [PATCH] refactor: clarify load-balanced dispatch reservations --- api/orchestration/dispatch.py | 100 ++++--- api/queue_orchestration.py | 390 ++++++++++++++++++---------- tests/test_dispatch_selection.py | 58 +++++ tests/test_load_balance_queueing.py | 221 ++++++++++++++++ 4 files changed, 595 insertions(+), 174 deletions(-) create mode 100644 tests/test_load_balance_queueing.py diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index 79fe59f..3cfcf7d 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -1,5 +1,6 @@ import asyncio import json +import time import uuid import aiohttp @@ -207,9 +208,13 @@ async def _probe_worker_queue(worker, semaphore, probe_timeout): payload = await probe_worker(worker_url, timeout=probe_timeout) if payload is None: return None + try: + reserved_slots = int(worker.get("reserved_slots", 0) or 0) + except (TypeError, ValueError): + reserved_slots = 0 return { "worker": worker, - "queue_remaining": _extract_queue_remaining(payload), + "queue_remaining": _extract_queue_remaining(payload) + max(reserved_slots, 0), } @@ -222,17 +227,7 @@ def _select_idle_round_robin(statuses): return statuses[index] -async def select_least_busy_worker( - workers, - trace_execution_id=None, - probe_concurrency=8, - probe_timeout=3.0, -): - """Select one worker by queue depth, round-robin among idle workers.""" - if not workers: - return None - - probe_limit = parse_positive_int(probe_concurrency, 8) +async def _probe_worker_queues(workers, probe_limit, probe_timeout): probe_semaphore = asyncio.Semaphore(probe_limit) statuses = await asyncio.gather( *[ @@ -240,29 +235,66 @@ async def select_least_busy_worker( for worker in workers ] ) - statuses = [status for status in statuses if status is not None] - if not statuses: - if trace_execution_id: - trace_info(trace_execution_id, "Least-busy selection failed: no worker queue probes succeeded.") - else: - log("[Distributed] Least-busy selection failed: no worker queue probes succeeded.") - return None + return [status for status in statuses if status is not None] + +def _select_worker_status(statuses, require_idle=False): idle_statuses = [status for status in statuses if status["queue_remaining"] == 0] if idle_statuses: - selected = _select_idle_round_robin(idle_statuses) - else: - selected = min(statuses, key=lambda status: status["queue_remaining"]) + return _select_idle_round_robin(idle_statuses) + if require_idle: + return None + return min(statuses, key=lambda status: status["queue_remaining"]) - worker = selected["worker"] - queue_remaining = selected["queue_remaining"] - if trace_execution_id: - trace_debug( - trace_execution_id, - f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}", - ) - else: - debug_log( - f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}" - ) - return worker + +async def select_least_busy_worker( + workers, + trace_execution_id=None, + probe_concurrency=8, + probe_timeout=3.0, + require_idle=False, + idle_poll_interval=1.0, + idle_wait_timeout=None, +): + """Select a worker by queue depth, optionally waiting until one is idle.""" + if not workers: + return None + + probe_limit = parse_positive_int(probe_concurrency, 8) + poll_interval = max(float(idle_poll_interval or 0), 0.0) + deadline = None + if idle_wait_timeout is not None: + deadline = time.monotonic() + max(float(idle_wait_timeout), 0.0) + + while True: + statuses = await _probe_worker_queues(workers, probe_limit, probe_timeout) + if not statuses: + if trace_execution_id: + trace_info(trace_execution_id, "Least-busy selection failed: no worker queue probes succeeded.") + else: + log("[Distributed] Least-busy selection failed: no worker queue probes succeeded.") + return None + + selected = _select_worker_status(statuses, require_idle=require_idle) + if selected is not None: + worker = selected["worker"] + queue_remaining = selected["queue_remaining"] + if trace_execution_id: + trace_debug( + trace_execution_id, + f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}", + ) + else: + debug_log( + f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}" + ) + return worker + + if deadline is not None and time.monotonic() >= deadline: + if trace_execution_id: + trace_info(trace_execution_id, "No idle load-balance worker became available before timeout.") + else: + debug_log("No idle load-balance worker became available before timeout.") + return None + + await asyncio.sleep(poll_interval) diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index ab18e7f..307df24 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -46,12 +46,159 @@ def ensure_distributed_state(server_instance=None): ps.distributed_pending_jobs = {} if not hasattr(ps, "distributed_jobs_lock"): ps.distributed_jobs_lock = asyncio.Lock() + if not hasattr(ps, "distributed_worker_reservations"): + ps.distributed_worker_reservations = {} # Initialize top-level distributed queue state at module import time. ensure_distributed_state() +def _worker_id(worker): + return str(worker.get("id")) + + +def _non_negative_int(value, default=0): + try: + parsed = int(value if value is not None else default) + except (TypeError, ValueError): + parsed = default + return max(parsed, 0) + + +async def _release_worker_reservation(worker_id): + if not worker_id: + return + ensure_distributed_state() + async with prompt_server.distributed_jobs_lock: + reservations = prompt_server.distributed_worker_reservations + current = _non_negative_int(reservations.get(worker_id, 0)) + if current <= 1: + reservations.pop(worker_id, None) + else: + reservations[worker_id] = current - 1 + + +async def _candidate_workers_with_reservations(candidate_workers): + ensure_distributed_state() + async with prompt_server.distributed_jobs_lock: + reservations = dict(prompt_server.distributed_worker_reservations) + return [ + { + **worker, + "reserved_slots": reservations.get(_worker_id(worker), 0), + } + for worker in candidate_workers + ] + + +async def _reserve_selected_worker(selected_worker): + selected_worker_id = _worker_id(selected_worker) + selected_reserved_slots = _non_negative_int(selected_worker.get("reserved_slots", 0)) + async with prompt_server.distributed_jobs_lock: + reservations = prompt_server.distributed_worker_reservations + current_reserved_slots = _non_negative_int(reservations.get(selected_worker_id, 0)) + if current_reserved_slots != selected_reserved_slots: + return None + reservations[selected_worker_id] = current_reserved_slots + 1 + return selected_worker_id + + +def _normalized_idle_poll_interval(value): + try: + poll_interval = float(value or 0) + except (TypeError, ValueError): + poll_interval = 1.0 + return max(poll_interval, 0.05) + + +async def _select_and_reserve_load_balance_worker( + candidate_workers, + execution_trace_id, + worker_probe_concurrency, + idle_poll_interval, +): + poll_interval = _normalized_idle_poll_interval(idle_poll_interval) + while candidate_workers: + reserved_candidates = await _candidate_workers_with_reservations(candidate_workers) + selected_worker = await select_least_busy_worker( + reserved_candidates, + trace_execution_id=execution_trace_id, + probe_concurrency=worker_probe_concurrency, + require_idle=True, + idle_poll_interval=0, + idle_wait_timeout=0, + ) + if selected_worker is None: + await asyncio.sleep(poll_interval) + continue + + reserved_worker_id = await _reserve_selected_worker(selected_worker) + if reserved_worker_id is not None: + return selected_worker + + await asyncio.sleep(0) + return None + + +def _load_balance_candidate_workers(active_workers, delegate_master, master_url): + candidate_workers = list(active_workers) + if not delegate_master: + candidate_workers.append( + { + "id": "master", + "name": "Master", + "host": master_url, + "type": "local", + } + ) + return candidate_workers + + +async def _resolve_load_balance_selection( + active_workers, + delegate_master, + master_url, + config, + execution_trace_id, + worker_probe_concurrency, +): + candidate_workers = _load_balance_candidate_workers(active_workers, delegate_master, master_url) + selected_worker = None + if candidate_workers: + settings = (config.get("settings", {}) or {}) + idle_poll_interval = parse_positive_float( + settings.get("load_balance_idle_poll_interval_seconds"), + 1.0, + ) + selected_worker = await _select_and_reserve_load_balance_worker( + candidate_workers, + execution_trace_id, + worker_probe_concurrency, + idle_poll_interval, + ) + + if selected_worker is not None and _worker_id(selected_worker) == "master": + trace_debug( + execution_trace_id, + "Load-balance selected master for execution (workers skipped).", + ) + return [], False, "master" + + if selected_worker is not None: + trace_debug( + execution_trace_id, + f"Load-balance selected worker {selected_worker.get('id')} (master set to delegate-only).", + ) + return [selected_worker], True, _worker_id(selected_worker) + + trace_debug( + execution_trace_id, + "Load-balance requested but no idle/probeable execution candidates were available.", + ) + return [], False, None + + async def _ensure_distributed_queue(job_id): """Ensure a queue exists for the given distributed job ID.""" ensure_distributed_state() @@ -259,160 +406,123 @@ async def orchestrate_distributed_execution( probe_concurrency=worker_probe_concurrency, ) + reserved_worker_id = None if load_balance_requested: - candidate_workers = list(active_workers) - if not delegate_master: - # Include master in load balancing only when master participation is enabled. - candidate_workers.append( - { - "id": "master", - "name": "Master", - "host": master_url, - "type": "local", - } + active_workers, delegate_master, reserved_worker_id = await _resolve_load_balance_selection( + active_workers, + delegate_master, + master_url, + config, + execution_trace_id, + worker_probe_concurrency, + ) + + try: + enabled_ids = [worker["id"] for worker in active_workers] + + discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}" + job_id_map = generate_job_id_map(prompt_index, discovery_prefix) + + if not job_id_map: + trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.") + queue_result = await queue_prompt_payload( + prompt_obj, + workflow_meta, + client_id, + include_queue_metadata=True, + ) + return ( + queue_result["prompt_id"], + queue_result["number"], + 0, + queue_result.get("node_errors", {}), ) - selected_worker = None - if candidate_workers: - selected_worker = await select_least_busy_worker( - candidate_workers, - trace_execution_id=execution_trace_id, - probe_concurrency=worker_probe_concurrency, - ) - if selected_worker is None and candidate_workers: + for job_id in job_id_map.values(): + await _ensure_distributed_queue(job_id) + + master_prompt = prompt_index.copy_prompt() + master_prompt = apply_participant_overrides( + master_prompt, + "master", + enabled_ids, + job_id_map, + master_url, + delegate_master, + prompt_index, + ) + + if delegate_master: + collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector") + upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed") + if upscale_nodes: + debug_log( + "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master." + ) + elif not collector_ids: + debug_log( + "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master." + ) + else: + master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids) + + if active_workers: trace_debug( execution_trace_id, - "Load-balance selection probe failed; using first available candidate.", + "Active distributed workers: " + + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers), ) - selected_worker = candidate_workers[0] - - if selected_worker is not None and str(selected_worker.get("id")) == "master": - # Master selected as least busy; run master workload only. - active_workers = [] - delegate_master = False - trace_debug( - execution_trace_id, - "Load-balance selected master for execution (workers skipped).", + worker_payloads = [] + if active_workers: + worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency) + media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency) + worker_payloads = await asyncio.gather( + *[ + _prepare_worker_payload( + worker, + prompt_index, + enabled_ids, + job_id_map, + master_url, + config, + delegate_master, + execution_trace_id, + worker_prep_semaphore, + media_sync_semaphore, + media_sync_timeout_seconds, + ) + for worker in active_workers + ] ) - elif selected_worker is not None: - active_workers = [selected_worker] - # Worker selected as least busy; keep master orchestrator-only for this run. - delegate_master = True - trace_debug( - execution_trace_id, - f"Load-balance selected worker {selected_worker.get('id')} (master set to delegate-only).", + + if worker_payloads: + await asyncio.gather( + *[ + dispatch_worker_prompt( + worker, + wprompt, + workflow_meta, + client_id, + use_websocket=use_websocket, + trace_execution_id=execution_trace_id, + ) + for worker, wprompt in worker_payloads + ] ) - else: - trace_debug( - execution_trace_id, - "Load-balance requested but no execution candidates were available.", - ) - active_workers = [] - delegate_master = False - enabled_ids = [worker["id"] for worker in active_workers] - - discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}" - job_id_map = generate_job_id_map(prompt_index, discovery_prefix) - - if not job_id_map: - trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.") queue_result = await queue_prompt_payload( - prompt_obj, + master_prompt, workflow_meta, client_id, include_queue_metadata=True, ) - return ( - queue_result["prompt_id"], - queue_result["number"], - 0, - queue_result.get("node_errors", {}), - ) - - for job_id in job_id_map.values(): - await _ensure_distributed_queue(job_id) - - master_prompt = prompt_index.copy_prompt() - master_prompt = apply_participant_overrides( - master_prompt, - "master", - enabled_ids, - job_id_map, - master_url, - delegate_master, - prompt_index, - ) - - if delegate_master: - collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector") - upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed") - if upscale_nodes: - debug_log( - "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master." - ) - elif not collector_ids: - debug_log( - "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master." - ) - else: - master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids) - - if active_workers: + prompt_id = queue_result["prompt_id"] + prompt_number = queue_result["number"] + node_errors = queue_result.get("node_errors", {}) trace_debug( execution_trace_id, - "Active distributed workers: " - + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers), + f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}", ) - worker_payloads = [] - if active_workers: - worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency) - media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency) - worker_payloads = await asyncio.gather( - *[ - _prepare_worker_payload( - worker, - prompt_index, - enabled_ids, - job_id_map, - master_url, - config, - delegate_master, - execution_trace_id, - worker_prep_semaphore, - media_sync_semaphore, - media_sync_timeout_seconds, - ) - for worker in active_workers - ] - ) - - if worker_payloads: - await asyncio.gather( - *[ - dispatch_worker_prompt( - worker, - wprompt, - workflow_meta, - client_id, - use_websocket=use_websocket, - trace_execution_id=execution_trace_id, - ) - for worker, wprompt in worker_payloads - ] - ) - - queue_result = await queue_prompt_payload( - master_prompt, - workflow_meta, - client_id, - include_queue_metadata=True, - ) - prompt_id = queue_result["prompt_id"] - prompt_number = queue_result["number"] - node_errors = queue_result.get("node_errors", {}) - trace_debug( - execution_trace_id, - f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}", - ) - return prompt_id, prompt_number, len(worker_payloads), node_errors + return prompt_id, prompt_number, len(worker_payloads), node_errors + finally: + await _release_worker_reservation(reserved_worker_id) diff --git a/tests/test_dispatch_selection.py b/tests/test_dispatch_selection.py index 8792e72..4e578b2 100644 --- a/tests/test_dispatch_selection.py +++ b/tests/test_dispatch_selection.py @@ -250,5 +250,63 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase): self.assertIsNone(selected) + async def test_select_least_busy_worker_waits_for_idle_when_required(self): + workers = [ + {"id": "w1", "name": "Worker 1"}, + {"id": "w2", "name": "Worker 2"}, + ] + queue_sequences = { + "w1": [1, 1], + "w2": [1, 0], + } + + async def fake_probe(worker_url, timeout=3.0): + worker_id = worker_url.rsplit("/", 1)[-1] + sequence = queue_sequences[worker_id] + value = sequence.pop(0) if sequence else 0 + return {"exec_info": {"queue_remaining": value}} + + with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object( + dispatch, + "probe_worker", + side_effect=fake_probe, + ): + dispatch._least_busy_rr_index = 0 + selected = await dispatch.select_least_busy_worker( + workers, + probe_concurrency=2, + require_idle=True, + idle_poll_interval=0, + ) + + self.assertEqual(selected["id"], "w2") + + + async def test_select_least_busy_worker_counts_reserved_slots_as_busy(self): + workers = [ + {"id": "w1", "name": "Worker 1", "reserved_slots": 1}, + {"id": "w2", "name": "Worker 2"}, + ] + + async def fake_probe(worker_url, timeout=3.0): + return {"exec_info": {"queue_remaining": 0}} + + with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object( + dispatch, + "probe_worker", + side_effect=fake_probe, + ): + dispatch._least_busy_rr_index = 0 + selected = await dispatch.select_least_busy_worker( + workers, + probe_concurrency=2, + require_idle=True, + idle_poll_interval=0, + idle_wait_timeout=0, + ) + + self.assertEqual(selected["id"], "w2") + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_load_balance_queueing.py b/tests/test_load_balance_queueing.py new file mode 100644 index 0000000..cb96161 --- /dev/null +++ b/tests/test_load_balance_queueing.py @@ -0,0 +1,221 @@ +import importlib.util +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import AsyncMock, patch + + +def _load_queue_orchestration_module(): + module_path = Path(__file__).resolve().parents[1] / "api" / "queue_orchestration.py" + package_name = "dist_queue_orchestration_testpkg" + + for mod_name in list(sys.modules): + if mod_name == package_name or mod_name.startswith(f"{package_name}."): + del sys.modules[mod_name] + + root_pkg = types.ModuleType(package_name) + root_pkg.__path__ = [] + sys.modules[package_name] = root_pkg + + api_pkg = types.ModuleType(f"{package_name}.api") + api_pkg.__path__ = [] + sys.modules[f"{package_name}.api"] = api_pkg + + orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration") + orchestration_pkg.__path__ = [] + sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg + + utils_pkg = types.ModuleType(f"{package_name}.utils") + utils_pkg.__path__ = [] + sys.modules[f"{package_name}.utils"] = utils_pkg + + prompt_server_instance = types.SimpleNamespace( + distributed_pending_jobs={}, + distributed_worker_reservations={}, + ) + server_module = types.ModuleType("server") + server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance) + sys.modules["server"] = server_module + + async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers") + async_helpers_module.queue_prompt_payload = AsyncMock( + return_value={"prompt_id": "prompt-master", "number": 1, "node_errors": {}} + ) + sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module + + config_module = types.ModuleType(f"{package_name}.utils.config") + config_module.load_config = lambda: { + "workers": [{"id": "w1", "name": "Worker 1", "enabled": True, "host": "worker-1", "port": 8188}], + "settings": {"load_balance_idle_poll_interval_seconds": 0}, + } + sys.modules[f"{package_name}.utils.config"] = config_module + + constants_module = types.ModuleType(f"{package_name}.utils.constants") + constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 4 + constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 5.0 + constants_module.ORCHESTRATION_WORKER_PROBE_CONCURRENCY = 8 + constants_module.ORCHESTRATION_WORKER_PREP_CONCURRENCY = 4 + sys.modules[f"{package_name}.utils.constants"] = constants_module + + logging_module = types.ModuleType(f"{package_name}.utils.logging") + logging_module.debug_log = lambda *_args, **_kwargs: None + logging_module.log = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.utils.logging"] = logging_module + + network_module = types.ModuleType(f"{package_name}.utils.network") + network_module.build_master_url = lambda *_args, **_kwargs: "http://master:8188" + network_module.build_master_callback_url = lambda *_args, **_kwargs: "http://master:8188" + sys.modules[f"{package_name}.utils.network"] = network_module + + trace_module = types.ModuleType(f"{package_name}.utils.trace_logger") + trace_module.trace_debug = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.utils.trace_logger"] = trace_module + + schemas_module = types.ModuleType(f"{package_name}.api.schemas") + schemas_module.parse_positive_int = lambda value, default: int(value if value is not None else default) + schemas_module.parse_positive_float = lambda value, default: float(value if value is not None else default) + sys.modules[f"{package_name}.api.schemas"] = schemas_module + + dispatch_module = types.ModuleType(f"{package_name}.api.orchestration.dispatch") + dispatch_module.dispatch_worker_prompt = AsyncMock() + dispatch_module.select_active_workers = AsyncMock(return_value=([{"id": "w1", "name": "Worker 1"}], False)) + dispatch_module.select_least_busy_worker = AsyncMock() + sys.modules[f"{package_name}.api.orchestration.dispatch"] = dispatch_module + + media_sync_module = types.ModuleType(f"{package_name}.api.orchestration.media_sync") + media_sync_module.convert_paths_for_platform = lambda prompt, _sep: prompt + media_sync_module.fetch_worker_path_separator = AsyncMock(return_value=None) + media_sync_module.sync_worker_media = AsyncMock() + sys.modules[f"{package_name}.api.orchestration.media_sync"] = media_sync_module + + prompt_transform_module = types.ModuleType(f"{package_name}.api.orchestration.prompt_transform") + + class _PromptIndex: + def __init__(self, prompt): + self.prompt = prompt + self.inputs_by_node = {node_id: node.get("inputs", {}) for node_id, node in prompt.items()} + + def nodes_for_class(self, class_type): + return [ + node_id + for node_id, node in self.prompt.items() + if node.get("class_type") == class_type + ] + + def copy_prompt(self): + return {node_id: dict(node) for node_id, node in self.prompt.items()} + + prompt_transform_module.PromptIndex = _PromptIndex + prompt_transform_module.apply_participant_overrides = lambda prompt, *_args, **_kwargs: prompt + prompt_transform_module.find_nodes_by_class = lambda prompt, class_type: [ + node_id for node_id, node in prompt.items() if node.get("class_type") == class_type + ] + prompt_transform_module.generate_job_id_map = lambda _prompt_index, _prefix: {"10": "job-10"} + prompt_transform_module.prepare_delegate_master_prompt = lambda prompt, _collector_ids: prompt + prompt_transform_module.prune_prompt_for_worker = lambda prompt: prompt + sys.modules[f"{package_name}.api.orchestration.prompt_transform"] = prompt_transform_module + + spec = importlib.util.spec_from_file_location( + f"{package_name}.api.queue_orchestration", + module_path, + ) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + spec.loader.exec_module(module) + return module + + +queue_orchestration = _load_queue_orchestration_module() + + +class LoadBalanceQueueingTests(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + queue_orchestration.ensure_distributed_state() + queue_orchestration.prompt_server.distributed_worker_reservations.clear() + queue_orchestration.prompt_server.distributed_pending_jobs.clear() + + async def test_load_balance_reservation_selection_refreshes_reservations_between_polls(self): + queue_orchestration.prompt_server.distributed_worker_reservations["w1"] = 1 + calls = [] + + async def fake_select(candidates, **_kwargs): + reserved_slots = candidates[0].get("reserved_slots") + calls.append(reserved_slots) + if reserved_slots: + await queue_orchestration._release_worker_reservation("w1") + return None + return candidates[0] + + with patch.object(queue_orchestration, "select_least_busy_worker", side_effect=fake_select): + selected = await queue_orchestration._select_and_reserve_load_balance_worker( + [{"id": "w1", "name": "Worker 1"}], + "exec-test", + worker_probe_concurrency=1, + idle_poll_interval=0, + ) + + self.assertEqual(selected["id"], "w1") + self.assertEqual(calls, [1, 0]) + self.assertEqual(queue_orchestration.prompt_server.distributed_worker_reservations, {"w1": 1}) + + async def test_load_balance_does_not_fallback_to_first_candidate_when_idle_selection_fails(self): + prompt = { + "10": { + "class_type": "DistributedCollector", + "inputs": {"load_balance": True}, + } + } + + with patch.object( + queue_orchestration, + "_select_and_reserve_load_balance_worker", + new=AsyncMock(return_value=None), + ), patch.object( + queue_orchestration, + "dispatch_worker_prompt", + new=AsyncMock(), + ) as dispatch_mock: + prompt_id, prompt_number, worker_count, node_errors = await queue_orchestration.orchestrate_distributed_execution( + prompt, + workflow_meta={}, + client_id="client-1", + trace_execution_id="exec-test", + ) + + self.assertEqual((prompt_id, prompt_number, worker_count, node_errors), ("prompt-master", 1, 0, {})) + dispatch_mock.assert_not_awaited() + + async def test_load_balance_releases_worker_reservation_when_dispatch_fails(self): + prompt = { + "10": { + "class_type": "DistributedCollector", + "inputs": {"load_balance": True}, + } + } + + async def select_first_candidate(candidates, **_kwargs): + return candidates[0] + + with patch.object( + queue_orchestration, + "select_least_busy_worker", + side_effect=select_first_candidate, + ), patch.object( + queue_orchestration, + "dispatch_worker_prompt", + new=AsyncMock(side_effect=RuntimeError("dispatch failed")), + ): + with self.assertRaisesRegex(RuntimeError, "dispatch failed"): + await queue_orchestration.orchestrate_distributed_execution( + prompt, + workflow_meta={}, + client_id="client-1", + trace_execution_id="exec-test", + ) + + self.assertEqual(queue_orchestration.prompt_server.distributed_worker_reservations, {}) + + +if __name__ == "__main__": + unittest.main()