robertvoy-ComfyUI-Distributed/tests/api/test_queue_orchestration.py

import asyncio
import importlib.util
import sys
import types
import unittest
from pathlib import Path

from tests.api.harness import (
    bootstrap_test_package,
    install_server_stub,
)


def _bootstrap_package(package_name):
    bootstrap_test_package(
        package_name,
        with_api=True,
        with_utils=True,
        with_orchestration=True,
    )

    prompt_server_instance = types.SimpleNamespace(address="127.0.0.1", port=8188)
    install_server_stub(prompt_server_instance)

    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")

    async def _queue_prompt_payload(prompt_obj, workflow_meta, client_id):
        _ = (prompt_obj, workflow_meta, client_id)
        return "prompt-id"

    async_helpers_module.queue_prompt_payload = _queue_prompt_payload
    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: {"settings": {}, "workers": []}
    sys.modules[f"{package_name}.utils.config"] = config_module

    constants_module = types.ModuleType(f"{package_name}.utils.constants")
    constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 2
    constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 120.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 **_kwargs: "http://127.0.0.1:8188"
    sys.modules[f"{package_name}.utils.network"] = network_module

    trace_logger_module = types.ModuleType(f"{package_name}.utils.trace_logger")
    trace_logger_module.trace_debug = lambda *_args, **_kwargs: None
    sys.modules[f"{package_name}.utils.trace_logger"] = trace_logger_module

    runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")

    def _ensure_distributed_runtime_state(server_instance=None):
        ps = server_instance or prompt_server_instance
        if not hasattr(ps, "distributed_pending_jobs"):
            ps.distributed_pending_jobs = {}
        if not hasattr(ps, "distributed_jobs_lock"):
            ps.distributed_jobs_lock = asyncio.Lock()
        if not hasattr(ps, "distributed_job_allowed_workers"):
            ps.distributed_job_allowed_workers = {}
        if not hasattr(ps, "distributed_pending_tile_jobs"):
            ps.distributed_pending_tile_jobs = {}
        if not hasattr(ps, "distributed_tile_jobs_lock"):
            ps.distributed_tile_jobs_lock = asyncio.Lock()
        return types.SimpleNamespace(
            distributed_pending_jobs=ps.distributed_pending_jobs,
            distributed_jobs_lock=ps.distributed_jobs_lock,
            distributed_job_allowed_workers=ps.distributed_job_allowed_workers,
            distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs,
            distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock,
        )

    runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state
    runtime_state_module.get_prompt_server_instance = lambda: prompt_server_instance
    sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module

    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
    schemas_module.coerce_positive_int = (
        lambda value, default: int(value) if str(value).isdigit() and int(value) > 0 else default
    )
    schemas_module.coerce_positive_float = (
        lambda value, default: float(value) if isinstance(value, (int, float, str)) and float(value) > 0 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 = lambda *args, **kwargs: None
    dispatch_module.rank_workers_by_load = lambda workers, **kwargs: workers
    dispatch_module.select_active_workers = lambda workers, *_args, **_kwargs: (workers, False)
    dispatch_module.select_least_busy_worker = lambda workers, **_kwargs: workers[0] if workers else None
    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 = lambda *_args, **_kwargs: "/"
    media_sync_module.sync_worker_media = lambda *_args, **_kwargs: None
    sys.modules[f"{package_name}.api.orchestration.media_sync"] = media_sync_module

    prompt_transform_module = types.ModuleType(f"{package_name}.api.orchestration.prompt_transform")
    prompt_transform_module.PromptIndex = object
    prompt_transform_module.apply_participant_overrides = lambda prompt, *_args, **_kwargs: prompt
    prompt_transform_module.find_nodes_by_class = lambda _prompt, _class_name: []
    prompt_transform_module.generate_job_id_map = lambda _index, _prefix: {}
    prompt_transform_module.prepare_delegate_master_prompt = lambda prompt, _ids: prompt
    prompt_transform_module.prune_prompt_for_worker = lambda prompt: prompt
    sys.modules[f"{package_name}.api.orchestration.prompt_transform"] = prompt_transform_module


def _load_module(package_name, module_rel_path, module_name):
    module_path = Path(__file__).resolve().parents[2] / module_rel_path
    spec = importlib.util.spec_from_file_location(
        f"{package_name}.{module_name}",
        module_path,
    )
    module = importlib.util.module_from_spec(spec)
    assert spec is not None and spec.loader is not None
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def _load_queue_orchestration_module():
    package_name = "dist_queue_orch_testpkg"
    _bootstrap_package(package_name)
    return _load_module(package_name, "api/queue_orchestration.py", "api.queue_orchestration")


queue_orchestration = _load_queue_orchestration_module()


class QueueOrchestrationHelpersTests(unittest.TestCase):
    def test_resolve_enabled_workers_filters_disabled_and_requested(self):
        config = {
            "workers": [
                {"id": "1", "enabled": True, "host": "a", "port": 8188},
                {"id": "2", "enabled": False, "host": "b", "port": 8189},
                {"id": "3", "enabled": True, "host": "c", "port": "bad-port"},
            ]
        }

        enabled_only = queue_orchestration._resolve_enabled_workers(config)
        requested = queue_orchestration._resolve_enabled_workers(config, requested_ids={"2", "3"})

        self.assertEqual([worker["id"] for worker in enabled_only], ["1", "3"])
        self.assertEqual([worker["id"] for worker in requested], ["2", "3"])
        self.assertEqual(requested[1]["port"], 8188)

    def test_resolve_orchestration_limits_uses_defaults_for_invalid_values(self):
        config = {
            "settings": {
                "worker_probe_concurrency": "x",
                "worker_prep_concurrency": "3",
                "media_sync_concurrency": "-1",
                "media_sync_timeout_seconds": "42.5",
            }
        }

        limits = queue_orchestration._resolve_orchestration_limits(config)

        self.assertEqual(limits[0], 8)   # default
        self.assertEqual(limits[1], 3)   # parsed
        self.assertEqual(limits[2], 2)   # default
        self.assertEqual(limits[3], 42.5)

    def test_is_load_balance_enabled_accepts_common_forms(self):
        self.assertTrue(queue_orchestration._is_load_balance_enabled(True))
        self.assertTrue(queue_orchestration._is_load_balance_enabled(1))
        self.assertTrue(queue_orchestration._is_load_balance_enabled(" yes "))
        self.assertFalse(queue_orchestration._is_load_balance_enabled(0))
        self.assertFalse(queue_orchestration._is_load_balance_enabled("off"))

    def test_prompt_requests_load_balance_detects_collector_flag(self):
        prompt_index = types.SimpleNamespace(
            nodes_for_class=lambda class_name: ["10"] if class_name == "DistributedCollector" else [],
            inputs_by_node={"10": {"load_balance": "true"}},
        )
        self.assertTrue(queue_orchestration._prompt_requests_load_balance(prompt_index))

    def test_ensure_distributed_state_initializes_attrs(self):
        prompt_server = types.SimpleNamespace()
        queue_orchestration.ensure_distributed_state(prompt_server)
        self.assertTrue(hasattr(prompt_server, "distributed_pending_jobs"))
        self.assertTrue(hasattr(prompt_server, "distributed_jobs_lock"))
        self.assertIsInstance(prompt_server.distributed_jobs_lock, asyncio.Lock)


if __name__ == "__main__":
    unittest.main()