Leonid Pershin
a563ae06c4
- Updated version number in pyproject.toml and __init__.py to 0.2.0. - Revised README.md to reflect the current state of the project, including usage instructions and setup steps. - Improved CLI documentation in cli.md, adding details about new commands and their functionalities. - Enhanced the quick start section in README.md for better clarity on initial setup. - Updated local folder documentation to clarify file handling and commands. - Added a new command for listing GPU flavors and improved error handling in the CLI. - Implemented a watchdog feature in the tunnel to manage server states effectively.
112 lines
3.0 KiB
Python
112 lines
3.0 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
|