diff --git a/fastvideo/pipelines/basic/pan2/guardrails.py b/fastvideo/pipelines/basic/pan2/guardrails.py index 4a2ccfe29..1fcdb7b95 100644 --- a/fastvideo/pipelines/basic/pan2/guardrails.py +++ b/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: diff --git a/fastvideo/tests/stages/test_pan2_guardrails.py b/fastvideo/tests/stages/test_pan2_guardrails.py index 76a6b6012..7a79d2ca8 100644 --- a/fastvideo/tests/stages/test_pan2_guardrails.py +++ b/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)])