diff --git a/py/api/images.py b/py/api/images.py index 45196f0f..bfafcca0 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -2272,7 +2272,8 @@ class ImageRoutesMixin: ) image_path = Path(image["image_path"]) - allowed_dirs = list(self._get_all_output_dirs()) + self._lora_image_dirs() + lora_dirs = await self._run_in_executor(self._lora_image_dirs) + allowed_dirs = list(self._get_all_output_dirs()) + lora_dirs denied = _media_access_error(image_path, allowed_dirs) if denied is not None: @@ -2320,7 +2321,10 @@ class ImageRoutesMixin: LoraManager's guarded LoRA roots, sidecar root and example-images root (``lora_manager_layout.media_roots``, #355) and PromptManager's CivitAI download cache. An unusable LoraManager layout gives the - default LoRA roots and the cache. + default LoRA roots and the cache. The layout comes from the layout + cache (``lora_utils.cached_lora_manager_layout``, #362), which reloads + when LoraManager's settings files change; blocking, so the routes + call it in the executor. """ try: from ..config import IntegrationConfig @@ -2328,25 +2332,18 @@ class ImageRoutesMixin: if not IntegrationConfig.LORA_MANAGER_ENABLED: return [] from ..lora_utils import ( + cached_lora_manager_layout, get_lora_image_cache_dir, - layout_for_install, - lora_manager_layout, ) - lora_root = IntegrationConfig.LORA_MANAGER_PATH or "" - if lora_root.strip(): + configured = IntegrationConfig.LORA_MANAGER_PATH or "" + install, layout = cached_lora_manager_layout(configured) + if configured.strip() and install is None: # A configured path comes from config.json, which can be hand # edited: only a LoraManager install under custom_nodes may # widen the serving allow-list. An empty path keeps the # auto-detection, which itself only scans custom_nodes. - from ..lora_utils import resolve_lora_manager_path - - lora_root = resolve_lora_manager_path(lora_root) - if not lora_root: - return [] - layout = layout_for_install(lora_root) - else: - layout = lora_manager_layout("") + return [] dirs = [Path(d) for d in layout.media_roots] # Named, not made: listing and serving never create the cache. diff --git a/py/api/lora_import_routes.py b/py/api/lora_import_routes.py index cb31341f..a09ec689 100644 --- a/py/api/lora_import_routes.py +++ b/py/api/lora_import_routes.py @@ -194,8 +194,13 @@ class LoraImportRoutesMixin: return False async def replace_lora_imports(self, db, staged, *, expected_ids, allow_empty): - """Swap the collection for ``staged.records`` in one transaction.""" - return await self._run_in_executor( + """Swap the collection for ``staged.records`` in one transaction. + + Then the image routes' layout cache is dropped (#362): the import may + have linked files under a folder LoraManager created after the cache + stored its layout (its first example-image download creates one). + """ + result = await self._run_in_executor( lambda: db.replace_category_prompts( CATEGORY, staged.records, @@ -203,6 +208,8 @@ class LoraImportRoutesMixin: allow_empty=allow_empty, ) ) + lora_utils.clear_layout_cache() + return result # ── preview ─────────────────────────────────────────────────────── diff --git a/py/lora_utils.py b/py/lora_utils.py index 5d334898..0345d417 100644 --- a/py/lora_utils.py +++ b/py/lora_utils.py @@ -30,6 +30,7 @@ except ImportError: from dataclasses import replace +from .lora_layout_cache import LayoutCache from .lora_manager_layout import ( MODE_ALONGSIDE, METADATA_SUFFIX, @@ -295,10 +296,35 @@ def layout_for_install(install: Optional[str]) -> LoraManagerLayout: def lora_manager_layout(lora_manager_path: Optional[str]) -> LoraManagerLayout: - """Where LoraManager keeps its files (``py/lora_manager_layout.py``).""" + """Where LoraManager keeps its files (``py/lora_manager_layout.py``). + + A fresh read: the import review, Start, the legacy scan and the trigger + word cache use it. The image hot paths use ``cached_lora_manager_layout``. + """ return layout_for_install(lora_manager_install(lora_manager_path)) +# The image hot paths' layout cache (#362). The lambdas look the functions up +# when called, so a test that patches them here is seen by the cache. +_layout_cache = LayoutCache( + lambda path: lora_manager_install(path), + lambda install: layout_for_install(install), +) + + +def cached_lora_manager_layout( + configured: Optional[str], +) -> Tuple[Optional[str], LoraManagerLayout]: + """``(install, layout)`` for the configured path, from the layout cache + while LoraManager's settings files are unchanged (``py/lora_layout_cache.py``).""" + return _layout_cache.get(configured or "") + + +def clear_layout_cache() -> None: + """Drop the cached layout (after an import linked files in new folders).""" + _layout_cache.clear() + + def _get_extra_lora_paths(comfyui_root: Path) -> List[Path]: """Parse extra_model_paths.yaml for additional LoRA directories.""" results = [] diff --git a/tests/test_api_images.py b/tests/test_api_images.py index 5b0fa3e3..896416e3 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -9,9 +9,11 @@ import asyncio import json import os import pathlib +import shutil import sys import tempfile import threading +import time import types import unittest from contextlib import ExitStack @@ -36,6 +38,8 @@ from PIL import Image # noqa: E402 from database.operations import PromptDatabase # noqa: E402 import py.api.images as images_module from py.api import gallery_listing # noqa: E402 +from py import lora_utils # noqa: E402 +from py import lora_manager_layout as layout_module # noqa: E402 from py.api.thumbnail_settings import ThumbnailSettings # noqa: E402 from py.api import PromptManagerAPI # noqa: E402 from py.api.prompts import MAX_PAGE_LIMIT # noqa: E402 @@ -111,13 +115,25 @@ class ImageAPITestCase(AioHTTPTestCase): self.assertGreater(image_id, 0, "fixture: image link failed") return image_id - def _make_symlink(self, target, link): + def _make_symlink(self, target, link, target_is_directory=False): """Create a symlink or skip the test when the platform refuses.""" try: - os.symlink(str(target), str(link)) + os.symlink(str(target), str(link), target_is_directory=target_is_directory) except (OSError, NotImplementedError, AttributeError) as exc: self.skipTest(f"symlinks unavailable: {exc}") + def _spy_executor(self): + """Record the __name__ of every callable sent to _run_in_executor.""" + original = self.api._run_in_executor + calls = [] + + async def spy(func, *args, **kwargs): + calls.append(getattr(func, "__name__", repr(func))) + return await original(func, *args, **kwargs) + + self.api._run_in_executor = spy + return calls + class TestServeOutputImage(ImageAPITestCase): """GET /prompt_manager/images/serve/{filepath}""" @@ -388,18 +404,6 @@ class ThumbnailTestCase(ImageAPITestCase): PROGRESS_URL = "/prompt_manager/images/generate-thumbnails/progress" - def _spy_executor(self): - """Record the __name__ of every callable sent to _run_in_executor.""" - original = self.api._run_in_executor - calls = [] - - async def spy(func, *args, **kwargs): - calls.append(getattr(func, "__name__", repr(func))) - return await original(func, *args, **kwargs) - - self.api._run_in_executor = spy - return calls - def _thumb_path(self, name, suffix=".webp"): return self.output_dir / "thumbnails" / f"{name}_thumb{suffix}" @@ -1216,6 +1220,13 @@ def _raise(*_args, **_kwargs): class ImageRouteCoverageCase(ImageAPITestCase): """Helpers for driving success and failure branches of each route.""" + def setUp(self): + # #362: the LoraManager layout cache is per process; every test here + # starts and ends with it empty (runs before the subclasses' setUpAsync). + super().setUp() + lora_utils.clear_layout_cache() + self.addCleanup(lora_utils.clear_layout_cache) + def _break_db(self, db_method): setattr(self.api.db, db_method, _raise) @@ -1327,6 +1338,9 @@ class TestServeErrorBranches(ImageRouteCoverageCase): ) self.assertEqual(resp.status, 200) + # The cache does not watch default_lora_roots (#362): without a + # clear the next request would reuse the folders above. + lora_utils.clear_layout_cache() with ( patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True), patch.object(IntegrationConfig, "LORA_MANAGER_PATH", ""), @@ -1419,32 +1433,39 @@ class TestLoraManagerMediaServing(ImageRouteCoverageCase): self.cache = root / "cache" for folder in (self.sidecars, self.examples, self.cache): folder.mkdir() + self.root = root self.write_settings(str(self.sidecars)) self.enabled = patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True) + self.default_roots = patch("py.lora_utils.default_lora_roots", return_value=[]) for item in ( self.enabled, patch.object(IntegrationConfig, "LORA_MANAGER_PATH", str(self.lm)), patch( "py.lora_utils.custom_nodes_directories", return_value=[custom_nodes] ), - patch("py.lora_utils.default_lora_roots", return_value=[]), + self.default_roots, patch("py.lora_utils.get_lora_image_cache_dir", return_value=self.cache), patch.dict(os.environ, {"LORA_MANAGER_SETTINGS_DIR": str(self.settings)}), ): item.start() self.addCleanup(item.stop) - def write_settings(self, sidecar_root): - (self.settings / "settings.json").write_text( + def write_settings(self, sidecar_root, examples=None): + path = self.settings / "settings.json" + path.write_text( json.dumps( { "sidecar_storage_mode": "centralized", "sidecar_storage_path": sidecar_root, - "example_images_path": str(self.examples), + "example_images_path": str(examples or self.examples), "civitai_api_key": "never-served", } ) ) + # Ten seconds old, so the layout cache (real clock here) stores the + # first load instead of treating the file as being written (#362). + stamp = time.time_ns() - 10 * 10**9 + os.utime(path, ns=(stamp, stamp)) async def _status(self, image_id): resp = await self.client.request( @@ -1464,11 +1485,62 @@ class TestLoraManagerMediaServing(ImageRouteCoverageCase): for image_id in ids: self.assertEqual(await self._status(image_id), 200) + # A warm layout cache does not outlive the integration (#362). + self.assertIsNotNone(lora_utils._layout_cache._entry) self.enabled.stop() for image_id in ids: self.assertEqual(await self._status(image_id), 403) self.enabled.start() + def _example(self, folder=None, name="1.mp4"): + path = (folder or self.examples) / "ab12" / name + path.parent.mkdir(exist_ok=True) + path.write_bytes(b"video") + return path + + async def test_unchanged_settings_reuse_the_cached_layout(self): + image_id = self._link_image(self._example()) + self.assertEqual(await self._status(image_id), 200) + self.assertIsNotNone(lora_utils._layout_cache._entry) + self.assertEqual(await self._status(image_id), 200) + self.assertEqual(lora_utils.default_lora_roots.call_count, 1) + + async def test_a_settings_change_is_seen_without_a_clear(self): + image_id = self._link_image(self._example()) + self.assertEqual(await self._status(image_id), 200) + self.assertIsNotNone(lora_utils._layout_cache._entry) + moved = self.root / "examples-moved" + moved.mkdir() + self.write_settings(str(self.sidecars), examples=moved) + self.assertEqual(await self._status(image_id), 403) + + async def test_a_cached_folder_swapped_for_a_symlink_to_home_is_refused(self): + # CUJ 7: the example-images folder becomes a symlink to a folder that + # holds the home folder; the same recorded file is refused. + outer = self.root / "outer" + (outer / "home").mkdir(parents=True) + self._example(outer) + image_id = self._link_image(self._example()) + with patch.object(layout_module.Path, "home", return_value=outer / "home"): + self.assertEqual(await self._status(image_id), 200) + self.assertIsNotNone(lora_utils._layout_cache._entry) + shutil.rmtree(self.examples) + self._make_symlink(outer, self.examples, target_is_directory=True) + self.assertEqual(await self._status(image_id), 403) + + async def test_a_configured_path_outside_custom_nodes_adds_nothing(self): + from py.config import IntegrationConfig + + self.assertEqual(await self._status(self._link_image(self._example())), 200) + self.assertIsNotNone(lora_utils._layout_cache._entry) + with patch.object(IntegrationConfig, "LORA_MANAGER_PATH", str(self.sidecars)): + self.assertEqual(self.api._lora_image_dirs(), []) + + async def test_serve_image_builds_the_allow_list_in_the_executor(self): + calls = self._spy_executor() + self.assertEqual(await self._status(self._link_image(self._example())), 200) + self.assertIn("_lora_image_dirs", calls) + async def test_a_filesystem_root_as_sidecar_folder_adds_nothing(self): self.write_settings(os.path.abspath(os.sep)) dirs = self.api._lora_image_dirs() diff --git a/tests/test_api_lora.py b/tests/test_api_lora.py index 6b7ec497..8fda0562 100644 --- a/tests/test_api_lora.py +++ b/tests/test_api_lora.py @@ -876,6 +876,18 @@ class TestLoraScanRoute(LoraAPITestCase): resp = await self.client.request("POST", "/prompt_manager/lora/scan") self.assertEqual(resp.status, 400) + async def test_a_finished_scan_clears_the_layout_cache(self): + # #362: the legacy scan replaces through the same method. + self._enable() + loras = self.tmp_root / "loras" + loras.mkdir() + _write_metadata(loras, "plain") + (loras / "plain.png").write_bytes(_png_bytes()) + with patch("py.lora_utils.clear_layout_cache") as clear: + resp, events, _calls = await self._scan(loras) + self.assertEqual(events[-1]["type"], "complete", events) + clear.assert_called_once_with() + async def test_scan_imports_loras_with_images_and_links_existing(self): self._enable() loras = self.tmp_root / "loras" diff --git a/tests/test_lora_import.py b/tests/test_lora_import.py index f5ff24d6..e3af65ed 100644 --- a/tests/test_lora_import.py +++ b/tests/test_lora_import.py @@ -529,6 +529,16 @@ class LoraImportRouteTests(LoraImportAPITestCase): resp, body = await self._post(APPLY, {"token": plan["token"]}) self.assertEqual((resp.status, body["code"]), (409, "plan-used")) + async def test_a_finished_import_clears_the_layout_cache(self): + # #362: the import may link files under a folder LoraManager created + # after the image routes cached its layout. + self._two_loras() + with patch("py.lora_utils.clear_layout_cache") as clear: + _plan, job = await self._apply_reviewed() + await job.task + self.assertEqual(job.state, "finished", job.log) + clear.assert_called_once_with() + async def test_discovery_failure_keeps_previous_imports(self): self._two_loras() before = _snapshot(self.api.db) @@ -936,6 +946,19 @@ class LoraImportRefusalTests(LoraImportAPITestCase): self.assertEqual((job.state, job.scope["failure"]), ("failed", "plan-changed")) self.assertEqual(_snapshot(self.api.db), before) + async def test_a_refused_import_keeps_the_layout_cache(self): + self._two_loras() + resp, plan = await self._post(PREVIEW) + self.api.db.save_prompt( + text="late", category=CATEGORY, prompt_hash=generate_prompt_hash("late") + ) + with patch("py.lora_utils.clear_layout_cache") as clear: + resp, body = await self._post(APPLY, {"token": plan["token"]}) + job = self.api._jobs.get(body["job"]["id"]) + await job.task + self.assertEqual(job.scope["failure"], "plan-changed") + clear.assert_not_called() + async def test_a_pair_saved_after_the_review_blocks_the_job(self): self._two_loras() resp, plan = await self._post(PREVIEW)