diff options
| author | tanaka takashi <isazaka@gmail.com> | 2026-10-14 02:21:50 +0900 |
|---|---|---|
| committer | tanaka takashi <isazaka@gmail.com> | 2026-10-14 02:21:50 +0900 |
| commit | 583a56f1ac3d9be9ce2f8c1715d77423347d4a4e (patch) | |
| tree | ec2789b91bfa4cfded83fa64adf044347926ecbf /tests | |
| parent | 213e738140859dbcc12c2540ce40d6b6d2875aa7 (diff) | |
stdio MCPのツール一覧取得でも、受付開始前にModelManagerがSD本体・
外部VAE・LoRAをロードし、CUDAへ転送していた。画像生成しない通常の
質問でもこの起動コストが発生するため、有効な生成要求の実行時まで
モデル初期化を遅らせ、一覧取得時の不要な遅延とVRAM消費を除去する。
実装:
- ModelManagerの構築時は未ロード状態とし、冪等なload()を追加する。
- 初期化ロックで二重ロードを防ぎ、成功時だけモデルを保持する。
失敗時は未ロード状態を維持し、次の要求で再試行できるようにする。
- 3種類の生成メソッドからload()を呼び、既存サービスの入力検証・
生成ロックを経由したワーカースレッド内で初期化する。
- ControlNetと前処理モデルの既存の遅延ロード・モデル共有を維持する。
- 起動・設定読み込みと実際のモデルロードのログを分離し、ロードの
開始・完了・失敗をstderrへ出力する。
- benchmarkの計測区間で明示的にload()を呼び、ロード時間とVRAMの
計測の意味を維持する。READMEにロードタイミングと確認手順を追記する。
変更範囲:
- stdioのプロセス起動・終了方式を維持し、常駐化やHTTP切り替え、
チャット側の接続管理・一覧キャッシュ・アイドル解放は追加しない。
- MCPスキーマ、設定形式、CLI引数、プロンプト・seed処理、scheduler、
生成サイズ、PNGメタデータ・保存仕様は変更しない。
- 設定とモデルファイルの存在確認は起動時に残し、モデル互換性・
CUDA・推論依存関係のエラー検出は初回生成時へ移す。
- 同一プロセス内ではモデルを再利用し、終了時に解放する。
別の生成プロセスでの再ロードは引き続き許容する。
回帰テスト:
- 修正前の既存69件に対し、修正後は追加18件を含む87件すべて成功。
既存テストと期待値は変更していない。
- CLI/stdioでの一覧取得時にモデルロードと推論ライブラリのimportを
禁止しても成功することを確認する。
- 全生成方式の初回ロード、ワーカースレッド実行、同時要求、冪等性、
不正入力・無効なControlNet指定、失敗後の再試行、ログを検証する。
- ruff check .は成功、pyrightは0 errors / 0 warnings。
実GPU・stdioでの確認:
- 現在の設定(外部VAE、2件のLoRA、LCM、6 steps)を維持して計測した。
- 起動から一覧取得までの時間は8.233秒から0.671秒に短縮した。
- 一覧取得前後のGPU全体の使用量は修正前808→3233 MiBに対し、
修正後682→682 MiBで増加しなかった。
- 同じ接続で2枚生成してもベースモデルのロードは1回で、一覧取得では
0回だった。ロード時間は7.812秒で、初回生成要求に移った。
- 同じ設定・プロンプト・seedによる修正前後の画像は画素が完全一致。
PNG形式・寸法・メタデータと、返却PNGおよび保存PNGの一致を確認した。
- 最終版スクリプトでの再確認でも一覧取得0.655秒、初回生成14.089秒、
2枚目1.221秒で、画素・保存結果の一致と生成プロセス終了を確認した。
- 通常質問の親アプリの処理経路でも、LLMのみをスタブ化し、実際の
設定済みuvコマンドで一覧取得・3ツールの登録を確認した。
- 計測スクリプトと詳細な検証報告を追加する。生のJSON・stderr・PNGは
gitignore対象のdata/verification-lazy-loadに保存する。
確認の限界:
- Windowsのプロセス別VRAMはN/Aのため、GPU全体の使用量・プロセス一覧・
tasklistによる終了確認で評価した。全体値には他プロセスの変動がある。
- 計測値は単一セッションの観測であり、統計的な速度保証ではない。
- ブラウザ操作と実LLMを含むチャット全体、および現在の設定で無効な
ControlNetの実GPU生成は未実施。誘導生成・Inpaintのロード順は
スタブを用いた自動テストで確認した。
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" + } |
