ability to apply is_flaky at a class-level

Author: sayakpaul

src/diffusers/utils/testing_utils.py

@@ -990,37 +990,44 @@ def summary_failures_short(tr):
     config.option.tbstyle = orig_tbstyle
 
 
-# Copied from https://github.com/huggingface/transformers/blob/000e52aec8850d3fe2f360adc6fd256e5b47fe4c/src/transformers/testing_utils.py#L1905
-def is_flaky(max_attempts: int = 5, wait_before_retry: Optional[float] = None, description: Optional[str] = None):
+# Adapted from https://github.com/huggingface/transformers/blob/000e52aec8850d3fe2f360adc6fd256e5b47fe4c/src/transformers/testing_utils.py#L1905
+def is_flaky(
+    max_attempts: int = 5,
+    wait_before_retry: Optional[float] = None,
+    description: Optional[str] = None,
+):
     """
-    To decorate flaky tests. They will be retried on failures.
-
-    Args:
-        max_attempts (`int`, *optional*, defaults to 5):
-            The maximum number of attempts to retry the flaky test.
-        wait_before_retry (`float`, *optional*):
-            If provided, will wait that number of seconds before retrying the test.
-        description (`str`, *optional*):
-            A string to describe the situation (what / where / why is flaky, link to GH issue/PR comments, errors,
-            etc.)
+    To decorate flaky tests (methods or entire classes). They will be retried on failures.
     """
 
-    def decorator(test_func_ref):
-        @functools.wraps(test_func_ref)
+    def decorator(obj):
+        # If decorating a class, wrap each test method on it
+        if inspect.isclass(obj):
+            for attr_name, attr_value in list(obj.__dict__.items()):
+                if callable(attr_value) and attr_name.startswith("test"):
+                    # recursively decorate the method
+                    setattr(obj, attr_name, decorator(attr_value))
+            return obj
+
+        # Otherwise we're decorating a single test function / method
+        @functools.wraps(obj)
         def wrapper(*args, **kwargs):
             retry_count = 1
-
             while retry_count < max_attempts:
                 try:
-                    return test_func_ref(*args, **kwargs)
-
+                    return obj(*args, **kwargs)
                 except Exception as err:
-                    print(f"Test failed with {err} at try {retry_count}/{max_attempts}.", file=sys.stderr)
+                    msg = (
+                        f"[FLAKY] {description or obj.__name__!r} "
+                        f"failed on attempt {retry_count}/{max_attempts}: {err}"
+                    )
+                    print(msg, file=sys.stderr)
                     if wait_before_retry is not None:
                         time.sleep(wait_before_retry)
                     retry_count += 1
 
-            return test_func_ref(*args, **kwargs)
+            # final attempt, let any exception bubble up
+            return obj(*args, **kwargs)
 
         return wrapper
 

tests/lora/test_lora_layers_wanvace.py

@@ -46,7 +46,7 @@
 
 @require_peft_backend
 @skip_mps
-@is_flaky
+@is_flaky(max_attempts=3, wait_before_retry=1, description="very flaky class")
 class WanVACELoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):
     pipeline_class = WanVACEPipeline
     scheduler_cls = FlowMatchEulerDiscreteScheduler