diff --git a/src/gpu_rent/civitai.py b/src/gpu_rent/civitai.py index c7e5649..b05153d 100644 --- a/src/gpu_rent/civitai.py +++ b/src/gpu_rent/civitai.py @@ -80,6 +80,19 @@ def pick_primary_file(version: dict) -> dict | None: return None +def file_size_bytes(info: dict) -> int | None: + """Civitai ``sizeKB`` (float) → bytes, or None.""" + raw = info.get("sizeKB") + if raw is None: + raw = info.get("sizeKb") + if raw is None: + return None + try: + return int(round(float(raw) * 1024.0)) + except (TypeError, ValueError): + return None + + def pick_preview_image(version: dict) -> tuple[str, str] | None: """First usable Civitai preview → (url, sidecar_suffix) for SwarmUI. diff --git a/src/gpu_rent/provision.py b/src/gpu_rent/provision.py index 6b980e1..879280d 100644 --- a/src/gpu_rent/provision.py +++ b/src/gpu_rent/provision.py @@ -11,7 +11,12 @@ from pathlib import Path import httpx -from gpu_rent.civitai import fetch_model_version, pick_preview_image, pick_primary_file +from gpu_rent.civitai import ( + fetch_model_version, + file_size_bytes, + pick_preview_image, + pick_primary_file, +) from gpu_rent.config import Config from gpu_rent.errors import CloudError, GpuRentError from gpu_rent.manifests import ( @@ -480,6 +485,9 @@ def seed_civitai(cfg: Config, host: str, log: Log) -> None: f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2), }, } + size_b = file_size_bytes(info) + if size_b: + job["size"] = size_b preview = pick_preview_image(version) if preview: preview_url, preview_suffix = preview diff --git a/src/gpu_rent/remote/civitai_fetch.py b/src/gpu_rent/remote/civitai_fetch.py index 1afd18b..baea0ef 100644 --- a/src/gpu_rent/remote/civitai_fetch.py +++ b/src/gpu_rent/remote/civitai_fetch.py @@ -4,6 +4,7 @@ from __future__ import annotations import hashlib import json +import os import sys import time import urllib.parse @@ -91,15 +92,65 @@ def sha256_path(path: Path) -> str: return digest.hexdigest() -def should_skip(dest: Path, expect_sha: str) -> tuple[bool, str]: - """Return (skip, reason). Existing file with matching sha — or any non-empty if no sha.""" +def sha_sidecar_path(dest: Path) -> Path: + """Beside weight: name.safetensors.sha256 (avoids stem collisions).""" + return Path(str(dest) + ".sha256") + + +def write_sha_sidecar(dest: Path, digest: str) -> None: + try: + sha_sidecar_path(dest).write_text(digest.lower().strip() + "\n", encoding="utf-8") + except OSError: + pass + + +def read_sha_sidecar(dest: Path) -> str: + side = sha_sidecar_path(dest) + if not side.is_file(): + return "" + try: + return side.read_text(encoding="utf-8").strip().split()[0].lower() + except (OSError, IndexError): + return "" + + +def should_skip( + dest: Path, + expect_sha: str, + *, + expect_size: int | None = None, + verify_sha: bool = False, +) -> tuple[bool, str]: + """Return (skip, reason). + + Re-up must stay fast: prefer size + sha sidecar over hashing 10+ GB weights. + Full SHA only when ``verify_sha`` or size unknown. + """ if not dest.is_file() or dest.stat().st_size <= 0: return False, "" + size = dest.stat().st_size expect = (expect_sha or "").lower() + # sizeKB from Civitai is float → allow small rounding slack + if expect_size is not None and expect_size > 0: + if abs(size - int(expect_size)) > 4096: + return False, f"размер {size} != {expect_size} — перекачиваю" if not expect: return True, "уже есть" + recorded = read_sha_sidecar(dest) + if recorded and recorded == expect: + return True, "уже есть (sha sidecar)" + if ( + expect_size is not None + and expect_size > 0 + and abs(size - int(expect_size)) <= 4096 + and not verify_sha + ): + # Trust size match; persist expected hash so next skip is sidecar-fast. + write_sha_sidecar(dest, expect) + return True, "уже есть (size ok)" got = sha256_path(dest).lower() if got == expect: + write_sha_sidecar(dest, got) return True, "уже есть (sha ok)" return False, "sha не совпал — перекачиваю" @@ -186,6 +237,12 @@ def main() -> int: if HF_TOKEN_PATH.is_file() else "" ) + verify_sha = (os.environ.get("GPU_RENT_VERIFY_SHA") or "").strip().lower() in { + "1", + "true", + "yes", + "on", + } jobs = json.loads(JOBS_PATH.read_text(encoding="utf-8")) total = len(jobs) failed = 0 @@ -194,6 +251,11 @@ def main() -> int: for index, job in enumerate(jobs, start=1): dest = Path(job["dest"]) expect = (job.get("sha256") or "").lower() + expect_size = job.get("size") + try: + expect_size_n = int(expect_size) if expect_size is not None else None + except (TypeError, ValueError): + expect_size_n = None prefix = f"[{index}/{total}]" auth = str(job.get("auth") or "civitai").lower() if auth == "hf": @@ -202,7 +264,9 @@ def main() -> int: else: token = civitai_token auth_host = "civitai" - skip, reason = should_skip(dest, expect) + skip, reason = should_skip( + dest, expect, expect_size=expect_size_n, verify_sha=verify_sha + ) if skip: skipped += 1 size = dest.stat().st_size if dest.is_file() else 0 @@ -234,6 +298,7 @@ def main() -> int: if got != expect: dest.unlink(missing_ok=True) raise RuntimeError(f"sha256 {got} != {expect}") + write_sha_sidecar(dest, got) for extra_name, extra_text in (job.get("sidecars") or {}).items(): extra = dest.parent / extra_name extra.write_text(extra_text, encoding="utf-8") diff --git a/src/gpu_rent/remote/install_swarm_comfy.py b/src/gpu_rent/remote/install_swarm_comfy.py index 636f3d2..b5c9711 100644 --- a/src/gpu_rent/remote/install_swarm_comfy.py +++ b/src/gpu_rent/remote/install_swarm_comfy.py @@ -11,6 +11,7 @@ import base64 import hashlib import json import os +import re import socket import struct import subprocess @@ -28,6 +29,7 @@ DLBACKEND = DATA / "dlbackend" COMFY_ROOT = DLBACKEND / "ComfyUI" COMFY_VENV = COMFY_ROOT / "venv" / "bin" / "python" SETTINGS = DATA / "Data" / "Settings.fds" +BACKENDS_FDS = DATA / "Data" / "Backends.fds" # Comfy clone + torch can take 20–40+ min on a cold disk. INSTALL_TIMEOUT = float(os.environ.get("GPU_RENT_COMFY_INSTALL_TIMEOUT") or "3600") @@ -255,6 +257,87 @@ def recover_errored_backends(*, wait_sec: float = 180.0) -> str: return last +def sanitize_backends_fds() -> bool: + """Repair ExtraArgs: \\x --flag (FDS empty + appended flag). Return True if changed.""" + path = DATA / "Data" / "Backends.fds" + if not path.is_file(): + return False + try: + text = path.read_text(encoding="utf-8") + except OSError: + return False + new, n = re.subn(r"^(\s*ExtraArgs:\s*)\\x(\s+)", r"\1", text, flags=re.M) + if not n: + return False + path.write_text(new, encoding="utf-8") + print(f"sanitized Backends.fds ExtraArgs \\x corruption ({n} line(s))", flush=True) + return True + + +def restart_swarmui_local() -> None: + print("systemctl restart swarmui (после fix Backends.fds)", flush=True) + try: + subprocess.run( + ["sudo", "-n", "systemctl", "restart", "swarmui"], + check=False, + timeout=120, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + print(f"WARN restart swarmui: {exc}", flush=True) + wait_http(time.time() + 180) + + +def list_backends(sid: str) -> dict: + try: + return post( + "/API/ListBackends", + {"session_id": sid, "nonreal": False, "full_data": False}, + timeout=30.0, + ) + except Exception as exc: + return {"error": str(exc)} + + +def add_comfy_selfstart(sid: str) -> dict: + return post( + "/API/AddNewBackend", + {"session_id": sid, "type_id": "comfyui_selfstart"}, + timeout=60.0, + ) + + +def recover_empty_backends() -> str: + """When API says empty but venv/IsInstalled exist — fix FDS and/or AddNewBackend.""" + changed = sanitize_backends_fds() + if changed: + restart_swarmui_local() + bstat, _ = backend_status_detail() + if bstat not in ("empty", "unknown") and not bstat.startswith("error:"): + print(f"after sanitize: backend_status={bstat}", flush=True) + return bstat + + sid = get_session(time.time() + 60) + backends = list_backends(sid) + if isinstance(backends, dict) and any( + isinstance(v, dict) and v.get("type") for v in backends.values() + ): + bstat, _ = backend_status_detail() + return bstat + + print("ListBackends empty — AddNewBackend comfyui_selfstart", flush=True) + try: + result = add_comfy_selfstart(sid) + print(f"AddNewBackend: {result}", flush=True) + except Exception as exc: + print(f"AddNewBackend FAIL: {exc}", flush=True) + return "empty" + + # Default StartScript may be dlbackend/comfy/...; our tree is dlbackend/ComfyUI. + time.sleep(2) + bstat, _ = backend_status_detail() + return bstat + + def run_diagnostics() -> None: """Best-effort: prefer uploaded swarm_diag.py, else journalctl snippet.""" diag = Path("/tmp/gpu-rent-swarm_diag.py") diff --git a/src/gpu_rent/remote/tune_swarm_perf.py b/src/gpu_rent/remote/tune_swarm_perf.py index 6b41145..f2d6d02 100644 --- a/src/gpu_rent/remote/tune_swarm_perf.py +++ b/src/gpu_rent/remote/tune_swarm_perf.py @@ -84,25 +84,42 @@ def find_pip() -> Path | None: def patch_backends_extra_args(extra: str) -> bool: - """Ensure ExtraArgs contains sage flag for Comfy self-start backends.""" + """Ensure ExtraArgs contains sage flag for Comfy self-start backends. + + FreneticDataSyntax encodes an empty string as ``\\x``. Appending to that + line corrupts the file so SwarmUI loads zero backends (ListBackends={}). + """ if not extra: return False if not BACKENDS.is_file(): print(f"no {BACKENDS} yet — skip ExtraArgs (Comfy not registered)") return False text = BACKENDS.read_text(encoding="utf-8") - if "--use-sage-attention" in text: + if "--use-sage-attention" in text and not re.search( + r"ExtraArgs:\s*\\x\s+--use-sage-attention", text + ): print("Backends.fds already has --use-sage-attention") return False lines = text.splitlines() changed = False out = [] for line in lines: - if re.match(r"^(\s*)ExtraArgs:\s*$", line) or re.match(r"^(\s*)ExtraArgs:\s*\"\"\s*$", line): - indent = re.match(r"^(\s*)", line).group(1) + m_empty = re.match(r"^(\s*)ExtraArgs:\s*(?:\\x)?\s*$", line) + m_corrupt = re.match(r"^(\s*)ExtraArgs:\s*\\x\s+(.*)$", line) + m_val = re.match(r"^(\s*)ExtraArgs:\s+(\S.*)$", line) + if m_empty: + indent = m_empty.group(1) out.append(f"{indent}ExtraArgs: {extra}") changed = True - elif re.match(r"^(\s*)ExtraArgs:\s+", line) and "--use-sage-attention" not in line: + elif m_corrupt: + indent = m_corrupt.group(1) + rest = m_corrupt.group(2).strip() + if "--use-sage-attention" in rest: + out.append(f"{indent}ExtraArgs: {rest}") + else: + out.append(f"{indent}ExtraArgs: {rest} {extra}".strip()) + changed = True + elif m_val and "--use-sage-attention" not in line: out.append(line.rstrip() + f" {extra}") changed = True else: @@ -115,6 +132,24 @@ def patch_backends_extra_args(extra: str) -> bool: 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"], diff --git a/tests/test_civitai_fetch.py b/tests/test_civitai_fetch.py index 1f377ea..ba55034 100644 --- a/tests/test_civitai_fetch.py +++ b/tests/test_civitai_fetch.py @@ -1,7 +1,12 @@ from pathlib import Path -from gpu_rent.civitai import pick_preview_image -from gpu_rent.remote.civitai_fetch import fmt_bytes, progress_line, should_skip +from gpu_rent.civitai import file_size_bytes, pick_preview_image +from gpu_rent.remote.civitai_fetch import ( + fmt_bytes, + progress_line, + should_skip, + write_sha_sidecar, +) def test_should_skip_missing(tmp_path: Path): @@ -17,21 +22,52 @@ def test_should_skip_exists_no_sha(tmp_path: Path): assert "уже есть" in reason -def test_should_skip_sha_match(tmp_path: Path): +def test_should_skip_size_ok_fast(tmp_path: Path): + """Re-up must not hash multi-GB weights when size matches.""" + dest = tmp_path / "big.safetensors" + data = b"x" * 10_000 + dest.write_bytes(data) + expect = "deadbeef" * 8 + skip, reason = should_skip(dest, expect, expect_size=len(data)) + assert skip is True + assert "size ok" in reason + # sidecar written for next run + assert (tmp_path / "big.safetensors.sha256").is_file() + + +def test_should_skip_sha_sidecar(tmp_path: Path): + dest = tmp_path / "a.safetensors" + dest.write_bytes(b"weights") + expect = "abcd" * 16 + write_sha_sidecar(dest, expect) + skip, reason = should_skip(dest, expect, expect_size=7) + assert skip is True + assert "sidecar" in reason + + +def test_should_skip_sha_match_slow(tmp_path: Path): dest = tmp_path / "a.safetensors" dest.write_bytes(b"weights") import hashlib digest = hashlib.sha256(b"weights").hexdigest() - skip, reason = should_skip(dest, digest) + skip, reason = should_skip(dest, digest, verify_sha=True) assert skip is True assert "sha ok" in reason +def test_should_not_skip_bad_size(tmp_path: Path): + dest = tmp_path / "a.safetensors" + dest.write_bytes(b"weights") + skip, reason = should_skip(dest, "deadbeef" * 8, expect_size=99999) + assert skip is False + assert "размер" in reason + + def test_should_not_skip_bad_sha(tmp_path: Path): dest = tmp_path / "a.safetensors" dest.write_bytes(b"weights") - skip, reason = should_skip(dest, "deadbeef" * 8) + skip, reason = should_skip(dest, "deadbeef" * 8, verify_sha=True) assert skip is False assert "перекачиваю" in reason @@ -80,3 +116,9 @@ def test_pick_preview_image_png_and_default(): ] == ".preview.jpg" assert pick_preview_image({"images": []}) is None assert pick_preview_image({}) is None + + +def test_file_size_bytes(): + assert file_size_bytes({"sizeKB": 1.0}) == 1024 + assert file_size_bytes({"sizeKB": 6775430.35}) == int(round(6775430.35 * 1024)) + assert file_size_bytes({}) is None