Add file size conversion and enhance Civitai job processing
- Introduced the `file_size_bytes` function to convert Civitai model sizes from kilobytes to bytes, improving data handling. - Updated the `seed_civitai` function to include file size in job definitions, enhancing model processing efficiency. - Enhanced the `should_skip` function to utilize expected size for faster decision-making during job processing. - Added tests for new functionality, ensuring robustness in file size handling and job processing logic.
This commit is contained in:
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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"],
|
||||
|
||||
Reference in New Issue
Block a user