From 491816b6793ec4610f349b6c4873e460c9ced208 Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Fri, 21 Aug 2026 11:56:53 +0300 Subject: [PATCH] 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. --- src/gpu_rent/remote/install_swarm_comfy.py | 149 ++++++++++++++++----- src/gpu_rent/remote/tune_swarm_perf.py | 45 ++++++- tests/test_tune_swarm_perf.py | 16 +++ 3 files changed, 169 insertions(+), 41 deletions(-) diff --git a/src/gpu_rent/remote/install_swarm_comfy.py b/src/gpu_rent/remote/install_swarm_comfy.py index f5a8186..f8de729 100644 --- a/src/gpu_rent/remote/install_swarm_comfy.py +++ b/src/gpu_rent/remote/install_swarm_comfy.py @@ -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: return post( "/API/DeleteBackend", @@ -246,7 +308,8 @@ def recover_errored_backends(*, wait_sec: float = 180.0) -> str: flush=True, ) ensure_dlbackend_bind() - if sanitize_backends_fds(): + disk_fix = sanitize_backends_fds() | repair_extra_args_on_disk() + if disk_fix: restart_swarmui_local() sid = get_session(time.time() + 60) @@ -256,7 +319,7 @@ def recover_errored_backends(*, wait_sec: float = 180.0) -> str: try: result = add_comfy_selfstart(sid) bid = int(result.get("id", 0)) - configure_comfy_backend(sid, bid) + configure_comfy_backend(sid, bid, extra_args="") except Exception as exc: print(f"AddNewBackend FAIL: {exc}", flush=True) run_diagnostics() @@ -265,8 +328,11 @@ def recover_errored_backends(*, wait_sec: float = 180.0) -> str: for bid, meta in entries: settings = meta.get("settings") or {} script = str(settings.get("StartScript") or "") - extra = str(settings.get("ExtraArgs") or "") - print(f"reconfigure errored id={bid} old StartScript={script!r}", flush=True) + extra = clean_extra_args(str(settings.get("ExtraArgs") or "")) + print( + f"reconfigure errored id={bid} old StartScript={script!r} ExtraArgs={extra!r}", + flush=True, + ) try: configure_comfy_backend(sid, bid, extra_args=extra) except Exception as exc: @@ -321,7 +387,7 @@ def recover_errored_backends(*, wait_sec: float = 180.0) -> str: result = add_comfy_selfstart(sid) print(f"AddNewBackend (recreate): {result}", flush=True) bid = int(result.get("id", 0)) - configure_comfy_backend(sid, bid) + configure_comfy_backend(sid, bid, extra_args="") except Exception as exc: print(f"recreate FAIL: {exc}", flush=True) return "errored" @@ -404,30 +470,43 @@ def comfy_start_script() -> str: def ensure_dlbackend_bind() -> None: - """Re-bind /opt/swarmui/dlbackend → data if mount dropped (common after reboot).""" - src = str(DATA / "dlbackend") - dst = str(SWARM_ROOT / "dlbackend") - if not (DATA / "dlbackend").is_dir(): - return - try: - mounted = subprocess.run( - ["findmnt", dst], - check=False, - capture_output=True, - text=True, - timeout=10, - ) - if mounted.returncode == 0 and src in (mounted.stdout or ""): - return - print(f"dlbackend bind missing — mount --bind {src} → {dst}", flush=True) - subprocess.run(["sudo", "-n", "mkdir", "-p", src, 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 dlbackend bind: {exc}", flush=True) + """Re-bind SwarmUI data mounts if dropped (git reset / reboot).""" + pairs = ( + ("dlbackend", DATA / "dlbackend", SWARM_ROOT / "dlbackend"), + ("Data", DATA / "Data", SWARM_ROOT / "Data"), + ("Models", DATA / "Models", SWARM_ROOT / "Models"), + ("Output", DATA / "Output", SWARM_ROOT / "Output"), + ) + for name, src_p, dst_p in pairs: + src, dst = str(src_p), str(dst_p) + if not src_p.is_dir(): + continue + try: + mounted = subprocess.run( + ["findmnt", "-n", "-o", "SOURCE", "--target", dst], + check=False, + capture_output=True, + text=True, + timeout=10, + ) + cur = (mounted.stdout or "").strip() + if mounted.returncode == 0 and cur == src: + continue + print(f"{name} bind missing — mount --bind {src} → {dst}", 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: @@ -460,10 +539,7 @@ def configure_comfy_backend( ) -> None: """Set StartScript after AddNewBackend (defaults are empty → stays disabled).""" script = comfy_start_script() - # Preserve ExtraArgs (e.g. --use-sage-attention); "" wipes sage every recover. - extra = "" if extra_args is None else str(extra_args).strip() - if extra in {"\\x", "x"}: - extra = "" + extra = clean_extra_args("" if extra_args is None else str(extra_args)) settings = { "StartScript": script, "ExtraArgs": extra, @@ -476,7 +552,10 @@ def configure_comfy_backend( "OverQueue": 1, "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: result = edit_backend( sid, backend_id, title="ComfyUI Self-Starting", settings=settings @@ -563,7 +642,7 @@ def recover_empty_backends() -> str: flush=True, ) try: - extra = str(settings.get("ExtraArgs") or "") + extra = clean_extra_args(str(settings.get("ExtraArgs") or "")) configure_comfy_backend(sid, bid, extra_args=extra) except Exception as exc: print(f"configure FAIL id={bid}: {exc}", flush=True) diff --git a/src/gpu_rent/remote/tune_swarm_perf.py b/src/gpu_rent/remote/tune_swarm_perf.py index 0bab583..ce393c1 100644 --- a/src/gpu_rent/remote/tune_swarm_perf.py +++ b/src/gpu_rent/remote/tune_swarm_perf.py @@ -84,24 +84,57 @@ def find_pip() -> Path | None: 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 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(): return False text = BACKENDS.read_text(encoding="utf-8") - # Only horizontal whitespace after \x (spaces/tabs), never newline. - new, n = re.subn( + new = text + new, n1 = re.subn( r"^(\s*ExtraArgs:\s*)\\x[ \t]+", r"\1", - text, + new, 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: - BACKENDS.write_text(new, encoding="utf-8") - print(f"sanitized Backends.fds ExtraArgs \\x corruption ({n} line(s))") + BACKENDS.write_text(new if new.endswith("\n") else new + "\n", encoding="utf-8") + print(f"sanitized Backends.fds ExtraArgs ({n} change(s))") return True return False diff --git a/tests/test_tune_swarm_perf.py b/tests/test_tune_swarm_perf.py index ca72bd7..5e7f97b 100644 --- a/tests/test_tune_swarm_perf.py +++ b/tests/test_tune_swarm_perf.py @@ -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 +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): mod = _load() data = tmp_path