Implement AI Host functionality with model caching and job management
- Added ai_host.py to handle AI tasks with FastAPI integration. - Introduced HostState and ModelCache classes for managing client connections and model loading. - Implemented job submission and cancellation endpoints in the AI host. - Enhanced tests for AI host functionality, including client tracking and job cancellation. - Updated existing tests to accommodate changes in the operator and command structure. - Refactored runtime commands to use request-based architecture instead of command-line arguments. - Improved model loading logic to handle CUDA memory management and reuse models efficiently.
This commit is contained in:
223
tests/test_ai_host.py
Normal file
223
tests/test_ai_host.py
Normal file
@@ -0,0 +1,223 @@
|
||||
import importlib.util
|
||||
import inspect
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def load_ai_host():
|
||||
specification = importlib.util.spec_from_file_location(
|
||||
"snapform_ai_host",
|
||||
PROJECT_ROOT / "src" / "snapform" / "ai_host.py",
|
||||
)
|
||||
module = importlib.util.module_from_spec(specification)
|
||||
specification.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class FakeCuda:
|
||||
class OutOfMemoryError(RuntimeError):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def is_available():
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def empty_cache():
|
||||
return None
|
||||
|
||||
|
||||
class HostTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.host = load_ai_host()
|
||||
|
||||
def create_cache(self):
|
||||
cache = self.host.ModelCache.__new__(self.host.ModelCache)
|
||||
cache.torch = SimpleNamespace(cuda=FakeCuda())
|
||||
cache.ben2_model = None
|
||||
cache.ben2_key = None
|
||||
cache.da3_model = None
|
||||
cache.da3_key = None
|
||||
return cache
|
||||
|
||||
def test_ben2_model_is_reused_for_matching_runtime_key(self):
|
||||
loads = []
|
||||
|
||||
class FakeModel:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_dir):
|
||||
loads.append(model_dir)
|
||||
return cls()
|
||||
|
||||
def to(self, device):
|
||||
self.device = device
|
||||
return self
|
||||
|
||||
def eval(self):
|
||||
return None
|
||||
|
||||
ben2 = ModuleType("ben2")
|
||||
ben2.BEN_Base = FakeModel
|
||||
cache = self.create_cache()
|
||||
device = SimpleNamespace(type="cuda")
|
||||
with patch.dict(sys.modules, {"ben2": ben2}):
|
||||
first, _load_ms = cache.get_ben2("model", device)
|
||||
second, cached_load_ms = cache.get_ben2("model", device)
|
||||
|
||||
self.assertIs(first, second)
|
||||
self.assertEqual(loads, ["model"])
|
||||
self.assertEqual(cache.status()["ben2"], "model")
|
||||
self.assertEqual(cached_load_ms, 0.0)
|
||||
|
||||
def test_da3_model_is_reused_for_matching_runtime_key(self):
|
||||
loads = []
|
||||
|
||||
class FakeModel:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_dir):
|
||||
loads.append(model_dir)
|
||||
return cls()
|
||||
|
||||
def to(self, device):
|
||||
self.device = device
|
||||
return self
|
||||
|
||||
def eval(self):
|
||||
return None
|
||||
|
||||
package = ModuleType("depth_anything_3")
|
||||
api = ModuleType("depth_anything_3.api")
|
||||
api.DepthAnything3 = FakeModel
|
||||
cache = self.create_cache()
|
||||
device = SimpleNamespace(type="cuda")
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"depth_anything_3": package,
|
||||
"depth_anything_3.api": api,
|
||||
},
|
||||
):
|
||||
first = cache.get_da3("model", device)
|
||||
second = cache.get_da3("model", device)
|
||||
|
||||
self.assertIs(first, second)
|
||||
self.assertEqual(loads, ["model"])
|
||||
self.assertEqual(cache.status()["da3"], "da3_model")
|
||||
|
||||
def test_http_host_has_health_jobs_auth_and_single_task_guard(self):
|
||||
source = inspect.getsource(self.host.create_app)
|
||||
|
||||
self.assertIn('@app.get("/health")', source)
|
||||
self.assertIn('@app.post("/jobs", status_code=202)', source)
|
||||
self.assertIn("x_snapform_token", source)
|
||||
self.assertIn("hmac.compare_digest", source)
|
||||
self.assertIn("parameters[\"status\"]", source)
|
||||
self.assertIn('@app.post("/clients", status_code=201)', source)
|
||||
self.assertIn('@app.post("/clients/{client_id}/heartbeat")', source)
|
||||
self.assertIn('@app.delete("/jobs/{job_id}", status_code=202)', source)
|
||||
self.assertIn('@app.post("/models/clear")', source)
|
||||
self.assertIn('@app.post("/shutdown", status_code=202)', source)
|
||||
self.assertNotIn("os._exit", source)
|
||||
|
||||
def test_host_tracks_clients_and_rejects_shared_shutdown(self):
|
||||
fake_models = SimpleNamespace(
|
||||
status=lambda: {"ben2": "", "da3": ""},
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(self.host, "ModelCache", return_value=fake_models):
|
||||
state = self.host.HostState("token", "instance")
|
||||
|
||||
state.register_client("first", "Blender A")
|
||||
state.register_client("second", "Blender B")
|
||||
|
||||
self.assertEqual(state.snapshot()["clients"], 2)
|
||||
self.assertEqual(state.request_shutdown("first"), "shared")
|
||||
self.assertTrue(
|
||||
state.reserve("job", "first", "reconstruct", "status.json")
|
||||
)
|
||||
self.assertFalse(
|
||||
state.reserve("other", "second", "reconstruct", "other.json")
|
||||
)
|
||||
|
||||
def test_cancelling_job_keeps_host_and_model_cache_alive(self):
|
||||
fake_models = SimpleNamespace(
|
||||
status=lambda: {"ben2": "BEN2_BASE", "da3": ""},
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(self.host, "ModelCache", return_value=fake_models):
|
||||
state = self.host.HostState("token", "instance")
|
||||
state.register_client("client", "Blender")
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
state.reserve(
|
||||
"job",
|
||||
"client",
|
||||
"remove-background",
|
||||
str(status_path),
|
||||
)
|
||||
|
||||
active = state.cancel_job("job", "client")
|
||||
self.assertTrue(active["cancel_requested"])
|
||||
self.assertTrue(state.is_job_cancelled("job"))
|
||||
state.run_job(
|
||||
{
|
||||
"job_id": "job",
|
||||
"command": "remove-background",
|
||||
"parameters": {"status": str(status_path)},
|
||||
}
|
||||
)
|
||||
status = self.host.json.loads(
|
||||
status_path.read_text(encoding="utf-8")
|
||||
)
|
||||
|
||||
self.assertEqual(status["state"], "cancelled")
|
||||
self.assertFalse(state.shutdown_event.is_set())
|
||||
self.assertEqual(state.snapshot()["cached_models"]["ben2"], "BEN2_BASE")
|
||||
|
||||
def test_expired_clients_are_removed(self):
|
||||
fake_models = SimpleNamespace(
|
||||
status=lambda: {"ben2": "", "da3": ""},
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(self.host, "ModelCache", return_value=fake_models):
|
||||
state = self.host.HostState("token", "instance")
|
||||
state.register_client("client", "Blender")
|
||||
state.clients["client"]["last_seen"] -= (
|
||||
self.host.CLIENT_TIMEOUT_SECONDS + 1.0
|
||||
)
|
||||
|
||||
state.expire_clients()
|
||||
|
||||
self.assertEqual(state.snapshot()["clients"], 0)
|
||||
|
||||
def test_discovery_is_only_removed_by_matching_host(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "host.json"
|
||||
args = SimpleNamespace(
|
||||
discovery=str(path),
|
||||
port=51234,
|
||||
token="token",
|
||||
instance_id="owner",
|
||||
)
|
||||
with patch.object(self.host.os, "getpid", return_value=42):
|
||||
self.host.write_discovery(args)
|
||||
|
||||
payload = self.host.json.loads(path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(payload["port"], 51234)
|
||||
self.assertEqual(payload["pid"], 42)
|
||||
self.host.remove_discovery(path, "other")
|
||||
self.assertTrue(path.is_file())
|
||||
self.host.remove_discovery(path, "owner")
|
||||
self.assertFalse(path.exists())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user