diff options
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/test_lazy_loading.py | 245 |
1 files changed, 245 insertions, 0 deletions
diff --git a/tests/test_lazy_loading.py b/tests/test_lazy_loading.py new file mode 100644 index 0000000..bd4b341 --- /dev/null +++ b/tests/test_lazy_loading.py @@ -0,0 +1,245 @@ +from __future__ import annotations + +import asyncio +import logging +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest +from mcp import Client, ClientSession +from mcp.client.stdio import StdioServerParameters, stdio_client +from PIL import Image + +from mcp_sd15.errors import ConfigError +from mcp_sd15.model_manager import ModelManager +from mcp_sd15.pipeline_factory import PipelineFactory +from mcp_sd15.schema import TextToImageRequest +from mcp_sd15.server import create_server +from mcp_sd15.service import ImageGenerationService + + +@pytest.fixture +def lazy_backend(app_config, monkeypatch): + loads: list[int] = [] + calls: list[dict[str, object]] = [] + + def pipeline(**kwargs): + calls.append(kwargs) + return SimpleNamespace(images=[Image.new("RGB", (4, 4), "red")]) + + def create_base(self): + loads.append(threading.get_ident()) + return pipeline + + monkeypatch.setattr(PipelineFactory, "create_base", create_base) + backend = ModelManager(app_config) + monkeypatch.setattr(backend, "_embed_prompts", lambda *args: ("positive", "negative")) + monkeypatch.setattr(backend, "_generator", lambda seed: seed) + + def preprocess(image, **kwargs): + assert backend.base is pipeline # Base must load before any GPU/preprocessor work. + return image + + monkeypatch.setattr(backend, "_depth", preprocess) + monkeypatch.setattr(backend, "_openpose", preprocess) + monkeypatch.setattr(backend, "_inpaint_condition", lambda image, mask: preprocess(image)) + monkeypatch.setattr(backend, "_pipeline_for", lambda *args, **kwargs: pipeline) + return backend, loads, calls + + +async def test_discovery_does_not_load_model(app_config, lazy_backend) -> None: + backend, loads, _ = lazy_backend + assert backend.base is None + server = create_server(ImageGenerationService(app_config, backend)) + assert loads == [] + async with Client(server, mode="legacy") as client: + tools = (await client.list_tools()).tools + assert len(tools) == 3 + assert loads == [] + assert backend.base is None + + +@pytest.mark.parametrize( + ("tool", "arguments"), + [ + ("generate_image", {}), + ("generate_guided_image", {"control_mode": "depth"}), + ("generate_guided_image", {"control_mode": "openpose"}), + ("generate_guided_image", {"control_mode": "openpose_full"}), + ("inpaint_image", {}), + ], +) +async def test_first_generation_loads_once_in_worker( + app_config, lazy_backend, image_base64, tool, arguments, caplog +) -> None: + caplog.set_level(logging.INFO, logger="mcp_sd15.model_manager") + backend, loads, calls = lazy_backend + arguments = {"positive_prompt": "cat", **arguments} + if tool != "generate_image": + arguments["input_image"] = image_base64 + if tool == "inpaint_image": + arguments["mask_image"] = image_base64 + server = create_server(ImageGenerationService(app_config, backend)) + async with Client(server) as client: + for _ in range(2): + result = await client.call_tool(tool, arguments) + assert result.is_error is False + assert result.content[0].type == "image" + assert len(loads) == 1 + assert loads[0] != threading.get_ident() + assert len(calls) == 2 + assert calls[0]["num_inference_steps"] == app_config.generation.steps + assert calls[0]["guidance_scale"] == app_config.generation.cfg_scale + assert caplog.text.count("model load started") == 1 + assert caplog.text.count("model load completed") == 1 + + +async def test_simultaneous_requests_serialize_load_and_generation( + app_config, lazy_backend +) -> None: + backend, loads, calls = lazy_backend + service = ImageGenerationService(app_config, backend) + await asyncio.gather( + *(service.generate_image(TextToImageRequest.create("cat")) for _ in range(4)) + ) + assert len(loads) == 1 + assert len(calls) == 4 + + +def test_load_is_thread_safe_and_idempotent(app_config, monkeypatch) -> None: + loads: list[int] = [] + pipeline = object() + + def create_base(self): + loads.append(threading.get_ident()) + time.sleep(0.03) + return pipeline + + monkeypatch.setattr(PipelineFactory, "create_base", create_base) + backend = ModelManager(app_config) + with ThreadPoolExecutor(max_workers=8) as executor: + list(executor.map(lambda _: backend.load(), range(8))) + backend.load() + assert len(loads) == 1 + assert backend.base is pipeline + + +@pytest.mark.parametrize( + ("tool", "arguments"), + [ + ("generate_image", {"positive_prompt": ""}), + ("generate_image", {"size": "invalid"}), + ("generate_guided_image", {"control_mode": "invalid", "input_image": "invalid"}), + ("generate_guided_image", {"control_mode": "depth", "input_image": "invalid"}), + ("inpaint_image", {"input_image": "invalid", "mask_image": "invalid"}), + ], +) +async def test_invalid_requests_do_not_load(app_config, lazy_backend, tool, arguments) -> None: + backend, loads, _ = lazy_backend + server = create_server(ImageGenerationService(app_config, backend)) + async with Client(server) as client: + result = await client.call_tool(tool, {"positive_prompt": "cat", **arguments}) + assert result.is_error is True + assert loads == [] + + +async def test_invalid_mask_does_not_load(app_config, lazy_backend, image_base64) -> None: + backend, loads, _ = lazy_backend + async with Client(create_server(ImageGenerationService(app_config, backend))) as client: + result = await client.call_tool( + "inpaint_image", + {"positive_prompt": "cat", "input_image": image_base64, "mask_image": "invalid"}, + ) + assert result.is_error is True + assert loads == [] + + +@pytest.mark.parametrize("tool", ["generate_guided_image", "inpaint_image"]) +async def test_disabled_controlnet_does_not_load(app_config, lazy_backend, image_base64, tool): + backend, loads, _ = lazy_backend + config = replace( + app_config, + controlnet=replace( + app_config.controlnet, + depth=replace(app_config.controlnet.depth, enabled=False), + inpaint=replace(app_config.controlnet.inpaint, enabled=False), + ), + ) + arguments = {"positive_prompt": "cat", "input_image": image_base64} + if tool == "generate_guided_image": + arguments["control_mode"] = "depth" + else: + arguments["mask_image"] = image_base64 + async with Client(create_server(ImageGenerationService(config, backend))) as client: + result = await client.call_tool(tool, arguments) + assert result.is_error is True + assert loads == [] + + +async def test_load_failure_is_tool_error_and_next_call_retries( + app_config, lazy_backend, monkeypatch, caplog +): + backend, loads, _ = lazy_backend + create_base = backend.factory.create_base + attempts = 0 + + def fail_once(): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise ConfigError("test model could not load") + return create_base() + + monkeypatch.setattr(backend.factory, "create_base", fail_once) + async with Client(create_server(ImageGenerationService(app_config, backend))) as client: + result = await client.call_tool("generate_image", {"positive_prompt": "cat"}) + assert result.is_error is True + assert "test model could not load" in caplog.text + assert backend.base is None + assert len((await client.list_tools()).tools) == 3 + result = await client.call_tool("generate_image", {"positive_prompt": "cat"}) + assert result.is_error is False + assert attempts == 2 + assert len(loads) == 1 + + +async def test_cli_stdio_discovery_forbids_inference_runtime(tmp_path: Path) -> None: + checkpoint = tmp_path / "checkpoints" / "placeholder.safetensors" + checkpoint.parent.mkdir() + checkpoint.touch() + config = tmp_path / "config.toml" + config.write_text('[model]\nroot = "."\ncheckpoint = "placeholder.safetensors"\n') + # Exercise the real CLI while making any inference import or model load fatal. + bootstrap = """ +import importlib.abc +import sys +class NoInference(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.split('.')[0] in {'torch', 'diffusers', 'controlnet_aux', 'sd_embed'}: + raise AssertionError('inference imported during discovery: ' + fullname) +sys.meta_path.insert(0, NoInference()) +from mcp_sd15.pipeline_factory import PipelineFactory +def forbidden(self): + raise AssertionError('model loaded during discovery') +PipelineFactory.create_base = forbidden +from mcp_sd15.cli import main +raise SystemExit(main()) +""" + params = StdioServerParameters( + command=sys.executable, + args=["-c", bootstrap, "--config", str(config)], + ) + async with ( + stdio_client(params) as (read, write), + ClientSession(read, write, read_timeout_seconds=15) as client, + ): + await client.initialize() + tools = (await client.list_tools()).tools + assert {tool.name for tool in tools} == { + "generate_image", "generate_guided_image", "inpaint_image" + } |