"""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_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