- Updated version number in pyproject.toml and __init__.py to 0.2.0. - Revised README.md to reflect the current state of the project, including usage instructions and setup steps. - Improved CLI documentation in cli.md, adding details about new commands and their functionalities. - Enhanced the quick start section in README.md for better clarity on initial setup. - Updated local folder documentation to clarify file handling and commands. - Added a new command for listing GPU flavors and improved error handling in the CLI. - Implemented a watchdog feature in the tunnel to manage server states effectively.
96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
"""Wait until SwarmUI HTTP is up and backend is Idle (on the VM)."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
from collections.abc import Callable
|
||
|
||
from gpu_rent.config import Config
|
||
from gpu_rent.errors import CloudError
|
||
from gpu_rent.ssh_ops import run_ssh
|
||
|
||
Log = Callable[[str], None]
|
||
|
||
_REMOTE_POLL = r"""
|
||
import json, time, urllib.request
|
||
url = "http://127.0.0.1:7801"
|
||
deadline = time.time() + 25
|
||
while time.time() < deadline:
|
||
try:
|
||
req = urllib.request.Request(
|
||
url + "/API/GetNewSession",
|
||
data=b"{}",
|
||
headers={"Content-Type": "application/json"},
|
||
method="POST",
|
||
)
|
||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||
session = json.loads(resp.read().decode())
|
||
sid = session.get("session_id")
|
||
if not sid:
|
||
time.sleep(2)
|
||
continue
|
||
body = json.dumps({"session_id": sid}).encode()
|
||
req2 = urllib.request.Request(
|
||
url + "/API/GetCurrentStatus",
|
||
data=body,
|
||
headers={"Content-Type": "application/json"},
|
||
method="POST",
|
||
)
|
||
with urllib.request.urlopen(req2, timeout=5) as resp:
|
||
data = json.loads(resp.read().decode())
|
||
st = data.get("status") or {}
|
||
be = data.get("backend_status") or {}
|
||
waiting = int(st.get("waiting_gens") or 0)
|
||
live = int(st.get("live_gens") or 0)
|
||
loading = int(st.get("loading_models") or 0)
|
||
bstat = str(be.get("status") or "unknown").lower()
|
||
if waiting or live or loading:
|
||
print(f"BUSY queue w={waiting} live={live} load={loading}")
|
||
elif bstat not in ("idle", "disabled", "all_disabled", "empty"):
|
||
print(f"BUSY backend={bstat}")
|
||
else:
|
||
print(f"READY backend={bstat}")
|
||
raise SystemExit(0)
|
||
except Exception as exc:
|
||
print(f"WAIT {exc}")
|
||
time.sleep(3)
|
||
print("WAIT timeout-slice")
|
||
"""
|
||
|
||
|
||
def wait_backend_idle(
|
||
cfg: Config,
|
||
host: str,
|
||
log: Log,
|
||
*,
|
||
timeout: float = 2400.0,
|
||
poll_every: float = 15.0,
|
||
) -> None:
|
||
"""Block until SwarmUI on the VM reports Idle backend (or timeout)."""
|
||
deadline = time.time() + timeout
|
||
log("жду HTTP :7801 и Idle backend на VM…")
|
||
last = ""
|
||
while time.time() < deadline:
|
||
try:
|
||
out = run_ssh(
|
||
cfg,
|
||
host,
|
||
"python3 - <<'PY'\n" + _REMOTE_POLL + "\nPY",
|
||
check=False,
|
||
timeout=40,
|
||
).strip()
|
||
except Exception as exc:
|
||
out = f"WAIT ssh: {exc}"
|
||
line = out.splitlines()[-1] if out else "WAIT empty"
|
||
if line != last:
|
||
log(line)
|
||
last = line
|
||
if line.startswith("READY"):
|
||
log("backend Idle")
|
||
return
|
||
time.sleep(poll_every)
|
||
raise CloudError(
|
||
f"backend не стал Idle за {int(timeout)} с. "
|
||
"Проверь journalctl -u swarmui на VM; GPU всё ещё жив."
|
||
)
|