[fix] PAN2: fail every rank at once when the guardrails fail on the first rank

The first sequence-parallel rank loads the guard models and runs each check. When either step raised there, the
other ranks waited at their next collective until the 10-minute store timeout. The first rank now broadcasts its
error as well as its result, and the other ranks raise with it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

fastvideo/pipelines/basic/pan2/guardrails.py

@@ -11,6 +11,7 @@ from __future__ import annotations
 import importlib.util
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
+from typing import TypeVar
 
 import numpy as np
 import torch
@@ -24,6 +25,8 @@ from fastvideo.utils import pixels_to_uint8
 
 logger = init_logger(__name__)
 
+T = TypeVar("T")
+
 _DISABLE_HINT = "or turn the guardrails off with `enable_guardrails=False` at your own risk."
 
 
@@ -38,6 +41,29 @@ def _sp_rank_and_size() -> tuple[int, int]:
     return group.rank_in_group, group.world_size
 
 
+def _on_first_rank(fn: Callable[[], T], what: str) -> T:
+    """Runs ``fn`` on the first rank of the sequence-parallel group and returns its result on every rank.
+
+    A failure on the first rank is broadcast too, so the other ranks raise with it at once instead of waiting at their
+    next collective until it times out.
+    """
+    rank, world_size = _sp_rank_and_size()
+    outcome: tuple[bool, object] | None = None
+    if rank == 0:
+        try:
+            outcome = (True, fn())
+        except BaseException as e:
+            if world_size > 1:
+                get_sp_group().broadcast_object((False, f"{type(e).__name__}: {e}"), src=0)
+            raise
+    if world_size > 1:
+        outcome = get_sp_group().broadcast_object(outcome, src=0)
+    ok, value = outcome
+    if not ok:
+        raise RuntimeError(f"The first sequence-parallel rank failed to {what}: {value}")
+    return value
+
+
 class PAN2Guardrails:
     """The checker, loaded on the first rank of each sequence-parallel group only.
 
@@ -52,8 +78,8 @@ class PAN2Guardrails:
                               _DISABLE_HINT)
         self.offload_to_cpu = offload_to_cpu
         self.checker = None
-        rank, _ = _sp_rank_and_size()
-        if rank == 0:
+
+        def load() -> None:
             from pan2_guardrail import PAN2SafetyChecker
 
             logger.info("Loading the PAN 2 guardrails (offload_to_cpu=%s)", offload_to_cpu)
@@ -62,6 +88,8 @@ class PAN2Guardrails:
             if not offload_to_cpu:
                 self.checker.to(get_local_torch_device())
 
+        _on_first_rank(load, "load the guardrails")
+
     @contextmanager
     def _on_device(self, module: torch.nn.Module) -> Iterator[None]:
         if not self.offload_to_cpu:
@@ -77,11 +105,7 @@ class PAN2Guardrails:
                 torch.cuda.empty_cache()
 
     def _passes(self, check: Callable[[], bool]) -> bool:
-        rank, world_size = _sp_rank_and_size()
-        safe = bool(check()) if rank == 0 else None
-        if world_size > 1:
-            safe = get_sp_group().broadcast_object(safe, src=0)
-        return bool(safe)
+        return bool(_on_first_rank(lambda: bool(check()), "run the guardrail check"))
 
     def check_prompt(self, prompt: str) -> None:
 

fastvideo/tests/stages/test_pan2_guardrails.py

@@ -128,24 +128,60 @@ def test_offloaded_guards_visit_the_device_only_for_their_check(fake_package, mo
     assert moves == [("text_guard", "cpu"), ("text_guard", "cpu"), ("video_guard", "cpu"), ("video_guard", "cpu")]
 
 
-def test_ranks_other_than_the_first_take_the_broadcast_verdict(fake_package, monkeypatch):
-    broadcasts = []
+class FakeGroup:
+    """A two-rank sequence-parallel group seen from ``rank``; a peer receives ``replies`` in order."""
 
-    class Group:
-        rank_in_group, world_size = 1, 2
+    def __init__(self, rank, replies=()):
+        self.rank_in_group, self.world_size = rank, 2
+        self.replies, self.broadcasts = list(replies), []
 
-        def broadcast_object(self, obj, src=0):
-            broadcasts.append(obj)
-            return False
+    def broadcast_object(self, obj, src=0):
+        self.broadcasts.append(obj)
+        return obj if self.rank_in_group == src else self.replies.pop(0)
 
+
+def _sp_group(monkeypatch, group):
     monkeypatch.setattr(pan2_guardrails, "model_parallel_is_initialized", lambda: True)
-    monkeypatch.setattr(pan2_guardrails, "get_sp_group", lambda: Group())
+    monkeypatch.setattr(pan2_guardrails, "get_sp_group", lambda: group)
+    return group
+
+
+def test_ranks_other_than_the_first_take_the_broadcast_verdict(fake_package, monkeypatch):
+    group = _sp_group(monkeypatch, FakeGroup(rank=1, replies=[(True, None), (True, False)]))
     guardrails = PAN2Guardrails()
     # This rank loads no checker and runs no check: it raises on rank 0's verdict.
     assert guardrails.checker is None
     with pytest.raises(PAN2GuardrailViolationError):
         guardrails.check_prompt("A red fox.")
-    assert broadcasts == [None]
+    assert group.broadcasts == [None, None]
+
+
+def test_a_failed_load_on_the_first_rank_reaches_the_other_ranks_at_once(fake_package, monkeypatch):
+    def fail():
+        raise OSError("no guard weights")
+
+    monkeypatch.setattr(FakeSafetyChecker, "from_pretrained", staticmethod(fail))
+    first = _sp_group(monkeypatch, FakeGroup(rank=0))
+    with pytest.raises(OSError, match="no guard weights"):
+        PAN2Guardrails()
+    assert first.broadcasts == [(False, "OSError: no guard weights")]
+
+    _sp_group(monkeypatch, FakeGroup(rank=1, replies=first.broadcasts))
+    with pytest.raises(RuntimeError, match="failed to load the guardrails: OSError: no guard weights"):
+        PAN2Guardrails()
+
+
+def test_a_failed_check_on_the_first_rank_reaches_the_other_ranks_at_once(fake_package, monkeypatch):
+    first = _sp_group(monkeypatch, FakeGroup(rank=0))
+    guardrails = PAN2Guardrails()
+    monkeypatch.setattr(guardrails.checker, "check_text_safety", lambda prompt: 1 / 0)
+    with pytest.raises(ZeroDivisionError):
+        guardrails.check_prompt("A red fox.")
+    assert first.broadcasts[-1] == (False, "ZeroDivisionError: division by zero")
+
+    _sp_group(monkeypatch, FakeGroup(rank=1, replies=[(True, None), first.broadcasts[-1]]))
+    with pytest.raises(RuntimeError, match="failed to run the guardrail check"):
+        PAN2Guardrails().check_prompt("A red fox.")
 
 
 @pytest.mark.parametrize("config", [None, SimpleNamespace(enable_guardrails=True, offload_guardrail_models=True)])