From fbfb411402b70060579868ef6ea1225bbc3d62b7 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 10 Oct 2026 21:59:18 +0000 Subject: [PATCH] fix: isolate local worker databases and avoid port conflicts --- README.md | 4 + api/worker_routes.py | 15 +++- tests/api/test_worker_routes.py | 35 +++++++- tests/test_worker_ports.py | 72 +++++++++++++++ tests/test_worker_process_runtime.py | 127 ++++++++++++++++++++++++++- web/masterDetection.js | 11 ++- web/tests/masterDetection.test.js | 65 ++++++++++++++ workers/ports.py | 88 +++++++++++++++++++ workers/process/launch_builder.py | 37 ++++++++ workers/process/lifecycle.py | 4 + 10 files changed, 454 insertions(+), 4 deletions(-) create mode 100644 tests/test_worker_ports.py create mode 100644 web/tests/masterDetection.test.js create mode 100644 workers/ports.py diff --git a/README.md b/README.md index f5725c5..b025f93 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,10 @@ ComfyUI Distributed supports three types of workers: - **Remote Workers** - GPUs on other computers in your network - **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels +Local auto-setup chooses available ports above the actual master port, skipping ports already in use. Existing worker ports are never silently reassigned: a launch conflict reports which port needs changing. + +On ComfyUI versions that support `--database-url`, each managed local worker automatically uses its own asset database under `/distributed/workers/`. The filename is stable for the worker ID, even after renaming it or changing its GPU/port. Models and custom nodes remain shared. Explicit worker `--database-url` and `--disable-assets` arguments are preserved; older ComfyUI versions receive no automatic database flag. Set the worker port in the UI/config, rather than overriding it with `--port` in extra arguments. + > For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md) --- diff --git a/api/worker_routes.py b/api/worker_routes.py index 336eec7..dfdf767 100644 --- a/api/worker_routes.py +++ b/api/worker_routes.py @@ -16,12 +16,14 @@ from ..utils.logging import debug_log, log from ..utils.network import ( build_worker_url, get_client_session, + get_server_port, handle_api_error, normalize_host, probe_worker, ) from ..utils.constants import CHUNK_SIZE from ..workers import get_worker_manager +from ..workers.ports import allocate_worker_ports from .schemas import require_fields, validate_worker_id from ..workers.detection import ( get_machine_id, @@ -277,6 +279,13 @@ def _get_cuda_info(): def _collect_network_info_sync(): """Collect network/cuda info in a worker thread to avoid blocking route handlers.""" cuda_device, cuda_device_count, physical_device_count = _get_cuda_info() + device_count = physical_device_count if physical_device_count > 0 else cuda_device_count + master_port = get_server_port() + config = load_config() + worker_ports = [] + if not config.get("settings", {}).get("has_auto_populated_workers") and not config.get("workers"): + worker_count = device_count - (1 if cuda_device is not None and 0 <= cuda_device < device_count else 0) + worker_ports = allocate_worker_ports(master_port, config.get("workers", []), worker_count) hostname = socket.gethostname() all_ips = get_network_ips() recommended_ip = get_recommended_ip(all_ips) @@ -285,7 +294,9 @@ def _collect_network_info_sync(): "all_ips": all_ips, "recommended_ip": recommended_ip, "cuda_device": cuda_device, - "cuda_device_count": physical_device_count if physical_device_count > 0 else cuda_device_count, + "cuda_device_count": device_count, + "master_port": master_port, + "local_worker_ports": worker_ports, } @@ -483,6 +494,8 @@ async def launch_worker_endpoint(request): "message": f"Worker {worker['name']} launched", "log_file": log_file }) + except ValueError as e: + return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 400) except Exception as e: return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 500) diff --git a/tests/api/test_worker_routes.py b/tests/api/test_worker_routes.py index 9a4c562..56aabb7 100644 --- a/tests/api/test_worker_routes.py +++ b/tests/api/test_worker_routes.py @@ -106,7 +106,7 @@ def _load_worker_routes_module(): sys.modules[f"{package_name}.utils"] = utils_pkg workers_pkg = types.ModuleType(f"{package_name}.workers") - workers_pkg.__path__ = [] + workers_pkg.__path__ = [str(module_path.parents[1] / "workers")] workers_pkg.get_worker_manager = lambda: _DummyWorkerManager() sys.modules[f"{package_name}.workers"] = workers_pkg @@ -213,6 +213,7 @@ def _load_worker_routes_module(): raise RuntimeError("not used in these tests") network_module.get_client_session = _get_client_session + network_module.get_server_port = lambda: 8189 sys.modules[f"{package_name}.utils.network"] = network_module constants_module = types.ModuleType(f"{package_name}.utils.constants") @@ -269,6 +270,38 @@ worker_routes = _load_worker_routes_module() class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase): + async def test_network_info_reports_actual_master_and_available_ports(self): + with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \ + patch.object(worker_routes, "get_network_ips", return_value=["127.0.0.1"]), \ + patch.object(worker_routes, "load_config", return_value={"workers": []}), \ + patch.object(worker_routes, "get_server_port", return_value=8189), \ + patch.object(worker_routes, "allocate_worker_ports", return_value=[8191, 8192, 8193]) as allocate: + response = await worker_routes.get_network_info_endpoint(_FakeRequest()) + self.assertEqual(response.payload["master_port"], 8189) + self.assertEqual(response.payload["local_worker_ports"], [8191, 8192, 8193]) + self.assertEqual(response.payload["cuda_device_count"], 4) + allocate.assert_called_once_with(8189, [], 3) + + async def test_network_info_does_not_reallocate_existing_configuration(self): + with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 4, 4)), \ + patch.object(worker_routes, "get_network_ips", return_value=[]), \ + patch.object(worker_routes, "load_config", return_value={"workers": [{"id": "existing", "port": 8189}]}), \ + patch.object(worker_routes, "allocate_worker_ports") as allocate: + response = await worker_routes.get_network_info_endpoint(_FakeRequest()) + self.assertEqual(response.payload["local_worker_ports"], []) + allocate.assert_not_called() + + async def test_launch_conflict_is_a_clear_client_error(self): + manager = _DummyWorkerManager() + config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8189}]} + with patch.object(worker_routes, "get_worker_manager", return_value=manager), \ + patch.object(worker_routes, "load_config", return_value=config), \ + patch.object(manager, "launch_worker", side_effect=ValueError("Worker port 8189 conflicts with the master port 8189")): + response = await worker_routes.launch_worker_endpoint(_FakeRequest({"worker_id": "worker-a"})) + self.assertEqual(response.status, 400) + self.assertIn("master port 8189", response.payload["message"]) + self.assertEqual(manager.processes, {}) + async def test_launch_worker_valid_id_returns_200(self): manager = _DummyWorkerManager() config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]} diff --git a/tests/test_worker_ports.py b/tests/test_worker_ports.py new file mode 100644 index 0000000..2995879 --- /dev/null +++ b/tests/test_worker_ports.py @@ -0,0 +1,72 @@ +import importlib.util +import socket +import unittest +from pathlib import Path +from unittest.mock import patch + + +def load_ports(): + path = Path(__file__).resolve().parents[1] / "workers" / "ports.py" + spec = importlib.util.spec_from_file_location("worker_ports_test", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class WorkerPortTests(unittest.TestCase): + def setUp(self): + self.ports = load_ports() + + def test_allocation_starts_above_actual_master_and_skips_assigned_and_occupied(self): + workers = [ + {"id": "existing", "host": "localhost", "port": 8190, "enabled": False}, + {"id": "remote", "host": "remote.example", "port": 8192}, + ] + with patch.object(self.ports, "is_port_available", side_effect=lambda port: port != 8191): + self.assertEqual(self.ports.allocate_worker_ports(8189, workers, 3), [8192, 8193, 8194]) + + def test_local_host_forms_reserve_ports(self): + for host in ["::1", "[::1]", "http://localhost/", "HTTPS://LOCALHOST", "0.0.0.0", None]: + with self.subTest(host=host), patch.object(self.ports, "is_port_available", return_value=True): + self.assertEqual(self.ports.allocate_worker_ports(8189, [{"id": "local", "host": host, "port": 8190}], 1), [8191]) + + def test_exhaustion_is_explicit_not_a_partial_allocation(self): + with patch.object(self.ports, "is_port_available", return_value=True): + with self.assertRaisesRegex(ValueError, "available.*ports"): + self.ports.allocate_worker_ports(65534, [], 2) + + def test_launch_rejects_master_port_without_reassigning_configuration(self): + worker = {"id": "manual", "port": 8189} + with self.assertRaisesRegex(ValueError, "master.*8189"): + self.ports.validate_worker_port(worker, 8189, []) + self.assertEqual(worker["port"], 8189) + + def test_launch_rejects_another_local_workers_reserved_port(self): + worker = {"id": "manual", "port": 8190} + other = {"id": "other", "host": None, "port": 8190, "enabled": False} + with self.assertRaisesRegex(ValueError, "other"): + self.ports.validate_worker_port(worker, 8189, [worker, other]) + + def test_launch_ignores_itself_and_remote_workers_with_same_port(self): + worker = {"id": "manual", "port": 8190} + other = {"id": "remote", "host": "example.com", "port": 8190} + with patch.object(self.ports, "is_port_available", return_value=True): + self.ports.validate_worker_port(worker, 8189, [worker, other]) + + def test_launch_rejects_an_occupied_port(self): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: + listener.bind(("127.0.0.1", 0)) + listener.listen() + port = listener.getsockname()[1] + self.assertFalse(self.ports.is_port_available(port)) + with self.assertRaisesRegex(ValueError, "already in use"): + self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, []) + + def test_invalid_ports_fail_clearly(self): + for port in [0, 65536, "not-a-port", None, True, 8190.5]: + with self.subTest(port=port), self.assertRaisesRegex(ValueError, "port"): + self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_worker_process_runtime.py b/tests/test_worker_process_runtime.py index dec0659..66b219f 100644 --- a/tests/test_worker_process_runtime.py +++ b/tests/test_worker_process_runtime.py @@ -1,5 +1,6 @@ import importlib.util import sys +import tempfile import types import unittest from argparse import Namespace @@ -21,7 +22,7 @@ def _load_process_module(module_filename: str): sys.modules[package_name] = root_pkg workers_pkg = types.ModuleType(f"{package_name}.workers") - workers_pkg.__path__ = [] + workers_pkg.__path__ = [str(module_path.parents[1])] sys.modules[f"{package_name}.workers"] = workers_pkg process_pkg = types.ModuleType(f"{package_name}.workers.process") @@ -39,8 +40,23 @@ def _load_process_module(module_filename: str): process_module = types.ModuleType(f"{package_name}.utils.process") process_module.get_python_executable = lambda: "/usr/bin/test-python" + process_module.is_process_alive = lambda _pid: False + process_module.terminate_process = lambda *_args, **_kwargs: None sys.modules[f"{package_name}.utils.process"] = process_module + config_module = types.ModuleType(f"{package_name}.utils.config") + config_module.load_config = lambda: {"workers": [], "settings": {"stop_workers_on_master_exit": False}} + config_module.save_config = lambda _config: None + sys.modules[f"{package_name}.utils.config"] = config_module + constants_module = types.ModuleType(f"{package_name}.utils.constants") + constants_module.PROCESS_TERMINATION_TIMEOUT = 1 + constants_module.PROCESS_WAIT_TIMEOUT = 1 + constants_module.WORKER_CHECK_INTERVAL = 0.01 + sys.modules[f"{package_name}.utils.constants"] = constants_module + network_module = types.ModuleType(f"{package_name}.utils.network") + network_module.get_server_port = lambda: 8189 + sys.modules[f"{package_name}.utils.network"] = network_module + spec = importlib.util.spec_from_file_location( f"{package_name}.workers.process.{module_name}", module_path, @@ -70,6 +86,88 @@ class ComfyRootDiscoveryTests(unittest.TestCase): class LaunchCommandBuilderTests(unittest.TestCase): + def build_command(self, root, worker=None, runtime=None): + builder = launch_builder_module.LaunchCommandBuilder() + worker = worker or {"id": "worker-a", "port": 8190} + runtime = runtime if runtime is not None else Namespace(database_url=None) + with patch.object(builder, "_get_runtime_args", return_value=runtime): + return builder.build_launch_command(worker, str(root)) + + def test_database_is_unique_and_stable_by_worker_id(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "main.py").touch() + first = self.build_command(root, {"id": "worker-a", "port": 8190}) + renamed = self.build_command(root, { + "id": "worker-a", "port": 9001, "name": "Renamed", "cuda_device": 3, + }) + second = self.build_command(root, {"id": "worker-b", "port": 8191}) + database = first[first.index("--database-url") + 1] + self.assertEqual(database, renamed[renamed.index("--database-url") + 1]) + self.assertNotEqual(database, second[second.index("--database-url") + 1]) + self.assertTrue(database.startswith("sqlite:///" + root.as_posix() + "/")) + self.assertTrue(Path(database.removeprefix("sqlite:///")).parent.is_dir()) + self.assertNotIn("--disable-assets", first) + + def test_database_uses_effective_user_or_base_directory(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "main.py").touch() + for extra_args, expected in [ + (f'--user-directory "{root / "custom user"}"', root / "custom user"), + (f'--base-directory "{root / "custom base"}"', root / "custom base" / "user"), + ]: + with self.subTest(extra_args=extra_args): + cmd = self.build_command(root, { + "id": "worker-a", "port": 8190, "extra_args": extra_args, + }) + database = cmd[cmd.index("--database-url") + 1] + self.assertTrue(database.startswith("sqlite:///" + expected.as_posix() + "/")) + cmd = self.build_command(root, runtime=Namespace( + database_url="sqlite:///master.db", user_directory=str(root / "runtime user"), + )) + database = cmd[cmd.index("--database-url") + 1] + self.assertTrue(database.startswith("sqlite:///" + (root / "runtime user").as_posix() + "/")) + self.assertNotEqual(database, "sqlite:///master.db") + + def test_explicit_database_and_disabled_assets_are_preserved(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "main.py").touch() + for extra in ["--database-url sqlite:///explicit.db", "--database-url=sqlite:///explicit.db", "--disable-assets"]: + with self.subTest(extra=extra): + cmd = self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra}) + self.assertEqual(sum(arg.split("=")[0] == "--database-url" for arg in cmd), 0 if extra == "--disable-assets" else 1) + self.assertFalse((root / "user").exists()) + + def test_older_comfyui_without_database_flag_keeps_existing_command(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "main.py").touch() + cmd = self.build_command(root, runtime=Namespace()) + self.assertNotIn("--database-url", cmd) + self.assertFalse((root / "user").exists()) + + def test_worker_id_cannot_escape_database_directory(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "main.py").touch() + databases = [] + for worker_id in ["../outside", "a/b", "a_b"]: + cmd = self.build_command(root, {"id": worker_id, "port": 8190}) + path = Path(cmd[cmd.index("--database-url") + 1].removeprefix("sqlite:///")) + self.assertTrue(path.is_relative_to(root / "user")) + databases.append(path) + self.assertEqual(len(set(databases)), 3) + + def test_extra_args_cannot_silently_override_the_configured_port(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "main.py").touch() + for extra in ["--port 8189", "--port=8189"]: + with self.subTest(extra=extra), self.assertRaisesRegex(ValueError, "configured worker port"): + self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra}) + def test_inherits_runtime_layout_args_for_desktop(self): builder = launch_builder_module.LaunchCommandBuilder() runtime_args = Namespace( @@ -130,5 +228,32 @@ class LaunchCommandBuilderTests(unittest.TestCase): self.assertNotIn("--auto-launch", cmd) +class ProcessLaunchPortTests(unittest.TestCase): + def test_conflict_is_rejected_before_building_or_spawning(self): + lifecycle_module = _load_process_module("lifecycle.py") + manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={}) + with patch.object(lifecycle_module.subprocess, "Popen") as spawn: + with self.assertRaisesRegex(ValueError, "master port 8189"): + lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189}) + spawn.assert_not_called() + self.assertEqual(manager.processes, {}) + + def test_available_port_reaches_existing_launch_path(self): + lifecycle_module = _load_process_module("lifecycle.py") + with tempfile.TemporaryDirectory() as directory: + manager = types.SimpleNamespace( + find_comfy_root=lambda: directory, + build_launch_command=lambda _worker, _root: ["python", "main.py", "--port", "8190"], + processes={}, save_processes=lambda: None, + ) + with patch.object(lifecycle_module, "validate_worker_port") as validate, \ + patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn: + pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190}) + validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, []) + spawn.assert_called_once() + self.assertEqual(pid, 1234) + self.assertEqual(manager.processes["manual"]["pid"], 1234) + + if __name__ == "__main__": unittest.main() diff --git a/web/masterDetection.js b/web/masterDetection.js index 40508bd..395993d 100644 --- a/web/masterDetection.js +++ b/web/masterDetection.js @@ -42,6 +42,15 @@ export async function detectMasterIP(extension) { if (shouldAutoPopulate) { extension.log(`Auto-populating workers based on ${extension.cudaDeviceCount} CUDA devices (excluding master on CUDA ${extension.masterCudaDevice})`, "info"); + const workerCount = extension.cudaDeviceCount - ( + Number.isInteger(extension.masterCudaDevice) && + extension.masterCudaDevice >= 0 && + extension.masterCudaDevice < extension.cudaDeviceCount ? 1 : 0 + ); + const workerPorts = data.local_worker_ports; + if (!Array.isArray(workerPorts) || workerPorts.length !== workerCount) { + throw new Error("Could not allocate enough available local worker ports; check the server log"); + } const newWorkers = []; let workerNum = 1; let portOffset = 0; @@ -56,7 +65,7 @@ export async function detectMasterIP(extension) { id: generateUUID(), name: `Worker ${workerNum}`, host: isRunpod ? null : "localhost", - port: 8189 + portOffset, + port: workerPorts[portOffset], cuda_device: i, enabled: true, extra_args: isRunpod ? "--listen" : "", diff --git a/web/tests/masterDetection.test.js b/web/tests/masterDetection.test.js new file mode 100644 index 0000000..e7a4b8a --- /dev/null +++ b/web/tests/masterDetection.test.js @@ -0,0 +1,65 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { detectMasterIP } from "../masterDetection.js"; + +describe("automatic local worker ports", () => { + let originalWindow; + beforeEach(() => { + originalWindow = globalThis.window; + globalThis.window = { location: { hostname: "localhost", port: "443" } }; + }); + afterEach(() => { globalThis.window = originalWindow; }); + + function makeExtension(overrides = {}) { + return { + config: { workers: [], settings: {}, master: {} }, + log: vi.fn(), + api: { + getNetworkInfo: vi.fn().mockResolvedValue({ + cuda_device: 0, cuda_device_count: 3, + master_port: 8189, local_worker_ports: [8191, 8193], + ...overrides, + }), + updateMaster: vi.fn().mockResolvedValue({}), + updateWorker: vi.fn().mockResolvedValue({}), + updateSetting: vi.fn().mockResolvedValue({}), + }, + ui: { updateMasterDisplay: vi.fn() }, + app: {}, + loadConfig: vi.fn().mockResolvedValue(), + }; + } + + it("uses server-allocated ports, not hardcoded ports or the browser proxy port", async () => { + const extension = makeExtension(); + await detectMasterIP(extension); + expect(extension.config.workers.map(worker => worker.port)).toEqual([8191, 8193]); + expect(extension.config.workers.map(worker => worker.cuda_device)).toEqual([1, 2]); + expect(extension.api.updateWorker).toHaveBeenCalledTimes(2); + }); + + it("does not overwrite an existing manually configured worker", async () => { + const extension = makeExtension(); + const worker = { id: "manual", host: "localhost", port: 8189 }; + extension.config.workers = [worker]; + await detectMasterIP(extension); + expect(extension.config.workers).toEqual([worker]); + expect(extension.api.updateWorker).not.toHaveBeenCalled(); + }); + + it("fails without partially saving workers if not enough ports are available", async () => { + const extension = makeExtension({ local_worker_ports: [8191] }); + await detectMasterIP(extension); + expect(extension.config.workers).toEqual([]); + expect(extension.api.updateWorker).not.toHaveBeenCalled(); + expect(extension.api.updateSetting).not.toHaveBeenCalled(); + expect(extension.log).toHaveBeenCalledWith(expect.stringContaining("port"), "error"); + }); + + it("preserves Runpod worker host and launch settings", async () => { + globalThis.window.location.hostname = "pod-8189.proxy.runpod.net"; + const extension = makeExtension(); + await detectMasterIP(extension); + expect(extension.config.workers.map(worker => worker.port)).toEqual([8191, 8193]); + expect(extension.config.workers.every(worker => worker.host === null && worker.extra_args === "--listen")).toBe(true); + }); +}); diff --git a/workers/ports.py b/workers/ports.py new file mode 100644 index 0000000..febfbfa --- /dev/null +++ b/workers/ports.py @@ -0,0 +1,88 @@ +"""Host-local worker port selection and launch preflight (no ComfyUI imports).""" +import errno +import socket +from urllib.parse import urlsplit + + +def _port_number(value): + try: + port = int(value) + except (TypeError, ValueError): + raise ValueError("Worker port must be an integer between 1 and 65535") from None + if isinstance(value, (bool, float)) or not 1 <= port <= 65535: + raise ValueError("Worker port must be an integer between 1 and 65535") + return port + + +def _local_config(worker): + local_hosts = ("localhost", "127.0.0.1", "0.0.0.0", "::1", "::") + host = (worker.get("host") or "localhost").strip().lower() + if worker.get("type") == "local" or host in local_hosts: + return True + try: + host = urlsplit(host if "://" in host else "//" + host).hostname + except ValueError: + return False + return host in local_hosts + + +def _reserved_ports(workers, exclude_id=None): + reserved = {} + for worker in workers: + if exclude_id is not None and str(worker.get("id")) == str(exclude_id): + continue + if not _local_config(worker): + continue + try: + port = _port_number(worker.get("port")) + except ValueError: + continue + reserved[port] = worker.get("name") or worker.get("id") or "another local worker" + return reserved + + +def is_port_available(port): + """Probe bindability without SO_REUSEADDR; include IPv6-only listeners.""" + addresses = [(socket.AF_INET, "0.0.0.0")] + if socket.has_ipv6: + addresses.append((socket.AF_INET6, "::")) + for family, address in addresses: + try: + with socket.socket(family, socket.SOCK_STREAM) as probe: + probe.bind((address, port)) + except OSError as exc: + if family == socket.AF_INET6 and exc.errno in (errno.EAFNOSUPPORT, errno.EADDRNOTAVAIL): + continue + return False + return True + + +def allocate_worker_ports(master_port, workers, count): + """Select ports above the actual master, skipping reservations/listeners. + + This is a snapshot, not a reservation. Launch must validate again because + another process can claim a suggested port before the user starts a worker. + """ + master_port = _port_number(master_port) + reserved = _reserved_ports(workers) + selected = [] + if count == 0: + return selected + for port in range(master_port + 1, 65536): + if port not in reserved and is_port_available(port): + selected.append(port) + if len(selected) == count: + return selected + raise ValueError(f"Not enough available worker ports above master port {master_port}; need {count}") + + +def validate_worker_port(worker, master_port, workers): + """Reject conflicts without silently changing a configured worker port.""" + port = _port_number(worker.get("port")) + if port == _port_number(master_port): + raise ValueError(f"Worker port {port} conflicts with the master port {master_port}; choose another worker port") + reserved = _reserved_ports(workers, exclude_id=worker.get("id")) + if port in reserved: + raise ValueError(f"Worker port {port} is assigned to local worker {reserved[port]}; choose another worker port") + if not is_port_available(port): + raise ValueError(f"Worker port {port} is already in use on this host; choose another worker port") diff --git a/workers/process/launch_builder.py b/workers/process/launch_builder.py index c4d510d..962f719 100644 --- a/workers/process/launch_builder.py +++ b/workers/process/launch_builder.py @@ -1,4 +1,5 @@ import glob +import hashlib import os import shlex import shutil @@ -87,6 +88,39 @@ class LaunchCommandBuilder: return wt_path return None + @staticmethod + def _arg_value(cmd, flag): + """Return the last value, matching argparse's override semantics.""" + value = None + for index, arg in enumerate(cmd): + if arg == flag and index + 1 < len(cmd): + value = cmd[index + 1] + elif arg.startswith(flag + "="): + value = arg.split("=", 1)[1] + return value + + def _add_worker_database(self, cmd, worker_config, comfy_root): + args = self._get_runtime_args() + # The parsed namespace advertises support without launching ComfyUI + # (or importing another checkout's argument parser). + if args is None or not hasattr(args, "database_url"): + return + if any(arg.split("=", 1)[0] in ("--database-url", "--disable-assets") for arg in cmd): + return + + user_directory = self._arg_value(cmd, "--user-directory") + base_directory = self._arg_value(cmd, "--base-directory") + if not user_directory: + user_directory = os.path.join(base_directory or comfy_root, "user") + database_directory = os.path.abspath(os.path.join(user_directory, "distributed", "workers")) + os.makedirs(database_directory, exist_ok=True) + # Hash the immutable ID: names/GPU/ports may change, and IDs must not + # become filesystem paths or collide after filename sanitization. + worker_id = str(worker_config["id"]) + filename = hashlib.sha256(worker_id.encode("utf-8")).hexdigest() + ".db" + database_path = os.path.join(database_directory, filename).replace(os.sep, "/") + cmd.extend(["--database-url", "sqlite:///" + database_path]) + def build_launch_command(self, worker_config, comfy_root): """Build the command to launch a worker.""" main_py = os.path.join(comfy_root, "main.py") @@ -141,4 +175,7 @@ class LaunchCommandBuilder: raise ValueError(f"Invalid characters in extra_args: {arg}. Forbidden: {forbidden}") cmd.extend(extra_args_list) + if int(self._arg_value(cmd, "--port")) != int(worker_config["port"]): + raise ValueError("extra_args must not override the configured worker port; set the port in the UI/config instead") + self._add_worker_database(cmd, worker_config, comfy_root) return cmd diff --git a/workers/process/lifecycle.py b/workers/process/lifecycle.py index 62a3653..f5e0c8e 100644 --- a/workers/process/lifecycle.py +++ b/workers/process/lifecycle.py @@ -7,7 +7,9 @@ import time from ...utils.config import load_config, save_config from ...utils.constants import PROCESS_TERMINATION_TIMEOUT, PROCESS_WAIT_TIMEOUT, WORKER_CHECK_INTERVAL from ...utils.logging import debug_log, log +from ...utils.network import get_server_port from ...utils.process import get_python_executable, is_process_alive, terminate_process +from ..ports import validate_worker_port try: import psutil @@ -28,6 +30,8 @@ class ProcessLifecycle: """Launch a worker process with logging.""" _ = show_window # Kept for API compatibility. comfy_root = self._manager.find_comfy_root() + config = load_config() + validate_worker_port(worker_config, get_server_port(), config.get("workers", [])) env = os.environ.copy() env["CUDA_VISIBLE_DEVICES"] = str(worker_config.get("cuda_device", 0))