77 lines
3.3 KiB
Python
77 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import os
|
|
import unittest
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
|
|
ROOT_DIR = Path(__file__).resolve().parents[1]
|
|
SPEC = importlib.util.spec_from_file_location("modelhub_agent_entrypoint", ROOT_DIR / "main.py")
|
|
if SPEC is None or SPEC.loader is None:
|
|
raise RuntimeError("Unable to load the hosted agent entrypoint")
|
|
ENTRYPOINT = importlib.util.module_from_spec(SPEC)
|
|
SPEC.loader.exec_module(ENTRYPOINT)
|
|
|
|
|
|
class HostedAgentEntrypointTests(unittest.TestCase):
|
|
def test_hosted_worker_always_uses_unlimited_cycle_submissions(self) -> None:
|
|
with patch.dict(
|
|
os.environ,
|
|
{
|
|
"MODELHUB_AGENT_MAX_SUBMITS_PER_RUN": "5",
|
|
"MODELHUB_AGENT_EXTRA_ARGS": "--max-submits-per-run 3",
|
|
},
|
|
clear=False,
|
|
):
|
|
command = ENTRYPOINT._worker_command()
|
|
|
|
self.assertEqual(["--max-submits-per-run", "0"], command[-2:])
|
|
self.assertIn("--state-sync", command)
|
|
|
|
def test_readiness_file_is_separate_from_liveness(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
path = Path(temporary_dir) / "readiness.json"
|
|
path.write_text('{"ready": false, "reason": "state_sync_unhealthy"}', encoding="utf-8")
|
|
with patch.object(ENTRYPOINT, "READINESS_PATH", path):
|
|
readiness = ENTRYPOINT._readiness()
|
|
self.assertFalse(readiness["ready"])
|
|
self.assertEqual("state_sync_unhealthy", readiness["reason"])
|
|
|
|
def test_runtime_image_uses_python_git_client_without_apt_layer(self) -> None:
|
|
dockerfile = (ROOT_DIR / "Dockerfile").read_text(encoding="utf-8")
|
|
requirements = (ROOT_DIR / "requirements.txt").read_text(encoding="utf-8")
|
|
self.assertNotIn("apt-get", dockerfile)
|
|
self.assertIn("dulwich", requirements.casefold())
|
|
|
|
def test_exited_worker_is_backed_off_without_exiting_supervisor(self) -> None:
|
|
class ExitedWorker:
|
|
def poll(self) -> int:
|
|
return 137
|
|
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
root = Path(temporary_dir)
|
|
with (
|
|
patch.object(ENTRYPOINT, "worker", ExitedWorker()),
|
|
patch.object(ENTRYPOINT, "worker_start_enabled", True),
|
|
patch.object(ENTRYPOINT, "shutdown_requested", False),
|
|
patch.object(ENTRYPOINT, "config_error", None),
|
|
patch.object(ENTRYPOINT, "worker_started_at", 90.0),
|
|
patch.object(ENTRYPOINT, "worker_restart_count", 0),
|
|
patch.object(ENTRYPOINT, "worker_next_restart_at", 0.0),
|
|
patch.object(ENTRYPOINT, "WORKER_CRASH_PATH", root / "worker_crashes.jsonl"),
|
|
patch.object(ENTRYPOINT, "READINESS_PATH", root / "readiness.json"),
|
|
):
|
|
ENTRYPOINT._supervise_worker(now=100.0)
|
|
self.assertIsNone(ENTRYPOINT.worker)
|
|
self.assertEqual(1, ENTRYPOINT.worker_restart_count)
|
|
self.assertEqual(105.0, ENTRYPOINT.worker_next_restart_at)
|
|
self.assertEqual("worker_restarting", ENTRYPOINT._readiness()["reason"])
|
|
self.assertIn('"exitCode":137', (root / "worker_crashes.jsonl").read_text(encoding="utf-8"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|