diff --git a/README.md b/README.md index 2b20f91..84bb35c 100644 --- a/README.md +++ b/README.md @@ -558,7 +558,7 @@ During the parameter sweep, each candidate recomputes raw LoRA diffs (A@B matmul | `ram` | All diffs in RAM. Fastest, but uses ~1.5 GB (SDXL) to ~6 GB (Flux) | | `disk` | All diffs to temp files with memory-mapping. Slowest cache mode, but minimal RAM | -When `auto` mode runs out of disk space, it falls back to RAM automatically. +**Disk bound:** the disk spill (`auto`/`disk`) is capped by free space — it keeps ~5 GB free on the temp volume, then stops caching and **recomputes the overflow diffs on demand** instead of filling the disk. Large models with several LoRAs (e.g. LTX-2 audio-video × 3) can otherwise spill tens of GB of full-rank diffs to `ComfyUI/temp/`. If you hit the cap a lot, point ComfyUI's temp dir at a larger drive or use `diff_cache_mode=disabled`. | Setting | Default | Effect | |---------|---------|--------| diff --git a/lora_optimizer.py b/lora_optimizer.py index 4e518c9..8669711 100644 --- a/lora_optimizer.py +++ b/lora_optimizer.py @@ -715,6 +715,13 @@ class _DiffCache: self._ram_limit = None self._cache_dir = None self._disk_failed = False + self._disk_bytes = 0 + self._disk_full = False + # Never let the diff cache fill the temp volume — keep this much free. + # When writing would drop below it, stop disk-caching and recompute the + # overflow on demand (a large model x several LoRAs can otherwise spill + # tens of GB of full-rank diffs). + self._disk_min_free = 5 * 1024 * 1024 * 1024 # 5 GB if mode in ("disk", "auto"): import tempfile, atexit, weakref self._cache_dir = tempfile.mkdtemp(prefix="lora_diff_cache_", @@ -745,6 +752,25 @@ class _DiffCache: # auto: spill to disk if RAM limit would be exceeded return (self._ram_bytes + tensor_bytes) > self._ram_limit + def _disk_has_space(self, tensor_bytes): + """True if writing tensor_bytes keeps the temp volume above the reserve.""" + if self._disk_full: + return False + try: + import shutil + free = shutil.disk_usage(self._cache_dir).free + except Exception: + return True # can't tell — don't block + if free - tensor_bytes < self._disk_min_free: + self._disk_full = True + logging.warning( + f"[DiffCache] Temp volume low on space ({free / 1024**3:.1f}GB free) — " + f"stopping disk caching ({self._disk_bytes / 1024**3:.1f}GB cached so far); " + f"remaining diffs will be recomputed on demand. Free up space, point " + f"ComfyUI's temp dir at a larger drive, or set diff_cache_mode=disabled.") + return False + return True + def prefetch(self, keys): """Load disk entries into a prefetch buffer in a background thread.""" self._wait_prefetch() @@ -799,6 +825,10 @@ class _DiffCache: cached = cached.cpu() tensor_bytes = cached.nelement() * cached.element_size() if self._use_disk(tensor_bytes) and not self._disk_failed: + if not self._disk_has_space(tensor_bytes): + # Volume nearly full: skip caching this diff entirely (it will be + # recomputed on the get() miss) rather than fill the disk or blow RAM. + return try: import hashlib, numpy as np name_hash = hashlib.sha256(f"{key[0]}_{key[1]}".encode()).hexdigest()[:16] @@ -807,6 +837,7 @@ class _DiffCache: disk_t = cached.half() if cached.dtype == torch.bfloat16 else cached np.save(path, disk_t.numpy()) self._disk_store[key] = path + self._disk_bytes += tensor_bytes return except Exception as e: self._disk_failed = True diff --git a/tests/test_lora_optimizer.py b/tests/test_lora_optimizer.py index f87baf8..7050c90 100644 --- a/tests/test_lora_optimizer.py +++ b/tests/test_lora_optimizer.py @@ -4129,6 +4129,25 @@ class TestDiffCacheRamLimit(unittest.TestCase): cache = lora_optimizer._DiffCache(mode="auto", ram_pct=0.5) self.assertEqual(cache._ram_limit, 50 * 1024 ** 3) + @unittest.skipIf(torch is None, "torch is not installed") + def test_disk_guard_skips_caching_when_volume_low(self): + cache = lora_optimizer._DiffCache(mode="disk") + low = types.SimpleNamespace(free=1 * 1024 ** 3, total=100 * 1024 ** 3, + used=99 * 1024 ** 3) + with mock.patch("shutil.disk_usage", return_value=low): + cache.put(("alias_a", 0), torch.zeros(4, 4)) + self.assertNotIn(("alias_a", 0), cache) # skipped -> caller recomputes + self.assertTrue(cache._disk_full) + + @unittest.skipIf(torch is None, "torch is not installed") + def test_disk_caches_when_space_available(self): + cache = lora_optimizer._DiffCache(mode="disk") + ok = types.SimpleNamespace(free=50 * 1024 ** 3, total=100 * 1024 ** 3, + used=50 * 1024 ** 3) + with mock.patch("shutil.disk_usage", return_value=ok): + cache.put(("alias_a", 0), torch.zeros(4, 4)) + self.assertIn(("alias_a", 0), cache) + @unittest.skipIf(torch is None, "torch is not installed in this environment") class TestDiffCacheWarming(unittest.TestCase):