From c8453a139bb8d8f3dd4104c200d375719c583330 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 4 Apr 2026 00:15:00 +0000 Subject: [PATCH] Use loopback callback URLs for local workers --- api/queue_orchestration.py | 11 +++++++++-- tests/test_network_helpers.py | 18 ++++++++++++++++++ utils/network.py | 19 +++++++++++++++++++ 3 files changed, 46 insertions(+), 2 deletions(-) diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 06f3a7f..226c645 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -13,7 +13,7 @@ from ..utils.constants import ( ORCHESTRATION_WORKER_PREP_CONCURRENCY, ) from ..utils.logging import debug_log, log -from ..utils.network import build_master_url +from ..utils.network import build_master_url, build_master_callback_url from ..utils.trace_logger import trace_debug from .schemas import parse_positive_float, parse_positive_int from .orchestration.dispatch import ( @@ -144,6 +144,7 @@ async def _prepare_worker_payload( enabled_ids, job_id_map, master_url, + config, delegate_master, trace_execution_id, worker_prep_semaphore, @@ -153,6 +154,11 @@ async def _prepare_worker_payload( """Prepare one worker prompt payload with bounded concurrency and media-sync timeout.""" async with worker_prep_semaphore: worker_prompt = prompt_index.copy_prompt() + worker_master_url = build_master_callback_url( + worker, + config=config, + prompt_server_instance=prompt_server, + ) worker_type = str(worker.get("type") or "local").strip().lower() is_remote_like = bool(worker.get("host")) and worker_type != "local" @@ -167,7 +173,7 @@ async def _prepare_worker_payload( worker["id"], enabled_ids, job_id_map, - master_url, + worker_master_url, delegate_master, prompt_index, ) @@ -360,6 +366,7 @@ async def orchestrate_distributed_execution( enabled_ids, job_id_map, master_url, + config, delegate_master, execution_trace_id, worker_prep_semaphore, diff --git a/tests/test_network_helpers.py b/tests/test_network_helpers.py index a9d6954..ffac6fd 100644 --- a/tests/test_network_helpers.py +++ b/tests/test_network_helpers.py @@ -114,6 +114,24 @@ class NetworkHelpersTests(unittest.TestCase): "http://127.0.0.1:8190", ) + def test_build_master_callback_url_uses_loopback_for_local_worker(self): + cfg = {"master": {"host": "192.168.68.56"}} + prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001) + worker = {"id": "w1", "type": "local", "host": "localhost", "port": 8189} + self.assertEqual( + network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server), + "http://127.0.0.1:8001", + ) + + def test_build_master_callback_url_keeps_public_master_url_for_remote_worker(self): + cfg = {"master": {"host": "192.168.68.56"}} + prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001) + worker = {"id": "w2", "type": "remote", "host": "192.168.68.99", "port": 8189} + self.assertEqual( + network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server), + "http://192.168.68.56:8001", + ) + if __name__ == "__main__": unittest.main() diff --git a/utils/network.py b/utils/network.py index 4668881..6d7f33f 100644 --- a/utils/network.py +++ b/utils/network.py @@ -180,3 +180,22 @@ def build_master_url(config=None, prompt_server_instance=None): default_port_for_scheme = 443 if scheme == "https" else 80 port_part = "" if port == default_port_for_scheme else f":{port}" return f"{scheme}://{address}{port_part}" + + +def build_master_callback_url(worker, config=None, prompt_server_instance=None): + """Build the callback URL a specific worker should use to reach the master.""" + prompt_server_instance = prompt_server_instance or server.PromptServer.instance + + worker_type = str((worker or {}).get("type") or "").strip().lower() + worker_host = normalize_host((worker or {}).get("host")) + local_hosts = {"", "localhost", "127.0.0.1", "::1", "[::1]", "0.0.0.0"} + + is_local_worker = worker_type == "local" or worker_host in local_hosts + if is_local_worker: + port = int(getattr(prompt_server_instance, "port", 8188) or 8188) + scheme = "https" if port == 443 else "http" + default_port_for_scheme = 443 if scheme == "https" else 80 + port_part = "" if port == default_port_for_scheme else f":{port}" + return f"{scheme}://127.0.0.1{port_part}" + + return build_master_url(config=config, prompt_server_instance=prompt_server_instance)