Enhance backend management and performance tuning
- Introduced remounting logic for data directories to ensure correct bindings before restarting services, improving reliability during performance tuning. - Added functions to sanitize and ensure absolute paths for backend scripts, preventing issues with relative paths in bind mounts. - Enhanced the `patch_backends_extra_args` function to clean up corrupted entries and ensure proper configuration of backend parameters. - Updated tests to validate the new sanitization and path handling functionalities, ensuring robustness in backend management.
This commit is contained in:
Leonid Pershin
@@ -68,9 +68,17 @@ ensure_data_mount() {
|
||||
ensure_bind() {
|
||||
local src="$1" dst="$2"
|
||||
mkdir -p "$src" "$dst"
|
||||
if ! findmnt "$dst" >/dev/null 2>&1; then
|
||||
mount --bind "$src" "$dst"
|
||||
if findmnt "$dst" >/dev/null 2>&1; then
|
||||
# git reset --hard can leave a stale/wrong bind; remount if SOURCE mismatch.
|
||||
local src_mnt
|
||||
src_mnt="$(findmnt -n -o SOURCE --target "$dst" 2>/dev/null || true)"
|
||||
if [[ "$src_mnt" == "$src" ]]; then
|
||||
return 0
|
||||
fi
|
||||
log "bind remount ${dst} (было: ${src_mnt:-?})"
|
||||
umount "$dst" 2>/dev/null || umount -l "$dst" 2>/dev/null || true
|
||||
fi
|
||||
mount --bind "$src" "$dst" || log "WARN: mount --bind ${src} → ${dst} failed"
|
||||
if ! grep -Fq " ${dst} " /etc/fstab; then
|
||||
echo "${src} ${dst} none bind,nofail 0 0" >> /etc/fstab
|
||||
fi
|
||||
|
||||
@@ -83,6 +83,77 @@ def find_pip() -> Path | None:
|
||||
return find_comfy_python()
|
||||
|
||||
|
||||
def sanitize_backends_fds() -> bool:
|
||||
"""Repair ``ExtraArgs: \\x --flag`` and bare ``ExtraArgs: \\x`` → clean values."""
|
||||
if not BACKENDS.is_file():
|
||||
return False
|
||||
text = BACKENDS.read_text(encoding="utf-8")
|
||||
new = text
|
||||
# ``ExtraArgs: \x --flags`` → ``ExtraArgs: --flags``
|
||||
new, n1 = re.subn(
|
||||
r"^(\s*ExtraArgs:\s*)\\x(\s+)",
|
||||
r"\1",
|
||||
new,
|
||||
flags=re.M,
|
||||
)
|
||||
# Bare empty marker left alone is OK for Swarm, but normalize trailing junk.
|
||||
new, n2 = re.subn(
|
||||
r"^(\s*ExtraArgs:\s*)\\x\s*$",
|
||||
r"\1",
|
||||
new,
|
||||
flags=re.M,
|
||||
)
|
||||
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))")
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def ensure_absolute_start_script() -> bool:
|
||||
"""Rewrite relative StartScript to absolute data-volume path (bind-safe)."""
|
||||
if not BACKENDS.is_file():
|
||||
return False
|
||||
candidates = (
|
||||
DATA / "dlbackend" / "ComfyUI" / "main.py",
|
||||
DATA / "dlbackend" / "comfy" / "ComfyUI" / "main.py",
|
||||
Path("/opt/swarmui/dlbackend/ComfyUI/main.py"),
|
||||
Path("/opt/swarmui/dlbackend/comfy/ComfyUI/main.py"),
|
||||
)
|
||||
abs_script = None
|
||||
for p in candidates:
|
||||
if p.is_file():
|
||||
abs_script = str(p.resolve())
|
||||
break
|
||||
if not abs_script:
|
||||
return False
|
||||
text = BACKENDS.read_text(encoding="utf-8")
|
||||
old_m = re.search(r"(?m)^\s*StartScript:\s*(.*)$", text)
|
||||
old = (old_m.group(1).strip() if old_m else "")
|
||||
if old == abs_script:
|
||||
return False
|
||||
# Relative paths are not bind-safe even if they resolve from cwd.
|
||||
if old:
|
||||
op = Path(old)
|
||||
if op.is_absolute() and op.is_file():
|
||||
try:
|
||||
if op.resolve() == Path(abs_script):
|
||||
return False
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def _repl(m: re.Match[str]) -> str:
|
||||
return f"{m.group(1)}{abs_script}"
|
||||
|
||||
new, n = re.subn(r"^(\s*StartScript:\s*).*$", _repl, text, flags=re.M)
|
||||
if not n or new == text:
|
||||
return False
|
||||
BACKENDS.write_text(new, encoding="utf-8")
|
||||
print(f"patched StartScript -> {abs_script} (was {old!r})")
|
||||
return True
|
||||
|
||||
|
||||
def patch_backends_extra_args(extra: str) -> bool:
|
||||
"""Ensure ExtraArgs contains sage flag for Comfy self-start backends.
|
||||
|
||||
@@ -94,12 +165,18 @@ def patch_backends_extra_args(extra: str) -> bool:
|
||||
if not BACKENDS.is_file():
|
||||
print(f"no {BACKENDS} yet — skip ExtraArgs (Comfy not registered)")
|
||||
return False
|
||||
# Always sanitize first so we never append onto ``\x``.
|
||||
sanitize_backends_fds()
|
||||
text = BACKENDS.read_text(encoding="utf-8")
|
||||
if "--use-sage-attention" in text and not re.search(
|
||||
r"ExtraArgs:\s*\\x\s+--use-sage-attention", text
|
||||
if re.search(
|
||||
rf"(?m)^\s*ExtraArgs:\s*{re.escape(extra)}\s*$", text
|
||||
) or re.search(
|
||||
rf"(?m)^\s*ExtraArgs:\s*.*\b{re.escape(extra)}\b", text
|
||||
):
|
||||
print("Backends.fds already has --use-sage-attention")
|
||||
return False
|
||||
# Already present without leading \x
|
||||
if not re.search(r"(?m)^\s*ExtraArgs:\s*\\x", text):
|
||||
print("Backends.fds already has --use-sage-attention")
|
||||
return False
|
||||
lines = text.splitlines()
|
||||
changed = False
|
||||
out = []
|
||||
@@ -114,13 +191,17 @@ def patch_backends_extra_args(extra: str) -> bool:
|
||||
elif m_corrupt:
|
||||
indent = m_corrupt.group(1)
|
||||
rest = m_corrupt.group(2).strip()
|
||||
if "--use-sage-attention" in rest:
|
||||
if extra in rest:
|
||||
out.append(f"{indent}ExtraArgs: {rest}")
|
||||
else:
|
||||
out.append(f"{indent}ExtraArgs: {rest} {extra}".strip())
|
||||
parts = f"{rest} {extra}".strip()
|
||||
out.append(f"{indent}ExtraArgs: {parts}")
|
||||
changed = True
|
||||
elif m_val and "--use-sage-attention" not in line:
|
||||
out.append(line.rstrip() + f" {extra}")
|
||||
elif m_val and extra not in line:
|
||||
indent = m_val.group(1)
|
||||
rest = m_val.group(2).strip()
|
||||
# Replace entirely if rest is only whitespace-ish flags we manage
|
||||
out.append(f"{indent}ExtraArgs: {rest} {extra}".strip())
|
||||
changed = True
|
||||
else:
|
||||
out.append(line)
|
||||
@@ -128,28 +209,12 @@ def patch_backends_extra_args(extra: str) -> bool:
|
||||
print("Backends.fds: no ExtraArgs field patched")
|
||||
return False
|
||||
BACKENDS.write_text("\n".join(out) + "\n", encoding="utf-8")
|
||||
# Final pass: never leave \x prefix
|
||||
sanitize_backends_fds()
|
||||
print(f"patched {BACKENDS} ExtraArgs += {extra}")
|
||||
return True
|
||||
|
||||
|
||||
def sanitize_backends_fds() -> bool:
|
||||
"""Repair ``ExtraArgs: \\x --flag`` corruption; return True if file changed."""
|
||||
if not BACKENDS.is_file():
|
||||
return False
|
||||
text = BACKENDS.read_text(encoding="utf-8")
|
||||
new, n = re.subn(
|
||||
r"^(\s*ExtraArgs:\s*)\\x(\s+)",
|
||||
r"\1",
|
||||
text,
|
||||
flags=re.M,
|
||||
)
|
||||
if n:
|
||||
BACKENDS.write_text(new, encoding="utf-8")
|
||||
print(f"sanitized Backends.fds ExtraArgs \\x corruption ({n} line(s))")
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def sage_already_importable(py: Path) -> bool:
|
||||
code, _ = _run(
|
||||
[str(py), "-c", "import triton, sageattention"],
|
||||
@@ -206,8 +271,10 @@ def main() -> int:
|
||||
except json.JSONDecodeError:
|
||||
prev = {}
|
||||
same_gpu = bool(prev.get("uuid") and prev.get("uuid") == plan["uuid"])
|
||||
# Always repair FDS corruption even when tune is otherwise a no-op.
|
||||
# Always repair FDS corruption / relative StartScript even when tune is a no-op.
|
||||
fds_fixed = sanitize_backends_fds()
|
||||
if ensure_absolute_start_script():
|
||||
fds_fixed = True
|
||||
if same_gpu and prev.get("extra_args") == plan["extra_args"]:
|
||||
if prev.get("pip_ok") or not plan["use_sage"]:
|
||||
print(f"perf tune already applied for {plan['name']} ({plan['tier']})")
|
||||
@@ -231,6 +298,8 @@ def main() -> int:
|
||||
# Only patch ExtraArgs when wheels installed — otherwise Comfy may break.
|
||||
if pip_ok and patch_backends_extra_args(plan["extra_args"]):
|
||||
restarted_needed = True
|
||||
if ensure_absolute_start_script():
|
||||
restarted_needed = True
|
||||
|
||||
marker = {
|
||||
"uuid": plan["uuid"],
|
||||
|
||||
Reference in New Issue
Block a user