MAINT GCG: make optimization iteration state explicit and typed

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

@@ -53,6 +53,52 @@
 ]
 
 
+class StopReason(str, Enum):
+    """Why an optimization run stopped iterating."""
+
+    MAX_STEPS_REACHED = "max_steps_reached"
+    ALL_PROMPTS_JAILBROKEN = "all_prompts_jailbroken"
+
+
+@dataclass
+class OptimizationRunState:
+    """
+    Typed iteration state for a single optimization run.
+
+    Captures the current suffix, losses, best result, counters, and stop reason
+    explicitly instead of leaving them as loose loop locals, so each phase of
+    the optimization loop has a stable contract that can be asserted under
+    seeded tests. Exposed as ``MultiPromptAttack.last_run_state`` after a call
+    to :meth:`MultiPromptAttack.run`.
+    """
+
+    control: str
+    best_control: str
+    loss: float
+    best_loss: float
+    steps_completed: int = 0
+    runtime: float = 0.0
+    stop_reason: StopReason | None = None
+
+
+@dataclass
+class ProgressiveScheduleState:
+    """
+    Typed schedule state for :class:`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 :meth:`ProgressiveMultiPromptAttack.run`.
+    """
+
+    goals_admitted: int
+    workers_admitted: int
+    steps_completed: int = 0
+    loss: float = float("inf")
+    stop_inner_on_success: bool = False
+
+
 class NpEncoder(json.JSONEncoder):
     """Encode NumPy scalar and array values for JSON output."""
 
@@ -900,6 +946,52 @@ def step(self, *args: Any, **kwargs: Any) -> tuple[str, float]:
         """Execute one attack optimization step."""
         raise NotImplementedError("Attack step function not yet implemented")
 
+    def _all_training_prompts_jailbroken(self) -> bool:
+        """
+        Check whether every worker jailbreaks every training prompt.
+
+        This is the stopping phase of the optimization loop.
+
+        Returns:
+            bool: True when all jailbreak tests pass for every worker.
+        """
+        model_tests_jb, _, _ = self.test(self.workers, self.prompts)
+        return all(all(tests for tests in model_test) for model_test in model_tests_jb)
+
+    def _log_best_checkpoint(
+        self,
+        *,
+        global_step: int,
+        n_steps_total: int,
+        runtime: float,
+        verbose: bool,
+        state: OptimizationRunState,
+    ) -> None:
+        """
+        Test the best-known suffix and write one periodic log entry.
+
+        This is the logging phase of the optimization loop.
+
+        Temporarily swaps ``self.control_str`` to the best-known suffix so the
+        held-out evaluation reflects it, then restores the active suffix.
+
+        Args:
+            global_step (int): The step number used for logging (including ``anneal_from``).
+            n_steps_total (int): The total step budget used for logging.
+            runtime (float): Runtime of the most recent optimization step, in seconds.
+            verbose (bool): Whether the log entry should print progress output.
+            state (OptimizationRunState): The current run state to read from.
+        """
+        last_control = self.control_str
+        try:
+            self.control_str = state.best_control
+            model_tests = self.test_all()
+            self.log(
+                global_step, n_steps_total, self.control_str, state.best_loss, runtime, model_tests, verbose=verbose
+            )
+        finally:
+            self.control_str = last_control
+
     def run(
         self,
         n_steps: int = 100,
@@ -949,22 +1041,31 @@ def control_weight_fn(_: int) -> float:
             def control_weight_fn(_: int) -> float:
                 return control_weight
 
-        steps = 0
-        loss = best_loss = 1e6
-        best_control = self.control_str
-        runtime = 0.0
+        state = OptimizationRunState(
+            control=self.control_str,
+            best_control=self.control_str,
+            loss=1e6,
+            best_loss=1e6,
+        )
 
         if self.logfile is not None and log_first:
             model_tests = self.test_all()
-            self.log(anneal_from, n_steps + anneal_from, self.control_str, loss, runtime, model_tests, verbose=verbose)
+            self.log(
+                anneal_from,
+                n_steps + anneal_from,
+                self.control_str,
+                state.loss,
+                state.runtime,
+                model_tests,
+                verbose=verbose,
+            )
 
         for i in range(n_steps):
-            if stop_on_success:
-                model_tests_jb, model_tests_mb, _ = self.test(self.workers, self.prompts)
-                if all(all(tests for tests in model_test) for model_test in model_tests_jb):
-                    break
+            if stop_on_success and self._all_training_prompts_jailbroken():
+                state.stop_reason = StopReason.ALL_PROMPTS_JAILBROKEN
+                break
 
-            steps += 1
+            state.steps_completed += 1
             start = time.time()
             control, loss = self.step(
                 batch_size=batch_size,
@@ -976,35 +1077,34 @@ def control_weight_fn(_: int) -> float:
                 filter_cand=filter_cand,
                 verbose=verbose,
             )
-            runtime = time.time() - start
+            state.runtime = time.time() - start
             keep_control = True if not anneal else acceptance_probability(prev_loss, loss, i + anneal_from)
             if keep_control:
                 self.control_str = control
+                state.control = control
 
             prev_loss = loss
-            if loss < best_loss:
-                best_loss = loss
-                best_control = control
-            logger.info(f"Current Loss: {loss}, Best Loss: {best_loss}")
+            state.loss = loss
+            if loss < state.best_loss:
+                state.best_loss = loss
+                state.best_control = control
+            logger.info(f"Current Loss: {loss}, Best Loss: {state.best_loss}")
 
             if self.logfile is not None and (i + 1 + anneal_from) % test_steps == 0:
-                last_control = self.control_str
-                self.control_str = best_control
-
-                model_tests = self.test_all()
-                self.log(
-                    i + 1 + anneal_from,
-                    n_steps + anneal_from,
-                    self.control_str,
-                    best_loss,
-                    runtime,
-                    model_tests,
+                self._log_best_checkpoint(
+                    global_step=i + 1 + anneal_from,
+                    n_steps_total=n_steps + anneal_from,
+                    runtime=state.runtime,
                     verbose=verbose,
+                    state=state,
                 )
 
-                self.control_str = last_control
+        if state.stop_reason is None:
+            state.stop_reason = StopReason.MAX_STEPS_REACHED
+
+        self.last_run_state = state
 
-        return self.control_str, loss, steps
+        return self.control_str, state.loss, state.steps_completed
 
     def test(
         self, workers: list[ModelWorker], prompts: list[PromptManager], include_loss: bool = False
@@ -1246,6 +1346,22 @@ def filter_mpa_kwargs(**kwargs: Any) -> dict[str, Any]:
         """Return options whose names use the ``mpa_`` prefix."""
         return {key[4:]: value for key, value in kwargs.items() if key.startswith("mpa_")}
 
+    def _finalize_progressive_run(
+        self, *, attack: MultiPromptAttack, step: int, n_steps: int, loss: float, verbose: bool
+    ) -> None:
+        """
+        Result-construction phase: run the final held-out evaluation and record the closing log entry.
+
+        Args:
+            attack (MultiPromptAttack): The fully-admitted inner attack that just finished.
+            step (int): The global step count reached by the progressive schedule.
+            n_steps (int): The total step budget of the progressive run.
+            loss (float): The final loss reported by the inner attack.
+            verbose (bool): Whether the closing log entry should print progress output.
+        """
+        model_tests = attack.test_all()
+        attack.log(step, n_steps, self.control, loss, 0.0, model_tests, verbose=verbose)
+
     def run(
         self,
         n_steps: int = 1000,
@@ -1313,17 +1429,18 @@ def run(
             },
         )
 
-        num_goals = 1 if self.progressive_goals else len(self.goals)
-        num_workers = 1 if self.progressive_models else len(self.workers)
-        step = 0
-        stop_inner_on_success = self.progressive_goals
+        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,
+        )
         loss = np.inf
 
-        while step < n_steps:
+        while schedule.steps_completed < n_steps:
             attack = self.managers["MPA"](
-                self.goals[:num_goals],
-                self.targets[:num_goals],
-                self.workers[:num_workers],
+                self.goals[: schedule.goals_admitted],
+                self.targets[: schedule.goals_admitted],
+                self.workers[: schedule.workers_admitted],
                 self.control,
                 self.test_prefixes,
                 self.logfile,
@@ -1332,39 +1449,40 @@ def run(
                 self.test_targets,
                 self.test_workers,
             )
-            if num_goals == len(self.goals) and num_workers == len(self.workers):
-                stop_inner_on_success = False
+            if schedule.goals_admitted == len(self.goals) and schedule.workers_admitted == len(self.workers):
+                schedule.stop_inner_on_success = False
             inner_result: tuple[str, float, int] = attack.run(
-                n_steps=n_steps - step,
+                n_steps=n_steps - schedule.steps_completed,
                 batch_size=batch_size,
                 topk=topk,
                 temp=temp,
                 allow_non_ascii=allow_non_ascii,
                 target_weight=target_weight,
                 control_weight=control_weight,
                 anneal=anneal,
-                anneal_from=step,
+                anneal_from=schedule.steps_completed,
                 prev_loss=loss,
-                stop_on_success=stop_inner_on_success,
+                stop_on_success=schedule.stop_inner_on_success,
                 test_steps=test_steps,
                 filter_cand=filter_cand,
                 verbose=verbose,
             )
             control, loss, inner_steps = inner_result
 
-            step += inner_steps
+            schedule.steps_completed += inner_steps
             self.control = control
 
-            if num_goals < len(self.goals):
-                num_goals += 1
+            if schedule.goals_admitted < len(self.goals):
+                schedule.goals_admitted += 1
                 loss = np.inf
-            elif num_goals == len(self.goals):
-                if num_workers < len(self.workers):
-                    num_workers += 1
+            elif schedule.goals_admitted == len(self.goals):
+                if schedule.workers_admitted < len(self.workers):
+                    schedule.workers_admitted += 1
                     loss = np.inf
-                elif num_workers == len(self.workers) and stop_on_success:
-                    model_tests = attack.test_all()
-                    attack.log(step, n_steps, self.control, loss, 0.0, model_tests, verbose=verbose)
+                elif schedule.workers_admitted == len(self.workers) and stop_on_success:
+                    self._finalize_progressive_run(
+                        attack=attack, step=schedule.steps_completed, n_steps=n_steps, loss=loss, verbose=verbose
+                    )
                     break
                 else:
                     if isinstance(control_weight, (int, float)) and incr_control:
@@ -1374,9 +1492,11 @@ def run(
                             if verbose:
                                 logger.info(f"Control weight increased to {control_weight:.5}")
                         else:
-                            stop_inner_on_success = False
+                            schedule.stop_inner_on_success = False
+
+        self.last_schedule_state = schedule
 
-        return self.control, step
+        return self.control, schedule.steps_completed
 
 
 class IndividualPromptAttack:

pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py

@@ -229,6 +229,34 @@ def _get_control_length(self, *, control: str) -> int | None:
         except (AttributeError, TypeError, ValueError):
             return None
 
+    def _select_best_candidate(
+        self,
+        *,
+        control_cands: list[list[str]],
+        losses: torch.Tensor,
+        batch_size: int,
+    ) -> tuple[str, torch.Tensor]:
+        """
+        Return the candidate with the minimal aggregate loss.
+
+        This is the selection phase of a GCG optimization step.
+
+        Candidate batches from multiple worker-model groups are concatenated, so
+        the flat argmin index is decomposed back into group and in-batch slots.
+
+        Args:
+            control_cands (list[list[str]]): Filtered candidate suffixes per worker-model group.
+            losses (torch.Tensor): Aggregate loss per candidate across all groups.
+            batch_size (int): Number of candidates sampled per group.
+
+        Returns:
+            tuple[str, torch.Tensor]: The best candidate suffix and its loss.
+        """
+        min_idx = losses.argmin()
+        model_idx = min_idx // batch_size
+        batch_idx = min_idx % batch_size
+        return control_cands[model_idx][batch_idx], losses[min_idx]
+
     def step(
         self,
         *,
@@ -352,10 +380,9 @@ def step(
                             f"loss={loss[j * batch_size : (j + 1) * batch_size].min().item() / (i + 1):.4f}"
                         )
 
-            min_idx = loss.argmin()
-            model_idx = min_idx // batch_size
-            batch_idx = min_idx % batch_size
-            next_control, cand_loss = control_cands[model_idx][batch_idx], loss[min_idx]
+            next_control, cand_loss = self._select_best_candidate(
+                control_cands=control_cands, losses=loss, batch_size=batch_size
+            )
         del control_cands, loss
 
         current_length = self._get_control_length(control=next_control)