MAINT GCG: model progressive admission transitions (#2665)

pyrit/executor/promptgen/gcg/attack/base/attack_manager.py

@@ -24,6 +24,11 @@
 from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM
 from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
 
+from pyrit.executor.promptgen.gcg.attack.base.progressive_schedule import (
+    ProgressiveScheduleController,
+    ProgressiveScheduleState,
+    ScheduleTransitionAction,
+)
 from pyrit.executor.promptgen.gcg.experiments.log import (
     log_gpu_memory,
     log_loss,
@@ -80,24 +85,6 @@ class OptimizationRunState:
     stop_reason: StopReason | None = None
 
 
-@dataclass
-class ProgressiveScheduleState:
-    """
-    Typed schedule state for ``ProgressiveMultiPromptAttack``.
-
-    Tracks how many goals and workers have been admitted so far, together with
-    the shared step counter and the loss carried between progressive rounds.
-    Exposed as ``ProgressiveMultiPromptAttack.last_schedule_state`` after a call
-    to ``ProgressiveMultiPromptAttack.run``.
-    """
-
-    goals_admitted: int
-    workers_admitted: int
-    steps_completed: int = 0
-    loss: float = float("inf")
-    stop_inner_on_success: bool = False
-
-
 @dataclass
 class RngBundle:
     """Per-run RNG state bundle for deterministic GCG execution."""
@@ -1507,23 +1494,25 @@ def run(
             },
         )
 
-        schedule = ProgressiveScheduleState(
-            goals_admitted=1 if self.progressive_goals else len(self.goals),
-            workers_admitted=1 if self.progressive_models else len(self.workers),
-            stop_inner_on_success=self.progressive_goals,
+        controller = ProgressiveScheduleController(
+            total_goals=len(self.goals),
+            total_workers=len(self.workers),
+            progressive_goals=self.progressive_goals,
+            progressive_models=self.progressive_models,
+            n_steps=n_steps,
+            control_weight=control_weight,
+            incr_control=incr_control,
+            stop_on_success=stop_on_success,
+            verbose=verbose,
         )
-        # Whether ``schedule.loss`` currently reflects an inner run's measured
-        # loss, as opposed to the ``inf`` sentinel written when a new round is
-        # admitted. Tracked explicitly so a legitimately non-finite inner loss
-        # (non-finite model loss or numeric overflow) is not mistaken for an
-        # unupdated sentinel value.
-        loss_is_measured = False
-
-        while schedule.steps_completed < n_steps:
+
+        while not controller.is_complete:
+            controller.before_inner_run()
+            schedule = controller.state
             attack = self.managers["MPA"](
-                self.goals[: schedule.goals_admitted],
-                self.targets[: schedule.goals_admitted],
-                self.workers[: schedule.workers_admitted],
+                self.goals[: controller.active_goal_count],
+                self.targets[: controller.active_goal_count],
+                self.workers[: controller.active_worker_count],
                 self.control,
                 self.test_prefixes,
                 self.logfile,
@@ -1532,17 +1521,15 @@ def run(
                 self.test_targets,
                 self.test_workers,
             )
-            if schedule.goals_admitted == len(self.goals) and schedule.workers_admitted == len(self.workers):
-                schedule.stop_inner_on_success = False
             attack._rng_bundle = rng_bundle
             inner_result: tuple[str, float, int] = attack.run(
-                n_steps=n_steps - schedule.steps_completed,
+                n_steps=controller.remaining_steps,
                 batch_size=batch_size,
                 topk=topk,
                 temp=temp,
                 allow_non_ascii=allow_non_ascii,
                 target_weight=target_weight,
-                control_weight=control_weight,
+                control_weight=controller.control_weight,
                 anneal=anneal,
                 anneal_from=schedule.steps_completed,
                 prev_loss=schedule.loss,
@@ -1553,28 +1540,13 @@ def run(
                 random_seed=random_seed,
             )
             control, inner_loss, inner_steps = inner_result
-            schedule.loss = inner_loss
-            loss_is_measured = True
-
-            schedule.steps_completed += inner_steps
             self.control = control
 
-            # Once the step budget is spent, stop preparing further rounds:
-            # admissions and their sentinel resets would strand ``inf`` on
-            # ``schedule.loss`` for a run that legitimately ends right here.
-            prepare_next_round = schedule.steps_completed < n_steps
-
-            if schedule.goals_admitted < len(self.goals):
-                if prepare_next_round:
-                    schedule.goals_admitted += 1
-                    schedule.loss = np.inf
-                    loss_is_measured = False
-            elif schedule.workers_admitted < len(self.workers):
-                if prepare_next_round:
-                    schedule.workers_admitted += 1
-                    schedule.loss = np.inf
-                    loss_is_measured = False
-            elif schedule.workers_admitted == len(self.workers) and stop_on_success:
+            action = controller.advance_after_inner_run(
+                inner_loss=inner_loss,
+                inner_steps=inner_steps,
+            )
+            if action == ScheduleTransitionAction.FINALIZE_AND_STOP:
                 self._finalize_progressive_run(
                     attack=attack,
                     step=schedule.steps_completed,
@@ -1583,27 +1555,11 @@ def run(
                     verbose=verbose,
                 )
                 break
-            elif prepare_next_round and isinstance(control_weight, (int, float)) and incr_control:
-                if control_weight <= 0.09:
-                    control_weight += 0.01
-                    schedule.loss = np.inf
-                    loss_is_measured = False
-                    if verbose:
-                        logger.info(f"Control weight increased to {control_weight:.5}")
-                else:
-                    schedule.stop_inner_on_success = False
-
-        # The inner run must have produced a measured loss whenever any
-        # optimization happened; guards against silent carry-over regressions.
-        # Whether the loss was measured is tracked explicitly (a completed
-        # inner run may legitimately report a non-finite loss), never inferred
-        # from the numeric value.
-        if schedule.steps_completed > 0:
-            assert loss_is_measured, "schedule.loss was never updated by the inner run"
 
-        self.last_schedule_state = schedule
+        controller.validate_post_run()
+        self.last_schedule_state = controller.state
 
-        return self.control, schedule.steps_completed
+        return self.control, controller.state.steps_completed
 
 
 class IndividualPromptAttack:

pyrit/executor/promptgen/gcg/attack/base/progressive_schedule.py

@@ -0,0 +1,190 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""Progressive schedule controller and state models for Greedy Coordinate Gradient (GCG) attacks."""
+
+import logging
+from dataclasses import dataclass
+from enum import Enum, auto
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ProgressiveScheduleState:
+    """
+    Typed schedule state for ``ProgressiveMultiPromptAttack``.
+
+    Tracks how many goals and workers have been admitted so far, together with
+    the shared step counter and the loss carried between progressive rounds.
+    Exposed as ``ProgressiveMultiPromptAttack.last_schedule_state`` after a call
+    to ``ProgressiveMultiPromptAttack.run``.
+    """
+
+    goals_admitted: int
+    workers_admitted: int
+    steps_completed: int = 0
+    loss: float = float("inf")
+    stop_inner_on_success: bool = False
+
+
+class ScheduleTransitionAction(Enum):
+    """Action to take following an inner attack result in progressive scheduling."""
+
+    CONTINUE = auto()
+    FINALIZE_AND_STOP = auto()
+
+
+class ProgressiveScheduleController:
+    """Encapsulates progressive admission, step budget scheduling, and state transitions."""
+
+    def __init__(
+        self,
+        *,
+        total_goals: int,
+        total_workers: int,
+        progressive_goals: bool = True,
+        progressive_models: bool = True,
+        n_steps: int = 1000,
+        control_weight: float | None = None,
+        incr_control: bool = True,
+        stop_on_success: bool = True,
+        verbose: bool = True,
+    ) -> None:
+        """
+        Initialize the progressive schedule controller.
+
+        Args:
+            total_goals: Total number of attack goals available for admission.
+            total_workers: Total number of model workers available for admission.
+            progressive_goals: Whether goals are admitted progressively one at a time.
+            progressive_models: Whether models/workers are admitted progressively one at a time.
+            n_steps: Total step budget across all progressive rounds.
+            control_weight: Current control weight, or None.
+            incr_control: Whether to increment control weight when all goals/models are admitted.
+            stop_on_success: Whether to finalize and stop when fully admitted and success is achieved.
+            verbose: Whether to log control weight updates.
+
+        Raises:
+            ValueError: If total_goals or total_workers is not positive.
+        """
+        if total_goals <= 0:
+            raise ValueError(f"total_goals must be positive, got {total_goals}")
+        if total_workers <= 0:
+            raise ValueError(f"total_workers must be positive, got {total_workers}")
+
+        self._total_goals = total_goals
+        self._total_workers = total_workers
+        self._n_steps = n_steps
+        self._control_weight = control_weight
+        self._incr_control = incr_control
+        self._stop_on_success = stop_on_success
+        self._verbose = verbose
+
+        self._state = ProgressiveScheduleState(
+            goals_admitted=1 if progressive_goals else total_goals,
+            workers_admitted=1 if progressive_models else total_workers,
+            stop_inner_on_success=progressive_goals,
+        )
+        self._loss_is_measured = False
+
+    @property
+    def state(self) -> ProgressiveScheduleState:
+        """The current progressive schedule state."""
+        return self._state
+
+    @property
+    def is_complete(self) -> bool:
+        """Whether the overall step budget has been exhausted."""
+        return self._state.steps_completed >= self._n_steps
+
+    @property
+    def remaining_steps(self) -> int:
+        """The remaining number of optimization steps in the budget."""
+        return max(0, self._n_steps - self._state.steps_completed)
+
+    @property
+    def active_goal_count(self) -> int:
+        """The number of currently admitted goals."""
+        return self._state.goals_admitted
+
+    @property
+    def active_worker_count(self) -> int:
+        """The number of currently admitted workers."""
+        return self._state.workers_admitted
+
+    @property
+    def control_weight(self) -> float | None:
+        """The current control weight."""
+        return self._control_weight
+
+    @property
+    def is_fully_admitted(self) -> bool:
+        """Whether all goals and workers have been admitted."""
+        return self._state.goals_admitted == self._total_goals and self._state.workers_admitted == self._total_workers
+
+    def before_inner_run(self) -> None:
+        """Prepare schedule state immediately before launching an inner attack."""
+        if self.is_fully_admitted:
+            self._state.stop_inner_on_success = False
+
+    def advance_after_inner_run(
+        self,
+        *,
+        inner_loss: float,
+        inner_steps: int,
+    ) -> ScheduleTransitionAction:
+        """
+        Update schedule state and determine the next action after an inner attack round completes.
+
+        Args:
+            inner_loss: Final loss reported by the inner attack.
+            inner_steps: Number of steps completed by the inner attack.
+
+        Returns:
+            ScheduleTransitionAction indicating whether to continue or finalize and stop.
+        """
+        self._state.loss = inner_loss
+        self._loss_is_measured = True
+        self._state.steps_completed += inner_steps
+
+        prepare_next_round = self._state.steps_completed < self._n_steps
+
+        if self._state.goals_admitted < self._total_goals:
+            if prepare_next_round:
+                self._state.goals_admitted += 1
+                self._state.loss = float("inf")
+                self._loss_is_measured = False
+            return ScheduleTransitionAction.CONTINUE
+
+        if self._state.workers_admitted < self._total_workers:
+            if prepare_next_round:
+                self._state.workers_admitted += 1
+                self._state.loss = float("inf")
+                self._loss_is_measured = False
+            return ScheduleTransitionAction.CONTINUE
+
+        if self._state.workers_admitted == self._total_workers and self._stop_on_success:
+            return ScheduleTransitionAction.FINALIZE_AND_STOP
+
+        if prepare_next_round and isinstance(self._control_weight, (int, float)) and self._incr_control:
+            if self._control_weight <= 0.09:
+                self._control_weight += 0.01
+                self._state.loss = float("inf")
+                self._loss_is_measured = False
+                if self._verbose:
+                    logger.info(f"Control weight increased to {self._control_weight:.5}")
+            else:
+                self._state.stop_inner_on_success = False
+
+        return ScheduleTransitionAction.CONTINUE
+
+    def validate_post_run(self) -> None:
+        """
+        Validate post-run invariants.
+
+        Raises:
+            AssertionError: If steps were completed but schedule.loss was never updated by an inner run.
+        """
+        if self._state.steps_completed > 0:
+            assert self._loss_is_measured, "schedule.loss was never updated by the inner run"