Fix restart losing quotes around arguments that contain spaces (#3488)
os.execv on Windows joins its arguments with spaces and quotes nothing, so only the interpreter and script path (quoted by hand) survived a restart. Any other argument with a space, e.g. --browser-profile "Profile 1", was split in two and ComfyUI failed to start with 'unrecognized arguments'. Build the command in one helper, manager_util.build_restart_command, used by both the reboot endpoint and the prestartup restart. On Windows every entry is quoted with subprocess.list2cmdline; other platforms are unchanged. Fixes #3339
This commit is contained in:
Pablo Apiolazza
+1
-11
@@ -1967,17 +1967,7 @@ def restart(self):
|
||||
|
||||
print("\nRestarting... [Legacy Mode]\n\n") # This printing should not be logging - that will be ugly
|
||||
|
||||
sys_argv = sys.argv.copy()
|
||||
if '--windows-standalone-build' in sys_argv:
|
||||
sys_argv.remove('--windows-standalone-build')
|
||||
|
||||
if sys_argv[0].endswith("__main__.py"): # this is a python module
|
||||
module_name = os.path.basename(os.path.dirname(sys_argv[0]))
|
||||
cmds = [sys.executable, '-m', module_name] + sys_argv[1:]
|
||||
elif sys.platform.startswith('win32'):
|
||||
cmds = ['"' + sys.executable + '"', '"' + sys_argv[0] + '"'] + sys_argv[1:]
|
||||
else:
|
||||
cmds = [sys.executable] + sys_argv
|
||||
cmds = manager_util.build_restart_command(sys.argv, drop_args=('--windows-standalone-build',))
|
||||
|
||||
print(f"Command: {cmds}", flush=True)
|
||||
|
||||
|
||||
@@ -693,6 +693,32 @@ def robust_readlines(fullpath):
|
||||
return []
|
||||
|
||||
|
||||
def build_restart_command(argv, drop_args=()):
|
||||
"""
|
||||
Rebuild the command line used to re-exec the current ComfyUI process.
|
||||
|
||||
`argv` is `sys.argv` of the running process; entries listed in `drop_args` are left out.
|
||||
The result is meant to be passed to `os.execv(sys.executable, ...)`.
|
||||
|
||||
On Windows, `os.execv` does not quote anything: it joins the list with spaces into one
|
||||
command line and the new process splits it again. Every entry is therefore quoted with
|
||||
the MSVC rules (`subprocess.list2cmdline`), otherwise an argument such as `Profile 1`
|
||||
comes back as two arguments. Other platforms pass the list through untouched.
|
||||
"""
|
||||
argv = [x for x in argv if x not in drop_args]
|
||||
|
||||
if argv[0].endswith("__main__.py"): # this is a python module
|
||||
module_name = os.path.basename(os.path.dirname(argv[0]))
|
||||
cmds = [sys.executable, '-m', module_name] + argv[1:]
|
||||
else:
|
||||
cmds = [sys.executable] + argv
|
||||
|
||||
if sys.platform.startswith('win32'):
|
||||
cmds = [subprocess.list2cmdline([x]) for x in cmds]
|
||||
|
||||
return cmds
|
||||
|
||||
|
||||
def restore_pip_snapshot(pips, options):
|
||||
non_url = []
|
||||
local_url = []
|
||||
|
||||
@@ -864,15 +864,7 @@ if script_executed:
|
||||
print("--------------------------------------------------------------------------\n")
|
||||
exit(0)
|
||||
else:
|
||||
sys_argv = sys.argv.copy()
|
||||
|
||||
if sys_argv[0].endswith("__main__.py"): # this is a python module
|
||||
module_name = os.path.basename(os.path.dirname(sys_argv[0]))
|
||||
cmds = [sys.executable, '-m', module_name] + sys_argv[1:]
|
||||
elif sys.platform.startswith('win32'):
|
||||
cmds = ['"' + sys.executable + '"', '"' + sys_argv[0] + '"'] + sys_argv[1:]
|
||||
else:
|
||||
cmds = [sys.executable] + sys_argv
|
||||
cmds = manager_util.build_restart_command(sys.argv)
|
||||
|
||||
print(f"Command: {cmds}", flush=True)
|
||||
print("--------------------------------------------------------------------------\n")
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
"""Regression tests for the restart command line (#3339).
|
||||
|
||||
Both restart paths (the reboot endpoint and the prestartup restart after a
|
||||
dependency install) re-exec ComfyUI with os.execv. On Windows os.execv joins
|
||||
its arguments with spaces and quotes nothing, so `--browser-profile "Profile 1"`
|
||||
used to reach the new process as `--browser-profile Profile 1`.
|
||||
"""
|
||||
import ast
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from manager_test_utils import GLOB_DIR, load_manager_util
|
||||
|
||||
MANAGER_UTIL = load_manager_util()
|
||||
|
||||
BACKSLASH = chr(92)
|
||||
TRICKY_ARGS = [
|
||||
"--browser-profile", "Profile 1",
|
||||
"--output-directory", "C:" + BACKSLASH + "my dir" + BACKSLASH,
|
||||
"--note", 'say "hi"',
|
||||
"--empty", "",
|
||||
"--port", "8188",
|
||||
]
|
||||
|
||||
|
||||
def test_windows_quotes_only_what_needs_quoting(monkeypatch):
|
||||
monkeypatch.setattr(sys, "platform", "win32")
|
||||
monkeypatch.setattr(sys, "executable", "C:" + BACKSLASH + "Program Files" + BACKSLASH + "python.exe")
|
||||
|
||||
cmds = MANAGER_UTIL.build_restart_command(["C:" + BACKSLASH + "Comfy UI" + BACKSLASH + "main.py"] + TRICKY_ARGS)
|
||||
|
||||
assert cmds == [
|
||||
'"C:' + BACKSLASH + 'Program Files' + BACKSLASH + 'python.exe"',
|
||||
'"C:' + BACKSLASH + 'Comfy UI' + BACKSLASH + 'main.py"',
|
||||
"--browser-profile", '"Profile 1"',
|
||||
"--output-directory", '"C:' + BACKSLASH + 'my dir' + BACKSLASH * 2 + '"',
|
||||
"--note", '"say ' + BACKSLASH + '"hi' + BACKSLASH + '""',
|
||||
"--empty", '""',
|
||||
"--port", "8188",
|
||||
]
|
||||
|
||||
|
||||
def test_posix_passes_arguments_through(monkeypatch):
|
||||
monkeypatch.setattr(sys, "platform", "linux")
|
||||
monkeypatch.setattr(sys, "executable", "/usr/bin/python3")
|
||||
|
||||
argv = ["/opt/Comfy UI/main.py"] + TRICKY_ARGS
|
||||
|
||||
assert MANAGER_UTIL.build_restart_command(argv) == ["/usr/bin/python3"] + argv
|
||||
|
||||
|
||||
def test_module_entry_point_uses_dash_m(monkeypatch):
|
||||
monkeypatch.setattr(sys, "platform", "linux")
|
||||
monkeypatch.setattr(sys, "executable", "/usr/bin/python3")
|
||||
|
||||
cmds = MANAGER_UTIL.build_restart_command(["/opt/site-packages/comfy/__main__.py", "--port", "8188"])
|
||||
|
||||
assert cmds == ["/usr/bin/python3", "-m", "comfy", "--port", "8188"]
|
||||
|
||||
|
||||
def test_module_entry_point_is_quoted_on_windows(monkeypatch):
|
||||
monkeypatch.setattr(sys, "platform", "win32")
|
||||
monkeypatch.setattr(sys, "executable", "C:" + BACKSLASH + "Program Files" + BACKSLASH + "python.exe")
|
||||
|
||||
# Forward slashes: the module name is then found by both the POSIX and the Windows os.path
|
||||
cmds = MANAGER_UTIL.build_restart_command(
|
||||
["C:/site packages/comfy/__main__.py", "--browser-profile", "Profile 1"])
|
||||
|
||||
assert cmds == [
|
||||
'"C:' + BACKSLASH + 'Program Files' + BACKSLASH + 'python.exe"',
|
||||
"-m", "comfy",
|
||||
"--browser-profile", '"Profile 1"',
|
||||
]
|
||||
|
||||
|
||||
def test_drop_args_removes_only_listed_arguments(monkeypatch):
|
||||
monkeypatch.setattr(sys, "platform", "linux")
|
||||
monkeypatch.setattr(sys, "executable", "/usr/bin/python3")
|
||||
argv = ["main.py", "--windows-standalone-build", "--cpu"]
|
||||
|
||||
assert MANAGER_UTIL.build_restart_command(argv) == ["/usr/bin/python3"] + argv
|
||||
assert MANAGER_UTIL.build_restart_command(argv, drop_args=("--windows-standalone-build",)) == \
|
||||
["/usr/bin/python3", "main.py", "--cpu"]
|
||||
assert argv == ["main.py", "--windows-standalone-build", "--cpu"] # caller's list is untouched
|
||||
|
||||
|
||||
def restart_call_drop_args(path):
|
||||
"""The `drop_args` that the one restart site in `path` passes for `sys.argv`."""
|
||||
calls = [node for node in ast.walk(ast.parse(path.read_text(encoding="utf-8")))
|
||||
if isinstance(node, ast.Call) and ast.unparse(node.func) == "manager_util.build_restart_command"]
|
||||
assert len(calls) == 1, f"expected one restart site in {path.name}"
|
||||
call = calls[0]
|
||||
assert ast.unparse(call.args[0]) == "sys.argv"
|
||||
|
||||
drop_args = call.args[1:] + [keyword.value for keyword in call.keywords if keyword.arg == "drop_args"]
|
||||
assert len(call.args) + len(call.keywords) == 1 + len(drop_args), "unexpected argument"
|
||||
return tuple(ast.literal_eval(drop_args[0])) if drop_args else ()
|
||||
|
||||
|
||||
def test_only_the_reboot_endpoint_drops_the_standalone_build_flag():
|
||||
# ComfyUI opens a browser when it starts with this flag. The reboot endpoint restarts a
|
||||
# ComfyUI whose page is already open, so it drops the flag (#1160).
|
||||
assert restart_call_drop_args(GLOB_DIR / "manager_server.py") == ("--windows-standalone-build",)
|
||||
|
||||
prestartup_script = GLOB_DIR.resolve().parent / "prestartup_script.py"
|
||||
if os.environ.get("MANAGER_GLOB_DIR") and not prestartup_script.exists():
|
||||
pytest.skip("MANAGER_GLOB_DIR points at a glob directory without prestartup_script.py beside it")
|
||||
# The prestartup restart happens before ComfyUI has opened anything, so the flag must stay.
|
||||
assert restart_call_drop_args(prestartup_script) == ()
|
||||
|
||||
|
||||
LAUNCHER = """
|
||||
import importlib.util, os, sys
|
||||
spec = importlib.util.spec_from_file_location("_manager_util_under_test", sys.argv[1])
|
||||
manager_util = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(manager_util)
|
||||
os.execv(sys.executable, manager_util.build_restart_command(sys.argv[2:]))
|
||||
"""
|
||||
|
||||
CHILD = """
|
||||
import json, sys
|
||||
with open(sys.argv[1] + ".tmp", "w", encoding="utf-8") as f:
|
||||
json.dump(sys.argv[2:], f)
|
||||
import os
|
||||
os.replace(sys.argv[1] + ".tmp", sys.argv[1])
|
||||
"""
|
||||
|
||||
|
||||
def test_arguments_survive_a_real_execv(tmp_path):
|
||||
"""End to end on the running platform: what the re-exec'ed process sees must equal what was sent."""
|
||||
work_dir = tmp_path / "dir with spaces"
|
||||
work_dir.mkdir()
|
||||
child = work_dir / "main entry.py"
|
||||
child.write_text(CHILD, encoding="utf-8")
|
||||
result_file = work_dir / "argv seen.json"
|
||||
|
||||
subprocess.run(
|
||||
[sys.executable, "-c", LAUNCHER, str(GLOB_DIR / "manager_util.py"), str(child), str(result_file)] + TRICKY_ARGS,
|
||||
check=True, timeout=60,
|
||||
)
|
||||
|
||||
# On Windows os.execv returns control to the caller before the new process is done.
|
||||
deadline = time.time() + 30
|
||||
while not result_file.exists():
|
||||
if time.time() > deadline:
|
||||
pytest.fail("re-exec'ed process never reported its arguments")
|
||||
time.sleep(0.1)
|
||||
|
||||
assert json.loads(result_file.read_text(encoding="utf-8")) == TRICKY_ARGS
|
||||
Reference in New Issue
Block a user