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.
2026-08-21 11:44:25 +03:00
parent fc60fef279
commit 021feaa839
4 changed files with 174 additions and 33 deletions
@@ -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"],
@@ -178,12 +178,36 @@ def _bind_access(
        clock.mark("Idle", log)
        try:
            if tune_swarm_perf(cfg, ip, log):
+                # git reset / prior boots can drop binds; remount before restart.
+                run_ssh(
+                    cfg,
+                    ip,
+                    "sudo -n bash -c '"
+                    "for p in dlbackend Data Models Output; do "
+                    "src=/mnt/swarm_data/$p; dst=/opt/swarmui/$p; "
+                    "mkdir -p \"$src\" \"$dst\"; "
+                    "cur=$(findmnt -n -o SOURCE --target \"$dst\" 2>/dev/null || true); "
+                    "if [[ \"$cur\" != \"$src\" ]]; then "
+                    "umount \"$dst\" 2>/dev/null || umount -l \"$dst\" 2>/dev/null || true; "
+                    "mount --bind \"$src\" \"$dst\" || true; "
+                    "echo remounted $dst; "
+                    "fi; "
+                    "done'",
+                    check=False,
+                    timeout=60,
+                )
                log("systemctl restart swarmui (perf ExtraArgs)")
                run_ssh(cfg, ip, "sudo -n systemctl restart swarmui", timeout=120)
                try:
                    wait_backend_idle(cfg, ip, log)
                except CloudError as exc:
-                    log(f"ready after perf: {exc}")
+                    log(f"ready after perf: {exc} — recover backends")
+                    try:
+                        ensure_swarm_comfy_installed(cfg, ip, log)
+                        wait_backend_idle(cfg, ip, log)
+                    except Exception as exc2:
+                        log(f"recover after perf: {exc2}")
+                        raise exc from exc2
            else:
                log("perf tune: restart не нужен")
        except Exception as exc:
@@ -314,8 +338,26 @@ def cmd_up(
            state.server_id = existing.id
            state.server_name = getattr(existing, "name", None)
            if status == "ACTIVE":
-                if state.bootstrapped and state.floating_ip:
-                    # Warm re-up: skip git pull/stop cascade unless --update.
+                # Prefer FIP from state; may still need ensure below.
+                fip = state.floating_ip
+                # Warm if local state says so, OR VM still has bootstrap marker
+                # (local bootstrapped often cleared on stop while disk/VM remain).
+                vm_warm = False
+                if fip and not (state.bootstrapped and state.floating_ip):
+                    try:
+                        if probe_ssh(cfg, fip, attempts=1) == "ok":
+                            probe = run_ssh(
+                                cfg,
+                                fip,
+                                "test -f /opt/swarmui/.gpu-rent-bootstrapped && echo yes || echo no",
+                                check=False,
+                                timeout=20,
+                            ).strip()
+                            vm_warm = probe.endswith("yes") or probe == "yes"
+                    except Exception:
+                        vm_warm = False
+
+                if (state.bootstrapped and state.floating_ip) or vm_warm:
                    if update is None:
                        do_update = False
                        log(
@@ -324,6 +366,9 @@ def cmd_up(
                        )
                    else:
                        do_update = update
+                    if vm_warm and not state.bootstrapped:
+                        log("warm: маркер на VM — восстанавливаю local bootstrapped")
+                        state.bootstrapped = True
                    log("сервер уже ACTIVE — второй GPU не создаём")
                    state.phase = "bootstrapping"
                    save_state(state)
@@ -331,7 +376,6 @@ def cmd_up(
                    return state

                # Bootstrap не завершён: почти всегда VM без authorized_keys.
-                fip = state.floating_ip
                if not fip:
                    try:
                        fip, fip_id = ensure_floating_ip(