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,10 +165,16 @@ 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
|
||||
):
|
||||
# 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()
|
||||
@@ -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"],
|
||||
|
||||
+48
-4
@@ -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(
|
||||
|
||||
@@ -176,3 +176,23 @@ def test_sanitize_backends_fds_corruption(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(mod, "BACKENDS", backends)
|
||||
assert mod.sanitize_backends_fds() is True
|
||||
assert backends.read_text(encoding="utf-8") == "ExtraArgs: --use-sage-attention\n"
|
||||
|
||||
|
||||
def test_ensure_absolute_start_script(tmp_path, monkeypatch):
|
||||
mod = _load()
|
||||
data = tmp_path
|
||||
main_py = data / "dlbackend" / "ComfyUI" / "main.py"
|
||||
main_py.parent.mkdir(parents=True)
|
||||
main_py.write_text("# comfy\n", encoding="utf-8")
|
||||
backends = data / "Data" / "Backends.fds"
|
||||
backends.parent.mkdir(parents=True)
|
||||
backends.write_text(
|
||||
"0:\n\ttype: comfyui_selfstart\n\tStartScript: dlbackend/ComfyUI/main.py\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.setattr(mod, "DATA", data)
|
||||
monkeypatch.setattr(mod, "BACKENDS", backends)
|
||||
assert mod.ensure_absolute_start_script() is True
|
||||
text = backends.read_text(encoding="utf-8")
|
||||
assert str(main_py.resolve()) in text
|
||||
assert mod.ensure_absolute_start_script() is False
|
||||
|
||||
Reference in New Issue
Block a user