Leonid Pershin
ccba40a228
- Updated the `_print_checks` function to replace console prints with logging functions for better traceability. - Introduced timing functionality in the `doctor`, `dry_run`, and `up` functions to log the duration of preflight checks. - Modified the `wait_ssh` function to accept a logging callback, improving SSH wait feedback. - Enhanced the `mark` method in `PhaseTimes` to log phase durations, aiding in performance analysis. - Updated various remote scripts to ensure error messages are printed to stderr for better error handling.
91 lines
2.8 KiB
Python
91 lines
2.8 KiB
Python
#!/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}")
|
|
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}")
|
|
return 1
|
|
if resp.get("error"):
|
|
print(f"SetAPIKey {key_type}: {resp['error']}")
|
|
return 1
|
|
if not resp.get("success"):
|
|
print(f"SetAPIKey {key_type}: unexpected {resp}")
|
|
return 1
|
|
print(f"SetAPIKey {key_type}=ok")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|