Enhance ExtraArgs sanitization and backend recovery logic
- Introduced the `clean_extra_args` function to filter out unsafe CLI flags and comment garbage from `ExtraArgs`, improving backend configuration integrity. - Updated the `repair_extra_args_on_disk` function to rewrite corrupted `ExtraArgs` lines in the FDS file, ensuring cleaner backend settings. - Enhanced the `recover_errored_backends` function to utilize the new sanitization logic, preserving valid `ExtraArgs` during backend reconfiguration. - Improved the `sanitize_backends_fds` function to drop comment garbage and retain valid entries, enhancing overall backend management. - Added tests to validate the new sanitization behavior, ensuring robustness in handling `ExtraArgs` during backend operations.
This commit is contained in:
@@ -229,6 +229,68 @@ def restart_backends(sid: str, *, which: str = "all") -> dict:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def clean_extra_args(raw: str | None) -> str:
|
||||||
|
"""Keep only safe Comfy CLI flags; drop FDS/comment garbage."""
|
||||||
|
s = (raw or "").strip()
|
||||||
|
if not s or s in {"\\x", "x"}:
|
||||||
|
return ""
|
||||||
|
# Corrupted FDS often leaves ConfigComment text in ExtraArgs.
|
||||||
|
low = s.lower()
|
||||||
|
if (
|
||||||
|
s.startswith("#")
|
||||||
|
or "unchecked" in low
|
||||||
|
or "automatic args" in low
|
||||||
|
or "disableinternal" in low.replace(" ", "")
|
||||||
|
or len(s) > 180
|
||||||
|
):
|
||||||
|
print(f"ExtraArgs garbage dropped: {s[:80]!r}...", flush=True)
|
||||||
|
return ""
|
||||||
|
keep: list[str] = []
|
||||||
|
if "--use-sage-attention" in s:
|
||||||
|
keep.append("--use-sage-attention")
|
||||||
|
if keep:
|
||||||
|
return " ".join(keep)
|
||||||
|
if re.fullmatch(r"[-\w.\s=/]+", s) and s.startswith("-"):
|
||||||
|
return s
|
||||||
|
print(f"ExtraArgs unrecognized dropped: {s[:80]!r}", flush=True)
|
||||||
|
return ""
|
||||||
|
|
||||||
|
|
||||||
|
def repair_extra_args_on_disk() -> bool:
|
||||||
|
"""Rewrite ExtraArgs lines that contain comment/garbage (post-sanitize merge)."""
|
||||||
|
path = DATA / "Data" / "Backends.fds"
|
||||||
|
if not path.is_file():
|
||||||
|
return False
|
||||||
|
try:
|
||||||
|
text = path.read_text(encoding="utf-8")
|
||||||
|
except OSError:
|
||||||
|
return False
|
||||||
|
lines = text.splitlines()
|
||||||
|
out: list[str] = []
|
||||||
|
changed = False
|
||||||
|
for line in lines:
|
||||||
|
m = re.match(r"^(\s*ExtraArgs:\s*)(.*)$", line)
|
||||||
|
if not m:
|
||||||
|
out.append(line)
|
||||||
|
continue
|
||||||
|
indent, val = m.group(1), m.group(2).strip()
|
||||||
|
if val in {"\\x", ""}:
|
||||||
|
out.append(line)
|
||||||
|
continue
|
||||||
|
cleaned = clean_extra_args(val)
|
||||||
|
if cleaned == val:
|
||||||
|
out.append(line)
|
||||||
|
continue
|
||||||
|
# Empty → leave blank value (Swarm will normalize to \x on save)
|
||||||
|
out.append(f"{indent}{cleaned}" if cleaned else f"{indent}")
|
||||||
|
changed = True
|
||||||
|
print(f"repaired ExtraArgs on disk: {val[:60]!r} -> {cleaned!r}", flush=True)
|
||||||
|
if not changed:
|
||||||
|
return False
|
||||||
|
path.write_text("\n".join(out) + "\n", encoding="utf-8")
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
def delete_backend(sid: str, backend_id: int) -> dict:
|
def delete_backend(sid: str, backend_id: int) -> dict:
|
||||||
return post(
|
return post(
|
||||||
"/API/DeleteBackend",
|
"/API/DeleteBackend",
|
||||||
@@ -246,7 +308,8 @@ def recover_errored_backends(*, wait_sec: float = 180.0) -> str:
|
|||||||
flush=True,
|
flush=True,
|
||||||
)
|
)
|
||||||
ensure_dlbackend_bind()
|
ensure_dlbackend_bind()
|
||||||
if sanitize_backends_fds():
|
disk_fix = sanitize_backends_fds() | repair_extra_args_on_disk()
|
||||||
|
if disk_fix:
|
||||||
restart_swarmui_local()
|
restart_swarmui_local()
|
||||||
|
|
||||||
sid = get_session(time.time() + 60)
|
sid = get_session(time.time() + 60)
|
||||||
@@ -256,7 +319,7 @@ def recover_errored_backends(*, wait_sec: float = 180.0) -> str:
|
|||||||
try:
|
try:
|
||||||
result = add_comfy_selfstart(sid)
|
result = add_comfy_selfstart(sid)
|
||||||
bid = int(result.get("id", 0))
|
bid = int(result.get("id", 0))
|
||||||
configure_comfy_backend(sid, bid)
|
configure_comfy_backend(sid, bid, extra_args="")
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
print(f"AddNewBackend FAIL: {exc}", flush=True)
|
print(f"AddNewBackend FAIL: {exc}", flush=True)
|
||||||
run_diagnostics()
|
run_diagnostics()
|
||||||
@@ -265,8 +328,11 @@ def recover_errored_backends(*, wait_sec: float = 180.0) -> str:
|
|||||||
for bid, meta in entries:
|
for bid, meta in entries:
|
||||||
settings = meta.get("settings") or {}
|
settings = meta.get("settings") or {}
|
||||||
script = str(settings.get("StartScript") or "")
|
script = str(settings.get("StartScript") or "")
|
||||||
extra = str(settings.get("ExtraArgs") or "")
|
extra = clean_extra_args(str(settings.get("ExtraArgs") or ""))
|
||||||
print(f"reconfigure errored id={bid} old StartScript={script!r}", flush=True)
|
print(
|
||||||
|
f"reconfigure errored id={bid} old StartScript={script!r} ExtraArgs={extra!r}",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
try:
|
try:
|
||||||
configure_comfy_backend(sid, bid, extra_args=extra)
|
configure_comfy_backend(sid, bid, extra_args=extra)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
@@ -321,7 +387,7 @@ def recover_errored_backends(*, wait_sec: float = 180.0) -> str:
|
|||||||
result = add_comfy_selfstart(sid)
|
result = add_comfy_selfstart(sid)
|
||||||
print(f"AddNewBackend (recreate): {result}", flush=True)
|
print(f"AddNewBackend (recreate): {result}", flush=True)
|
||||||
bid = int(result.get("id", 0))
|
bid = int(result.get("id", 0))
|
||||||
configure_comfy_backend(sid, bid)
|
configure_comfy_backend(sid, bid, extra_args="")
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
print(f"recreate FAIL: {exc}", flush=True)
|
print(f"recreate FAIL: {exc}", flush=True)
|
||||||
return "errored"
|
return "errored"
|
||||||
@@ -404,30 +470,43 @@ def comfy_start_script() -> str:
|
|||||||
|
|
||||||
|
|
||||||
def ensure_dlbackend_bind() -> None:
|
def ensure_dlbackend_bind() -> None:
|
||||||
"""Re-bind /opt/swarmui/dlbackend → data if mount dropped (common after reboot)."""
|
"""Re-bind SwarmUI data mounts if dropped (git reset / reboot)."""
|
||||||
src = str(DATA / "dlbackend")
|
pairs = (
|
||||||
dst = str(SWARM_ROOT / "dlbackend")
|
("dlbackend", DATA / "dlbackend", SWARM_ROOT / "dlbackend"),
|
||||||
if not (DATA / "dlbackend").is_dir():
|
("Data", DATA / "Data", SWARM_ROOT / "Data"),
|
||||||
return
|
("Models", DATA / "Models", SWARM_ROOT / "Models"),
|
||||||
try:
|
("Output", DATA / "Output", SWARM_ROOT / "Output"),
|
||||||
mounted = subprocess.run(
|
)
|
||||||
["findmnt", dst],
|
for name, src_p, dst_p in pairs:
|
||||||
check=False,
|
src, dst = str(src_p), str(dst_p)
|
||||||
capture_output=True,
|
if not src_p.is_dir():
|
||||||
text=True,
|
continue
|
||||||
timeout=10,
|
try:
|
||||||
)
|
mounted = subprocess.run(
|
||||||
if mounted.returncode == 0 and src in (mounted.stdout or ""):
|
["findmnt", "-n", "-o", "SOURCE", "--target", dst],
|
||||||
return
|
check=False,
|
||||||
print(f"dlbackend bind missing — mount --bind {src} → {dst}", flush=True)
|
capture_output=True,
|
||||||
subprocess.run(["sudo", "-n", "mkdir", "-p", src, dst], check=False, timeout=15)
|
text=True,
|
||||||
subprocess.run(
|
timeout=10,
|
||||||
["sudo", "-n", "mount", "--bind", src, dst],
|
)
|
||||||
check=False,
|
cur = (mounted.stdout or "").strip()
|
||||||
timeout=30,
|
if mounted.returncode == 0 and cur == src:
|
||||||
)
|
continue
|
||||||
except (OSError, subprocess.TimeoutExpired) as exc:
|
print(f"{name} bind missing — mount --bind {src} → {dst}", flush=True)
|
||||||
print(f"WARN dlbackend bind: {exc}", flush=True)
|
subprocess.run(["sudo", "-n", "mkdir", "-p", src, dst], check=False, timeout=15)
|
||||||
|
if mounted.returncode == 0:
|
||||||
|
subprocess.run(
|
||||||
|
["sudo", "-n", "umount", dst],
|
||||||
|
check=False,
|
||||||
|
timeout=15,
|
||||||
|
)
|
||||||
|
subprocess.run(
|
||||||
|
["sudo", "-n", "mount", "--bind", src, dst],
|
||||||
|
check=False,
|
||||||
|
timeout=30,
|
||||||
|
)
|
||||||
|
except (OSError, subprocess.TimeoutExpired) as exc:
|
||||||
|
print(f"WARN {name} bind: {exc}", flush=True)
|
||||||
|
|
||||||
|
|
||||||
def edit_backend(sid: str, backend_id: int, *, title: str, settings: dict) -> dict:
|
def edit_backend(sid: str, backend_id: int, *, title: str, settings: dict) -> dict:
|
||||||
@@ -460,10 +539,7 @@ def configure_comfy_backend(
|
|||||||
) -> None:
|
) -> None:
|
||||||
"""Set StartScript after AddNewBackend (defaults are empty → stays disabled)."""
|
"""Set StartScript after AddNewBackend (defaults are empty → stays disabled)."""
|
||||||
script = comfy_start_script()
|
script = comfy_start_script()
|
||||||
# Preserve ExtraArgs (e.g. --use-sage-attention); "" wipes sage every recover.
|
extra = clean_extra_args("" if extra_args is None else str(extra_args))
|
||||||
extra = "" if extra_args is None else str(extra_args).strip()
|
|
||||||
if extra in {"\\x", "x"}:
|
|
||||||
extra = ""
|
|
||||||
settings = {
|
settings = {
|
||||||
"StartScript": script,
|
"StartScript": script,
|
||||||
"ExtraArgs": extra,
|
"ExtraArgs": extra,
|
||||||
@@ -476,7 +552,10 @@ def configure_comfy_backend(
|
|||||||
"OverQueue": 1,
|
"OverQueue": 1,
|
||||||
"AutoRestart": True,
|
"AutoRestart": True,
|
||||||
}
|
}
|
||||||
print(f"EditBackend id={backend_id} StartScript={script}", flush=True)
|
print(
|
||||||
|
f"EditBackend id={backend_id} StartScript={script} ExtraArgs={extra!r}",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
try:
|
try:
|
||||||
result = edit_backend(
|
result = edit_backend(
|
||||||
sid, backend_id, title="ComfyUI Self-Starting", settings=settings
|
sid, backend_id, title="ComfyUI Self-Starting", settings=settings
|
||||||
@@ -563,7 +642,7 @@ def recover_empty_backends() -> str:
|
|||||||
flush=True,
|
flush=True,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
extra = str(settings.get("ExtraArgs") or "")
|
extra = clean_extra_args(str(settings.get("ExtraArgs") or ""))
|
||||||
configure_comfy_backend(sid, bid, extra_args=extra)
|
configure_comfy_backend(sid, bid, extra_args=extra)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
print(f"configure FAIL id={bid}: {exc}", flush=True)
|
print(f"configure FAIL id={bid}: {exc}", flush=True)
|
||||||
|
|||||||
@@ -84,24 +84,57 @@ def find_pip() -> Path | None:
|
|||||||
|
|
||||||
|
|
||||||
def sanitize_backends_fds() -> bool:
|
def sanitize_backends_fds() -> bool:
|
||||||
"""Repair ``ExtraArgs: \\x --flag`` only. Bare ``\\x`` is valid FDS empty.
|
"""Repair ``ExtraArgs: \\x --flag`` and comment/garbage ExtraArgs values.
|
||||||
|
|
||||||
Important: do not use ``\\s+`` after ``\\x`` — that eats the newline and
|
Important: do not use ``\\s+`` after ``\\x`` — that eats the newline and
|
||||||
merges the next FDS key onto the ExtraArgs line (breaks SwarmUI load).
|
merges the next FDS key onto the ExtraArgs line (breaks SwarmUI load).
|
||||||
|
Bare ``\\x`` is valid FDS empty and must be left alone.
|
||||||
"""
|
"""
|
||||||
if not BACKENDS.is_file():
|
if not BACKENDS.is_file():
|
||||||
return False
|
return False
|
||||||
text = BACKENDS.read_text(encoding="utf-8")
|
text = BACKENDS.read_text(encoding="utf-8")
|
||||||
# Only horizontal whitespace after \x (spaces/tabs), never newline.
|
new = text
|
||||||
new, n = re.subn(
|
new, n1 = re.subn(
|
||||||
r"^(\s*ExtraArgs:\s*)\\x[ \t]+",
|
r"^(\s*ExtraArgs:\s*)\\x[ \t]+",
|
||||||
r"\1",
|
r"\1",
|
||||||
text,
|
new,
|
||||||
flags=re.M,
|
flags=re.M,
|
||||||
)
|
)
|
||||||
|
# Drop ConfigComment garbage left in ExtraArgs after a bad merge.
|
||||||
|
lines_out: list[str] = []
|
||||||
|
n2 = 0
|
||||||
|
for line in new.splitlines():
|
||||||
|
m = re.match(r"^(\s*ExtraArgs:\s*)(.*)$", line)
|
||||||
|
if not m:
|
||||||
|
lines_out.append(line)
|
||||||
|
continue
|
||||||
|
indent, val = m.group(1), m.group(2).strip()
|
||||||
|
if val in {"", "\\x"}:
|
||||||
|
lines_out.append(line)
|
||||||
|
continue
|
||||||
|
low = val.lower()
|
||||||
|
bad = (
|
||||||
|
val.startswith("#")
|
||||||
|
or "unchecked" in low
|
||||||
|
or "automatic args" in low
|
||||||
|
or len(val) > 180
|
||||||
|
)
|
||||||
|
if bad:
|
||||||
|
lines_out.append(f"{indent}")
|
||||||
|
n2 += 1
|
||||||
|
print(f"sanitized ExtraArgs garbage: {val[:60]!r}")
|
||||||
|
elif "--use-sage-attention" in val and not val.startswith("--"):
|
||||||
|
# Keep only the sage flag if mixed with junk
|
||||||
|
lines_out.append(f"{indent}--use-sage-attention")
|
||||||
|
n2 += 1
|
||||||
|
else:
|
||||||
|
lines_out.append(line)
|
||||||
|
if n2:
|
||||||
|
new = "\n".join(lines_out) + "\n"
|
||||||
|
n = n1 + n2
|
||||||
if n and new != text:
|
if n and new != text:
|
||||||
BACKENDS.write_text(new, encoding="utf-8")
|
BACKENDS.write_text(new if new.endswith("\n") else new + "\n", encoding="utf-8")
|
||||||
print(f"sanitized Backends.fds ExtraArgs \\x corruption ({n} line(s))")
|
print(f"sanitized Backends.fds ExtraArgs ({n} change(s))")
|
||||||
return True
|
return True
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
|||||||
@@ -189,6 +189,22 @@ def test_sanitize_does_not_eat_newline_on_bare_empty(tmp_path, monkeypatch):
|
|||||||
assert backends.read_text(encoding="utf-8") == original
|
assert backends.read_text(encoding="utf-8") == original
|
||||||
|
|
||||||
|
|
||||||
|
def test_sanitize_drops_comment_garbage_extra_args(tmp_path, monkeypatch):
|
||||||
|
mod = _load()
|
||||||
|
backends = tmp_path / "Backends.fds"
|
||||||
|
junk = (
|
||||||
|
'\tExtraArgs: "#If unchecked, the system will automatically add some '
|
||||||
|
'relevant arguments to the comfy launch."\n'
|
||||||
|
"\tStartScript: /mnt/x/main.py\n"
|
||||||
|
)
|
||||||
|
backends.write_text(junk, encoding="utf-8")
|
||||||
|
monkeypatch.setattr(mod, "BACKENDS", backends)
|
||||||
|
assert mod.sanitize_backends_fds() is True
|
||||||
|
text = backends.read_text(encoding="utf-8")
|
||||||
|
assert "unchecked" not in text
|
||||||
|
assert "StartScript: /mnt/x/main.py" in text
|
||||||
|
|
||||||
|
|
||||||
def test_ensure_absolute_start_script(tmp_path, monkeypatch):
|
def test_ensure_absolute_start_script(tmp_path, monkeypatch):
|
||||||
mod = _load()
|
mod = _load()
|
||||||
data = tmp_path
|
data = tmp_path
|
||||||
|
|||||||
Reference in New Issue
Block a user