Merge pull request #4691 from awsr/errorlimiter-update

Update Errorlimiter to be thread-safe

modules/errorlimiter.py

@@ -1,44 +1,53 @@
 from __future__ import annotations
+
 from contextlib import contextmanager
-from typing import TYPE_CHECKING
+from threading import Lock
+from typing import TYPE_CHECKING, ClassVar
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
+_instance_id = 0
+_lock = Lock()
 
-class ErrorLimiterTrigger(BaseException):  # Use BaseException to avoid being caught by "except Exception:".
+
+def _make_unique(name: str):
+    global _instance_id
+    with _lock:  # Guard against race conditions
+        new_name = f"{name}__{_instance_id}"
+        _instance_id += 1
+        return new_name
+
+
+class _ErrorLimiterTrigger(BaseException):  # Use BaseException to avoid being caught by "except Exception:".
     def __init__(self, name: str, *args):
         super().__init__(*args)
-        self.name = name
+        self.name = name.rsplit("__", 1)[0]
 
 
-class ErrorLimiterAbort(RuntimeError):
-    def __init__(self, msg: str):
-        super().__init__(msg)
-
-
-class ErrorLimiter:
-    _store: dict[str, int] = {}
+class _ErrorLimiter:
+    _store: ClassVar[dict[str, int]] = {}
 
     @classmethod
     def start(cls, name: str, limit: int = 5):
         cls._store[name] = limit
 
     @classmethod
-    def notify(cls, name: str | Iterable[str]):  # Can be manually triggered if execution is spread across multiple files
-        if isinstance(name, str):
-            name = (name,)
-        for key in name:
-            if key in cls._store.keys():
-                cls._store[key] = cls._store[key] - 1
-                if cls._store[key] <= 0:
-                    raise ErrorLimiterTrigger(key)
+    def notify(cls, key: str):
+        cls._store[key] = cls._store[key] - 1
+        if cls._store[key] <= 0:
+            raise _ErrorLimiterTrigger(key)
 
     @classmethod
     def end(cls, name: str):
         cls._store.pop(name)
 
 
+class ErrorLimiterAbort(RuntimeError):
+    def __init__(self, msg: str):
+        super().__init__(msg)
+
+
 @contextmanager
 def limit_errors(name: str, limit: int = 5):
     """Limiter for aborting execution after being triggered a specified number of times (default 5).
@@ -64,10 +73,11 @@ def limit_errors(name: str, limit: int = 5):
     Yields:
         Callable: Notification function to indicate that an error occurred.
     """
+    name_id = _make_unique(name)
     try:
-        ErrorLimiter.start(name, limit)
-        yield lambda: ErrorLimiter.notify(name)
-    except ErrorLimiterTrigger as e:
+        _ErrorLimiter.start(name_id, limit)
+        yield lambda: _ErrorLimiter.notify(name_id)
+    except _ErrorLimiterTrigger as e:
         raise ErrorLimiterAbort(f"HALTING. Too many errors during '{e.name}'") from None
     finally:
-        ErrorLimiter.end(name)
+        _ErrorLimiter.end(name_id)

modules/lora/lora_apply.py

@@ -1,12 +1,17 @@
+from __future__ import annotations
+
 import re
 import time
+from typing import TYPE_CHECKING
 import torch
 import diffusers.models.lora
-from modules.errorlimiter import ErrorLimiter
 from modules.lora import lora_common as l
 from modules import shared, devices, errors, model_quant
 from modules.logger import log
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
 
 bnb = None
 re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)")
@@ -76,7 +81,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr
     return backup_size
 
 
-def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, use_previous: bool = False):
+def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, use_previous: bool = False, *, elimit: Callable[[], None] | None = None):
     if shared.opts.diffusers_offload_mode == "none":
         try:
             self.to(devices.device)
@@ -142,7 +147,8 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou
             if l.debug:
                 errors.display(e, 'LoRA')
                 raise RuntimeError('LoRA apply weight') from e
-            ErrorLimiter.notify(("network_activate", "network_deactivate"))
+            if elimit is not None:
+                elimit()
         continue
     return batch_updown, batch_ex_bias
 

modules/lora/networks.py

@@ -18,7 +18,7 @@ def network_activate(include=None, exclude=None):
     if include is None:
         include = []
     t0 = time.time()
-    with limit_errors("network_activate"):
+    with limit_errors("network_activate") as elimit:
         sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
         if shared.opts.diffusers_offload_mode == "sequential":
             sd_models.disable_offload(sd_model)
@@ -56,7 +56,7 @@ def network_activate(include=None, exclude=None):
                             pbar.update(task, advance=1)
                         continue
                     backup_size += network_backup_weights(module, network_layer_name, wanted_names)
-                    batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name)
+                    batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit)
                     if shared.opts.lora_fuse_native:
                         network_apply_direct(module, batch_updown, batch_ex_bias, device=device)
                     else:
@@ -92,7 +92,7 @@ def network_deactivate(include=None, exclude=None):
     if len(l.previously_loaded_networks) == 0:
         return
     t0 = time.time()
-    with limit_errors("network_deactivate"):
+    with limit_errors("network_deactivate") as elimit:
         sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
         if shared.opts.diffusers_offload_mode == "sequential":
             sd_models.disable_offload(sd_model)
@@ -124,7 +124,7 @@ def network_deactivate(include=None, exclude=None):
                         if task is not None:
                             pbar.update(task, advance=1)
                         continue
-                    batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True)
+                    batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True, elimit=elimit)
                     if shared.opts.lora_fuse_native:
                         network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True)
                     else: