gpu-rent/tests/test_tune_swarm_perf.py

"""Tests for remote tune_swarm_perf pip_ok / ExtraArgs gating."""

from __future__ import annotations

import importlib.util
import json
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
REMOTE = ROOT / "src" / "gpu_rent" / "remote" / "tune_swarm_perf.py"


def _load():
    spec = importlib.util.spec_from_file_location("tune_swarm_perf", REMOTE)
    assert spec and spec.loader
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def test_pip_fail_skips_extra_args(tmp_path, monkeypatch):
    mod = _load()
    data = tmp_path
    backends = data / "Data" / "Backends.fds"
    backends.parent.mkdir(parents=True)
    backends.write_text("ExtraArgs: \n", encoding="utf-8")
    gpu_json = data / ".gpu-rent-gpu.json"
    gpu_json.write_text(
        json.dumps(
            {
                "vram_mib": 24576,
                "compute_cap": "8.9",
                "uuid": "gpu-1",
                "name": "RTX",
            }
        ),
        encoding="utf-8",
    )
    pip = data / "fake-pip"
    pip.write_text("#!/bin/sh\n", encoding="utf-8")

    monkeypatch.setattr(mod, "DATA", data)
    monkeypatch.setattr(mod, "GPU_JSON", gpu_json)
    monkeypatch.setattr(mod, "MARKER", data / ".gpu-rent-perf-tuned")
    monkeypatch.setattr(mod, "BACKENDS", backends)
    monkeypatch.setattr(mod, "find_pip", lambda: pip)
    monkeypatch.setattr(mod, "find_comfy_python", lambda: pip)
    monkeypatch.setattr(mod, "pip_install_sage", lambda _p: False)

    assert mod.main() == 0
    marker = json.loads((data / ".gpu-rent-perf-tuned").read_text(encoding="utf-8"))
    assert marker["pip_ok"] is False
    assert marker["extra_args"] == ""
    assert "--use-sage-attention" not in backends.read_text(encoding="utf-8")


def test_pip_ok_patches_extra_args(tmp_path, monkeypatch):
    mod = _load()
    data = tmp_path
    backends = data / "Data" / "Backends.fds"
    backends.parent.mkdir(parents=True)
    backends.write_text("ExtraArgs: \n", encoding="utf-8")
    gpu_json = data / ".gpu-rent-gpu.json"
    gpu_json.write_text(
        json.dumps(
            {
                "vram_mib": 24576,
                "compute_cap": "8.9",
                "uuid": "gpu-1",
                "name": "RTX",
            }
        ),
        encoding="utf-8",
    )
    pip = data / "fake-pip"
    pip.write_text("#!/bin/sh\n", encoding="utf-8")

    monkeypatch.setattr(mod, "DATA", data)
    monkeypatch.setattr(mod, "GPU_JSON", gpu_json)
    monkeypatch.setattr(mod, "MARKER", data / ".gpu-rent-perf-tuned")
    monkeypatch.setattr(mod, "BACKENDS", backends)
    monkeypatch.setattr(mod, "find_pip", lambda: pip)
    monkeypatch.setattr(mod, "find_comfy_python", lambda: pip)
    monkeypatch.setattr(mod, "pip_install_sage", lambda _p: True)

    assert mod.main() == 0
    marker = json.loads((data / ".gpu-rent-perf-tuned").read_text(encoding="utf-8"))
    assert marker["pip_ok"] is True
    assert "--use-sage-attention" in marker["extra_args"]
    assert "--use-sage-attention" in backends.read_text(encoding="utf-8")


def test_pip_fail_retries_next_run(tmp_path, monkeypatch):
    mod = _load()
    data = tmp_path
    backends = data / "Data" / "Backends.fds"
    backends.parent.mkdir(parents=True)
    backends.write_text("ExtraArgs: \n", encoding="utf-8")
    gpu_json = data / ".gpu-rent-gpu.json"
    gpu_json.write_text(
        json.dumps(
            {
                "vram_mib": 24576,
                "compute_cap": "8.9",
                "uuid": "gpu-1",
                "name": "RTX",
            }
        ),
        encoding="utf-8",
    )
    marker = data / ".gpu-rent-perf-tuned"
    marker.write_text(
        json.dumps(
            {
                "uuid": "gpu-1",
                "extra_args": "",
                "pip_ok": False,
                "tier": "high",
                "name": "RTX",
            }
        ),
        encoding="utf-8",
    )
    pip = data / "fake-pip"
    pip.write_text("#!/bin/sh\n", encoding="utf-8")

    monkeypatch.setattr(mod, "DATA", data)
    monkeypatch.setattr(mod, "GPU_JSON", gpu_json)
    monkeypatch.setattr(mod, "MARKER", marker)
    monkeypatch.setattr(mod, "BACKENDS", backends)
    monkeypatch.setattr(mod, "find_pip", lambda: pip)
    monkeypatch.setattr(mod, "find_comfy_python", lambda: pip)
    monkeypatch.setattr(mod, "pip_install_sage", lambda _p: True)

    assert mod.main() == 0
    new_m = json.loads(marker.read_text(encoding="utf-8"))
    assert new_m["pip_ok"] is True
    assert "--use-sage-attention" in backends.read_text(encoding="utf-8")


def test_pip_install_uses_python_dash_m(tmp_path, monkeypatch):
    mod = _load()
    py = tmp_path / "python"
    py.write_text("", encoding="utf-8")
    seen: list[list[str]] = []

    def fake_call(cmd, env=None):
        seen.append(list(cmd))
        return 0

    monkeypatch.setattr(mod, "sage_already_importable", lambda _p: False)
    monkeypatch.setattr(mod.subprocess, "check_call", fake_call)
    assert mod.pip_install_sage(py) is True
    assert seen and seen[0][:4] == [str(py), "-m", "pip", "install"]


def test_patch_extra_args_fds_empty_x(tmp_path, monkeypatch):
    """FDS empty ExtraArgs is '\\x' — must replace, not append."""
    mod = _load()
    backends = tmp_path / "Backends.fds"
    backends.write_text("0:\n\ttype: comfyui_selfstart\n\tExtraArgs: \\x\n", encoding="utf-8")
    monkeypatch.setattr(mod, "BACKENDS", backends)
    assert mod.patch_backends_extra_args("--use-sage-attention") is True
    text = backends.read_text(encoding="utf-8")
    assert "ExtraArgs: --use-sage-attention" in text
    assert "\\x --use-sage" not in text


def test_sanitize_backends_fds_corruption(tmp_path, monkeypatch):
    mod = _load()
    backends = tmp_path / "Backends.fds"
    backends.write_text(
        "ExtraArgs: \\x --use-sage-attention\n",
        encoding="utf-8",
    )
    monkeypatch.setattr(mod, "BACKENDS", backends)
    assert mod.sanitize_backends_fds() is True
    assert backends.read_text(encoding="utf-8") == "ExtraArgs: --use-sage-attention\n"