horde-engine 3.2.9__tar.gz → 3.3.0__tar.gz

{horde_engine-3.2.9 → horde_engine-3.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: horde-engine
-Version: 3.2.9
+Version: 3.3.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.2.9 → horde_engine-3.3.0}/hordelib/_version.py

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

{horde_engine-3.2.9 → horde_engine-3.3.0}/hordelib/api.py

@@ -7,6 +7,8 @@ nothing here triggers a ComfyUI import (``hordelib.initialise()`` does that expl
 The contract is enforced by ``tests/meta/test_public_api_contract.py``.
 """
 
+from typing import TYPE_CHECKING
+
 from hordelib import initialise, is_initialised
 from hordelib.execution.interface import (
     ExecutionBackend,
@@ -67,7 +69,6 @@ from hordelib.pipeline.payload_pp import (
     classify_post_processor,
 )
 from hordelib.preload import controlnet_annotators_present
-from hordelib.shared_model_manager import SharedModelManager
 from hordelib.utils.ioredirect import ComfyUIProgress, ComfyUIProgressUnit
 from hordelib.utils.logger import HordeLog
 from hordelib.utils.torch_memory import (
@@ -77,10 +78,22 @@ from hordelib.utils.torch_memory import (
     enumerate_accelerators,
     get_accelerator_utilization_percent,
     get_free_ram_mb,
+    get_loaded_weights_offloaded_mb,
+    get_torch_device_free_vram_mb,
     get_torch_free_vram_mb,
     get_torch_total_vram_mb,
     log_free_ram,
 )
+from hordelib.vram_planning import (
+    compute_extra_reserved_mb,
+    compute_inference_reserve_mb,
+    compute_weight_budget_mb,
+)
+
+if TYPE_CHECKING:
+    # Re-exported lazily (see ``__getattr__`` below) so the type checker still sees the name without the
+    # eager import that would drag torch into every consumer of this facade.
+    from hordelib.shared_model_manager import SharedModelManager
 
 __all__ = [
     "CONTROLNET_ANNOTATOR_DOWNLOAD_BYTES",
@@ -121,6 +134,9 @@ __all__ = [
     "available_features",
     "classify_post_processor",
     "clear_accelerator_cache",
+    "compute_extra_reserved_mb",
+    "compute_inference_reserve_mb",
+    "compute_weight_budget_mb",
     "controlnet_annotator_download_bytes",
     "controlnet_annotators_present",
     "enumerate_accelerators",
@@ -133,6 +149,8 @@ __all__ = [
     "missing_packages",
     "get_accelerator_utilization_percent",
     "get_free_ram_mb",
+    "get_loaded_weights_offloaded_mb",
+    "get_torch_device_free_vram_mb",
     "get_metrics_collector",
     "get_torch_free_vram_mb",
     "get_torch_total_vram_mb",
@@ -142,3 +160,21 @@ __all__ = [
     "log_free_ram",
     "set_gpu_sampling_lease",
 ]
+
+
+def __getattr__(name: str) -> object:
+    """Lazily resolve the one torch-heavy re-export so the facade's import cost stays honest.
+
+    Every other name above comes from a torch-free submodule, so ``import hordelib.api`` is cheap
+    (~75MB) and pulls no torch. ``SharedModelManager`` is the exception: it transitively imports the
+    concrete model managers, which import torch (~500MB RSS). Re-exporting it eagerly would tax every
+    consumer of this facade with a full torch load, including the worker's torch-free orchestrator and
+    benchmark planner that only want the pure-Python helpers. PEP 562 module ``__getattr__`` defers that
+    import until ``SharedModelManager`` is actually accessed, so a consumer pays for torch only when it
+    asks for something that genuinely needs it.
+    """
+    if name == "SharedModelManager":
+        from hordelib.shared_model_manager import SharedModelManager
+
+        return SharedModelManager
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

{horde_engine-3.2.9 → horde_engine-3.3.0}/hordelib/comfy_horde.py

@@ -511,12 +511,81 @@ def get_torch_free_vram_mb():
     return round(_comfy_get_free_memory() / (1024 * 1024))
 
 
+def get_torch_device_free_vram_mb():
+    """Device-wide free VRAM (mem_get_info), excluding this process's reclaimable torch cache.
+
+    See :func:`hordelib.utils.torch_memory.get_torch_device_free_vram_mb`. This is the honest figure for
+    over-commit/streaming decisions; ``get_torch_free_vram_mb`` (comfy's view) adds back a single
+    process's cache and so over-states cross-process headroom.
+    """
+    from hordelib.utils.torch_memory import get_torch_device_free_vram_mb as _device_free
+
+    return _device_free()
+
+
+_LAST_LOW_VRAM_WARN_TIME: float = 0.0
+"""Monotonic-ish wall time of the last low-free-VRAM warning, so a job that runs its whole length below the
+inference reserve warns periodically rather than on every (very frequent) log_free_ram call."""
+_LOW_VRAM_WARN_INTERVAL_SECONDS: float = 15.0
+
+
 def log_free_ram():
-    logger.debug(
-        f"Free VRAM: {get_torch_free_vram_mb():0.0f} MB, "
-        f"Free RAM: {psutil.virtual_memory().available / (1024 * 1024):0.0f} MB",
+    global _LAST_LOW_VRAM_WARN_TIME
+
+    # Report the *device-wide* free VRAM (mem_get_info), not comfy's get_free_memory: the latter adds back
+    # this process's reserved-but-inactive torch allocator cache, over-stating free by several GB and
+    # hiding VRAM held by other (including leaked) processes. The driver's system-memory fallback triggers
+    # on the device-wide figure, so that is the honest number for a streaming/over-commit warning.
+    free_vram_mb = get_torch_device_free_vram_mb()
+    comfy_free_vram_mb = get_torch_free_vram_mb()
+    reclaimable_cache_mb = max(0.0, comfy_free_vram_mb - free_vram_mb)
+    free_ram_mb = psutil.virtual_memory().available / (1024 * 1024)
+
+    offloaded_mb = 0
+    try:
+        from hordelib.utils.torch_memory import get_loaded_weights_offloaded_mb
+
+        offloaded_mb = get_loaded_weights_offloaded_mb()
+    except Exception:
+        offloaded_mb = 0
+    offloaded_note = f", Weights offloaded to RAM: {offloaded_mb:0.0f} MB" if offloaded_mb > 0 else ""
+    # Surface the gap between the device-wide free and comfy's view so a live run shows how much of the
+    # apparent headroom is merely this process's reclaimable cache (returned by empty_cache), not room a
+    # sibling model load could actually use.
+    cache_note = (
+        f" (+{reclaimable_cache_mb:0.0f} MB reclaimable torch cache; comfy sees {comfy_free_vram_mb:0.0f})"
+        if reclaimable_cache_mb >= 256
+        else ""
     )
 
+    logger.debug(f"Free VRAM: {free_vram_mb:0.0f} MB{cache_note}, Free RAM: {free_ram_mb:0.0f} MB{offloaded_note}")
+
+    # A driver-level system-memory fallback (e.g. the Windows WDDM shared-memory pool, on by default) spills
+    # device allocations to host RAM once free VRAM nears zero, collapsing the sampling rate the same way
+    # ComfyUI's own weight offloading does but *without* appearing in its offload accounting. ComfyUI can
+    # report a model as fully resident while the driver pages its per-step activations. So warn purely on
+    # measured free VRAM falling below the inference working-set reserve: below that floor a sampling step has
+    # no room for its activations and something (ComfyUI or the driver) must stream over the bus. Throttled so
+    # a job that stays under the floor for its whole run periodically reports rather than flooding the log.
+    try:
+        from hordelib.vram_planning import compute_inference_reserve_mb
+
+        reserve_mb = compute_inference_reserve_mb(get_torch_total_vram_mb())
+    except Exception:
+        return
+    if free_vram_mb < reserve_mb:
+        now = time.time()
+        if now - _LAST_LOW_VRAM_WARN_TIME >= _LOW_VRAM_WARN_INTERVAL_SECONDS:
+            _LAST_LOW_VRAM_WARN_TIME = now
+            if offloaded_mb > 0:
+                cause = f"ComfyUI reports {offloaded_mb:0.0f} MB of weights offloaded"
+            else:
+                cause = "ComfyUI reports no offload, so the GPU driver's system-memory fallback is the likely cause"
+            logger.warning(
+                f"Free VRAM {free_vram_mb:0.0f} MB is below the {reserve_mb:0.0f} MB inference reserve: "
+                f"sampling activations will stream from host RAM and run several times slower ({cause}).",
+            )
+
 
 def interrupt_comfyui_processing():
     logger.warning("Interrupting comfyui processing")

{horde_engine-3.2.9 → horde_engine-3.3.0}/hordelib/feature_impact.py

@@ -73,6 +73,18 @@ class BaselineBurden(BaseModel):
     max_recommended_batch: int
     typical_disk_gb: float
     """Typical checkpoint size on disk, used when no model record is available."""
+    vram_weights_mb: int = 0
+    """Resident weight footprint (MB) at the baseline's typical worker dtype, distinct from activations.
+
+    This is the figure ComfyUI compares against its weight budget when deciding whether to keep weights
+    resident or stream them from host RAM, so it (not the activation-inflated ``vram_base_mb``) is what a
+    consumer should use to forecast streaming and whole-card residency. Zero means unspecified, in which
+    case :meth:`resident_weight_estimate_mb` falls back to ``vram_base_mb``. The seeded figures reflect the
+    dtypes the worker typically runs on consumer cards (for example, a Flux fp8 checkpoint)."""
+
+    def resident_weight_estimate_mb(self) -> int:
+        """Return the resident weight footprint (MB), falling back to the conflated base when unseeded."""
+        return self.vram_weights_mb or self.vram_base_mb
 
 
 class BurdenEstimate(BaseModel):
@@ -153,6 +165,7 @@ _BASELINE_SEEDS: list[BaselineBurden] = [
         min_recommended_vram_mb=8000,
         max_recommended_batch=4,
         typical_disk_gb=6.9,
+        vram_weights_mb=4900,
     ),
     BaselineBurden(
         baseline="stable_cascade",
@@ -163,6 +176,7 @@ _BASELINE_SEEDS: list[BaselineBurden] = [
         min_recommended_vram_mb=10000,
         max_recommended_batch=2,
         typical_disk_gb=9.0,
+        vram_weights_mb=8000,
     ),
     BaselineBurden(
         baseline="flux_1",
@@ -173,6 +187,7 @@ _BASELINE_SEEDS: list[BaselineBurden] = [
         min_recommended_vram_mb=14000,
         max_recommended_batch=1,
         typical_disk_gb=17.0,
+        vram_weights_mb=11500,
     ),
     BaselineBurden(
         baseline="flux_schnell",
@@ -183,6 +198,7 @@ _BASELINE_SEEDS: list[BaselineBurden] = [
         min_recommended_vram_mb=14000,
         max_recommended_batch=1,
         typical_disk_gb=17.0,
+        vram_weights_mb=11500,
     ),
     BaselineBurden(
         baseline="flux_dev",
@@ -193,6 +209,7 @@ _BASELINE_SEEDS: list[BaselineBurden] = [
         min_recommended_vram_mb=14000,
         max_recommended_batch=1,
         typical_disk_gb=17.0,
+        vram_weights_mb=11500,
     ),
     BaselineBurden(
         baseline="qwen_image",
@@ -203,6 +220,7 @@ _BASELINE_SEEDS: list[BaselineBurden] = [
         min_recommended_vram_mb=16000,
         max_recommended_batch=1,
         typical_disk_gb=20.0,
+        vram_weights_mb=12000,
     ),
     BaselineBurden(
         baseline="z_image_turbo",
@@ -225,6 +243,7 @@ _FALLBACK_BASELINE_BURDEN = BaselineBurden(
     min_recommended_vram_mb=10000,
     max_recommended_batch=1,
     typical_disk_gb=8.0,
+    vram_weights_mb=7000,
 )
 """Used for baselines with no registry entry (including ``infer``), erring on the heavy side."""
 

{horde_engine-3.2.9 → horde_engine-3.3.0}/hordelib/utils/torch_memory.py

@@ -132,13 +132,32 @@ def get_torch_free_vram_mb() -> int:
     return round(_torch_fallback_vram_bytes(free=True) / _MB)
 
 
+def get_torch_device_free_vram_mb() -> int:
+    """Physically free VRAM on the device, in MB, from the driver's device-wide accounting.
+
+    Unlike :func:`get_torch_free_vram_mb` (which, when ComfyUI is loaded, returns comfy's
+    ``get_free_memory`` == device-wide free PLUS this process's reserved-but-inactive torch allocator
+    cache), this is the raw device-wide free from ``mem_get_info``. The distinction is load-bearing for
+    cross-process residency budgeting and for predicting the driver's system-memory fallback: the comfy
+    number adds back a single process's reclaimable cache, so it over-states what a *new* allocation can
+    use and hides VRAM held by *other* processes (including leaked/orphaned ones). The driver spills to
+    host RAM on device-wide free nearing zero, not on the inflated number, so the worker must budget
+    against this. Falls back to system RAM on CPU/MPS, matching the other helpers here.
+    """
+    return round(_torch_fallback_vram_bytes(free=True) / _MB)
+
+
 def get_free_ram_mb() -> int:
     """Available system RAM, in MB."""
     return round(psutil.virtual_memory().available / _MB)
 
 
 def log_free_ram() -> None:
-    logger.debug(f"Free VRAM: {get_torch_free_vram_mb():0.0f} MB, Free RAM: {get_free_ram_mb():0.0f} MB")
+    offloaded_mb = get_loaded_weights_offloaded_mb()
+    offloaded_note = f", Weights offloaded to RAM: {offloaded_mb:0.0f} MB" if offloaded_mb > 0 else ""
+    logger.debug(
+        f"Free VRAM: {get_torch_free_vram_mb():0.0f} MB, Free RAM: {get_free_ram_mb():0.0f} MB{offloaded_note}",
+    )
 
 
 def get_accelerator_utilization_percent(index: int = 0) -> int | None:
@@ -266,6 +285,34 @@ def enumerate_accelerators() -> list[AcceleratorInfo]:
     return _enumerate_via_torch()
 
 
+def get_loaded_weights_offloaded_mb() -> int:
+    """Return how much of the currently-loaded models' weights sit off the compute device, in MB.
+
+    ComfyUI keeps a model fully resident when it fits the weight budget; otherwise it offloads the
+    remainder to host RAM and streams it to VRAM during sampling (the slow path that drops device-free
+    to near zero and runs several times slower). A non-zero result here is the positive runtime signal
+    that weight streaming is happening for the loaded model(s), which the worker uses to confirm a
+    streaming forecast and to grant the over-budget step grace instead of killing a slow-but-progressing
+    job. Returns 0 when ComfyUI is not loaded in this process or nothing is resident.
+    """
+    mm = _comfy_model_management()
+    if mm is None:
+        return 0
+    loaded_models = getattr(mm, "current_loaded_models", None)
+    if not loaded_models:
+        return 0
+    offloaded_bytes = 0
+    for loaded_model in loaded_models:
+        read_offloaded = getattr(loaded_model, "model_offloaded_memory", None)
+        if read_offloaded is None:
+            continue
+        try:
+            offloaded_bytes += max(0, int(read_offloaded()))
+        except Exception as e:
+            logger.debug(f"offloaded-memory read failed for a loaded model: {e}")
+    return round(offloaded_bytes / _MB)
+
+
 def clear_accelerator_cache() -> None:
     """Release cached device memory on whatever backend is active (a no-op on CPU).
 

horde_engine-3.3.0/hordelib/vram_planning.py

@@ -0,0 +1,92 @@
+"""Pure-Python mirrors of ComfyUI's VRAM accounting, usable before (or without) torch/ComfyUI.
+
+The worker's orchestrator decides whether a model will fit a device *before* dispatching it to a
+torch child, and it must do so without importing torch (a ~500MB RSS cost it cannot pay; see the
+torch-free-orchestrator invariant in horde-worker-reGen). ComfyUI's decision to keep a model's
+weights resident or stream them from host RAM hinges on one quantity computed in
+``comfy/model_management.py``: ``minimum_inference_memory()``, the free VRAM it keeps for activations.
+A model whose weights leave less than that free at load time is offloaded and streamed during
+sampling, which collapses the step rate on a memory-constrained device. This module reproduces that
+formula exactly from the device's total VRAM and platform so the worker can predict streaming and gate
+accordingly.
+
+Keep these constants in lockstep with the vendored ``comfy/model_management.py``. For runtime-exact
+agreement (across ComfyUI version skew) prefer the live value a torch child reports; this module is
+the torch-free estimate the orchestrator uses for admission and the no-boot benchmark planner.
+
+Comfy-free and torch-free: safe to import in the orchestrator and the benchmark planner.
+"""
+
+from __future__ import annotations
+
+import sys
+
+# Mirror of ``comfy/model_management.py`` module-level constants. ComfyUI reserves a base amount of
+# VRAM for other applications, raised on Windows (shared-VRAM driver behaviour) and again on cards
+# above 15GB. An explicit ``--reserve-vram`` overrides all of it.
+_BASE_EXTRA_RESERVED_MB = 400
+_WINDOWS_EXTRA_RESERVED_MB = 600
+_WINDOWS_LARGE_CARD_BONUS_MB = 100
+_LARGE_CARD_THRESHOLD_MB = 15 * 1024
+# The fixed activation working set ComfyUI keeps free on top of the reserve (0.8 GB).
+_INFERENCE_WORKING_SET_MB = 0.8 * 1024
+# ComfyUI scales total VRAM by this before subtracting the reserve to get the static weight budget.
+_WEIGHT_BUDGET_FRACTION = 0.88
+
+
+def _resolve_is_windows(is_windows: bool | None) -> bool:
+    """Return whether to apply ComfyUI's Windows reserve rules (autodetect when not forced)."""
+    return sys.platform.startswith("win") if is_windows is None else is_windows
+
+
+def compute_extra_reserved_mb(
+    total_vram_mb: int,
+    *,
+    is_windows: bool | None = None,
+    reserve_vram_gb: float | None = None,
+) -> int:
+    """Return ComfyUI's ``EXTRA_RESERVED_VRAM`` (MB) for a device of ``total_vram_mb``.
+
+    Mirrors ``comfy/model_management.py``: a base reserve, higher on Windows, with a bonus on cards
+    above 15GB. An explicit ``--reserve-vram`` (``reserve_vram_gb``) overrides all of it.
+    """
+    if reserve_vram_gb is not None:
+        return round(reserve_vram_gb * 1024)
+    if _resolve_is_windows(is_windows):
+        reserved = _WINDOWS_EXTRA_RESERVED_MB
+        if total_vram_mb > _LARGE_CARD_THRESHOLD_MB:
+            reserved += _WINDOWS_LARGE_CARD_BONUS_MB
+        return reserved
+    return _BASE_EXTRA_RESERVED_MB
+
+
+def compute_inference_reserve_mb(
+    total_vram_mb: int,
+    *,
+    is_windows: bool | None = None,
+    reserve_vram_gb: float | None = None,
+) -> int:
+    """Return ComfyUI's ``minimum_inference_memory()`` (MB): the free VRAM it keeps for activations.
+
+    A model whose weights leave less than this free at load time is partially offloaded to host RAM
+    and streamed during sampling (the slow path). Equal to the 0.8GB working set plus the extra
+    reserve, so it is the headroom the worker must preserve to avoid weight streaming.
+    """
+    extra = compute_extra_reserved_mb(total_vram_mb, is_windows=is_windows, reserve_vram_gb=reserve_vram_gb)
+    return round(_INFERENCE_WORKING_SET_MB + extra)
+
+
+def compute_weight_budget_mb(
+    total_vram_mb: int,
+    *,
+    is_windows: bool | None = None,
+    reserve_vram_gb: float | None = None,
+) -> int:
+    """Return ComfyUI's static ``maximum_vram_for_weights()`` (MB): the most weights can occupy.
+
+    ``total * 0.88 - inference_reserve``. A model whose weights exceed this budget cannot be held
+    resident even alone and will always stream; the runtime per-load budget can be lower still when
+    other models are resident, which is why the worker also reasons about reclaimable siblings.
+    """
+    reserve = compute_inference_reserve_mb(total_vram_mb, is_windows=is_windows, reserve_vram_gb=reserve_vram_gb)
+    return round(total_vram_mb * _WEIGHT_BUDGET_FRACTION) - reserve