From b718775b16c913c24fbef3a95577f8db84394028 Mon Sep 17 00:00:00 2001 From: feiiiiii5 Date: Sun, 23 Aug 2026 04:13:50 +0800 Subject: [PATCH] MAINT GCG: make optimization iteration state explicit and typed Implements the structural half of #2416 while keeping public attack behavior and extension protocols unchanged. - add StopReason enum plus typed OptimizationRunState and ProgressiveScheduleState dataclasses capturing suffix, losses, best result, counters, and stop reason - MultiPromptAttack.run now tracks state through the typed object and exposes it as last_run_state; stopping and periodic logging phases are extracted into _all_training_prompts_jailbroken and _log_best_checkpoint with stable contracts - ProgressiveMultiPromptAttack.run tracks admission scheduling through ProgressiveScheduleState (exposed as last_schedule_state) and moves final evaluation into _finalize_progressive_run - GCGMultiPromptAttack extracts the candidate-selection phase into _select_best_candidate; candidate batches intentionally remain step-local to bound VRAM - add deterministic seeded regression tests covering stop reasons, best tracking under annealing rejection, checkpoint restore, argmin decomposition across worker groups, and progressive finalize path --- .../gcg/attack/base/attack_manager.py | 224 +++++++++++++---- .../promptgen/gcg/attack/gcg/gcg_attack.py | 35 ++- .../executor/promptgen/gcg/test_run_state.py | 230 ++++++++++++++++++ 3 files changed, 433 insertions(+), 56 deletions(-) create mode 100644 tests/unit/executor/promptgen/gcg/test_run_state.py diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index 6a00c8b3c6..b319d7ed6c 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/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,10 +1449,10 @@ 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, @@ -1343,28 +1460,29 @@ def run( 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: diff --git a/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py b/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py index d9287b69cd..4d02fac985 100644 --- a/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py +++ b/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) diff --git a/tests/unit/executor/promptgen/gcg/test_run_state.py b/tests/unit/executor/promptgen/gcg/test_run_state.py new file mode 100644 index 0000000000..b1f70634c0 --- /dev/null +++ b/tests/unit/executor/promptgen/gcg/test_run_state.py @@ -0,0 +1,230 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for typed optimization-iteration state in the GCG attack loop.""" + +import random +from typing import Any +from unittest.mock import MagicMock + +import pytest + +attack_manager_mod = pytest.importorskip( + "pyrit.executor.promptgen.gcg.attack.base.attack_manager", + reason="attack_manager module not importable", +) +torch = pytest.importorskip("torch", reason="torch not installed") + +MultiPromptAttack = attack_manager_mod.MultiPromptAttack +OptimizationRunState = attack_manager_mod.OptimizationRunState +ProgressiveMultiPromptAttack = attack_manager_mod.ProgressiveMultiPromptAttack +ProgressiveScheduleState = attack_manager_mod.ProgressiveScheduleState +StopReason = attack_manager_mod.StopReason + + +def _bare_multi_prompt_attack(step_results: list[tuple[str, float]]) -> MultiPromptAttack: + """Build a MultiPromptAttack without __init__ whose step() replays canned results.""" + attack = object.__new__(MultiPromptAttack) + prompt_manager = MagicMock() + prompt_manager.control_str = "initial" + attack.prompts = [prompt_manager] + attack.workers = [MagicMock()] + attack.control_str = "initial" + attack.logfile = None + attack.step = MagicMock(side_effect=list(step_results)) + return attack + + +class TestStopReason: + def test_has_expected_members(self) -> None: + assert StopReason.MAX_STEPS_REACHED == "max_steps_reached" + assert StopReason.ALL_PROMPTS_JAILBROKEN == "all_prompts_jailbroken" + + +class TestOptimizationRunState: + def test_counters_and_stop_reason_default(self) -> None: + state = OptimizationRunState(control="c", best_control="c", loss=1e6, best_loss=1e6) + + assert state.steps_completed == 0 + assert state.runtime == 0.0 + assert state.stop_reason is None + + +class TestProgressiveScheduleState: + def test_defaults(self) -> None: + schedule = ProgressiveScheduleState(goals_admitted=1, workers_admitted=2) + + assert schedule.steps_completed == 0 + assert schedule.loss == float("inf") + assert schedule.stop_inner_on_success is False + + +class TestMultiPromptRunStateTracking: + def test_run_sets_max_steps_reached_when_loop_exhausts(self) -> None: + attack = _bare_multi_prompt_attack([("better", 1.0)]) + + control, loss, steps = attack.run(n_steps=1, prev_loss=2.0, stop_on_success=False, anneal=True) + + assert (control, loss, steps) == ("better", 1.0, 1) + state: OptimizationRunState | None = getattr(attack, "last_run_state", None) + assert state is not None + assert state.steps_completed == 1 + assert state.stop_reason == StopReason.MAX_STEPS_REACHED + assert state.best_control == "better" + assert state.best_loss == 1.0 + assert state.control == "better" + + def test_run_records_jailbroken_stop_reason_without_counting_final_check(self) -> None: + attack = _bare_multi_prompt_attack([]) + attack.test = MagicMock(return_value=([[True]], [[1]], [[1.0]])) + + control, loss, steps = attack.run(n_steps=5, stop_on_success=True) + + assert (control, loss, steps) == ("initial", 1e6, 0) + state: OptimizationRunState = attack.last_run_state + assert state.steps_completed == 0 + assert state.stop_reason == StopReason.ALL_PROMPTS_JAILBROKEN + attack.step.assert_not_called() + + def test_rejected_candidate_keeps_active_suffix_but_updates_loss(self) -> None: + attack = _bare_multi_prompt_attack([("better", 1.0), ("worse", 5.0)]) + random.seed(2026) + + control, loss, steps = attack.run(n_steps=2, prev_loss=2.0, stop_on_success=False, anneal=True) + + # The worse candidate must be rejected by annealing with overwhelming + # probability under this seed; the active suffix stays "better". + assert control == "better" + assert steps == 2 + state: OptimizationRunState = attack.last_run_state + assert state.best_control == "better" + assert state.best_loss == 1.0 + assert state.loss == 5.0 + assert state.stop_reason == StopReason.MAX_STEPS_REACHED + + def test_periodic_checkpoint_restores_active_suffix(self) -> None: + attack = _bare_multi_prompt_attack([("better", 1.0), ("best-yet", 0.25)]) + attack.logfile = "unused-by-test.json" # gate for periodic checkpoints; log/test_all are mocked + attack.test_all = MagicMock(return_value=([[False]], [[0]], [[0.5]])) + attack.log = MagicMock() + + attack.run( + n_steps=2, + prev_loss=2.0, + stop_on_success=False, + anneal=True, + test_steps=1, + ) + + # Each periodic checkpoint evaluates the best-known suffix and then + # restores whatever suffix was active for optimization. + assert attack.control_str == "best-yet" + assert attack.log.call_count == 2 + first_log_args = attack.log.call_args_list[0].args + assert first_log_args[2] == "better" + second_log_args = attack.log.call_args_list[1].args + assert second_log_args[2] == "best-yet" + + def test_seeded_runs_produce_identical_trajectories(self) -> None: + results = [] + for _ in range(2): + random.seed(1234) + attack = _bare_multi_prompt_attack([("a", 3.0), ("b", 2.0), ("c", 1.5)]) + results.append(attack.run(n_steps=3, prev_loss=4.0, stop_on_success=False, anneal=True)) + + assert results[0] == results[1] + assert results[0] == ("c", 1.5, 3) + + +class TestGCGCandidateSelection: + def test_selects_minimum_within_single_group(self) -> None: + from pyrit.executor.promptgen.gcg.attack.gcg.gcg_attack import GCGMultiPromptAttack + + attack = object.__new__(GCGMultiPromptAttack) + next_control, cand_loss = attack._select_best_candidate( + control_cands=[["aa", "bb"]], + losses=torch.tensor([0.5, 9.0]), + batch_size=2, + ) + + assert next_control == "aa" + assert cand_loss.item() == pytest.approx(0.5) + + def test_decomposes_cross_group_argmin_index(self) -> None: + from pyrit.executor.promptgen.gcg.attack.gcg.gcg_attack import GCGMultiPromptAttack + + attack = object.__new__(GCGMultiPromptAttack) + next_control, cand_loss = attack._select_best_candidate( + control_cands=[["aa", "bb"], ["cc", "dd"]], + losses=torch.tensor([9.0, 8.0, 7.0, 6.0]), + batch_size=2, + ) + + assert next_control == "dd" + assert cand_loss.item() == pytest.approx(6.0) + + +class TestProgressiveRunScheduleState: + def _bare_progressive_attack(self, inner_attack: Any) -> ProgressiveMultiPromptAttack: + progressive = object.__new__(ProgressiveMultiPromptAttack) + progressive.goals = ["goal"] + progressive.targets = ["target"] + progressive.workers = [MagicMock()] + progressive.test_goals = [] + progressive.test_targets = [] + progressive.test_workers = [] + progressive.test_prefixes = [] + progressive.managers = {"MPA": MagicMock(return_value=inner_attack)} + progressive.control = "initial" + progressive.logfile = None + progressive.progressive_goals = True + progressive.progressive_models = True + return progressive + + def test_finalize_phase_logs_final_evaluation_and_stops(self) -> None: + inner_attack = MagicMock() + inner_attack.run.return_value = ("ctrl", 0.5, 2) + model_tests = ([[True]], [[1]], [[1.0]]) + inner_attack.test_all.return_value = model_tests + progressive = self._bare_progressive_attack(inner_attack) + + control, steps = progressive.run(n_steps=10, stop_on_success=True) + + assert (control, steps) == ("ctrl", 2) + schedule: ProgressiveScheduleState = progressive.last_schedule_state + assert schedule.steps_completed == 2 + assert schedule.goals_admitted == 1 + assert schedule.workers_admitted == 1 + inner_attack.test_all.assert_called_once() + inner_attack.log.assert_called_once_with(2, 10, "ctrl", 0.5, 0.0, model_tests, verbose=True) + + def test_schedule_exhaustion_continues_until_step_budget_spent(self) -> None: + inner_attack = MagicMock() + inner_attack.run.return_value = ("ctrl", 0.5, 2) + progressive = self._bare_progressive_attack(inner_attack) + + control, steps = progressive.run(n_steps=10, stop_on_success=False) + + assert (control, steps) == ("ctrl", 10) + schedule: ProgressiveScheduleState = progressive.last_schedule_state + assert schedule.steps_completed == 10 + assert schedule.stop_inner_on_success is False + inner_attack.run.assert_called_with( + n_steps=2, + batch_size=1024, + topk=256, + temp=1.0, + allow_non_ascii=False, + target_weight=None, + control_weight=None, + anneal=True, + anneal_from=8, + # The inner result's loss feeds back as the next phase's prev_loss + # so the annealing temperature schedule stays continuous across + # progressive admissions. + prev_loss=0.5, + stop_on_success=False, + test_steps=50, + filter_cand=True, + verbose=True, + )