diff --git a/lora_optimizer.py b/lora_optimizer.py index c81e9e9..0d27ffc 100644 --- a/lora_optimizer.py +++ b/lora_optimizer.py @@ -7687,6 +7687,66 @@ class LoRAAutoTuner(LoRAOptimizer): except OSError: pass + @staticmethod + def _lora_identity_hash(lora_item): + """16-char hex hash of a single LoRA's file identity (name+mtime+size).""" + name = lora_item["name"] + path = folder_paths.get_full_path("loras", name) + if path is not None: + try: + st = os.stat(path) + entry = (name, st.st_mtime, st.st_size) + except OSError: + entry = (name, 0, 0) + else: + entry = (name, 0, 0) + hash_input = json.dumps(entry, separators=(",", ":")) + return hashlib.sha256(hash_input.encode()).hexdigest()[:16] + + @staticmethod + def _lora_cache_path(lora_hash): + return os.path.join(AUTOTUNER_MEMORY_DIR, f"{lora_hash}.lora.json") + + @staticmethod + def _lora_cache_load(lora_hash): + """Load per-LoRA prefix stats. Returns per_prefix dict or None on miss/stale.""" + path = LoRAAutoTuner._lora_cache_path(lora_hash) + if not os.path.exists(path): + return None + try: + with open(path) as f: + data = json.load(f) + if data.get("algo_version") != AUTOTUNER_ALGO_VERSION: + logging.info("[AutoTuner Lora Cache] Stale algo version, ignoring") + return None + return data.get("per_prefix") + except Exception as e: + logging.warning(f"[AutoTuner Lora Cache] Failed to load: {e}") + return None + + @staticmethod + def _lora_cache_save(lora_hash, per_prefix): + """Atomic write of per-LoRA cache to disk.""" + from datetime import datetime + path = LoRAAutoTuner._lora_cache_path(lora_hash) + entry = { + "algo_version": AUTOTUNER_ALGO_VERSION, + "created_at": datetime.now().isoformat(timespec="seconds"), + "per_prefix": per_prefix, + } + tmp_path = path + ".tmp" + try: + with open(tmp_path, "w") as f: + json.dump(entry, f) + os.replace(tmp_path, path) + logging.info(f"[AutoTuner Lora Cache] Saved: {path}") + except Exception as e: + logging.warning(f"[AutoTuner Lora Cache] Failed to save: {e}") + try: + os.unlink(tmp_path) + except OSError: + pass + @staticmethod def _memory_load(lora_hash, settings_hash, requested_top_n): """Load and validate a memory entry. Returns tuner_data or None.""" diff --git a/tests/test_lora_optimizer.py b/tests/test_lora_optimizer.py index 14e259b..4eb4b7a 100644 --- a/tests/test_lora_optimizer.py +++ b/tests/test_lora_optimizer.py @@ -1749,5 +1749,97 @@ class TestAnalysisPartialLifecycle(unittest.TestCase): self.assertFalse(os.path.exists(partial_path)) +class TestLoraCacheIO(unittest.TestCase): + + def test_lora_identity_hash_returns_16char_hex(self): + with tempfile.TemporaryDirectory() as tmpdir: + # Create a fake lora file so os.stat works + lora_path = os.path.join(tmpdir, "test.safetensors") + with open(lora_path, "w") as f: + f.write("x") + with mock.patch("lora_optimizer.folder_paths.get_full_path", + return_value=lora_path): + h = lora_optimizer.LoRAAutoTuner._lora_identity_hash( + {"name": "test.safetensors", "strength": 1.0}) + self.assertIsInstance(h, str) + self.assertEqual(len(h), 16) + self.assertTrue(all(c in "0123456789abcdef" for c in h)) + + def test_lora_identity_hash_differs_for_different_files(self): + with tempfile.TemporaryDirectory() as tmpdir: + path_a = os.path.join(tmpdir, "a.safetensors") + path_b = os.path.join(tmpdir, "b.safetensors") + with open(path_a, "w") as f: f.write("x") + with open(path_b, "w") as f: f.write("y") + def fake_get_full_path(folder, name): + return path_a if name == "a.safetensors" else path_b + with mock.patch("lora_optimizer.folder_paths.get_full_path", + side_effect=fake_get_full_path): + ha = lora_optimizer.LoRAAutoTuner._lora_identity_hash( + {"name": "a.safetensors", "strength": 1.0}) + hb = lora_optimizer.LoRAAutoTuner._lora_identity_hash( + {"name": "b.safetensors", "strength": 1.0}) + self.assertNotEqual(ha, hb) + + def test_lora_identity_hash_ignores_strength(self): + with tempfile.TemporaryDirectory() as tmpdir: + lora_path = os.path.join(tmpdir, "test.safetensors") + with open(lora_path, "w") as f: f.write("x") + with mock.patch("lora_optimizer.folder_paths.get_full_path", + return_value=lora_path): + h1 = lora_optimizer.LoRAAutoTuner._lora_identity_hash( + {"name": "test.safetensors", "strength": 0.5}) + h2 = lora_optimizer.LoRAAutoTuner._lora_identity_hash( + {"name": "test.safetensors", "strength": 1.0}) + self.assertEqual(h1, h2) + + def test_lora_cache_path_uses_lora_suffix(self): + with tempfile.TemporaryDirectory() as tmpdir: + with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir): + path = lora_optimizer.LoRAAutoTuner._lora_cache_path("abc123") + self.assertTrue(path.endswith("abc123.lora.json")) + + def test_lora_cache_load_missing_returns_none(self): + with tempfile.TemporaryDirectory() as tmpdir: + with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir): + self.assertIsNone( + lora_optimizer.LoRAAutoTuner._lora_cache_load("nonexistent")) + + def test_lora_cache_load_stale_algo_version_returns_none(self): + with tempfile.TemporaryDirectory() as tmpdir: + with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir): + path = os.path.join(tmpdir, "stale.lora.json") + with open(path, "w") as f: + json.dump({"algo_version": "0.0.0", "per_prefix": {}}, f) + self.assertIsNone( + lora_optimizer.LoRAAutoTuner._lora_cache_load("stale")) + + def test_lora_cache_roundtrip(self): + with tempfile.TemporaryDirectory() as tmpdir: + with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir): + per_prefix = { + "prefix_a": { + "norm_sq": 1.5, "rank": 16, + "magnitude_samples_unscaled": [0.1, 0.2], + "strength_sign": 1, + "target_key": "layer.weight", + "is_clip": False, "skip_count": 0, "raw_n": 1, + } + } + lora_optimizer.LoRAAutoTuner._lora_cache_save("abc123", per_prefix) + loaded = lora_optimizer.LoRAAutoTuner._lora_cache_load("abc123") + self.assertIsNotNone(loaded) + self.assertIn("prefix_a", loaded) + self.assertEqual(loaded["prefix_a"]["norm_sq"], 1.5) + + def test_lora_cache_save_is_atomic(self): + with tempfile.TemporaryDirectory() as tmpdir: + with mock.patch("lora_optimizer.AUTOTUNER_MEMORY_DIR", tmpdir): + lora_optimizer.LoRAAutoTuner._lora_cache_save("abc123", {}) + path = lora_optimizer.LoRAAutoTuner._lora_cache_path("abc123") + self.assertTrue(os.path.exists(path)) + self.assertFalse(os.path.exists(path + ".tmp")) + + if __name__ == "__main__": unittest.main()