horde-engine 3.3.2__tar.gz → 3.4.0__tar.gz

{horde_engine-3.3.2 → horde_engine-3.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: horde-engine
-Version: 3.3.2
+Version: 3.4.0
 Summary: A wrapper around ComfyUI to allow use by AI Horde.
 Project-URL: Homepage, https://github.com/Haidra-Org/hordelib
 Project-URL: Bug Tracker, https://github.com/Haidra-Org/hordelib/issues

{horde_engine-3.3.2 → horde_engine-3.4.0}/hordelib/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '3.3.2'
-__version_tuple__ = version_tuple = (3, 3, 2)
+__version__ = version = '3.4.0'
+__version_tuple__ = version_tuple = (3, 4, 0)
 
 __commit_id__ = commit_id = None

{horde_engine-3.3.2 → horde_engine-3.4.0}/hordelib/comfy_horde.py

@@ -405,6 +405,10 @@ def pin_models_in_vram() -> bool:
 def unload_all_models_vram():
     global _comfy_current_loaded_models
 
+    from hordelib.metrics import get_metrics_collector
+
+    _unload_start = time.perf_counter()
+
     log_free_ram()
 
     from hordelib.shared_model_manager import SharedModelManager
@@ -448,6 +452,12 @@ def unload_all_models_vram():
         clear_gc_and_torch_cache()
     log_free_ram()
 
+    # Surface the VRAM-eviction cost to the in-process collector so an embedder can attribute the
+    # between-jobs reload churn (this is the "unload" half; the reload is already timed as
+    # disk_to_ram/ram_to_vram). Recorded under the job in progress when the unload happens, or the
+    # next job on this process when it is an idle eviction.
+    get_metrics_collector().record_phase("model_unload", time.perf_counter() - _unload_start)
+
 
 def unload_all_models_ram():
     global _comfy_current_loaded_models

{horde_engine-3.3.2 → horde_engine-3.4.0}/hordelib/model_manager/civitai_adhoc.py

@@ -18,11 +18,13 @@ is inherited. See :class:`~hordelib.model_manager.lora.LoraModelManager` and
 
 from __future__ import annotations
 
+import errno
 import glob
 import hashlib
 import json
 import os
 import re
+import shutil
 import threading
 import time
 import uuid
@@ -51,9 +53,27 @@ TESTS_ONGOING = os.getenv("TESTS_ONGOING", "0") == "1"
 TIMESTAMP_FORMAT = "%Y-%m-%d %H:%M:%S"
 """The strftime/strptime format used for all cache timestamps persisted to disk."""
 
+_TERMINAL_DISK_ERRNOS: frozenset[int] = frozenset(
+    e for e in (getattr(errno, name, None) for name in ("ENOSPC", "EDQUOT", "EROFS", "EFBIG")) if e is not None
+)
+"""Disk-write failures that retrying cannot fix: full disk, exceeded quota, read-only FS, file too large.
+
+Such a write fails identically on every attempt, yet each retry first re-downloads the entire (often
+hundreds-of-MB) weight file, and the worker thread holds the record for the whole retry budget. The
+ad-hoc downloader treats these as terminal so the queue drains promptly instead of thrashing.
+"""
+
 A1111_TRIGGER_PATTERN = re.compile(r"<(?:lora|ti):(.*):.*>")
 """Matches A1111-style inline trigger syntax so it can be reduced to the bare trigger word."""
 
+DEFAULT_MIN_FREE_DISK_MB = 1024
+"""Keep at least this many MB free on the cache volume; ad-hoc downloads evict (then refuse) below it.
+
+A full cache disk turns every weight write into an ENOSPC failure (and can take the whole host's
+working volume down with it). Treating a low-free-space floor as a hard constraint on ad-hoc growth
+keeps the cache self-limiting even if the configured byte budget is mis-set or the volume is shared.
+"""
+
 
 def now_timestamp() -> str:
     """Return the current local time formatted with :data:`TIMESTAMP_FORMAT`."""
@@ -181,6 +201,7 @@ class CivitaiAdhocModelManager[RecordT: GenericModelRecord](
         reference_backups: bool | None = None,
         max_top_disk: int,
         max_adhoc_disk: int,
+        min_free_disk_mb: int = DEFAULT_MIN_FREE_DISK_MB,
     ) -> None:
         """Create an ad-hoc CivitAI model manager.
 
@@ -195,9 +216,12 @@ class CivitaiAdhocModelManager[RecordT: GenericModelRecord](
                 backups follow whether a multiprocessing lock was supplied.
             max_top_disk: The default-set cache budget, in megabytes.
             max_adhoc_disk: The ad-hoc cache budget, in megabytes.
+            min_free_disk_mb: Keep at least this much free space (in megabytes) on the cache volume;
+                ad-hoc downloads evict to make room and refuse rather than cross the floor.
         """
         self._max_top_disk = max_top_disk
         self.max_adhoc_disk = max_adhoc_disk
+        self.min_free_disk_mb = max(0, min_free_disk_mb)
 
         self._data: dict | None = None
         self._next_page_url: str | None = None
@@ -503,6 +527,28 @@ class CivitaiAdhocModelManager[RecordT: GenericModelRecord](
                     self.save_reference_to_disk()
                     return
 
+                # Refuse to fetch a fresh weight when the volume can't hold it above the floor, even
+                # after evicting every ad-hoc entry. Writing anyway risks an ENOSPC that takes the
+                # whole cache (and any co-located worker data) down; skipping leaves the worker able to
+                # serve the job without this LoRA. Defaults are never evicted, so an over-full default
+                # set legitimately yields "no room" here.
+                if not self._ensure_room_for_download(target):
+                    download_logger.warning(
+                        "adhoc.download_skipped_disk_full",
+                        filename=target.filename,
+                        size_mb=target.size_mb,
+                        free_mb=round(self.disk_free_mb() or -1.0),
+                        floor_mb=self.min_free_disk_mb,
+                    )
+                    self._record_download_event(
+                        record,
+                        success=False,
+                        size_bytes=0,
+                        duration_seconds=time.perf_counter() - overall_start,
+                        retries=retries,
+                    )
+                    return
+
                 download_start = time.perf_counter()
                 download_url = target.url
                 if self._civitai_api_token and self.is_model_url_from_civitai(download_url):
@@ -578,6 +624,25 @@ class CivitaiAdhocModelManager[RecordT: GenericModelRecord](
                         retries=retries,
                     )
                     return
+            except OSError as disk_error:
+                if disk_error.errno in _TERMINAL_DISK_ERRNOS:
+                    # A full / read-only / over-quota disk will not heal between attempts, so retrying only
+                    # re-downloads the whole weight file to fail the write again while pinning this worker
+                    # thread (and thus the callers blocked in wait_for_downloads/reset_adhoc_cache). Record
+                    # the terminal failure and move on; the caller proceeds without this model.
+                    download_logger.error("adhoc.download_disk_error", errno=disk_error.errno, error=str(disk_error))
+                    self._metric_network_errors.add(1, {"error_type": "disk"})
+                    self._record_download_event(
+                        record,
+                        success=False,
+                        size_bytes=0,
+                        duration_seconds=time.perf_counter() - overall_start,
+                        retries=retries,
+                    )
+                    return
+                # Any other OS error may be transient; fall through to the shared retry path below.
+                self._metric_network_errors.add(1, {"error_type": "fatal"})
+                download_logger.error("adhoc.download_fatal_error", error=str(disk_error))
             except Exception as fatal_error:
                 self._metric_network_errors.add(1, {"error_type": "fatal"})
                 download_logger.error("adhoc.download_fatal_error", error=str(fatal_error))
@@ -617,6 +682,10 @@ class CivitaiAdhocModelManager[RecordT: GenericModelRecord](
     def wait_for_downloads(self, timeout: float | None = None) -> None:
         """Block until the download queue drains and all workers are idle.
 
+        ``timeout`` is the budget in seconds: ``None`` waits forever, and any number (including ``0``) is
+        honoured. Using a truthiness test here instead would silently turn ``timeout=0`` into "wait
+        forever", so a caller that asked for a bounded drain could hang on a wedged background download.
+
         Raises:
             TimeoutError: If *timeout* seconds elapse before downloads complete.
         """
@@ -624,7 +693,7 @@ class CivitaiAdhocModelManager[RecordT: GenericModelRecord](
         while not self.are_downloads_complete():
             time.sleep(self.THREAD_WAIT_TIME)
             waited += self.THREAD_WAIT_TIME
-            if timeout and waited > timeout:
+            if timeout is not None and waited > timeout:
                 raise TimeoutError(f"{self.METRIC_PREFIX} downloads exceeded specified timeout ({timeout})")
 
     def are_downloads_complete(self) -> bool:
@@ -692,17 +761,17 @@ class CivitaiAdhocModelManager[RecordT: GenericModelRecord](
 
     def is_default_cache_full(self) -> bool:
         """Return whether the default cache has reached its budget."""
-        return self.calculate_default_cache() >= self._max_top_disk * 1024
+        return self.calculate_default_cache() >= self._max_top_disk
 
     def is_adhoc_cache_full(self) -> bool:
         """Return whether the ad-hoc cache has reached its budget."""
-        return self.calculate_adhoc_cache() >= self.max_adhoc_disk * 1024
+        return self.calculate_adhoc_cache() >= self.max_adhoc_disk
 
     def amount_of_adhoc_to_delete(self) -> int:
         """Return how many ad-hoc entries to evict: one, plus one more per 4GB over budget."""
         if not self.is_adhoc_cache_full():
             return 0
-        return 1 + int((self.calculate_adhoc_cache() - self.max_adhoc_disk * 1024) / 4096)
+        return 1 + int((self.calculate_adhoc_cache() - self.max_adhoc_disk) / 4096)
 
     def find_oldest_adhoc_entry(self) -> CacheEntry | None:
         """Return the least-recently-used ad-hoc cache entry, or ``None`` if there are none."""
@@ -726,6 +795,72 @@ class CivitaiAdhocModelManager[RecordT: GenericModelRecord](
         for _eviction in range(self.amount_of_adhoc_to_delete()):
             self.delete_oldest()
 
+    def enforce_adhoc_budget(self) -> None:
+        """Evict ad-hoc entries back within :attr:`max_adhoc_disk` and persist, if over budget.
+
+        The public entry point for callers (e.g. the worker) that have just lowered
+        ``max_adhoc_disk`` to constrain the cache and want the eviction applied immediately.
+        """
+        if self.is_adhoc_cache_full():
+            self._evict_adhoc_over_limit()
+            self.save_reference_to_disk()
+
+    # ------------------------------------------------------------------
+    # Disk-space safety floor
+    # ------------------------------------------------------------------
+
+    def disk_free_mb(self) -> float | None:
+        """Return free space (in megabytes) on the cache volume, or ``None`` if it can't be read."""
+        try:
+            return shutil.disk_usage(self.model_folder_path).free / (1024 * 1024)
+        except OSError as usage_error:
+            logger.bind(manager=self.METRIC_PREFIX).warning("adhoc.disk_usage_failed", error=str(usage_error))
+            return None
+
+    def is_disk_below_floor(self) -> bool:
+        """Return whether free space on the cache volume is below :attr:`min_free_disk_mb`."""
+        free = self.disk_free_mb()
+        return free is not None and free < self.min_free_disk_mb
+
+    def evict_adhoc_for_free_space(self, *, required_mb: float = 0.0) -> bool:
+        """Evict oldest ad-hoc entries until the volume has room for the floor plus *required_mb*.
+
+        Re-samples free space after each deletion (only the OS knows the true effect of a delete on a
+        shared volume). Persists once if anything was evicted.
+
+        Returns:
+            ``True`` once free space meets the target (or the volume can't be sampled), ``False`` if
+            the ad-hoc entries are exhausted and the target still isn't met (an "unsolvable" disk).
+        """
+        target_mb = self.min_free_disk_mb + max(required_mb, 0.0)
+        evicted_any = False
+        while True:
+            free = self.disk_free_mb()
+            if free is None or free >= target_mb:
+                if evicted_any:
+                    self.save_reference_to_disk()
+                return True
+            oldest = self.find_oldest_adhoc_entry()
+            if oldest is None:
+                if evicted_any:
+                    self.save_reference_to_disk()
+                return False
+            self._delete_model_entry(oldest.model_key, oldest.version_key)
+            evicted_any = True
+
+    def _ensure_room_for_download(self, target: DownloadTarget) -> bool:
+        """Return whether *target* can be written without crossing the disk floor, evicting if needed.
+
+        Evicts oldest ad-hoc entries to make room. Returns ``False`` when even an empty ad-hoc cache
+        would leave the volume below the floor (a download that must be skipped).
+        """
+        free = self.disk_free_mb()
+        if free is None:
+            return True  # Can't sample the volume; defer to the normal ENOSPC handling on write.
+        if free - target.size_mb >= self.min_free_disk_mb:
+            return True
+        return self.evict_adhoc_for_free_space(required_mb=target.size_mb)
+
     # ------------------------------------------------------------------
     # Unused-file cleanup (shared)
     # ------------------------------------------------------------------

{horde_engine-3.3.2 → horde_engine-3.4.0}/hordelib/model_manager/lora.py

@@ -25,6 +25,7 @@ from loguru import logger
 
 import hordelib.exceptions as he
 from hordelib.model_manager.civitai_adhoc import (
+    DEFAULT_MIN_FREE_DISK_MB,
     TESTS_ONGOING,
     CacheEntry,
     CivitaiAdhocModelManager,
@@ -75,6 +76,8 @@ class LoraModelManager(CivitaiAdhocModelManager[HordeLoraModelRecord]):
             allowed_top_lora_storage: The default-set cache budget, in megabytes.
             allowed_adhoc_lora_storage: The ad-hoc cache budget, in megabytes. Overridden by the
                 ``AIWORKER_LORA_CACHE_SIZE`` environment variable when set.
+                The ``AIWORKER_LORA_MIN_DISK_FREE_MB`` environment variable likewise overrides the
+                free-space floor below which ad-hoc downloads evict (then refuse) to spare the disk.
             download_wait: Whether :meth:`download_default_models` blocks until downloads complete.
             multiprocessing_lock: Optional cross-process lock guarding on-disk reference writes.
             civitai_api_token: Optional CivitAI API token.
@@ -93,6 +96,15 @@ class LoraModelManager(CivitaiAdhocModelManager[HordeLoraModelRecord]):
                 logger.bind(manager="lora").warning("lora.env_cache_size_invalid", raw_value=env_cache_size)
                 adhoc_storage = AIWORKER_LORA_CACHE_SIZE_DEFAULT
 
+        min_free_disk_mb = DEFAULT_MIN_FREE_DISK_MB
+        env_min_free = os.getenv("AIWORKER_LORA_MIN_DISK_FREE_MB")
+        if env_min_free is not None:
+            try:
+                min_free_disk_mb = int(env_min_free)
+            except (ValueError, TypeError):
+                logger.bind(manager="lora").warning("lora.env_min_free_invalid", raw_value=env_min_free)
+                min_free_disk_mb = DEFAULT_MIN_FREE_DISK_MB
+
         models_db_path = horde_model_reference_paths.legacy_path.joinpath("lora.json").resolve()
         super().__init__(
             model_category=MODEL_REFERENCE_CATEGORY.lora,
@@ -103,6 +115,7 @@ class LoraModelManager(CivitaiAdhocModelManager[HordeLoraModelRecord]):
             reference_backups=reference_backups,
             max_top_disk=allowed_top_lora_storage,
             max_adhoc_disk=adhoc_storage,
+            min_free_disk_mb=min_free_disk_mb,
         )
 
     def ensure_is_version(self, lora_version: int | str) -> str | None: