Leonid Pershin
3c8225a69e
- Added a new function `verify_gpu_env` to check GPU stack readiness, including nvidia-smi, CUDA, and torch in the Comfy virtual environment when SwarmUI is enabled. - Updated the session management to call `verify_gpu_env`, capturing GPU environment status and errors in the state notes. - Enhanced documentation in `cli.md` to reflect the new GPU environment verification process. - Added tests for `verify_gpu_env` to ensure proper functionality and error handling during GPU checks.
142 lines
4.1 KiB
Python
142 lines
4.1 KiB
Python
from gpu_rent.errors import CloudError
|
|
from gpu_rent.ready import ServiceCheck, _expected_services, verify_gpu_env, verify_stack_local
|
|
|
|
|
|
class _Cfg:
|
|
enable_swarmui = True
|
|
llm_runtime = "none"
|
|
swarmui_local_port = 17801
|
|
ollama_local_port = 17811
|
|
llamacpp_local_port = 17812
|
|
|
|
|
|
def test_expected_services_swarm_only():
|
|
assert _expected_services(_Cfg()) == (True, False, False)
|
|
|
|
|
|
def test_expected_services_llm_only():
|
|
class C:
|
|
enable_swarmui = False
|
|
llm_runtime = "llamacpp"
|
|
|
|
assert _expected_services(C()) == (False, False, True)
|
|
|
|
|
|
def test_verify_stack_local_empty_when_nothing():
|
|
class C:
|
|
enable_swarmui = False
|
|
llm_runtime = "none"
|
|
swarmui_local_port = 17801
|
|
ollama_local_port = 17811
|
|
llamacpp_local_port = 17812
|
|
|
|
logs: list[str] = []
|
|
assert verify_stack_local(C(), logs.append, timeout=0.1) == []
|
|
|
|
|
|
def test_verify_stack_local_fails_closed_port(monkeypatch):
|
|
class C:
|
|
enable_swarmui = False
|
|
llm_runtime = "ollama"
|
|
swarmui_local_port = 17801
|
|
ollama_local_port = 17999
|
|
llamacpp_local_port = 17812
|
|
|
|
monkeypatch.setattr(
|
|
"gpu_rent.ready._tcp_ok", lambda port, host="127.0.0.1", timeout=0.8: False
|
|
)
|
|
logs: list[str] = []
|
|
try:
|
|
verify_stack_local(C(), logs.append, timeout=0.3, poll_every=0.1)
|
|
assert False, "expected CloudError"
|
|
except Exception as exc:
|
|
assert "ollama" in str(exc).lower() or "не отвечает" in str(exc)
|
|
|
|
|
|
def test_service_check_dataclass():
|
|
c = ServiceCheck("x", True, "ok", "vm")
|
|
assert c.ok and c.where == "vm"
|
|
|
|
|
|
def test_verify_gpu_env_ok(monkeypatch):
|
|
payload = {
|
|
"ok": True,
|
|
"checks": [
|
|
{"name": "nvidia-smi", "required": True, "ok": True, "detail": "A100"},
|
|
{"name": "cuda", "required": True, "ok": True, "detail": "libcuda"},
|
|
{
|
|
"name": "torch",
|
|
"required": True,
|
|
"ok": True,
|
|
"detail": "torch=2.0 cuda=12 available=True",
|
|
},
|
|
{
|
|
"name": "triton/sage",
|
|
"required": False,
|
|
"ok": True,
|
|
"detail": "есть: triton; нет: —",
|
|
},
|
|
],
|
|
}
|
|
|
|
def fake_run_python(cfg, host, script, **kw):
|
|
assert "GPU_RENT_CHECK_SWARM" in script
|
|
return json.dumps(payload) + "\n"
|
|
|
|
import json
|
|
|
|
monkeypatch.setattr("gpu_rent.ssh_ops.run_python", fake_run_python)
|
|
# patch where used
|
|
monkeypatch.setattr(
|
|
"gpu_rent.ssh_ops.run_python",
|
|
fake_run_python,
|
|
raising=False,
|
|
)
|
|
|
|
import gpu_rent.ready as ready_mod
|
|
|
|
monkeypatch.setattr(
|
|
ready_mod,
|
|
"run_python",
|
|
fake_run_python,
|
|
raising=False,
|
|
)
|
|
|
|
# verify_gpu_env imports run_python inside the function
|
|
import gpu_rent.ssh_ops as ssh_ops
|
|
|
|
monkeypatch.setattr(ssh_ops, "run_python", fake_run_python)
|
|
|
|
logs: list[str] = []
|
|
out = verify_gpu_env(_Cfg(), "1.2.3.4", logs.append, timeout=5.0, poll_every=0.1)
|
|
assert all(c.ok for c in out)
|
|
assert any(c.name == "torch" for c in out)
|
|
|
|
|
|
def test_verify_gpu_env_fails_without_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"},
|
|
],
|
|
}
|
|
monkeypatch.setattr(
|
|
ssh_ops,
|
|
"run_python",
|
|
lambda *a, **k: json.dumps(payload),
|
|
)
|
|
logs: list[str] = []
|
|
try:
|
|
verify_gpu_env(
|
|
_Cfg(), "1.2.3.4", logs.append, timeout=0.4, poll_every=0.1
|
|
)
|
|
assert False, "expected CloudError"
|
|
except CloudError as exc:
|
|
assert "cuda" in str(exc).lower() or "GPU-стек" in str(exc)
|