Leonid Pershin
1ec615c03e
- Marked critical bugs as resolved in the review documentation, including changes to the `arm_idle_killer` function to raise errors on credential creation failures and ensure proper file permissions for JSON credentials. - Introduced a new `_try_arm_idle_killer` function in `provision.py` to manage idle-killer state more effectively, ensuring it arms correctly during provisioning. - Updated the `swarm_busy` function in `remote/idle_killer.py` to allow idle state after a specified duration of Swarm unavailability, preventing unnecessary billing. - Enhanced performance tuning logic in `tune_swarm_perf.py` to ensure proper handling of pip installation success before applying extra arguments. - Added tests to validate the new idle-killer behavior and swarm management logic, ensuring robustness in handling idle states and error conditions.
143 lines
4.1 KiB
Python
143 lines
4.1 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_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
|