Leonid Pershin
1785ab369c
- Updated the idle-killer logic to treat SwarmUI `empty` and `disabled` states as busy, preventing unnecessary idle time during provisioning. - Enhanced the `wait_backend_idle` function to recognize suspended backends as ready, improving resource utilization and user feedback. - Refined the `install_swarm_comfy` script to skip installation when backends are already present, streamlining the setup process. - Improved the `resolve_llm_runtime` function to prioritize live configuration over stale state notes, ensuring accurate runtime detection. - Added tests to validate the new backend status handling and idle management logic, ensuring robustness and reliability.
250 lines
7.3 KiB
Python
250 lines
7.3 KiB
Python
"""Tests for hold parsing and idle-killer busy classification."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from gpu_rent.errors import GpuRentError
|
|
from gpu_rent.hold import _parse_until
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
REMOTE_KILLER = ROOT / "src" / "gpu_rent" / "remote" / "idle_killer.py"
|
|
|
|
|
|
def _load_remote():
|
|
spec = importlib.util.spec_from_file_location("idle_killer_remote", REMOTE_KILLER)
|
|
assert spec and spec.loader
|
|
mod = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(mod)
|
|
return mod
|
|
|
|
|
|
def test_parse_until_iso():
|
|
ts = _parse_until("2030-01-01T00:00:00+00:00")
|
|
assert ts == int(datetime(2030, 1, 1, tzinfo=timezone.utc).timestamp())
|
|
|
|
|
|
def test_parse_until_unix():
|
|
assert _parse_until("1700000000") == 1700000000
|
|
|
|
|
|
def test_parse_until_bad():
|
|
with pytest.raises(GpuRentError):
|
|
_parse_until("not-a-date")
|
|
|
|
|
|
def test_classify_busy_from_status(monkeypatch):
|
|
mod = _load_remote()
|
|
|
|
class FakeResp:
|
|
def __init__(self, payload):
|
|
self._payload = payload
|
|
|
|
def read(self):
|
|
import json
|
|
|
|
return json.dumps(self._payload).encode()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return False
|
|
|
|
calls = {"n": 0}
|
|
|
|
def fake_urlopen(req, timeout=0, context=None):
|
|
calls["n"] += 1
|
|
url = getattr(req, "full_url", None) or req.get_full_url()
|
|
if "GetNewSession" in url:
|
|
return FakeResp({"session_id": "abc"})
|
|
return FakeResp(
|
|
{
|
|
"status": {"waiting_gens": 0, "live_gens": 0, "loading_models": 0},
|
|
"backend_status": {"status": "idle"},
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(mod.urllib.request, "urlopen", fake_urlopen)
|
|
busy, detail = mod.swarm_busy("http://127.0.0.1:7801")
|
|
assert busy is False
|
|
assert "idle" in detail
|
|
|
|
|
|
def test_classify_running_without_queue_not_busy(monkeypatch):
|
|
"""SwarmUI 'running' = ready; empty queue → idle-killer may stop GPU."""
|
|
mod = _load_remote()
|
|
|
|
class FakeResp:
|
|
def __init__(self, payload):
|
|
self._payload = payload
|
|
|
|
def read(self):
|
|
import json
|
|
|
|
return json.dumps(self._payload).encode()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return False
|
|
|
|
def fake_urlopen(req, timeout=0, context=None):
|
|
url = getattr(req, "full_url", None) or req.get_full_url()
|
|
if "GetNewSession" in url:
|
|
return FakeResp({"session_id": "abc"})
|
|
return FakeResp(
|
|
{
|
|
"status": {"waiting_gens": 0, "live_gens": 0, "loading_models": 0},
|
|
"backend_status": {"status": "running", "any_loading": False},
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(mod.urllib.request, "urlopen", fake_urlopen)
|
|
busy, detail = mod.swarm_busy("http://127.0.0.1:7801")
|
|
assert busy is False
|
|
assert "running" in detail
|
|
|
|
|
|
def test_classify_loading_is_busy(monkeypatch):
|
|
mod = _load_remote()
|
|
|
|
class FakeResp:
|
|
def __init__(self, payload):
|
|
self._payload = payload
|
|
|
|
def read(self):
|
|
import json
|
|
|
|
return json.dumps(self._payload).encode()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return False
|
|
|
|
def fake_urlopen(req, timeout=0, context=None):
|
|
url = getattr(req, "full_url", None) or req.get_full_url()
|
|
if "GetNewSession" in url:
|
|
return FakeResp({"session_id": "abc"})
|
|
return FakeResp(
|
|
{
|
|
"status": {"waiting_gens": 0, "live_gens": 0, "loading_models": 0},
|
|
"backend_status": {"status": "loading", "any_loading": True},
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(mod.urllib.request, "urlopen", fake_urlopen)
|
|
busy, detail = mod.swarm_busy("http://127.0.0.1:7801")
|
|
assert busy is True
|
|
assert "loading" in detail
|
|
|
|
|
|
def test_classify_empty_is_busy(monkeypatch):
|
|
"""No backends yet (first Comfy install) must not start idle clock."""
|
|
mod = _load_remote()
|
|
|
|
class FakeResp:
|
|
def __init__(self, payload):
|
|
self._payload = payload
|
|
|
|
def read(self):
|
|
import json
|
|
|
|
return json.dumps(self._payload).encode()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return False
|
|
|
|
def fake_urlopen(req, timeout=0, context=None):
|
|
url = getattr(req, "full_url", None) or req.get_full_url()
|
|
if "GetNewSession" in url:
|
|
return FakeResp({"session_id": "abc"})
|
|
return FakeResp(
|
|
{
|
|
"status": {"waiting_gens": 0, "live_gens": 0, "loading_models": 0},
|
|
"backend_status": {"status": "empty"},
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(mod.urllib.request, "urlopen", fake_urlopen)
|
|
busy, detail = mod.swarm_busy("http://127.0.0.1:7801")
|
|
assert busy is True
|
|
assert "empty" in detail
|
|
|
|
|
|
def test_classify_busy_queue(monkeypatch):
|
|
mod = _load_remote()
|
|
|
|
class FakeResp:
|
|
def __init__(self, payload):
|
|
self._payload = payload
|
|
|
|
def read(self):
|
|
import json
|
|
|
|
return json.dumps(self._payload).encode()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return False
|
|
|
|
def fake_urlopen(req, timeout=0, context=None):
|
|
url = getattr(req, "full_url", None) or req.get_full_url()
|
|
if "GetNewSession" in url:
|
|
return FakeResp({"session_id": "abc"})
|
|
return FakeResp(
|
|
{
|
|
"status": {"waiting_gens": 2, "live_gens": 0, "loading_models": 0},
|
|
"backend_status": {"status": "idle"},
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(mod.urllib.request, "urlopen", fake_urlopen)
|
|
busy, detail = mod.swarm_busy("http://127.0.0.1:7801")
|
|
assert busy is True
|
|
assert "waiting=2" in detail
|
|
|
|
|
|
def test_swarm_unreachable_starts_busy_then_allows_idle(tmp_path, monkeypatch):
|
|
mod = _load_remote()
|
|
monkeypatch.setattr(mod, "DATA", tmp_path)
|
|
monkeypatch.setattr(mod, "SWARM_DOWN_SINCE", tmp_path / ".gpu-rent-swarm-down-since")
|
|
monkeypatch.setattr(mod, "SWARM_UNREACHABLE_IDLE_SEC", 100)
|
|
|
|
def boom(*_a, **_k):
|
|
raise mod.urllib.error.URLError("down")
|
|
|
|
monkeypatch.setattr(mod.urllib.request, "urlopen", boom)
|
|
|
|
busy1, d1 = mod.swarm_busy("http://127.0.0.1:7801")
|
|
assert busy1 is True
|
|
assert "clock start" in d1
|
|
assert mod.SWARM_DOWN_SINCE.is_file()
|
|
|
|
# Still within window
|
|
since = mod.read_ts(mod.SWARM_DOWN_SINCE)
|
|
assert since is not None
|
|
mod.write_ts(mod.SWARM_DOWN_SINCE, since - 50)
|
|
busy2, d2 = mod.swarm_busy("http://127.0.0.1:7801")
|
|
assert busy2 is True
|
|
assert "50s" in d2 or "/ 100s" in d2
|
|
|
|
# Past window → not busy so idle clock can run
|
|
mod.write_ts(mod.SWARM_DOWN_SINCE, since - 120)
|
|
busy3, d3 = mod.swarm_busy("http://127.0.0.1:7801")
|
|
assert busy3 is False
|
|
assert "allow idle" in d3
|