gpu-rent/tests/test_verify_stack.py

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):
    import json

    import gpu_rent.ssh_ops as ssh_ops

    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"

    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)
    assert any("GPU-стека" in line for line in logs)


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 "GPU-стек" in str(exc) or "cuda" in str(exc).lower()


def test_verify_gpu_env_llm_only_skips_torch_requirement(monkeypatch):
    import json

    import gpu_rent.ssh_ops as ssh_ops

    class C:
        enable_swarmui = False
        llm_runtime = "llamacpp"

    payload = {
        "ok": True,
        "checks": [
            {"name": "nvidia-smi", "required": True, "ok": True, "detail": "ok"},
            {"name": "cuda", "required": True, "ok": True, "detail": "ok"},
            {
                "name": "torch",
                "required": False,
                "ok": True,
                "detail": "skip (llm-only, без Comfy venv)",
            },
        ],
    }

    def fake_run_python(cfg, host, script, **kw):
        assert "GPU_RENT_CHECK_SWARM" in script
        assert "os.environ['GPU_RENT_CHECK_SWARM']='0'" in script
        return json.dumps(payload)

    monkeypatch.setattr(ssh_ops, "run_python", fake_run_python)
    logs: list[str] = []
    out = verify_gpu_env(C(), "1.2.3.4", logs.append, timeout=5.0)
    assert all(c.ok for c in out)