nan
b3a07cde32
- Implement JobClient for handling HTTP requests to the JobServer. - Create ModalJobOperatorBase for managing Blender job lifecycle. - Introduce JobContext and JobServer for job management and state handling. - Add ServiceConnection and ServiceController for managing service processes. - Implement job submission, cancellation, and resource management endpoints in the server. - Ensure proper error handling and job state updates throughout the service.
262 lines
8.8 KiB
Python
262 lines
8.8 KiB
Python
import importlib.util
|
|
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_service():
|
|
specification = importlib.util.spec_from_file_location(
|
|
"snapform_service",
|
|
PROJECT_ROOT / "src" / "snapform" / "ai_service" / "server.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 ServiceTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.service = load_service()
|
|
|
|
def setUp(self):
|
|
self.jobs_temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.jobs_temp.cleanup)
|
|
self.jobs_directory = Path(self.jobs_temp.name) / "jobs"
|
|
|
|
def create_cache(self):
|
|
cache = self.service.ModelSessionCache.__new__(
|
|
self.service.ModelSessionCache
|
|
)
|
|
cache.ben2_model = None
|
|
cache.ben2_key = None
|
|
cache.da3_model = None
|
|
cache.da3_key = None
|
|
cache.moge_model = None
|
|
cache.moge_key = None
|
|
return cache
|
|
|
|
def test_ben2_model_is_reused_for_matching_runtime_key(self):
|
|
loads = []
|
|
|
|
ben2 = ModuleType("ben2_onnx")
|
|
ben2.create_session = lambda model_dir, device: (
|
|
loads.append((model_dir, device)) or object()
|
|
)
|
|
cache = self.create_cache()
|
|
device = SimpleNamespace(type="cuda")
|
|
with patch.dict(sys.modules, {"ben2_onnx": 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", "cuda")])
|
|
self.assertEqual(cache.snapshot()["ben2"], "model")
|
|
self.assertEqual(cached_load_ms, 0.0)
|
|
|
|
def test_da3_model_is_reused_for_matching_runtime_key(self):
|
|
loads = []
|
|
|
|
da3 = ModuleType("da3_onnx")
|
|
da3.create_session = lambda model_dir, device: (
|
|
loads.append((model_dir, device)) or object()
|
|
)
|
|
cache = self.create_cache()
|
|
device = SimpleNamespace(type="cuda")
|
|
with patch.dict(
|
|
sys.modules,
|
|
{
|
|
"da3_onnx": da3,
|
|
},
|
|
):
|
|
first = cache.get_da3("model", device)
|
|
second = cache.get_da3("model", device)
|
|
|
|
self.assertIs(first, second)
|
|
self.assertEqual(loads, [("model", "cuda")])
|
|
self.assertEqual(cache.snapshot()["da3"], "model")
|
|
|
|
def test_depth_model_cache_keeps_only_da3_or_moge(self):
|
|
cache = self.create_cache()
|
|
cache.ben2_model = object()
|
|
cache.ben2_key = ("ben2", "cuda")
|
|
cache.moge_model = object()
|
|
cache.moge_key = ("moge", "cuda")
|
|
device = SimpleNamespace(type="cuda")
|
|
da3 = object()
|
|
|
|
def load_da3(_model_dir, _device, key):
|
|
cache.da3_model = da3
|
|
cache.da3_key = key
|
|
return da3
|
|
|
|
with patch.object(cache, "_load_da3", side_effect=load_da3):
|
|
self.assertIs(cache.get_da3("da3", device), da3)
|
|
|
|
self.assertIsNone(cache.moge_model)
|
|
self.assertIsNotNone(cache.ben2_model)
|
|
|
|
moge = object()
|
|
|
|
fake_moge = ModuleType("moge_onnx")
|
|
fake_moge.create_session = lambda _model_dir, _device: moge
|
|
with patch.dict(sys.modules, {"moge_onnx": fake_moge}):
|
|
self.assertIs(cache.get_moge_onnx("moge", "cuda"), moge)
|
|
|
|
self.assertIsNone(cache.da3_model)
|
|
self.assertIsNotNone(cache.ben2_model)
|
|
|
|
def test_dedicated_service_runs_only_one_job_at_a_time(self):
|
|
fake_models = SimpleNamespace(
|
|
snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}},
|
|
clear=lambda: None,
|
|
close=lambda: None,
|
|
)
|
|
with patch.object(self.service, "ModelManager", return_value=fake_models):
|
|
server = self.service.create_server(jobs_directory=self.jobs_directory)
|
|
|
|
parameters = {"model": "BASE"}
|
|
first = server.submit("debug-da3", parameters, job_id="job")
|
|
second = server.submit("debug-da3", parameters, job_id="other")
|
|
|
|
self.assertIsNotNone(first)
|
|
self.assertIsNone(second)
|
|
self.assertTrue(server.snapshot("instance")["busy"])
|
|
self.assertNotIn("clients", server.snapshot("instance"))
|
|
|
|
def test_service_registers_tasks_instead_of_model_names(self):
|
|
server = self.service.create_server(jobs_directory=self.jobs_directory)
|
|
self.assertEqual(
|
|
set(server.handlers),
|
|
{
|
|
"download-model",
|
|
"remove-background",
|
|
"extract-frames",
|
|
"generate-da3-depth",
|
|
"generate-moge2-depth",
|
|
"generate-da3-surface",
|
|
"generate-moge2-surface",
|
|
"debug-da3",
|
|
"debug-moge2",
|
|
"debug-ben2",
|
|
},
|
|
)
|
|
|
|
def test_cancelling_job_keeps_service_and_model_cache_alive(self):
|
|
fake_models = SimpleNamespace(
|
|
snapshot=lambda: {
|
|
"loaded": {"ben2": "BEN2_BASE", "da3": ""}
|
|
},
|
|
clear=lambda: None,
|
|
close=lambda: None,
|
|
)
|
|
with patch.object(self.service, "ModelManager", return_value=fake_models):
|
|
server = self.service.create_server(jobs_directory=self.jobs_directory)
|
|
context = server.submit("remove-background", {}, job_id="job")
|
|
self.assertTrue(context.request_cancel())
|
|
self.assertTrue(context.is_cancelled())
|
|
server.run(context)
|
|
status = context.snapshot()
|
|
|
|
self.assertEqual(status["state"], "cancelled")
|
|
self.assertFalse(server.shutdown_event.is_set())
|
|
self.assertEqual(
|
|
server.snapshot("instance")["resources"]["models"]["loaded"]["ben2"],
|
|
"BEN2_BASE",
|
|
)
|
|
|
|
def test_job_progress_is_kept_in_service_memory(self):
|
|
server = self.service.JobServer("Test", jobs_directory=self.jobs_directory)
|
|
|
|
@server.job("example")
|
|
def example(context, _parameters):
|
|
context.update(0.5, "Half way")
|
|
|
|
context = server.submit("example", {})
|
|
server.run(context)
|
|
|
|
self.assertRegex(context.job_id, r"^[0-9a-f]{32}$")
|
|
self.assertEqual(context.directory.name, context.job_id)
|
|
self.assertEqual(context.snapshot()["state"], "succeeded")
|
|
self.assertEqual(context.snapshot()["progress"], 1.0)
|
|
|
|
def test_job_server_owns_and_clears_shared_resources(self):
|
|
events = []
|
|
resource = SimpleNamespace(
|
|
snapshot=lambda: {"loaded": "model"},
|
|
clear=lambda: events.append("clear"),
|
|
close=lambda: events.append("close"),
|
|
)
|
|
server = self.service.JobServer("Test", jobs_directory=self.jobs_directory)
|
|
server.add_resource("models", resource)
|
|
|
|
@server.job("example")
|
|
def example(context, _parameters):
|
|
self.assertIs(context.resource("models"), resource)
|
|
|
|
context = server.submit("example", {}, job_id="job")
|
|
self.assertFalse(server.clear_resource("models"))
|
|
server.run(context)
|
|
self.assertTrue(server.clear_resource("models"))
|
|
server.close()
|
|
|
|
self.assertEqual(
|
|
server.snapshot("instance")["resources"]["models"]["loaded"],
|
|
"model",
|
|
)
|
|
self.assertEqual(events, ["clear", "close"])
|
|
|
|
def test_service_dispatches_frame_extraction_without_models(self):
|
|
fake_models = SimpleNamespace(
|
|
snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}},
|
|
clear=lambda: None,
|
|
close=lambda: None,
|
|
)
|
|
with patch.object(self.service, "ModelManager", return_value=fake_models):
|
|
server = self.service.create_server(jobs_directory=self.jobs_directory)
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
input_directory = root / "frames"
|
|
input_directory.mkdir()
|
|
(input_directory / "frame-0001.jpg").touch()
|
|
(input_directory / "frame-0002.jpg").touch()
|
|
context = server.submit(
|
|
"extract-frames",
|
|
{
|
|
"input": str(input_directory),
|
|
"video_frames": 16,
|
|
"start_frame": 1,
|
|
},
|
|
job_id="job",
|
|
)
|
|
server.run(context)
|
|
status = context.snapshot()
|
|
self.assertTrue(
|
|
(context.directory / "frame.0001.jpg").is_file()
|
|
)
|
|
|
|
self.assertEqual(status["state"], "succeeded")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|