feat(lora): cache LoraManager's layout behind file signatures (#362, phase 1)
A new py/lora_layout_cache.py keeps one (install, layout) entry and reuses it while the configured path, the plugin settings file's and the settings file's signatures (st_mtime_ns, st_size), the install and every media root's real path are unchanged and the entry is younger than 30 s. A load whose watched file changed within 3 s (git's racy rule), could not be checked, or whose settings were unreadable serves its request only. clear() bumps a generation so a load that started earlier stores nothing. Hits take no lock. Not wired yet; phase 2 puts it behind the image hot paths.
This commit is contained in:
Vito Sansevero
@@ -0,0 +1,442 @@
|
||||
"""The LoraManager layout cache for the image hot paths (#362).
|
||||
|
||||
``py/lora_layout_cache.py``. Every test injects the install resolver, the
|
||||
layout loader (both counting calls), the wall clock, the monotonic clock and,
|
||||
where needed, ``stat`` and ``pathmod``, over a temporary folder. No test
|
||||
sleeps: modification times are set with ``os.utime`` and the clocks are
|
||||
plain attributes the test moves.
|
||||
"""
|
||||
|
||||
import ntpath
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path, PureWindowsPath
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from py import lora_layout_cache as cache_module
|
||||
from py.lora_layout_cache import (
|
||||
MAX_AGE_SECONDS,
|
||||
MISSING,
|
||||
RACY_SECONDS,
|
||||
LayoutCache,
|
||||
)
|
||||
from py.lora_manager_layout import UNREADABLE, LoraManagerLayout, _canonical
|
||||
|
||||
SECOND_NS = 10**9
|
||||
|
||||
|
||||
def make_layout(settings_path=None, media_roots=(), error=None, source="user config"):
|
||||
return LoraManagerLayout(
|
||||
settings_path=settings_path,
|
||||
source=source,
|
||||
error=error,
|
||||
mode="alongside",
|
||||
lora_roots=tuple(media_roots),
|
||||
sidecar_root=None,
|
||||
example_images_root=None,
|
||||
library="default",
|
||||
media_roots=tuple(media_roots),
|
||||
)
|
||||
|
||||
|
||||
class CacheTestCase(unittest.TestCase):
|
||||
"""An install, its plugin settings file, a settings file and a LoRA root."""
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = Path(tempfile.mkdtemp()).resolve()
|
||||
self.addCleanup(shutil.rmtree, self.tmp, True)
|
||||
self.install_dir = self.tmp / "custom_nodes" / "ComfyUI-Lora-Manager"
|
||||
self.install_dir.mkdir(parents=True)
|
||||
self.plugin_settings = self.install_dir / "settings.json"
|
||||
self.settings = self.tmp / "user-config" / "settings.json"
|
||||
self.settings.parent.mkdir()
|
||||
self.root_dir = self.tmp / "loras"
|
||||
self.root_dir.mkdir()
|
||||
# The wall clock is an hour past every file written here, so nothing
|
||||
# is racy unless a test sets a recent modification time.
|
||||
self.now_ns = time.time_ns() + 3600 * SECOND_NS
|
||||
self.mono = 1000.0
|
||||
self.resolved = []
|
||||
self.loaded = []
|
||||
self.install = _canonical(str(self.install_dir))
|
||||
self.layout = make_layout(
|
||||
settings_path=self.settings,
|
||||
media_roots=(_canonical(str(self.root_dir)),),
|
||||
)
|
||||
self.write(self.settings, "{}")
|
||||
|
||||
def write(self, path, text, mtime_ns=None):
|
||||
Path(path).write_text(text, encoding="utf-8")
|
||||
stamp = mtime_ns if mtime_ns is not None else self.now_ns - 600 * SECOND_NS
|
||||
os.utime(path, ns=(stamp, stamp))
|
||||
|
||||
def resolve(self, key):
|
||||
self.resolved.append(key)
|
||||
return self.install
|
||||
|
||||
def load(self, install):
|
||||
self.loaded.append(install)
|
||||
return self.layout
|
||||
|
||||
def cache(self, **kwargs):
|
||||
kwargs.setdefault("wall_ns", lambda: self.now_ns)
|
||||
kwargs.setdefault("monotonic", lambda: self.mono)
|
||||
return LayoutCache(
|
||||
kwargs.pop("resolve", self.resolve), kwargs.pop("load", self.load), **kwargs
|
||||
)
|
||||
|
||||
|
||||
class ValidationTests(CacheTestCase):
|
||||
def test_hit_reuses_the_entry_without_resolving_or_loading(self):
|
||||
cache = self.cache()
|
||||
first = cache.get("")
|
||||
second = cache.get("")
|
||||
self.assertEqual(len(self.resolved), 1)
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
self.assertIs(first[1], second[1])
|
||||
self.assertIs(first[0], second[0])
|
||||
self.assertEqual(first, (self.install, self.layout))
|
||||
|
||||
def test_settings_size_change_reloads_once(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
stamp = self.settings.stat().st_mtime_ns
|
||||
self.write(self.settings, '{"a": 1}', mtime_ns=stamp)
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_settings_mtime_change_reloads_once(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
stamp = self.settings.stat().st_mtime_ns + 1
|
||||
os.utime(self.settings, ns=(stamp, stamp))
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_settings_deletion_reloads_once(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.settings.unlink()
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_settings_creation_reloads_once(self):
|
||||
self.settings.unlink()
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.write(self.settings, "{}")
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_plugin_settings_change_reloads_once(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.write(self.plugin_settings, '{"use_portable_settings": true}')
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_a_different_configured_path_reloads(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
cache.get("other")
|
||||
self.assertEqual(self.resolved, ["", "other"])
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_configured_path_is_stripped_for_the_key_and_the_resolver(self):
|
||||
cache = self.cache()
|
||||
cache.get(" x ")
|
||||
cache.get("x")
|
||||
self.assertEqual(self.resolved, ["x"])
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
|
||||
def test_a_recently_modified_file_is_not_stored(self):
|
||||
self.write(self.settings, "{}", mtime_ns=self.now_ns - SECOND_NS)
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
self.assertIsNone(cache._entry)
|
||||
self.now_ns += RACY_SECONDS * SECOND_NS
|
||||
cache.get("")
|
||||
self.assertIsNotNone(cache._entry)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 3)
|
||||
|
||||
def test_racy_boundary(self):
|
||||
# Exactly RACY_SECONDS old is trusted; one nanosecond newer is not.
|
||||
boundary = self.now_ns - RACY_SECONDS * SECOND_NS
|
||||
self.write(self.settings, "{}", mtime_ns=boundary + 1)
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.assertIsNone(cache._entry)
|
||||
os.utime(self.settings, ns=(boundary, boundary))
|
||||
cache.get("")
|
||||
self.assertIsNotNone(cache._entry)
|
||||
|
||||
def test_a_modification_time_in_the_future_is_racy(self):
|
||||
self.write(self.plugin_settings, "{}", mtime_ns=self.now_ns + 60 * SECOND_NS)
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
self.assertIsNone(cache._entry)
|
||||
|
||||
def test_a_missing_file_is_not_racy(self):
|
||||
self.assertFalse(self.plugin_settings.exists())
|
||||
self.settings.unlink()
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.assertIsNotNone(cache._entry)
|
||||
signatures = dict(cache._entry.signatures)
|
||||
self.assertTrue(all(value is MISSING for value in signatures.values()))
|
||||
self.assertEqual(len(signatures), 2)
|
||||
|
||||
def test_entry_is_reused_below_the_maximum_age_and_reloaded_at_it(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
stored_at = self.mono
|
||||
self.mono = stored_at + MAX_AGE_SECONDS - 0.5
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
self.mono = stored_at + MAX_AGE_SECONDS
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_a_wall_clock_moved_backward_does_not_extend_the_age(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.now_ns -= 3600 * SECOND_NS
|
||||
self.mono += MAX_AGE_SECONDS
|
||||
# The old entry is too old; the new load sees files "from the future"
|
||||
# (racy) and is not stored, but it does load.
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_stat_failure_on_a_hit_reloads(self):
|
||||
failing = {"on": False}
|
||||
|
||||
def stat(path):
|
||||
if failing["on"]:
|
||||
raise PermissionError(13, "sharing violation", path)
|
||||
return os.stat(path)
|
||||
|
||||
cache = self.cache(stat=stat)
|
||||
cache.get("")
|
||||
failing["on"] = True
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
self.assertIsNone(cache._entry)
|
||||
|
||||
def test_stat_failure_on_a_miss_leaves_the_layout_unstored(self):
|
||||
def stat(path):
|
||||
raise PermissionError(13, "sharing violation", path)
|
||||
|
||||
cache = self.cache(stat=stat)
|
||||
self.assertEqual(cache.get(""), (self.install, self.layout))
|
||||
self.assertIsNone(cache._entry)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_an_unreadable_settings_layout_is_not_stored(self):
|
||||
self.layout = make_layout(settings_path=self.settings, error=UNREADABLE)
|
||||
cache = self.cache()
|
||||
self.assertEqual(cache.get("")[1].error, UNREADABLE)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
self.assertIsNone(cache._entry)
|
||||
|
||||
def test_other_layout_errors_are_stored(self):
|
||||
self.layout = make_layout(
|
||||
settings_path=self.settings, error="the settings file is not valid JSON"
|
||||
)
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
|
||||
def _symlink(self, link, target):
|
||||
try:
|
||||
os.symlink(target, link, target_is_directory=True)
|
||||
except (OSError, NotImplementedError) as error:
|
||||
self.skipTest(f"symlinks are not available here: {error}")
|
||||
|
||||
def test_a_media_root_swapped_for_a_symlink_reloads(self):
|
||||
elsewhere = self.tmp / "elsewhere"
|
||||
elsewhere.mkdir()
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
shutil.rmtree(self.root_dir)
|
||||
self._symlink(self.root_dir, elsewhere)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_an_install_swapped_for_a_symlink_reloads(self):
|
||||
elsewhere = self.tmp / "elsewhere-lm"
|
||||
elsewhere.mkdir()
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
moved = self.tmp / "moved-lm"
|
||||
self.install_dir.rename(moved)
|
||||
self._symlink(self.install_dir, elsewhere)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_a_root_that_no_longer_resolves_reloads(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
with patch.object(
|
||||
cache_module, "_canonical", side_effect=OSError("unreachable share")
|
||||
):
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_a_load_failure_propagates_and_stores_nothing(self):
|
||||
calls = []
|
||||
|
||||
def load(install):
|
||||
calls.append(install)
|
||||
if len(calls) == 1:
|
||||
raise RuntimeError("boom")
|
||||
return self.layout
|
||||
|
||||
cache = self.cache(load=load)
|
||||
with self.assertRaises(RuntimeError):
|
||||
cache.get("")
|
||||
self.assertIsNone(cache._entry)
|
||||
self.assertEqual(cache.get(""), (self.install, self.layout))
|
||||
self.assertEqual(len(calls), 2)
|
||||
|
||||
def test_without_an_install_nothing_is_watched_and_the_entry_ages_out(self):
|
||||
stats = []
|
||||
|
||||
def stat(path):
|
||||
stats.append(path)
|
||||
return os.stat(path)
|
||||
|
||||
self.layout = make_layout(media_roots=self.layout.media_roots, source=None)
|
||||
cache = self.cache(resolve=lambda key: self.resolved.append(key), stat=stat)
|
||||
self.assertEqual(cache.get(""), (None, self.layout))
|
||||
self.assertEqual(cache._entry.signatures, ())
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
self.assertEqual(stats, [])
|
||||
self.mono += MAX_AGE_SECONDS
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_windows_paths_differing_only_in_case_are_watched_once(self):
|
||||
stats = []
|
||||
|
||||
def stat(path):
|
||||
stats.append(path)
|
||||
return SimpleNamespace(st_mtime_ns=self.now_ns - 600 * SECOND_NS, st_size=2)
|
||||
|
||||
install = _canonical("C:\\ComfyUI\\custom_nodes\\LM", ntpath)
|
||||
self.layout = make_layout(
|
||||
settings_path=PureWindowsPath(
|
||||
"C:\\ComfyUI\\custom_nodes\\lm\\SETTINGS.JSON"
|
||||
)
|
||||
)
|
||||
cache = self.cache(resolve=lambda key: install, stat=stat, pathmod=ntpath)
|
||||
cache.get("")
|
||||
self.assertEqual(len(cache._entry.signatures), 1)
|
||||
self.assertEqual(len(stats), 1)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
self.assertEqual(len(stats), 2)
|
||||
|
||||
|
||||
class ConcurrencyTests(CacheTestCase):
|
||||
def test_a_load_that_started_before_clear_stores_nothing(self):
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def load(install):
|
||||
self.loaded.append(install)
|
||||
if len(self.loaded) == 1:
|
||||
started.set()
|
||||
self.assertTrue(release.wait(10))
|
||||
return self.layout
|
||||
|
||||
cache = self.cache(load=load)
|
||||
results = []
|
||||
worker = threading.Thread(target=lambda: results.append(cache.get("")))
|
||||
worker.start()
|
||||
self.assertTrue(started.wait(10))
|
||||
cache.clear()
|
||||
release.set()
|
||||
worker.join(10)
|
||||
self.assertFalse(worker.is_alive())
|
||||
self.assertEqual(results, [(self.install, self.layout)])
|
||||
self.assertIsNone(cache._entry)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
self.assertIsNotNone(cache._entry)
|
||||
|
||||
def test_concurrent_misses_all_get_a_layout_and_one_entry_is_stored(self):
|
||||
cache = self.cache()
|
||||
count = 6
|
||||
barrier = threading.Barrier(count)
|
||||
results = []
|
||||
lock = threading.Lock()
|
||||
|
||||
def worker():
|
||||
barrier.wait(10)
|
||||
result = cache.get("")
|
||||
with lock:
|
||||
results.append(result)
|
||||
|
||||
threads = [threading.Thread(target=worker) for _ in range(count)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join(10)
|
||||
self.assertEqual(results, [(self.install, self.layout)] * count)
|
||||
self.assertIsNotNone(cache._entry)
|
||||
loads = len(self.loaded)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), loads)
|
||||
|
||||
def test_a_valid_hit_takes_no_lock(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
results = []
|
||||
with cache._lock:
|
||||
worker = threading.Thread(target=lambda: results.append(cache.get("")))
|
||||
worker.start()
|
||||
worker.join(10)
|
||||
alive = worker.is_alive()
|
||||
worker.join(10)
|
||||
self.assertFalse(alive)
|
||||
self.assertEqual(results, [(self.install, self.layout)])
|
||||
|
||||
def test_clear_forces_a_reload(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
cache.clear()
|
||||
self.assertIsNone(cache._entry)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
|
||||
class ModuleTests(unittest.TestCase):
|
||||
def test_constants(self):
|
||||
self.assertEqual(cache_module.RACY_SECONDS, 3)
|
||||
self.assertEqual(cache_module.MAX_AGE_SECONDS, 30)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user