AnyImage/tests/test_server.py

import importlib.util
import sys
import tempfile
import threading
import time
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_server_module():
    server_directory = PROJECT_ROOT / "src" / "snapform" / "server"
    specification = importlib.util.spec_from_file_location(
        "server", server_directory / "__init__.py",
        submodule_search_locations=[str(server_directory)],
    )
    module = importlib.util.module_from_spec(specification)
    sys.modules["server"] = module
    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 ServerTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.server_module = load_server_module()

    def setUp(self):
        self.jobs_temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.jobs_temp.cleanup)
        self.storage_root = Path(self.jobs_temp.name)

    def test_required_model_download_skips_ready_model(self):
        manager = self.server_module.ModelManager(Path("."))
        downloads = []
        context = SimpleNamespace(progress=lambda *_args: None)
        with patch.object(
            manager,
            "ready",
            side_effect=lambda key: key == "BEN2_BASE",
        ), patch.object(
            manager,
            "_download",
            side_effect=lambda _context, key, _progress: downloads.append(key),
        ):
            result = manager.download_missing(
                context,
                ("MOGE2_VITS_NORMAL", "BEN2_BASE"),
            )

        self.assertEqual(downloads, ["MOGE2_VITS_NORMAL"])
        self.assertEqual(result, {"downloaded": ["MOGE2_VITS_NORMAL"]})

    def test_models_are_owned_by_the_server_resource(self):
        self.assertFalse((PROJECT_ROOT / "src" / "snapform" / "models.py").exists())
        with tempfile.TemporaryDirectory() as directory:
            manager = self.server_module.ModelManager(Path(directory))
            snapshot = manager.snapshot()
        self.assertEqual(
            snapshot,
            {"loaded": {"ben2": "", "da3": "", "moge2": ""}},
        )
        self.assertIn("server.model_catalog", sys.modules)
        self.assertNotIn("model_catalog", sys.modules)

    def create_cache(self):
        cache = self.server_module.ModelSessionCache.__new__(
            self.server_module.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 create_server(self):
        source = self.server_module.server
        server = self.server_module.JobServer(
            "SnapForm Job Server", storage_root=self.storage_root
        )
        server.handlers.update(source.handlers)
        server.resource_factories.update(source.resource_factories)
        server.bind(self.storage_root)
        return server

    def wait_for_job(self, context, timeout=2.0):
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            if context._snapshot()["state"] in {"succeeded", "failed", "cancelled"}:
                return context._snapshot()
            time.sleep(0.01)
        self.fail(f"Job did not finish: {context._snapshot()}")

    def test_ben2_model_is_reused_for_matching_runtime_key(self):
        loads = []

        ben2 = ModuleType("server.models.onnx_ben2")
        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, {"server.models.onnx_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", "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("server.models.onnx_da3")
        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,
            {
                "server.models.onnx_da3": 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("server.models.onnx_moge2")
        fake_moge.create_session = lambda _model_dir, _device: moge
        with patch.dict(sys.modules, {"server.models.onnx_moge2": fake_moge}):
            self.assertIs(cache.get_moge2("moge", "cuda"), moge)

        self.assertIsNone(cache.da3_model)
        self.assertIsNotNone(cache.ben2_model)

    def test_dedicated_server_runs_only_one_job_at_a_time(self):
        server = self.server_module.JobServer("Test", storage_root=self.storage_root)
        first_started = threading.Event()
        release_first = threading.Event()
        order = []

        @server.job("example")
        def example(_context, parameters):
            order.append(parameters["name"])
            if parameters["name"] == "first":
                first_started.set()
                release_first.wait(1.0)

        first = server.submit("example", {"name": "first"}, job_id="job")
        self.assertTrue(first_started.wait(1.0))
        second = server.submit("example", {"name": "second"}, job_id="other")

        self.assertIsNotNone(first)
        self.assertIsNotNone(second)
        self.assertTrue(server.snapshot("instance")["busy"])
        self.assertEqual(server.snapshot("instance")["queued_jobs"], 1)
        self.assertNotIn("clients", server.snapshot("instance"))
        release_first.set()
        self.wait_for_job(first)
        self.wait_for_job(second)
        self.assertEqual(order, ["first", "second"])
        server.close()

    def test_server_registers_tasks_instead_of_model_names(self):
        server = self.create_server()
        self.assertEqual(
            set(server.handlers),
            {
                "download-model",
                "download-required-models",
                "remove-background",
                "extract-frames",
                "generate-depth",
                "generate-surface",
                "debug-da3",
                "debug-moge2",
                "debug-ben2",
            },
        )

    def test_queued_job_can_be_cancelled_before_it_runs(self):
        server = self.server_module.JobServer("Test", storage_root=self.storage_root)
        first_started = threading.Event()
        release_first = threading.Event()
        executed = []

        @server.job("example")
        def example(_context, parameters):
            executed.append(parameters["name"])
            if parameters["name"] == "first":
                first_started.set()
                release_first.wait(1.0)

        first = server.submit("example", {"name": "first"})
        self.assertTrue(first_started.wait(1.0))
        second = server.submit("example", {"name": "second"})

        self.assertIs(server.cancel(second.job_id), second)
        release_first.set()
        self.wait_for_job(first)
        second_status = self.wait_for_job(second)

        self.assertEqual(second_status["state"], "cancelled")
        self.assertEqual(executed, ["first"])
        server.close()

    def test_cancelling_job_keeps_server_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.server_module, "ModelManager", return_value=fake_models
        ):
            server = self.create_server()
        started = threading.Event()

        @server.job("cancellable")
        def cancellable(context, _parameters):
            started.set()
            while not context._is_cancelled():
                time.sleep(0.001)
            context.check_cancelled()

        context = server.submit("cancellable", {}, job_id="job")
        self.assertTrue(started.wait(1.0))
        self.assertIsNotNone(server.cancel("job"))
        self.assertTrue(context._is_cancelled())
        status = self.wait_for_job(context)

        self.assertEqual(status["state"], "cancelled")
        self.assertFalse(server.shutdown_event.is_set())
        self.assertEqual(
            server.snapshot("instance")["resources"]["model_manager"]["loaded"]["ben2"],
            "BEN2_BASE",
        )

    def test_job_progress_is_kept_in_server_memory(self):
        server = self.server_module.JobServer(
            "Test", storage_root=self.storage_root
        )

        @server.job("example")
        def example(context, _parameters):
            context.progress(0.5, "Half way")

        context = server.submit("example", {})
        self.wait_for_job(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_runtime_storage_root_owns_the_jobs_directory(self):
        storage_root = self.storage_root
        server = self.server_module.JobServer("Test")
        server.bind(storage_root)

        @server.job("example")
        def example(_context, _parameters):
            return {"value": 42}

        context = server.submit("example", {})
        self.wait_for_job(context)

        self.assertEqual(server.jobs_directory, storage_root / "jobs")
        self.assertEqual(context.storage_root, storage_root.resolve())
        self.assertEqual(context._snapshot()["result"], {"value": 42})

    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.server_module.JobServer(
            "Test", storage_root=self.storage_root
        )
        server.add_resource("model_manager", resource)

        @server.job("example")
        def example(context, _parameters):
            self.assertIs(context.resource("model_manager"), resource)
            release.wait(1.0)

        release = threading.Event()
        context = server.submit("example", {}, job_id="job")
        self.assertFalse(server.clear_resource("model_manager"))
        release.set()
        self.wait_for_job(context)
        self.assertTrue(server.clear_resource("model_manager"))
        server.close()

        self.assertEqual(
            server.snapshot("instance")["resources"]["model_manager"]["loaded"],
            "model",
        )
        self.assertEqual(events, ["clear", "close"])

    def test_server_dispatches_frame_extraction_without_models(self):
        fake_models = SimpleNamespace(
            snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}},
            clear=lambda: None,
            close=lambda: None,
        )
        with patch.object(
            self.server_module, "ModelManager", return_value=fake_models
        ):
            server = self.create_server()
        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",
            )
            status = self.wait_for_job(context)
            self.assertTrue(
                (context.directory / "frame.0001.jpg").is_file()
            )

        self.assertEqual(status["state"], "succeeded")


if __name__ == "__main__":
    unittest.main()