119 lines
5.1 KiB
Python
119 lines
5.1 KiB
Python
import importlib.util
|
|
import io
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
spec = importlib.util.spec_from_file_location("watchdog", Path(__file__).resolve().parents[1] / "scripts/llama_watchdog.py")
|
|
watchdog = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(watchdog)
|
|
|
|
|
|
def busy(task=1, decoded=20):
|
|
return [{"id": 0, "id_task": task, "is_processing": True,
|
|
"next_token": [{"n_decoded": decoded}]}]
|
|
|
|
|
|
class WatchdogTests(unittest.TestCase):
|
|
def test_busy_request_is_allowed_then_expires_despite_progress(self):
|
|
state, reasons = watchdog.observe({}, busy(), "server1", 100, 600, 2048)
|
|
self.assertFalse(reasons)
|
|
state, reasons = watchdog.observe(state, busy(decoded=1000), "server1", 699, 600, 2048)
|
|
self.assertFalse(reasons)
|
|
_, reasons = watchdog.observe(state, busy(decoded=1100), "server1", 700, 600, 2048)
|
|
self.assertIn("600s", reasons[0])
|
|
|
|
def test_new_request_and_restarted_server_do_not_inherit_deadline(self):
|
|
state, _ = watchdog.observe({}, busy(), "server1", 100, 600, 2048)
|
|
for slots, invocation in [(busy(2), "server1"), (busy(), "server2")]:
|
|
_, reasons = watchdog.observe(state, slots, invocation, 900, 600, 2048)
|
|
self.assertFalse(reasons)
|
|
|
|
def test_idle_slot_drops_old_deadline(self):
|
|
state, _ = watchdog.observe({}, busy(), "server1", 100, 600, 2048)
|
|
state, reasons = watchdog.observe(state, [{"id": 0, "is_processing": False}], "server1", 900, 600, 2048)
|
|
self.assertEqual(state["active"], {})
|
|
self.assertFalse(reasons)
|
|
|
|
def test_excess_output_is_stopped_even_when_client_overrides_default(self):
|
|
_, reasons = watchdog.observe({}, busy(decoded=2049), "server1", 100, 600, 2048)
|
|
self.assertIn("2048 tokens", reasons[0])
|
|
|
|
def test_prompt_and_generated_text_are_not_persisted(self):
|
|
slots = busy()
|
|
slots[0].update(prompt="private prompt", generated="private output")
|
|
state, _ = watchdog.observe({}, slots, "server1", 100, 600, 2048)
|
|
self.assertNotIn("private", str(state))
|
|
|
|
|
|
class WatchdogIntegrationTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp.cleanup)
|
|
self.environment = patch.dict(watchdog.os.environ, {"XDG_STATE_HOME": self.temp.name,
|
|
"LLAMA_MAX_REQUEST_SECONDS": "600", "LLAMA_MAX_OUTPUT_TOKENS": "2048"})
|
|
self.environment.start()
|
|
self.addCleanup(self.environment.stop)
|
|
self.path = Path(self.temp.name) / "hermes-stack-maintenance/llama-watchdog.json"
|
|
self.boot = Path("/proc/sys/kernel/random/boot_id").read_text().strip()
|
|
self.calls = []
|
|
|
|
def control(self, *args):
|
|
self.calls.append(args)
|
|
return "ActiveState=active\nInvocationID=test-server" if args[0] == "show" else ""
|
|
|
|
def tick(self, now, slots=None, error=False):
|
|
response = io.BytesIO(json.dumps(slots or []).encode())
|
|
with patch.object(watchdog, "systemctl", self.control), \
|
|
patch.object(watchdog.time, "monotonic", return_value=now), \
|
|
patch.object(watchdog, "urlopen", side_effect=OSError() if error else None, return_value=response):
|
|
watchdog.main()
|
|
|
|
def restarts(self):
|
|
return [call for call in self.calls if call[0] == "restart"]
|
|
|
|
def test_deadline_restarts_actual_configured_service_once(self):
|
|
self.tick(100, busy())
|
|
self.tick(699, busy())
|
|
self.assertFalse(self.restarts())
|
|
self.tick(700, busy())
|
|
self.assertEqual(self.restarts(), [("restart", "--no-block", "llama-qwen.service")])
|
|
self.tick(701, busy(decoded=3000))
|
|
self.assertEqual(len(self.restarts()), 1)
|
|
self.assertEqual(self.path.stat().st_mode & 0o777, 0o600)
|
|
|
|
def test_only_persistent_probe_failure_triggers_recovery(self):
|
|
self.tick(100, error=True)
|
|
self.tick(699, error=True)
|
|
self.assertFalse(self.restarts())
|
|
self.tick(700, error=True)
|
|
self.assertEqual(len(self.restarts()), 1)
|
|
|
|
def test_successful_probe_resets_failure_deadline(self):
|
|
self.tick(100, error=True)
|
|
self.tick(500)
|
|
self.tick(700, error=True)
|
|
self.assertFalse(self.restarts())
|
|
|
|
def test_inactive_service_never_started_or_probed(self):
|
|
with patch.object(watchdog, "systemctl", return_value="ActiveState=inactive\nInvocationID=") as control, \
|
|
patch.object(watchdog, "urlopen") as probe:
|
|
watchdog.main()
|
|
probe.assert_not_called()
|
|
self.assertEqual(control.call_count, 1)
|
|
|
|
def test_reboot_clears_old_monotonic_deadlines_and_cooldown(self):
|
|
self.path.parent.mkdir()
|
|
self.path.write_text(json.dumps({"boot": "previous-boot", "invocation": "test-server",
|
|
"active": {"0:1": 100}, "last_restart": 10000, "unreachable_since": 50}))
|
|
self.tick(100, busy())
|
|
self.tick(700, busy())
|
|
self.assertEqual(len(self.restarts()), 1)
|
|
self.assertEqual(json.loads(self.path.read_text())["boot"], self.boot)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|