mrleo1nid/gpu-rent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enhance SwarmUI integration and GPU environment verification

Browse Source 
- Updated CLI documentation to reflect the new handling of `CIVITAI_API_TOKEN`, which is now automatically passed to SwarmUI user settings during startup.
- Improved the `render_access_panel` function to include additional warnings for idle-killer failures and stack errors, enhancing user feedback.
- Introduced a new function `seed_swarmui_api_keys` to manage API key injection into SwarmUI, ensuring seamless integration with the Model Downloader.
- Enhanced GPU environment verification logic to include fail-fast checks for critical components like CUDA, improving error handling and user notifications.
- Updated tests to validate the new API key handling and access panel behavior, ensuring robustness in the integration process.
This commit is contained in:
Leonid Pershin
2026-08-21 07:09:20 +03:00
parent 1ec615c03e
commit adba4976ee
20 changed files with 657 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/cli.md
+6 -2
View File
@@ -233,8 +233,12 @@ Application credential для idle-killer CLI создаёт на `up` (узко
2. TCP 22 / SSH
3. Backend Idle (если SwarmUI) → toast (если `NOTIFY_READY`)
4. На VM: HTTP сервисов стека (SwarmUI `:7801` / Ollama `:11434` / llama.cpp `:8080`) — `verify_stack_on_vm`
5. На VM: **nvidia-smi / CUDA** (+ **torch+cuda** в Comfy venv, если SwarmUI) — `verify_gpu_env`
6. Туннель + проверка **localhost** тех же сервисов → access-card
5. На VM: **nvidia-smi / CUDA** (fail-fast) + **torch+cuda** в Comfy venv при SwarmUI (ждём) — `verify_gpu_env`
6. В логе: строка **`тайминг up:`** (SSH / bootstrap / Idle / verify / …)
7. Туннель + проверка **localhost** тех же сервисов → access-card (красная рамка, если killer failed)
`gpu-rent logs --unit swarm|ollama|llamacpp|killer` — фильтр journalctl.
`gpu-rent status` — killer/hold, последний стек/GPU-env, тайминг up.
Локальный порт UI: **17801** (на VM по-прежнему 7801 на loopback).
docs/models.md
+2
View File
@@ -67,6 +67,8 @@ copy models.example.yaml models.yaml
| В `.env` есть `CIVITAI_API_TOKEN` **и** в манифесте есть хотя бы одна запись | Скачать перечисленные модели с Civitai **на VM** (не через ноут). Рядом положить метаданные. **Не** качать дефолтный чекпоинт установщика SwarmUI |
| Токена нет, или манифест пуст/отсутствует | Обычный первый запуск SwarmUI: пусть ставит свою стандартную модель. В лог — почему Civitai-seed пропущен |
После старта SwarmUI тот же `CIVITAI_API_TOKEN` прокидывается в UI через `/API/SetAPIKey` (`civitai_api`) — Model Downloader и gated-модели в браузере. При наличии `HF_TOKEN` / `HUGGING_FACE_HUB_TOKEN` — ещё `huggingface_api`.
Токен без манифеста = предупреждение и ветка SwarmUI default, не «пустой диск без моделей».
ComfyUI в `dlbackend` качается всегда. Речь только о **модели весов** установщика SwarmUI, не о backend.
docs/setup.md
+1
View File
@@ -169,6 +169,7 @@ notepad .env
2. [Account settings](https://civitai.com/user/account) → **API Keys** → **Add**.
3. Токен показывают **один раз** → в `.env`: `CIVITAI_API_TOKEN=...`
4. Хост по умолчанию **`civitai.red`** (полный каталог). С `.com` NSFW часто 404.
5. На `up` токен уходит и в seed на VM, и в SwarmUI User Settings (Civitai API Key) — вручную в UI вводить не нужно.
5. Манифест:
```powershell
src/gpu_rent/access_card.py
+19 -3
View File
@@ -161,8 +161,12 @@ def render_access_panel(
notes = load_state().notes or {}
if notes.get("idle_killer") == "failed":
warn_bits.append(
"idle-killer НЕ вооружён — GPU может крутиться без авто-stop"
"idle-killer НЕ вооружён — GPU может крутиться без авто-stop → gpu-rent stop"
)
if notes.get("stack_vm_error"):
warn_bits.append(f"стек VM: {str(notes['stack_vm_error'])[:140]}")
if notes.get("gpu_env_error"):
warn_bits.append(f"GPU-стек: {str(notes['gpu_env_error'])[:140]}")
if notes.get("llm_error"):
warn_bits.append(f"LLM ошибка: {str(notes['llm_error'])[:120]}")
except Exception:
@@ -170,8 +174,9 @@ def render_access_panel(
parts: list = [subtitle, Text("")]
if warn_bits:
parts.append(Text("⚠ ВНИМАНИЕ — биллинг / готовность", style="bold white on red"))
for w in warn_bits:
parts.append(Text(f"⚠ {w}", style="bold red"))
parts.append(Text(f" • {w}", style="bold red"))
parts.append(Text(""))
parts.extend(
[
@@ -185,10 +190,11 @@ def render_access_panel(
]
)
body = Group(*parts)
border = "red" if warn_bits else "bright_blue"
return Panel(
body,
title=f"[bold]{title}[/bold]",
border_style="bright_blue",
border_style=border,
padding=(1, 2),
)
@@ -212,6 +218,16 @@ def print_access_card(
# Plain fallback for non-Rich loggers
log("")
log("══ gpu-rent · доступы ══")
try:
notes = load_state().notes or {}
if notes.get("idle_killer") == "failed":
log("⚠ idle-killer НЕ вооружён — GPU без авто-stop → gpu-rent stop")
if notes.get("stack_vm_error"):
log(f"⚠ стек VM: {notes['stack_vm_error']}")
if notes.get("gpu_env_error"):
log(f"⚠ GPU-стек: {notes['gpu_env_error']}")
except Exception:
pass
for link in collect_access_links(cfg, tunneled=tunneled):
extra = f" ({link.note})" if link.note else ""
log(f" {link.label:14} {link.url}{extra}")
src/gpu_rent/cli.py
+102 -19
View File
@@ -47,6 +47,13 @@ def _die(exc: BaseException) -> None:
if _DEBUG:
traceback.print_exc()
err(str(exc))
hint = (
"Дальше: gpu-rent status · gpu-rent logs · gpu-rent stop "
"(Ctrl+C на туннеле GPU не гасит)"
)
msg = str(exc)
if "gpu-rent status" not in msg and "Дальше:" not in msg:
err(hint)
raise typer.Exit(1)
@@ -210,6 +217,7 @@ def flavors(
def status() -> None:
"""Локальный state + OpenStack, если .env есть. Туннель не нужен."""
state = load_state()
notes = dict(state.notes or {})
table = Table(title="status")
table.add_column("поле")
table.add_column("значение")
@@ -249,13 +257,18 @@ def status() -> None:
table.add_row("диск used/free", df or "нет df")
from gpu_rent.idle_killer import killer_status_lines
table.add_row("idle-killer", "; ".join(killer_status_lines(cfg, state.floating_ip)))
note_k = (state.notes or {}).get("idle_killer")
killer_line = "; ".join(killer_status_lines(cfg, state.floating_ip))
note_k = notes.get("idle_killer")
if note_k == "failed":
err = (state.notes or {}).get("idle_killer_error") or ""
table.add_row("idle-killer arm", f"[red]FAILED[/red] {err}"[:120])
err_k = notes.get("idle_killer_error") or ""
table.add_row(
"idle-killer",
f"[red]FAILED arm[/red] · {killer_line} · {err_k}"[:160],
)
elif note_k == "armed":
table.add_row("idle-killer arm", "ok (в сессии)")
table.add_row("idle-killer", f"{killer_line} · arm ok (сессия)")
else:
table.add_row("idle-killer", killer_line)
except GpuRentError as exc:
table.add_row("диск used/free", f"SSH: {exc}")
table.add_row("idle-killer", "нет SSH")
@@ -263,20 +276,47 @@ def status() -> None:
table.add_row("диск used/free", "нужен живой FIP + SSH-ключ")
table.add_row("idle-killer", "нужен SSH на живую VM")
# Last verify snapshots (no new SSH)
if notes.get("stack_vm_error"):
table.add_row("стек VM", f"[red]FAIL[/red] {notes['stack_vm_error']}"[:140])
elif notes.get("stack_vm"):
bits = notes["stack_vm"]
if isinstance(bits, list):
ok_n = sum(1 for x in bits if isinstance(x, dict) and x.get("ok"))
table.add_row("стек VM", f"ok {ok_n}/{len(bits)} (последний up)")
else:
table.add_row("стек VM", str(bits)[:120])
if notes.get("gpu_env_error"):
table.add_row("GPU-стек", f"[red]FAIL[/red] {notes['gpu_env_error']}"[:140])
elif notes.get("gpu_env"):
bits = notes["gpu_env"]
if isinstance(bits, list):
summary = ", ".join(
f"{x.get('name')}={'ok' if x.get('ok') else 'FAIL'}"
for x in bits
if isinstance(x, dict)
)
table.add_row("GPU-стек", summary[:140] or "—")
if notes.get("up_timing"):
table.add_row("тайминг up", str(notes["up_timing"])[:140])
from gpu_rent.local_watchdog import watchdog_status_lines
table.add_row("local-watchdog", "; ".join(watchdog_status_lines()))
from gpu_rent.access_card import resolve_llm_runtime
rt = resolve_llm_runtime(cfg)
noted = (state.notes or {}).get("llm_runtime")
llm_err = (state.notes or {}).get("llm_error")
noted = notes.get("llm_runtime")
llm_err = notes.get("llm_error")
detail = f"{rt}; ollama :{cfg.ollama_local_port} / llamacpp :{cfg.llamacpp_local_port}"
if noted and noted != rt:
detail += f" (notes: {noted})"
if llm_err:
detail += f" [red]err: {llm_err[:80]}[/red]"
table.add_row("LLM", detail)
swarm_note = notes.get("enable_swarmui")
if swarm_note is False or not cfg.enable_swarmui:
detail += " · llm-only"
table.add_row("LLM / workload", detail)
if cfg.auth_ok:
try:
@@ -288,11 +328,11 @@ def status() -> None:
", ".join(f"{s.name} {s.status}" for s in servers),
)
else:
table.add_row("Nova", "нет сервера gpu-rent")
table.add_row("Nova", "нет tagged server")
snap = find_snapshot_by_name(conn, cfg.boot_snapshot_name)
table.add_row("snapshot", cfg.boot_snapshot_name if snap else "нет")
except GpuRentError as exc:
table.add_row("Nova", f"не достучались: {exc}")
except Exception as exc:
table.add_row("Nova", f"ошибка: {exc}"[:120])
else:
table.add_row("Nova", "нет .env — только локальный state")
@@ -695,22 +735,65 @@ def ssh() -> None:
@app.command()
def logs() -> None:
"""cloud-init / journalctl -u swarmui на VM."""
def logs(
unit: Optional[str] = typer.Option(
None,
"--unit",
"-u",
help="swarm|ollama|llamacpp|killer|cloud-init (по умолчанию — всё)",
),
lines: int = typer.Option(80, "--lines", "-n", help="Строк journalctl"),
) -> None:
"""cloud-init / journalctl юнитов на VM."""
try:
cfg = load_config(require_auth=True)
state = load_state()
if not state.floating_ip:
raise GpuRentError("нет IP — VM не поднята")
key = (unit or "all").strip().lower().replace("_", "-")
aliases = {
"all": "all",
"swarm": "swarmui",
"swarmui": "swarmui",
"ollama": "ollama",
"llamacpp": "llamacpp",
"llama": "llamacpp",
"killer": "gpu-rent-idle-killer",
"idle-killer": "gpu-rent-idle-killer",
"idle": "gpu-rent-idle-killer",
"cloud-init": "cloud-init",
"cloud": "cloud-init",
}
if key not in aliases:
raise GpuRentError(
f"неизвестный --unit={unit!r}; "
"ожидаю: swarm|ollama|llamacpp|killer|cloud-init|all"
)
target = aliases[key]
n = max(10, min(int(lines), 500))
parts: list[str] = []
if target in {"all", "cloud-init"}:
parts.append(
"echo '=== cloud-init (tail) ==='; "
"sudo -n tail -n 60 /var/log/cloud-init-output.log 2>/dev/null || true"
)
journal_units = []
if target == "all":
journal_units = ["swarmui", "ollama", "llamacpp", "gpu-rent-idle-killer"]
elif target != "cloud-init":
journal_units = [target]
for ju in journal_units:
parts.append(
f"echo; echo '=== systemctl {ju} ==='; "
f"systemctl is-active {ju} 2>/dev/null || true; "
f"echo; echo '=== journalctl -u {ju} ==='; "
f"sudo -n journalctl -u {ju} -n {n} --no-pager 2>/dev/null || true"
)
cmd = "; ".join(parts)
out = run_ssh(
cfg,
state.floating_ip,
"echo '=== cloud-init (tail) ==='; "
"sudo -n tail -n 60 /var/log/cloud-init-output.log 2>/dev/null || true; "
"echo; echo '=== systemctl swarmui ==='; "
"systemctl is-active swarmui 2>/dev/null || true; "
"echo; echo '=== journalctl -u swarmui ==='; "
"sudo -n journalctl -u swarmui -n 80 --no-pager 2>/dev/null || true",
cmd,
check=False,
timeout=60,
)
src/gpu_rent/provision.py
+44
View File
@@ -228,6 +228,50 @@ def _download_url(host: str, version_id: int, file_info: dict) -> str:
return f"https://{host}/api/download/models/{version_id}"
def seed_swarmui_api_keys(cfg: Config, host: str, log: Log) -> None:
"""Write CIVITAI_API_TOKEN (and HF if set) into SwarmUI user keys via SetAPIKey.
Swarm stores them in Users.ldb GenericData — needed for Model Downloader in the UI.
Call after SwarmUI HTTP is up (after wait_backend / verify).
"""
import os
keys: dict[str, str] = {}
if cfg.civitai_api_token:
keys["civitai_api"] = cfg.civitai_api_token
hf = (
os.environ.get("HF_TOKEN")
or os.environ.get("HUGGING_FACE_HUB_TOKEN")
or ""
).strip()
if hf:
keys["huggingface_api"] = hf
if not keys:
log("SwarmUI API keys: нет CIVITAI_API_TOKEN / HF_TOKEN — skip")
return
put_text(
cfg,
host,
"/tmp/gpu-rent-swarm-api-keys.json",
json.dumps(keys) + "\n",
mode=0o600,
)
names = ", ".join(keys)
log(f"SwarmUI: прокидываю API keys ({names})")
try:
run_python(
cfg,
host,
_pkg_text("swarmui_set_api_keys.py"),
remote_path="/tmp/gpu-rent-swarmui_set_api_keys.py",
timeout=180,
log=log,
)
except CloudError as exc:
log(f"⚠ SwarmUI API keys: {exc}")
run_ssh(cfg, host, "rm -f /tmp/gpu-rent-swarm-api-keys.json", check=False)
def seed_civitai(cfg: Config, host: str, log: Log) -> None:
entries = parse_models(cfg.models_manifest)
if not cfg.civitai_api_token:
src/gpu_rent/ready.py
+34 -5
View File
@@ -17,6 +17,7 @@ from gpu_rent.config import Config
from gpu_rent.errors import CloudError
from gpu_rent.llm_runtime import normalize_runtime
from gpu_rent.ssh_ops import run_ssh
from gpu_rent.timing import WaitLog
Log = Callable[[str], None]
@@ -266,6 +267,7 @@ def verify_stack_on_vm(
deadline = time.time() + timeout
last: list[ServiceCheck] = []
wait = WaitLog(log, every=30.0)
while time.time() < deadline:
try:
last = _probe_vm_once(cfg, host)
@@ -277,7 +279,7 @@ def verify_stack_on_vm(
log("проверка VM: всё отвечает")
return last
bad = ", ".join(f"{c.name}={c.detail}" for c in last if not c.ok) or "пусто"
log(f" … ещё нет: {bad}")
wait.tick(f" … ещё нет: {bad}")
time.sleep(poll_every)
for c in last:
@@ -320,11 +322,18 @@ def verify_gpu_env(
poll_every: float = 15.0,
raise_on_fail: bool = True,
) -> list[ServiceCheck]:
"""nvidia-smi / CUDA / torch(+cuda) in Comfy venv when SwarmUI is on."""
"""nvidia-smi / CUDA / torch(+cuda) in Comfy venv when SwarmUI is on.
Driver/CUDA missing → fail immediately (won't appear later).
Torch/Comfy venv → poll until timeout (first Comfy start installs them).
"""
from importlib.resources import files
from gpu_rent.ssh_ops import run_python
# These never "appear later" on a broken image — don't burn the poll budget.
instant_fail_names = {"nvidia-smi", "cuda"}
want_swarm = bool(getattr(cfg, "enable_swarmui", True))
script = files("gpu_rent.remote").joinpath("stack_env_probe.py").read_text(
encoding="utf-8"
@@ -338,6 +347,7 @@ def verify_gpu_env(
deadline = time.time() + timeout
last: list[ServiceCheck] = []
wait = WaitLog(log, every=30.0)
while time.time() < deadline:
try:
out = run_python(
@@ -350,7 +360,7 @@ def verify_gpu_env(
)
except Exception as exc:
last = [ServiceCheck("gpu-env", False, str(exc)[:200], "vm")]
log(f" … gpu-env: {exc}")
wait.tick(f" … gpu-env: {exc}")
time.sleep(poll_every)
continue
@@ -366,6 +376,7 @@ def verify_gpu_env(
checks_raw = data.get("checks") if isinstance(data, dict) else None
if not isinstance(checks_raw, list):
last = [ServiceCheck("gpu-env", False, f"нет JSON: {out[-180:]}", "vm")]
wait.tick(f" … gpu-env: нет JSON")
time.sleep(poll_every)
continue
@@ -395,8 +406,22 @@ def verify_gpu_env(
log("проверка GPU-стека: ок")
return last
instant = [c for c in hard if c.name in instant_fail_names]
if instant:
for c in last:
mark = "ok" if c.ok else "FAIL"
log(f" [{mark}] {c.name}: {c.detail}")
if raise_on_fail:
failed = [c.name for c in instant]
raise CloudError(
f"GPU-стек: нет {', '.join(failed)} (fail-fast). "
"Проверь образ Driver / nvidia на VM. "
"Дальше: gpu-rent status · gpu-rent logs · gpu-rent stop"
)
return last
bad = ", ".join(f"{c.name}={c.detail}" for c in hard)
log(f" … ещё нет: {bad}")
wait.tick(f" … ждём torch/Comfy: {bad}")
time.sleep(poll_every)
for c in last:
@@ -407,7 +432,8 @@ def verify_gpu_env(
raise CloudError(
f"GPU-стек не готов за {int(timeout)} с: {', '.join(failed)}. "
"Нужны nvidia-smi, CUDA; для SwarmUI — torch с cuda в Comfy venv "
"(journalctl -u swarmui / первый старт backend)."
"(journalctl -u swarmui / первый старт backend). "
"Дальше: gpu-rent status · gpu-rent logs · gpu-rent stop"
)
return last
@@ -446,6 +472,7 @@ def verify_stack_local(
)
deadline = time.time() + timeout
last: list[ServiceCheck] = []
wait = WaitLog(log, every=15.0)
while time.time() < deadline:
last = []
for name, port, url in targets:
@@ -479,6 +506,8 @@ def verify_stack_local(
log(f" [ok] localhost {c.name}: {c.detail}")
log("проверка туннеля: всё доступно")
return last
bad = ", ".join(f"{c.name}={c.detail}" for c in last if not c.ok) or "пусто"
wait.tick(f" … localhost ещё нет: {bad}")
time.sleep(poll_every)
for c in last:
src/gpu_rent/remote/bootstrap.sh
+14 -2
View File
@@ -77,9 +77,21 @@ ensure_bind() {
}
log "пакеты (без upgrade ядра)"
if [[ "${GPU_RENT_BOOTSTRAP_LIGHT:-0}" == "1" && -f "$MARKER_BOOT" ]]; then
log "light bootstrap — пропускаем apt-get"
# Light: skip apt when Swarm already bootstrapped, or llm-only data disk already ready.
_light_ok=0
if [[ "${GPU_RENT_BOOTSTRAP_LIGHT:-0}" == "1" ]]; then
if [[ -f "$MARKER_BOOT" ]]; then
_light_ok=1
elif [[ "${GPU_RENT_SKIP_SWARMUI:-0}" == "1" && -f "$MARKER_DATA" ]]; then
_light_ok=1
fi
fi
if [[ "$_light_ok" == "1" ]]; then
log "light bootstrap — пропускаем apt-get (маркер уже есть)"
else
if [[ "${GPU_RENT_BOOTSTRAP_LIGHT:-0}" == "1" ]]; then
log "light запрошен, но маркера нет — полный apt"
fi
apt-get update -qq
apt-get install -y -qq git python3 python3-venv python3-pip python3-full curl ca-certificates
fi
src/gpu_rent/remote/clone_ext.py
+4 -2
View File
@@ -179,14 +179,16 @@ def main() -> int:
token = TOKEN_PATH.read_text(encoding="utf-8").strip() if TOKEN_PATH.is_file() else ""
jobs = json.loads(JOBS_PATH.read_text(encoding="utf-8"))
update = do_update()
print(f"extensions update={'on' if update else 'off'}")
print(f"extensions update={'on' if update else 'off'} jobs={len(jobs)}")
failed = 0
known: set[str] = set()
try:
for job in jobs:
total = len(jobs)
for i, job in enumerate(jobs, start=1):
try:
dest = str(Path(job["dest"]))
known.add(dest)
print(f"extensions [{i}/{total}] {dest}")
clone_one(job, token, update)
except Exception as exc:
failed += 1
src/gpu_rent/remote/swarmui_set_api_keys.py
+90
View File
@@ -0,0 +1,90 @@
#!/usr/bin/env python3
"""Push upstream API keys into SwarmUI (user GenericData via SetAPIKey).
Stdlib only. Keys file: /tmp/gpu-rent-swarm-api-keys.json (mode 600), shape:
{"civitai_api": "...", "huggingface_api": "..."} # omit empty
"""
from __future__ import annotations
import json
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
KEYS_PATH = Path("/tmp/gpu-rent-swarm-api-keys.json")
SWARM = "http://127.0.0.1:7801"
ACCEPTED = ("civitai_api", "huggingface_api", "stability_api")
def post(path: str, payload: dict, timeout: float = 15.0) -> dict:
body = json.dumps(payload).encode("utf-8")
req = urllib.request.Request(
f"{SWARM}{path}",
data=body,
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(req, timeout=timeout) as resp:
return json.loads(resp.read().decode("utf-8"))
def wait_session(deadline: float) -> str:
last = ""
while time.time() < deadline:
try:
data = post("/API/GetNewSession", {})
sid = str(data.get("session_id") or "")
if sid:
return sid
last = "no session_id"
except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError, json.JSONDecodeError) as exc:
last = str(exc)[:160]
time.sleep(2)
raise SystemExit(f"SwarmUI session unavailable: {last}")
def main() -> int:
if not KEYS_PATH.is_file():
print("no keys file — skip")
return 0
try:
raw = json.loads(KEYS_PATH.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as exc:
print(f"bad keys file: {exc}", file=sys.stderr)
return 1
finally:
try:
KEYS_PATH.unlink(missing_ok=True)
except OSError:
pass
keys = {k: str(v).strip() for k, v in (raw or {}).items() if k in ACCEPTED and str(v).strip()}
if not keys:
print("no api keys to set")
return 0
sid = wait_session(time.time() + 120)
for key_type, value in keys.items():
try:
resp = post(
"/API/SetAPIKey",
{"session_id": sid, "keyType": key_type, "key": value},
)
except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError, json.JSONDecodeError) as exc:
print(f"SetAPIKey {key_type} failed: {exc}", file=sys.stderr)
return 1
if resp.get("error"):
print(f"SetAPIKey {key_type}: {resp['error']}", file=sys.stderr)
return 1
if not resp.get("success"):
print(f"SetAPIKey {key_type}: unexpected {resp}", file=sys.stderr)
return 1
print(f"SetAPIKey {key_type}=ok")
return 0
if __name__ == "__main__":
raise SystemExit(main())
src/gpu_rent/remote/tune_swarm_perf.py
+5 -1
View File
@@ -129,11 +129,15 @@ def main() -> int:
prev = json.loads(MARKER.read_text(encoding="utf-8"))
except json.JSONDecodeError:
prev = {}
if prev.get("uuid") and prev.get("uuid") == plan["uuid"] and prev.get("extra_args") == plan["extra_args"]:
same_gpu = bool(prev.get("uuid") and prev.get("uuid") == plan["uuid"])
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']})")
return 0
if same_gpu and plan["use_sage"] and prev.get("pip_ok") is False:
print("perf tune: retry (previous pip_ok=false — sage/triton ещё не встали)")
print(f"perf tune: {plan['name']} tier={plan['tier']} sage={plan['use_sage']}")
pip_ok = not plan["use_sage"]
restarted_needed = False
src/gpu_rent/session.py
+65 -12
View File
@@ -21,7 +21,7 @@ from gpu_rent.cloud import (
wait_volume,
)
from gpu_rent.bootstrap import run_bootstrap
from gpu_rent.provision import provision_vm, tune_swarm_perf
from gpu_rent.provision import provision_vm, seed_swarmui_api_keys, tune_swarm_perf
from gpu_rent.ready import verify_gpu_env, verify_stack_on_vm, wait_backend_idle
from gpu_rent.snapshot import ensure_boot_snapshot
from gpu_rent.notify import notify_ready
@@ -48,6 +48,7 @@ from gpu_rent.os_client import (
from gpu_rent.ssh_keys import ensure_ed25519
from gpu_rent.ssh_ops import probe_ssh, run_ssh, wait_ssh
from gpu_rent.state import SessionState, load_state, save_state, utc_now
from gpu_rent.timing import PhaseTimes
Log = Callable[[str], None]
@@ -74,7 +75,9 @@ def _bind_access(
log: Log,
*,
update: bool = True,
phases: PhaseTimes | None = None,
) -> SessionState:
clock = phases or PhaseTimes()
ip, fip_id = ensure_floating_ip(
conn, server, state.floating_ip_id, state.floating_ip, log
)
@@ -83,6 +86,7 @@ def _bind_access(
state.floating_ip_id = fip_id
save_state(state)
wait_ssh(cfg, ip)
clock.mark("SSH")
log(f"SSH {cfg.ssh_user}@{ip}")
state.phase = "bootstrapping"
save_state(state)
@@ -92,20 +96,49 @@ def _bind_access(
if active == "active":
log("systemctl stop swarmui перед git update")
run_ssh(cfg, ip, "sudo -n systemctl stop swarmui", timeout=120, check=False)
# Skip apt-heavy bootstrap when the VM already finished first-boot.
# Detect existing markers so light/full choice is explicit (swarm ↔ llm-only).
probe = run_ssh(
cfg,
ip,
"echo swarm=$(test -f /opt/swarmui/.gpu-rent-bootstrapped && echo yes || echo no); "
"echo data=$(test -f /mnt/swarm_data/.gpu-rent-ready && echo yes || echo no); "
"echo llm=$(test -f /mnt/swarm_data/.gpu-rent-llm-only && echo yes || echo no)",
check=False,
)
flags: dict[str, str] = {}
for line in probe.splitlines():
if "=" in line:
k, v = line.strip().split("=", 1)
flags[k] = v
has_swarm = flags.get("swarm") == "yes"
has_data = flags.get("data") == "yes"
was_llm_only = flags.get("llm") == "yes"
if swarm:
marker_cmd = "test -f /opt/swarmui/.gpu-rent-bootstrapped && echo yes || echo no"
else:
marker_cmd = "test -f /mnt/swarm_data/.gpu-rent-ready && echo yes || echo no"
marker = run_ssh(cfg, ip, marker_cmd, check=False).strip()
if marker == "yes":
if not state.bootstrapped:
log("маркер bootstrap на VM — лёгкий проход (локальный bootstrapped был сброшен)")
light = has_swarm
if was_llm_only and not has_swarm:
log("bootstrap: был llm-only → полный проход (ставим SwarmUI)")
elif light:
why = (
"локальный bootstrapped был сброшен, маркер на VM есть"
if not state.bootstrapped
else "маркер /opt/swarmui/.gpu-rent-bootstrapped"
)
log(f"bootstrap: LIGHT (без apt) — {why}")
else:
log("bootstrap уже на VM — лёгкий проход (без apt)")
run_bootstrap(cfg, ip, log, update=update and swarm, light=True)
log("bootstrap: FULL (apt + SwarmUI) — маркера bootstrapped нет")
else:
run_bootstrap(cfg, ip, log, update=update and swarm, light=False)
light = has_data
if has_swarm and not was_llm_only:
log("bootstrap: llm-only на диске со SwarmUI — LIGHT data, Swarm unit stop")
if light:
log("bootstrap: LIGHT llm-only (без apt) — есть .gpu-rent-ready")
else:
log("bootstrap: FULL llm-only — маркера data ready нет")
run_bootstrap(cfg, ip, log, update=update and swarm, light=light)
clock.mark("bootstrap")
provision_vm(
cfg,
ip,
@@ -114,17 +147,31 @@ def _bind_access(
server_id=getattr(server, "id", None) or state.server_id,
update=update,
)
clock.mark("provision")
if swarm:
try:
wait_backend_idle(cfg, ip, log)
except CloudError as exc:
log(f"ready: {exc}")
clock.mark("Idle")
try:
if tune_swarm_perf(cfg, ip, log):
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}")
else:
log("perf tune: restart не нужен")
except Exception as exc:
log(f"perf tune: {exc}")
clock.mark("perf")
try:
seed_swarmui_api_keys(cfg, ip, log)
except Exception as exc:
log(f"SwarmUI API keys: {exc}")
clock.mark("api-keys")
else:
log("ready: llm-only (без ожидания SwarmUI Idle)")
@@ -134,11 +181,13 @@ def _bind_access(
state.notes["stack_vm"] = [
{"name": c.name, "ok": c.ok, "detail": c.detail} for c in checks
]
state.notes.pop("stack_vm_error", None)
except CloudError as exc:
state.notes = dict(state.notes or {})
state.notes["stack_vm_error"] = str(exc)[:500]
save_state(state)
raise
clock.mark("verify")
try:
gpu_checks = verify_gpu_env(cfg, ip, log, timeout=600.0)
@@ -146,11 +195,13 @@ def _bind_access(
state.notes["gpu_env"] = [
{"name": c.name, "ok": c.ok, "detail": c.detail} for c in gpu_checks
]
state.notes.pop("gpu_env_error", None)
except CloudError as exc:
state.notes = dict(state.notes or {})
state.notes["gpu_env_error"] = str(exc)[:500]
save_state(state)
raise
clock.mark("gpu-env")
try:
ensure_boot_snapshot(
@@ -187,7 +238,9 @@ def _bind_access(
state.phase = "ready_cloud"
state.notes = dict(state.notes or {})
state.notes["enable_swarmui"] = swarm
state.notes["up_timing"] = clock.summary_line()
save_state(state)
log(f"тайминг up: {clock.summary_line()}")
return state
src/gpu_rent/sync_files.py
+25 -9
View File
@@ -31,25 +31,29 @@ def push_tree(
log(f"push {local_root.name}: пусто — skip")
return 0
files = model_push_set(local_root) if models else iter_payload_files(local_root)
file_list = list(files)
total = len(file_list)
sent = 0
skipped = 0
client = open_ssh(cfg, host)
try:
for path in files:
for i, path in enumerate(file_list, start=1):
rel = path.relative_to(local_root).as_posix()
remote = f"{remote_root.rstrip('/')}/{rel}"
local_hash = sha256_file(path)
remote_hash = remote_sha256_on(client, remote)
if remote_hash and remote_hash.lower() == local_hash.lower():
skipped += 1
continue
log(f"push {rel}")
log(f"push [{i}/{total}] {rel}")
put_file_on(client, path, remote)
sent += 1
finally:
client.close()
if sent == 0:
log(f"push {local_root.name}: всё уже на VM")
log(f"push {local_root.name}: всё уже на VM ({total} файл(ов), skip={skipped})")
else:
log(f"push {local_root.name}: {sent} файл(ов)")
log(f"push {local_root.name}: {sent}/{total} отправлено (skip={skipped})")
return sent
@@ -63,19 +67,31 @@ def pull_tree(cfg: Config, host: str, remote_root: str, local_root: Path, log: L
timeout=120,
)
names = [line.strip() for line in listing.splitlines() if line.strip()]
work = [
rel
for rel in names
if not (
rel.endswith("/.gitkeep")
or rel.rsplit("/", 1)[-1] in {".gitkeep", "README.md"}
)
]
total = len(work)
pulled = 0
for rel in names:
if rel.endswith("/.gitkeep") or rel.rsplit("/", 1)[-1] in {".gitkeep", "README.md"}:
continue
skipped = 0
for i, rel in enumerate(work, start=1):
remote = f"{remote_root.rstrip('/')}/{rel}"
local = local_root / rel
remote_hash = remote_sha256_on(client, remote)
if local.is_file() and remote_hash and sha256_file(local).lower() == remote_hash.lower():
skipped += 1
continue
log(f"pull {rel}")
log(f"pull [{i}/{total}] {rel}")
get_file_on(client, remote, local)
pulled += 1
finally:
client.close()
log(f"pull Output: {pulled} файл(ов)" if pulled else "pull Output: нечего забирать")
if pulled:
log(f"pull Output: {pulled}/{total} (skip={skipped})")
else:
log(f"pull Output: нечего забирать ({total} файл(ов), skip={skipped})")
return pulled
src/gpu_rent/timing.py
+65
View File
@@ -0,0 +1,65 @@
"""Elapsed-time helpers for long `up` flows."""
from __future__ import annotations
import time
from collections.abc import Callable
def format_duration(seconds: float) -> str:
sec = max(0, int(round(seconds)))
if sec < 60:
return f"{sec}s"
minutes, rem = divmod(sec, 60)
if minutes < 60:
return f"{minutes}m {rem}s" if rem else f"{minutes}m"
hours, rem_m = divmod(minutes, 60)
return f"{hours}h {rem_m}m" if rem_m else f"{hours}h"
class PhaseTimes:
"""Record named milestones from a shared start."""
def __init__(self) -> None:
self._t0 = time.monotonic()
self._marks: list[tuple[str, float]] = []
def mark(self, name: str) -> float:
elapsed = time.monotonic() - self._t0
self._marks.append((name, elapsed))
return elapsed
@property
def total(self) -> float:
return time.monotonic() - self._t0
def deltas(self) -> list[tuple[str, float]]:
"""Per-phase duration (from previous mark or start)."""
out: list[tuple[str, float]] = []
prev = 0.0
for name, at in self._marks:
out.append((name, at - prev))
prev = at
return out
def summary_line(self) -> str:
parts = [f"{name} {format_duration(dt)}" for name, dt in self.deltas()]
parts.append(f"всего {format_duration(self.total)}")
return " · ".join(parts)
class WaitLog:
"""Log first wait message, then at most every `every` seconds."""
def __init__(self, log: Callable[[str], None], *, every: float = 30.0) -> None:
self._log = log
self._every = every
self._last = 0.0
self._n = 0
def tick(self, msg: str) -> None:
self._n += 1
now = time.monotonic()
if self._n == 1 or (now - self._last) >= self._every:
self._log(msg)
self._last = now
tests/test_access_card.py
+15 -1
View File
@@ -1,4 +1,4 @@
from gpu_rent.access_card import collect_access_links, mcp_snippet_lines
from gpu_rent.access_card import collect_access_links, mcp_snippet_lines, render_access_panel
class _Cfg:
@@ -6,6 +6,7 @@ class _Cfg:
llm_runtime = "ollama"
ollama_local_port = 17811
llamacpp_local_port = 17812
enable_swarmui = True
def test_collect_links_swarm_and_ollama(monkeypatch):
@@ -29,3 +30,16 @@ def test_collect_links_no_tunnel():
def test_mcp_snippet_json():
lines = mcp_snippet_lines(_Cfg())
assert any("17801/mcp" in line for line in lines)
def test_access_panel_red_when_killer_failed(monkeypatch):
monkeypatch.setattr(
"gpu_rent.access_card.load_state",
lambda: type(
"S",
(),
{"notes": {"idle_killer": "failed", "idle_killer_error": "no cred"}},
)(),
)
panel = render_access_panel(_Cfg(), tunneled=True)
assert panel.border_style == "red"
tests/test_bootstrap.py
+5 -2
View File
@@ -14,5 +14,8 @@ def test_bootstrap_script_is_native_swarmui():
assert "Data/Autocompletions" in script
assert "mkfs.ext4" in script
assert "apt-get upgrade" not in script
assert ".gpu-rent-ready" in script
assert "src/BuiltinExtensions/ComfyUIBackend/DLNodes" in script
assert "GPU_RENT_BOOTSTRAP_LIGHT" in script
assert "GPU_RENT_SKIP_SWARMUI" in script
assert "light bootstrap — пропускаем apt-get" in script
# llm-only re-up can skip apt when data marker exists (not only Swarm boot marker)
assert 'MARKER_DATA' in script or ".gpu-rent-ready" in script
tests/test_swarmui_api_keys.py
+52
View File
@@ -0,0 +1,52 @@
"""Tests for remote swarmui_set_api_keys."""
from __future__ import annotations
import importlib.util
import json
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
REMOTE = ROOT / "src" / "gpu_rent" / "remote" / "swarmui_set_api_keys.py"
def _load():
spec = importlib.util.spec_from_file_location("swarmui_set_api_keys", REMOTE)
assert spec and spec.loader
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
return mod
def test_set_civitai_key(tmp_path, monkeypatch):
mod = _load()
keys = tmp_path / "keys.json"
keys.write_text(json.dumps({"civitai_api": "tok-123"}), encoding="utf-8")
monkeypatch.setattr(mod, "KEYS_PATH", keys)
calls = []
def fake_post(path, payload, timeout=15.0):
calls.append((path, payload))
if path == "/API/GetNewSession":
return {"session_id": "sid-1"}
if path == "/API/SetAPIKey":
assert payload["session_id"] == "sid-1"
assert payload["keyType"] == "civitai_api"
assert payload["key"] == "tok-123"
return {"success": True}
raise AssertionError(path)
monkeypatch.setattr(mod, "post", fake_post)
monkeypatch.setattr(mod, "wait_session", lambda _d: "sid-1")
assert mod.main() == 0
assert not keys.exists()
assert any(c[0] == "/API/SetAPIKey" for c in calls)
def test_skip_empty_keys(tmp_path, monkeypatch):
mod = _load()
keys = tmp_path / "keys.json"
keys.write_text("{}", encoding="utf-8")
monkeypatch.setattr(mod, "KEYS_PATH", keys)
assert mod.main() == 0
tests/test_timing.py
+33
View File
@@ -0,0 +1,33 @@
from gpu_rent.timing import PhaseTimes, WaitLog, format_duration
def test_format_duration():
assert format_duration(5) == "5s"
assert format_duration(65) == "1m 5s"
assert format_duration(60) == "1m"
assert format_duration(3661) == "1h 1m"
def test_phase_times_summary(monkeypatch):
times = iter([100.0, 110.0, 130.0, 190.0])
monkeypatch.setattr("gpu_rent.timing.time.monotonic", lambda: next(times))
clock = PhaseTimes()
clock.mark("SSH")
clock.mark("bootstrap")
line = clock.summary_line()
assert "SSH 10s" in line
assert "bootstrap 20s" in line
assert "всего" in line
def test_wait_log_throttles(monkeypatch):
logs: list[str] = []
t = {"now": 0.0}
monkeypatch.setattr("gpu_rent.timing.time.monotonic", lambda: t["now"])
w = WaitLog(logs.append, every=30.0)
w.tick("a")
t["now"] = 10.0
w.tick("b")
t["now"] = 31.0
w.tick("c")
assert logs == ["a", "c"]
tests/test_tune_swarm_perf.py
+47
View File
@@ -86,3 +86,50 @@ def test_pip_ok_patches_extra_args(tmp_path, monkeypatch):
assert marker["pip_ok"] is True
assert "--use-sage-attention" in marker["extra_args"]
assert "--use-sage-attention" in backends.read_text(encoding="utf-8")
def test_pip_fail_retries_next_run(tmp_path, monkeypatch):
mod = _load()
data = tmp_path
backends = data / "Data" / "Backends.fds"
backends.parent.mkdir(parents=True)
backends.write_text("ExtraArgs: \n", encoding="utf-8")
gpu_json = data / ".gpu-rent-gpu.json"
gpu_json.write_text(
json.dumps(
{
"vram_mib": 24576,
"compute_cap": "8.9",
"uuid": "gpu-1",
"name": "RTX",
}
),
encoding="utf-8",
)
marker = data / ".gpu-rent-perf-tuned"
marker.write_text(
json.dumps(
{
"uuid": "gpu-1",
"extra_args": "",
"pip_ok": False,
"tier": "high",
"name": "RTX",
}
),
encoding="utf-8",
)
pip = data / "fake-pip"
pip.write_text("#!/bin/sh\n", encoding="utf-8")
monkeypatch.setattr(mod, "DATA", data)
monkeypatch.setattr(mod, "GPU_JSON", gpu_json)
monkeypatch.setattr(mod, "MARKER", marker)
monkeypatch.setattr(mod, "BACKENDS", backends)
monkeypatch.setattr(mod, "find_pip", lambda: pip)
monkeypatch.setattr(mod, "pip_install_sage", lambda _p: True)
assert mod.main() == 0
new_m = json.loads(marker.read_text(encoding="utf-8"))
assert new_m["pip_ok"] is True
assert "--use-sage-attention" in backends.read_text(encoding="utf-8")
tests/test_verify_stack.py
+29
View File
@@ -95,6 +95,35 @@ def test_verify_gpu_env_ok(monkeypatch):
assert any("GPU-стека" in line for line in logs)
def test_verify_gpu_env_fail_fast_cuda(monkeypatch):
import json
import gpu_rent.ssh_ops as ssh_ops
payload = {
"ok": False,
"checks": [
{"name": "nvidia-smi", "required": True, "ok": True, "detail": "ok"},
{"name": "cuda", "required": True, "ok": False, "detail": "нет libcuda"},
{"name": "torch", "required": True, "ok": False, "detail": "no venv"},
],
}
calls = {"n": 0}
def fake(*a, **k):
calls["n"] += 1
return json.dumps(payload)
monkeypatch.setattr(ssh_ops, "run_python", fake)
logs: list[str] = []
try:
verify_gpu_env(_Cfg(), "1.2.3.4", logs.append, timeout=600.0, poll_every=0.1)
assert False, "expected CloudError"
except CloudError as exc:
assert "fail-fast" in str(exc).lower() or "cuda" in str(exc).lower()
assert calls["n"] == 1
def test_verify_gpu_env_fails_without_cuda(monkeypatch):
import json