Robert Wojciechowski
a1faf7e916
* fix: isolate local worker databases and avoid port conflicts * fix: harden local worker startup preflight and database paths Handle auto-launch before the master binds, allow TIME_WAIT restarts without accepting active listeners, resolve relative paths in the worker cwd, and round-trip SQLite URLs through SQLAlchemy. Add regressions and remove the README additions.
194 lines
8.3 KiB
Python
194 lines
8.3 KiB
Python
import glob
|
|
import hashlib
|
|
import os
|
|
import shlex
|
|
import shutil
|
|
|
|
from ...utils.logging import debug_log
|
|
from ...utils.process import get_python_executable
|
|
|
|
|
|
class LaunchCommandBuilder:
|
|
"""Build command-lines for launching worker ComfyUI processes."""
|
|
|
|
def _extend_arg(self, cmd, flag, value):
|
|
if value in (None, "", [], ()):
|
|
return
|
|
cmd.extend([flag, str(value)])
|
|
|
|
def _extend_grouped_args(self, cmd, flag, values):
|
|
for group in values or []:
|
|
flattened = [str(item) for item in group if item]
|
|
if flattened:
|
|
cmd.append(flag)
|
|
cmd.extend(flattened)
|
|
|
|
def _get_runtime_args(self):
|
|
try:
|
|
from comfy.cli_args import args
|
|
return args
|
|
except Exception as exc:
|
|
debug_log(f"Could not read current ComfyUI CLI args for worker launch: {exc}")
|
|
return None
|
|
|
|
def _build_runtime_launch_args(self):
|
|
args = self._get_runtime_args()
|
|
if args is None:
|
|
return []
|
|
|
|
inherited = []
|
|
self._extend_arg(inherited, "--listen", getattr(args, "listen", None))
|
|
self._extend_arg(inherited, "--base-directory", getattr(args, "base_directory", None))
|
|
self._extend_arg(inherited, "--temp-directory", getattr(args, "temp_directory", None))
|
|
self._extend_arg(inherited, "--input-directory", getattr(args, "input_directory", None))
|
|
self._extend_arg(inherited, "--output-directory", getattr(args, "output_directory", None))
|
|
self._extend_arg(inherited, "--user-directory", getattr(args, "user_directory", None))
|
|
self._extend_arg(inherited, "--front-end-root", getattr(args, "front_end_root", None))
|
|
self._extend_grouped_args(
|
|
inherited,
|
|
"--extra-model-paths-config",
|
|
getattr(args, "extra_model_paths_config", None),
|
|
)
|
|
|
|
if getattr(args, "enable_manager", False):
|
|
inherited.append("--enable-manager")
|
|
if getattr(args, "disable_manager_ui", False):
|
|
inherited.append("--disable-manager-ui")
|
|
if getattr(args, "enable_manager_legacy_ui", False):
|
|
inherited.append("--enable-manager-legacy-ui")
|
|
if getattr(args, "windows_standalone_build", False):
|
|
inherited.append("--windows-standalone-build")
|
|
if getattr(args, "log_stdout", False):
|
|
inherited.append("--log-stdout")
|
|
|
|
verbose = getattr(args, "verbose", None)
|
|
if verbose and verbose != "INFO":
|
|
inherited.extend(["--verbose", str(verbose)])
|
|
|
|
return inherited
|
|
|
|
def _find_windows_terminal(self):
|
|
"""Find Windows Terminal executable."""
|
|
possible_paths = [
|
|
os.path.expandvars(r"%LOCALAPPDATA%\Microsoft\WindowsApps\wt.exe"),
|
|
os.path.expandvars(r"%PROGRAMFILES%\WindowsApps\Microsoft.WindowsTerminal_*\wt.exe"),
|
|
"wt.exe",
|
|
]
|
|
|
|
for path in possible_paths:
|
|
if os.path.exists(path):
|
|
return path
|
|
if "*" in path:
|
|
matches = glob.glob(path)
|
|
if matches:
|
|
return matches[0]
|
|
|
|
wt_path = shutil.which("wt")
|
|
if wt_path:
|
|
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")
|
|
if not os.path.isabs(user_directory):
|
|
# Popen uses comfy_root as cwd, which can differ from the master's.
|
|
user_directory = os.path.join(comfy_root, user_directory)
|
|
database_directory = os.path.abspath(os.path.join(user_directory, "distributed", "workers"))
|
|
# 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, "/")
|
|
# Use the same URL codec as ComfyUI; raw '?' and '%xx' can otherwise
|
|
# truncate/decode the path and make distinct workers share a database.
|
|
from sqlalchemy.engine import URL, make_url
|
|
database_url = URL.create("sqlite", database=database_path).render_as_string()
|
|
if make_url(database_url).database != database_path:
|
|
raise ValueError(
|
|
"Installed SQLAlchemy cannot safely encode the worker database path; "
|
|
"use a user directory without '?' or an explicit --database-url with a safe path"
|
|
)
|
|
os.makedirs(database_directory, exist_ok=True)
|
|
cmd.extend(["--database-url", database_url])
|
|
|
|
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")
|
|
|
|
if os.path.exists(main_py):
|
|
cmd = [
|
|
get_python_executable(),
|
|
main_py,
|
|
]
|
|
cmd.extend(self._build_runtime_launch_args())
|
|
cmd.extend(["--port", str(worker_config["port"])])
|
|
|
|
current_args = self._get_runtime_args()
|
|
current_cors = getattr(current_args, "enable_cors_header", None) if current_args else None
|
|
cmd.append("--enable-cors-header")
|
|
if current_cors is not None:
|
|
cmd.append(str(current_cors))
|
|
|
|
if "--disable-auto-launch" not in cmd:
|
|
cmd.append("--disable-auto-launch")
|
|
|
|
debug_log(f"Using main.py: {main_py}")
|
|
else:
|
|
error_msg = f"Could not find main.py in {comfy_root}\n"
|
|
error_msg += f"Searched for: {main_py}\n"
|
|
error_msg += f"Directory contents of {comfy_root}:\n"
|
|
try:
|
|
if os.path.exists(comfy_root):
|
|
files = os.listdir(comfy_root)[:20]
|
|
error_msg += " " + "\n ".join(files)
|
|
if len(os.listdir(comfy_root)) > 20:
|
|
error_msg += f"\n ... and {len(os.listdir(comfy_root)) - 20} more files"
|
|
else:
|
|
error_msg += f" Directory {comfy_root} does not exist!"
|
|
except Exception as exc:
|
|
error_msg += f" Error listing directory: {exc}"
|
|
|
|
error_msg += "\n\nPossible solutions:\n"
|
|
error_msg += "1. Check if ComfyUI is installed in a different location\n"
|
|
error_msg += "2. For Docker: ComfyUI might be in /ComfyUI or /app\n"
|
|
error_msg += "3. Ensure the custom node is installed in the correct location\n"
|
|
raise RuntimeError(error_msg)
|
|
|
|
if worker_config.get("extra_args"):
|
|
raw_args = worker_config["extra_args"].strip()
|
|
if raw_args:
|
|
extra_args_list = shlex.split(raw_args)
|
|
forbidden_chars = set(";|>&<`$()[]{}*!?")
|
|
for arg in extra_args_list:
|
|
if any(char in forbidden_chars for char in arg):
|
|
forbidden = "".join(forbidden_chars)
|
|
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
|