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


def test_expected_services_swarm_only():
    assert _expected_services(_Cfg()) == (True, False)


def test_expected_services_llm_only():
    class C:
        enable_swarmui = False
        llm_runtime = "ollama"

    assert _expected_services(C()) == (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

    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

    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 "from __future__" in script
        assert not script.lstrip().startswith("import os")
        assert (kw.get("env") or {}).get("GPU_RENT_CHECK_SWARM") == "1"
        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_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

    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 = "ollama"

    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 "from __future__" in script
        assert (kw.get("env") or {}).get("GPU_RENT_CHECK_SWARM") == "0"
        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)


def test_stack_probe_requires_ollama_models():
    from gpu_rent.ready import _REMOTE_STACK_PROBE

    assert "/api/tags" in _REMOTE_STACK_PROBE
    assert "0 models" in _REMOTE_STACK_PROBE
    assert "WANT_OLLAMA_MODELS" in _REMOTE_STACK_PROBE
    assert "retry" in _REMOTE_STACK_PROBE
    assert ".gpu-rent-ollama-pulling" in _REMOTE_STACK_PROBE
    assert "WARN 0 models" in _REMOTE_STACK_PROBE
    assert "GPU не гасим" in _REMOTE_STACK_PROBE


def test_verify_stack_zero_models_does_not_fail_up(monkeypatch):
    from gpu_rent.ready import ServiceCheck, verify_stack_on_vm

    class C:
        enable_swarmui = True
        llm_runtime = "ollama"

    def fake_probe(_cfg, _host):
        return [
            ServiceCheck("swarmui", True, "ok", "vm", retry=False),
            ServiceCheck(
                "ollama",
                True,
                "WARN 0 models — Assistent empty (GPU не гасим)",
                "vm",
                retry=False,
            ),
        ]

    monkeypatch.setattr("gpu_rent.ready._probe_vm_once", fake_probe)
    out = verify_stack_on_vm(C(), "1.2.3.4", [].append, timeout=5.0, poll_every=0.01)
    assert all(c.ok for c in out)


def test_verify_stack_retries_while_pulling(monkeypatch):
    from gpu_rent.ready import ServiceCheck, verify_stack_on_vm

    class C:
        enable_swarmui = False
        llm_runtime = "ollama"

    n = {"i": 0}

    def fake_probe(_cfg, _host):
        n["i"] += 1
        if n["i"] == 1:
            return [
                ServiceCheck(
                    "ollama", False, "0 models — pull идёт (12s)", "vm", retry=True
                )
            ]
        return [ServiceCheck("ollama", True, "1 models (qwen)", "vm", retry=False)]

    monkeypatch.setattr("gpu_rent.ready._probe_vm_once", fake_probe)
    out = verify_stack_on_vm(
        C(), "1.2.3.4", [].append, timeout=5.0, poll_every=0.01
    )
    assert out[0].ok
    assert n["i"] == 2