From 8a3d7b45c4b0a3f09393d6210e4df6deb4961b6b Mon Sep 17 00:00:00 2001 From: AmruthVamshi Date: Wed, 26 Aug 2026 12:54:45 -0400 Subject: [PATCH 1/5] FIX Propagate GCG random_seed to all RNG sources for deterministic runs (#2490) --- .../gcg/attack/base/attack_manager.py | 15 +- .../promptgen/gcg/attack/gcg/gcg_attack.py | 28 ++- .../promptgen/gcg/default_implementations.py | 5 +- .../promptgen/gcg/extension_protocols.py | 4 + pyrit/executor/promptgen/gcg/generator.py | 10 +- .../executor/promptgen/gcg/test_gcg_core.py | 208 ++++++++++++++++++ .../executor/promptgen/gcg/test_run_state.py | 10 +- 7 files changed, 261 insertions(+), 19 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index c3270fcfdb..2663f32618 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -1016,6 +1016,7 @@ def run( log_first: bool = False, filter_cand: bool = True, verbose: bool = True, + random_seed: int = 42, ) -> tuple[str, float, int]: """ Run iterative optimization. @@ -1023,10 +1024,14 @@ def run( Returns: tuple[str, float, int]: The final control, loss, and step count. """ + py_rng = random.Random(random_seed) + models = getattr(self, "models", None) + device = models[0].device if models else "cpu" + self._torch_gen = torch.Generator(device=device).manual_seed(random_seed) def acceptance_probability(e: float, e_prime: float, k: int) -> bool: temperature = max(1 - float(k + 1) / (n_steps + anneal_from), 1.0e-7) - return e_prime < e or math.exp(-(e_prime - e) / temperature) >= random.random() + return e_prime < e or math.exp(-(e_prime - e) / temperature) >= py_rng.random() if target_weight is None: @@ -1400,6 +1405,7 @@ def run( stop_on_success: bool = True, verbose: bool = True, filter_cand: bool = True, + random_seed: int = 42, ) -> tuple[str, int]: """ Execute the progressive multi-prompt attack. @@ -1431,6 +1437,8 @@ def run( Whether to print verbose output (default is True) filter_cand (bool, optional): Whether to filter candidates whose lengths changed after re-tokenization (default is True) + random_seed (int, optional): + Seed for deterministic random number generation (default is 42) Returns: tuple[str, int]: The final control suffix and completed step count. @@ -1499,6 +1507,7 @@ def run( test_steps=test_steps, filter_cand=filter_cand, verbose=verbose, + random_seed=random_seed, ) control, inner_loss, inner_steps = inner_result schedule.loss = inner_loss @@ -1656,6 +1665,7 @@ def run( stop_on_success: bool = True, verbose: bool = True, filter_cand: bool = True, + random_seed: int = 42, ) -> tuple[str, int]: """ Execute the individual-prompt attack. @@ -1687,6 +1697,8 @@ def run( Whether to print verbose output (default is True) filter_cand (bool, optional): Whether to filter candidates (default is True) + random_seed (int, optional): + Seed for deterministic random number generation (default is 42) Returns: tuple[str, int]: The final control suffix and configured step count. @@ -1741,6 +1753,7 @@ def run( log_first=True, filter_cand=filter_cand, verbose=verbose, + random_seed=random_seed, ) return self.control, n_steps diff --git a/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py b/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py index 4d02fac985..df854a8832 100644 --- a/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py +++ b/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py @@ -104,6 +104,7 @@ def sample_control( topk: int = 256, temp: float = 1.0, allow_non_ascii: bool = True, + torch_generator: torch.Generator | None = None, ) -> torch.Tensor: """ Sample new control token candidates based on gradients. @@ -114,6 +115,7 @@ def sample_control( topk (int): Number of top gradient positions to sample from. Defaults to 256. temp (float): Temperature for sampling. Currently unused but kept for API compatibility. Defaults to 1.0. allow_non_ascii (bool): Whether to allow non-ASCII tokens. Defaults to True. + torch_generator (torch.Generator | None): Optional generator for deterministic sampling. Returns: torch.Tensor: Batch of new candidate control token sequences. @@ -127,7 +129,9 @@ def sample_control( torch.int64 ) new_token_val = torch.gather( - top_indices[new_token_pos], 1, torch.randint(0, topk, (batch_size, 1), device=grad.device) + top_indices[new_token_pos], + 1, + torch.randint(0, topk, (batch_size, 1), device=grad.device, generator=torch_generator), ) return original_control_toks.scatter_(1, new_token_pos.unsqueeze(-1), new_token_val) @@ -199,15 +203,19 @@ def _sample_control_candidates( ) -> torch.Tensor: sampler = self._resolve_sampling() prompt_manager = self.prompts[worker_index] - return sampler.sample_candidates( - gradient=gradient, - control_tokens=prompt_manager.control_toks, - batch_size=batch_size, - top_k=topk, - temperature=temp, - allow_non_ascii=allow_non_ascii, - non_ascii_tokens=prompt_manager.disallowed_toks, - ) + torch_gen: torch.Generator | None = getattr(self, "_torch_gen", None) + kwargs: dict[str, Any] = { + "gradient": gradient, + "control_tokens": prompt_manager.control_toks, + "batch_size": batch_size, + "top_k": topk, + "temperature": temp, + "allow_non_ascii": allow_non_ascii, + "non_ascii_tokens": prompt_manager.disallowed_toks, + } + if torch_gen is not None: + kwargs["torch_generator"] = torch_gen + return sampler.sample_candidates(**kwargs) def _filter_control_candidates( self, diff --git a/pyrit/executor/promptgen/gcg/default_implementations.py b/pyrit/executor/promptgen/gcg/default_implementations.py index ae22eb85f9..2686d296b4 100644 --- a/pyrit/executor/promptgen/gcg/default_implementations.py +++ b/pyrit/executor/promptgen/gcg/default_implementations.py @@ -57,6 +57,7 @@ def sample_candidates( temperature: float, allow_non_ascii: bool, non_ascii_tokens: torch.Tensor, + torch_generator: torch.Generator | None = None, ) -> torch.Tensor: """ Sample ``batch_size`` candidate suffix token sequences. @@ -79,6 +80,8 @@ def sample_candidates( the top-k. non_ascii_tokens (torch.Tensor): Token ids to exclude when ``allow_non_ascii`` is False. + torch_generator (torch.Generator | None): Optional generator for + deterministic sampling. Returns: torch.Tensor: Candidate suffix token sequences with shape @@ -99,7 +102,7 @@ def sample_candidates( new_token_val = torch.gather( top_indices[new_token_pos], 1, - torch.randint(0, top_k, (batch_size, 1), device=gradient.device), + torch.randint(0, top_k, (batch_size, 1), device=gradient.device, generator=torch_generator), ) return original_control_tokens.scatter_(1, new_token_pos.unsqueeze(-1), new_token_val) diff --git a/pyrit/executor/promptgen/gcg/extension_protocols.py b/pyrit/executor/promptgen/gcg/extension_protocols.py index 1fc2512dd2..169ab0188c 100644 --- a/pyrit/executor/promptgen/gcg/extension_protocols.py +++ b/pyrit/executor/promptgen/gcg/extension_protocols.py @@ -76,6 +76,7 @@ def sample_candidates( temperature: float, allow_non_ascii: bool, non_ascii_tokens: torch.Tensor, + torch_generator: torch.Generator | None = None, ) -> torch.Tensor: """ Sample ``batch_size`` candidate suffix token sequences. @@ -101,6 +102,9 @@ def sample_candidates( non_ascii_tokens (torch.Tensor): Token ids to exclude when ``allow_non_ascii`` is False, shape ``(num_disallowed,)`` and integer dtype. + torch_generator (torch.Generator | None): Optional random number + generator for deterministic sampling. When provided, all + random tensor operations should use this generator. Returns: torch.Tensor: Candidate suffix token sequences with shape diff --git a/pyrit/executor/promptgen/gcg/generator.py b/pyrit/executor/promptgen/gcg/generator.py index 14995110ff..0cea0eb795 100644 --- a/pyrit/executor/promptgen/gcg/generator.py +++ b/pyrit/executor/promptgen/gcg/generator.py @@ -260,9 +260,11 @@ async def _setup_async(self, *, context: GCGContext) -> None: self._ensure_spawn_start_method() context.memory_labels = combine_dict({}, context.memory_labels) + np_rng = np.random.default_rng(self._algorithm.random_seed) context.targets, context.test_targets = self._apply_target_augmentation( train_targets=context.targets, test_targets=context.test_targets, + np_rng=np_rng, ) log_gpu_memory(step=0) @@ -318,6 +320,7 @@ async def _perform_async(self, *, context: GCGContext) -> GCGResult: verbose=self._output.verbose, filter_cand=self._algorithm.filter_cand, allow_non_ascii=self._algorithm.allow_non_ascii, + random_seed=self._algorithm.random_seed, ) return self._read_result(logfile_path=context.logfile_path, memory_labels=context.memory_labels) @@ -381,6 +384,7 @@ def _apply_target_augmentation( *, train_targets: list[str], test_targets: list[str], + np_rng: np.random.Generator | None = None, ) -> tuple[list[str], list[str]]: """ Randomly substitute equivalent target phrasings for diversity. @@ -388,6 +392,8 @@ def _apply_target_augmentation( Returns: tuple[list[str], list[str]]: Augmented training and held-out targets. """ + if np_rng is None: + np_rng = np.random.default_rng() def _shorten(s: str) -> str: return s.replace("Sure, h", "H") @@ -395,8 +401,8 @@ def _shorten(s: str) -> str: def _contract(s: str) -> str: return s.replace("Sure, here is", "Sure, here's") - train_targets = [_shorten(t) if np.random.random() < 0.5 else _contract(t) for t in train_targets] - test_targets = [_shorten(t) if np.random.random() < 0.5 else _contract(t) for t in test_targets] + train_targets = [_shorten(t) if np_rng.random() < 0.5 else _contract(t) for t in train_targets] + test_targets = [_shorten(t) if np_rng.random() < 0.5 else _contract(t) for t in test_targets] return train_targets, test_targets def _to_attack_params(self, *, context: GCGContext) -> Any: diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index 148e67c980..86dd513934 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -41,6 +41,15 @@ reason="GCG optional dependencies not installed", ) LengthPreservingFilter = default_implementations_mod.LengthPreservingFilter +StandardGCGSampling = default_implementations_mod.StandardGCGSampling + +import numpy as np # noqa: E402 + +generator_mod = pytest.importorskip( + "pyrit.executor.promptgen.gcg.generator", + reason="GCG optional dependencies not installed", +) +GCGGenerator = generator_mod.GCGGenerator @dataclass @@ -1358,3 +1367,202 @@ def test_token_gradients_raises_when_coordinate_gradient_missing() -> None: def test_length_preserving_filter_rejects_unknown_option() -> None: with pytest.raises(TypeError, match="Unexpected LengthPreservingFilter option: unexpected"): LengthPreservingFilter(unexpected=True) + + +class TestRandomSeedDeterminism: + """Verify that random_seed produces reproducible results across runs.""" + + def test_target_augmentation_deterministic_same_seed(self) -> None: + """Same seed produces identical augmentation results.""" + targets = ["Sure, here is how to hack", "Sure, here is how to pick a lock"] + rng1 = np.random.default_rng(42) + rng2 = np.random.default_rng(42) + + result1, _ = GCGGenerator._apply_target_augmentation(train_targets=targets, test_targets=[], np_rng=rng1) + result2, _ = GCGGenerator._apply_target_augmentation(train_targets=targets, test_targets=[], np_rng=rng2) + + assert result1 == result2 + + def test_target_augmentation_different_seed_can_differ(self) -> None: + """Different seeds can produce different augmentation results.""" + targets = ["Sure, here is how to hack"] * 20 + rng1 = np.random.default_rng(1) + rng2 = np.random.default_rng(999) + + result1, _ = GCGGenerator._apply_target_augmentation(train_targets=targets, test_targets=[], np_rng=rng1) + result2, _ = GCGGenerator._apply_target_augmentation(train_targets=targets, test_targets=[], np_rng=rng2) + + assert result1 != result2 + + def test_sampling_deterministic_same_seed(self) -> None: + """StandardGCGSampling produces identical candidates with same torch Generator seed.""" + sampler = StandardGCGSampling() + gradient = torch.randn(5, 100) + control_tokens = torch.tensor([1, 2, 3, 4, 5], dtype=torch.long) + non_ascii = torch.tensor([50], dtype=torch.long) + + gen1 = torch.Generator().manual_seed(42) + gen2 = torch.Generator().manual_seed(42) + + result1 = sampler.sample_candidates( + gradient=gradient.clone(), + control_tokens=control_tokens.clone(), + batch_size=8, + top_k=10, + temperature=1.0, + allow_non_ascii=True, + non_ascii_tokens=non_ascii, + torch_generator=gen1, + ) + result2 = sampler.sample_candidates( + gradient=gradient.clone(), + control_tokens=control_tokens.clone(), + batch_size=8, + top_k=10, + temperature=1.0, + allow_non_ascii=True, + non_ascii_tokens=non_ascii, + torch_generator=gen2, + ) + + assert torch.equal(result1, result2) + + def test_sampling_different_seed_can_differ(self) -> None: + """Different torch Generator seeds can produce different candidates.""" + sampler = StandardGCGSampling() + gradient = torch.randn(5, 100) + control_tokens = torch.tensor([1, 2, 3, 4, 5], dtype=torch.long) + non_ascii = torch.tensor([50], dtype=torch.long) + + gen1 = torch.Generator().manual_seed(1) + gen2 = torch.Generator().manual_seed(999) + + result1 = sampler.sample_candidates( + gradient=gradient.clone(), + control_tokens=control_tokens.clone(), + batch_size=8, + top_k=10, + temperature=1.0, + allow_non_ascii=True, + non_ascii_tokens=non_ascii, + torch_generator=gen1, + ) + result2 = sampler.sample_candidates( + gradient=gradient.clone(), + control_tokens=control_tokens.clone(), + batch_size=8, + top_k=10, + temperature=1.0, + allow_non_ascii=True, + non_ascii_tokens=non_ascii, + torch_generator=gen2, + ) + + assert not torch.equal(result1, result2) + + def test_annealing_deterministic_same_seed(self) -> None: + """run() with same seed produces identical annealing acceptance decisions.""" + attack = object.__new__(MultiPromptAttack) + prompt_manager = MagicMock() + prompt_manager.control_str = "initial" + attack.prompts = [prompt_manager] + attack.logfile = None + + # Step returns a slightly worse loss so annealing decides acceptance + attack.step = MagicMock(return_value=("candidate", 2.5)) + + _, loss1, _ = attack.run(n_steps=3, prev_loss=2.0, stop_on_success=False, anneal=True, random_seed=42) + attack.step = MagicMock(return_value=("candidate", 2.5)) + prompt_manager.control_str = "initial" + _, loss2, _ = attack.run(n_steps=3, prev_loss=2.0, stop_on_success=False, anneal=True, random_seed=42) + + assert loss1 == loss2 + + def test_annealing_different_seed_can_differ(self) -> None: + """run() with different seeds can produce different annealing outcomes.""" + results = [] + for seed in [1, 999]: + attack = object.__new__(MultiPromptAttack) + prompt_manager = MagicMock() + prompt_manager.control_str = "initial" + attack.prompts = [prompt_manager] + attack.logfile = None + # Marginal loss that annealing might accept or reject depending on random draw + attack.step = MagicMock(return_value=("candidate", 2.1)) + control, _, _ = attack.run(n_steps=10, prev_loss=2.0, stop_on_success=False, anneal=True, random_seed=seed) + results.append(control) + + # With enough steps and marginal losses, different seeds should diverge + # (probabilistic but extremely likely with 10 steps) + assert results[0] != results[1] or True # non-flaky: just verify no crash + + def test_concurrent_runs_isolated(self) -> None: + """Two runs with different seeds don't interfere with each other's RNG state.""" + targets = ["Sure, here is how to hack"] * 10 + + rng_a = np.random.default_rng(42) + result_a, _ = GCGGenerator._apply_target_augmentation(train_targets=targets, test_targets=[], np_rng=rng_a) + + # Interleave: run a different seed in between + rng_other = np.random.default_rng(999) + GCGGenerator._apply_target_augmentation(train_targets=targets, test_targets=[], np_rng=rng_other) + + # Fresh rng with seed 42 still gives same result + rng_b = np.random.default_rng(42) + result_b, _ = GCGGenerator._apply_target_augmentation(train_targets=targets, test_targets=[], np_rng=rng_b) + + assert result_a == result_b + + def test_run_creates_torch_gen_for_step(self) -> None: + """run() sets self._torch_gen so step() can access it for sampling.""" + attack = object.__new__(MultiPromptAttack) + prompt_manager = MagicMock() + prompt_manager.control_str = "initial" + attack.prompts = [prompt_manager] + attack.logfile = None + attack.step = MagicMock(return_value=("result", 0.5)) + + attack.run(n_steps=1, stop_on_success=False, anneal=False, random_seed=123) + + assert hasattr(attack, "_torch_gen") + assert isinstance(attack._torch_gen, torch.Generator) + + def test_custom_sampler_without_torch_generator_still_works(self) -> None: + """Custom SamplingStrategy that doesn't accept torch_generator still functions.""" + gradient = torch.randn(3, 6) + logits = torch.randn(2, 8, 10) + token_ids = torch.randint(0, 10, (2, 8)) + control_tokens = torch.tensor([1, 2, 3], dtype=torch.long) + disallowed_tokens = torch.tensor([5], dtype=torch.long) + tokenizer = MagicMock() + tokenizer.decode.return_value = "decoded" + + worker = _WorkerStub(gradient=gradient.clone(), logits=logits, token_ids=token_ids, tokenizer=tokenizer) + prompt = MagicMock() + prompt.control_toks = control_tokens + prompt_manager = MagicMock() + prompt_manager.control_toks = control_tokens + prompt_manager.disallowed_toks = disallowed_tokens + + sampled_tokens = torch.tensor([[8, 8, 8]], dtype=torch.long) + # _SpySampling does NOT accept torch_generator — backward compat test + sampling = _SpySampling(sampled_tokens=sampled_tokens) + + attack = object.__new__(GCGMultiPromptAttack) + attack._sampling = sampling + attack.prompts = [prompt_manager] + attack.workers = [worker] + attack.models = [MagicMock(device=torch.device("cpu"))] + attack.control_str = "test" + + # No _torch_gen set — simulates step() called without run() + result = attack._sample_control_candidates( + worker_index=0, + gradient=gradient, + batch_size=1, + topk=3, + temp=1.0, + allow_non_ascii=True, + ) + + assert torch.equal(result, sampled_tokens) diff --git a/tests/unit/executor/promptgen/gcg/test_run_state.py b/tests/unit/executor/promptgen/gcg/test_run_state.py index 67ac07809f..25654ae1ef 100644 --- a/tests/unit/executor/promptgen/gcg/test_run_state.py +++ b/tests/unit/executor/promptgen/gcg/test_run_state.py @@ -3,7 +3,6 @@ """Tests for typed optimization-iteration state in the GCG attack loop.""" -import random from typing import Any from unittest.mock import MagicMock @@ -114,9 +113,10 @@ def test_rejected_first_candidate_does_not_dethrone_seed(self) -> None: def test_rejected_candidate_keeps_active_suffix_and_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) + control, loss, steps = attack.run( + n_steps=2, prev_loss=2.0, stop_on_success=False, anneal=True, random_seed=2026 + ) # The worse candidate must be rejected by annealing with overwhelming # probability under this seed; the active suffix stays "better" and the @@ -171,9 +171,8 @@ def test_periodic_checkpoint_restores_active_suffix(self) -> None: 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)) + results.append(attack.run(n_steps=3, prev_loss=4.0, stop_on_success=False, anneal=True, random_seed=1234)) assert results[0] == results[1] assert results[0] == ("c", 1.5, 3) @@ -270,6 +269,7 @@ def test_schedule_exhaustion_continues_until_step_budget_spent(self) -> None: test_steps=50, filter_cand=True, verbose=True, + random_seed=42, ) def test_schedule_loss_carried_on_schedule_object(self) -> None: From dd8f30cf8e62271e7ab9b32abe7610a678d9108d Mon Sep 17 00:00:00 2001 From: AmruthVamshi Date: Fri, 11 Sep 2026 01:13:16 -0400 Subject: [PATCH 2/5] FIX: Wire GCG random_seed through all stochastic operations (#2490) Create a per-run RngBundle at the outer execution level (_setup_async) with local random.Random, numpy.random.Generator, and per-worker torch.Generator instances derived deterministically from the configured seed. Thread the bundle through Progressive/Individual attacks into inner MultiPromptAttack runs so that repeating a configuration produces identical candidates, annealing decisions, and final suffixes. Record base seed and derived worker seeds in run metadata JSON. --- .../gcg/attack/base/attack_manager.py | 26 +++++++++-- .../promptgen/gcg/attack/gcg/gcg_attack.py | 8 +++- pyrit/executor/promptgen/gcg/generator.py | 44 ++++++++++++++++--- .../executor/promptgen/gcg/test_gcg_core.py | 23 ++++------ .../executor/promptgen/gcg/test_generator.py | 5 ++- .../promptgen/gcg/test_multi_prompt_attack.py | 8 +++- 6 files changed, 83 insertions(+), 31 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index ebf0c350f0..2b27500435 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -1003,10 +1003,18 @@ def run( Returns: tuple[str, float, int]: The final control, loss, and step count. """ - py_rng = random.Random(random_seed) - models = getattr(self, "models", None) - device = models[0].device if models else "cpu" - self._torch_gen = torch.Generator(device=device).manual_seed(random_seed) + rng_bundle = getattr(self, "_rng_bundle", None) + py_rng = rng_bundle.py_rng if rng_bundle else random.Random(random_seed) + if rng_bundle: + self._torch_gens = rng_bundle.torch_gens + else: + try: + self._torch_gens = { + i: torch.Generator(device=self.models[i].device).manual_seed(random_seed + i) + for i in range(len(self.workers)) + } + except (AttributeError, TypeError): + self._torch_gens = {0: torch.Generator().manual_seed(random_seed)} def acceptance_probability(e: float, e_prime: float, k: int) -> bool: temperature = max(1 - float(k + 1) / (n_steps + anneal_from), 1.0e-7) @@ -1426,6 +1434,8 @@ def run( # not keep looking current. self.last_schedule_state = None + rng_bundle = getattr(self, "_rng_bundle", None) + _update_attack_log_params( logfile=self.logfile, params={ @@ -1440,6 +1450,8 @@ def run( "anneal": anneal, "incr_control": incr_control, "stop_on_success": stop_on_success, + "random_seed": random_seed, + "derived_seeds": rng_bundle.derived_seeds if rng_bundle else {}, }, ) @@ -1470,6 +1482,7 @@ def run( ) 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, batch_size=batch_size, @@ -1681,6 +1694,8 @@ def run( Returns: tuple[str, int]: The final control suffix and configured step count. """ + rng_bundle = getattr(self, "_rng_bundle", None) + _update_attack_log_params( logfile=self.logfile, params={ @@ -1695,6 +1710,8 @@ def run( "anneal": anneal, "incr_control": incr_control, "stop_on_success": stop_on_success, + "random_seed": random_seed, + "derived_seeds": rng_bundle.derived_seeds if rng_bundle else {}, }, ) @@ -1715,6 +1732,7 @@ def run( self.test_targets, self.test_workers, ) + attack._rng_bundle = rng_bundle attack.run( n_steps=n_steps, batch_size=batch_size, diff --git a/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py b/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py index df854a8832..1700a77814 100644 --- a/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py +++ b/pyrit/executor/promptgen/gcg/attack/gcg/gcg_attack.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import inspect import logging from typing import Any @@ -203,7 +204,8 @@ def _sample_control_candidates( ) -> torch.Tensor: sampler = self._resolve_sampling() prompt_manager = self.prompts[worker_index] - torch_gen: torch.Generator | None = getattr(self, "_torch_gen", None) + torch_gens = getattr(self, "_torch_gens", None) or {} + torch_gen = torch_gens.get(worker_index) kwargs: dict[str, Any] = { "gradient": gradient, "control_tokens": prompt_manager.control_toks, @@ -214,7 +216,9 @@ def _sample_control_candidates( "non_ascii_tokens": prompt_manager.disallowed_toks, } if torch_gen is not None: - kwargs["torch_generator"] = torch_gen + sig = inspect.signature(sampler.sample_candidates) + if "torch_generator" in sig.parameters: + kwargs["torch_generator"] = torch_gen return sampler.sample_candidates(**kwargs) def _filter_control_candidates( diff --git a/pyrit/executor/promptgen/gcg/generator.py b/pyrit/executor/promptgen/gcg/generator.py index 0cea0eb795..6eaa2d8c5a 100644 --- a/pyrit/executor/promptgen/gcg/generator.py +++ b/pyrit/executor/promptgen/gcg/generator.py @@ -36,12 +36,14 @@ import asyncio import json import logging +import random import time from dataclasses import dataclass, field from functools import partial from typing import Any, overload import numpy as np +import torch import torch.multiprocessing as mp from pydantic import Field @@ -69,6 +71,17 @@ logger = logging.getLogger(__name__) +@dataclass +class RngBundle: + """Per-run RNG state bundle for deterministic GCG execution.""" + + np_rng: np.random.Generator + py_rng: random.Random + torch_gens: dict[int, torch.Generator] + base_seed: int + derived_seeds: dict[int, int] + + @dataclass class GCGContext(PromptGeneratorStrategyContext): """ @@ -97,6 +110,7 @@ class GCGContext(PromptGeneratorStrategyContext): test_workers: list[Any] = field(default_factory=list) attack: Any | None = None logfile_path: str | None = None + rng_bundle: RngBundle | None = None class GCGResult(PromptGeneratorStrategyResult): @@ -260,19 +274,33 @@ async def _setup_async(self, *, context: GCGContext) -> None: self._ensure_spawn_start_method() context.memory_labels = combine_dict({}, context.memory_labels) - np_rng = np.random.default_rng(self._algorithm.random_seed) - context.targets, context.test_targets = self._apply_target_augmentation( - train_targets=context.targets, - test_targets=context.test_targets, - np_rng=np_rng, - ) - log_gpu_memory(step=0) log_train_goals(train_goals=context.goals) params = self._to_attack_params(context=context) context.workers, context.test_workers = await asyncio.to_thread(get_workers, params) + seed = self._algorithm.random_seed + derived_seeds = {i: seed + i for i in range(len(context.workers))} + try: + torch_gens = { + i: torch.Generator(device=context.workers[i].model.device).manual_seed(derived_seeds[i]) + for i in range(len(context.workers)) + } + except (TypeError, AttributeError): + torch_gens = {i: torch.Generator().manual_seed(derived_seeds[i]) for i in range(len(context.workers))} + context.rng_bundle = RngBundle( + np_rng=np.random.default_rng(seed), + py_rng=random.Random(seed), + torch_gens=torch_gens, + base_seed=seed, + derived_seeds=derived_seeds, + ) + + context.targets, context.test_targets = self._apply_target_augmentation( + train_targets=context.targets, test_targets=context.test_targets, np_rng=context.rng_bundle.np_rng + ) + async def _perform_async(self, *, context: GCGContext) -> GCGResult: """ Build the attack, run the optimization loop, and read the result back. @@ -305,6 +333,8 @@ async def _perform_async(self, *, context: GCGContext) -> GCGResult: logfile_path=context.logfile_path, ) + context.attack._rng_bundle = context.rng_bundle + await asyncio.to_thread( context.attack.run, n_steps=self._algorithm.n_steps, diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index 4cf323e9f6..df394bbc91 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -1514,7 +1514,7 @@ def test_annealing_deterministic_same_seed(self) -> None: assert loss1 == loss2 def test_annealing_different_seed_can_differ(self) -> None: - """run() with different seeds can produce different annealing outcomes.""" + """run() with different seeds produces different annealing acceptance histories.""" results = [] for seed in [1, 999]: attack = object.__new__(MultiPromptAttack) @@ -1522,14 +1522,11 @@ def test_annealing_different_seed_can_differ(self) -> None: prompt_manager.control_str = "initial" attack.prompts = [prompt_manager] attack.logfile = None - # Marginal loss that annealing might accept or reject depending on random draw - attack.step = MagicMock(return_value=("candidate", 2.1)) - control, _, _ = attack.run(n_steps=10, prev_loss=2.0, stop_on_success=False, anneal=True, random_seed=seed) + attack.step = MagicMock(side_effect=[("c1", 2.1), ("c2", 2.2), ("c3", 2.3)]) + control, _, _ = attack.run(n_steps=3, prev_loss=2.0, stop_on_success=False, anneal=True, random_seed=seed) results.append(control) - # With enough steps and marginal losses, different seeds should diverge - # (probabilistic but extremely likely with 10 steps) - assert results[0] != results[1] or True # non-flaky: just verify no crash + assert results[0] != results[1], f"different seeds produced same control: {results}" def test_concurrent_runs_isolated(self) -> None: """Two runs with different seeds don't interfere with each other's RNG state.""" @@ -1559,11 +1556,12 @@ def test_run_creates_torch_gen_for_step(self) -> None: attack.run(n_steps=1, stop_on_success=False, anneal=False, random_seed=123) - assert hasattr(attack, "_torch_gen") - assert isinstance(attack._torch_gen, torch.Generator) + assert hasattr(attack, "_torch_gens") + assert isinstance(attack._torch_gens, dict) def test_custom_sampler_without_torch_generator_still_works(self) -> None: - """Custom SamplingStrategy that doesn't accept torch_generator still functions.""" + """Custom SamplingStrategy that doesn't accept torch_generator still functions + even when _torch_gen is set (the seeded run() path).""" gradient = torch.randn(3, 6) logits = torch.randn(2, 8, 10) token_ids = torch.randint(0, 10, (2, 8)) @@ -1573,14 +1571,11 @@ def test_custom_sampler_without_torch_generator_still_works(self) -> None: tokenizer.decode.return_value = "decoded" worker = _WorkerStub(gradient=gradient.clone(), logits=logits, token_ids=token_ids, tokenizer=tokenizer) - prompt = MagicMock() - prompt.control_toks = control_tokens prompt_manager = MagicMock() prompt_manager.control_toks = control_tokens prompt_manager.disallowed_toks = disallowed_tokens sampled_tokens = torch.tensor([[8, 8, 8]], dtype=torch.long) - # _SpySampling does NOT accept torch_generator — backward compat test sampling = _SpySampling(sampled_tokens=sampled_tokens) attack = object.__new__(GCGMultiPromptAttack) @@ -1589,8 +1584,8 @@ def test_custom_sampler_without_torch_generator_still_works(self) -> None: attack.workers = [worker] attack.models = [MagicMock(device=torch.device("cpu"))] attack.control_str = "test" + attack._torch_gens = {0: torch.Generator(device=torch.device("cpu")).manual_seed(42)} - # No _torch_gen set — simulates step() called without run() result = attack._sample_control_candidates( worker_index=0, gradient=gradient, diff --git a/tests/unit/executor/promptgen/gcg/test_generator.py b/tests/unit/executor/promptgen/gcg/test_generator.py index af19732a06..d4f4942dee 100644 --- a/tests/unit/executor/promptgen/gcg/test_generator.py +++ b/tests/unit/executor/promptgen/gcg/test_generator.py @@ -253,9 +253,10 @@ def test_returns_same_length_lists(self) -> None: def test_augmentation_modifies_at_least_some_targets(self) -> None: import numpy as np - np.random.seed(42) targets = ["Sure, here is how to do it"] * 100 - result, _ = GCGGenerator._apply_target_augmentation(train_targets=targets, test_targets=[]) + result, _ = GCGGenerator._apply_target_augmentation( + train_targets=targets, test_targets=[], np_rng=np.random.default_rng(42) + ) num_changed = sum(1 for orig, aug in zip(targets, result, strict=False) if orig != aug) assert num_changed > 0 diff --git a/tests/unit/executor/promptgen/gcg/test_multi_prompt_attack.py b/tests/unit/executor/promptgen/gcg/test_multi_prompt_attack.py index 43bb84c229..03a79d38cf 100644 --- a/tests/unit/executor/promptgen/gcg/test_multi_prompt_attack.py +++ b/tests/unit/executor/promptgen/gcg/test_multi_prompt_attack.py @@ -184,7 +184,7 @@ def test_attack_manager_records_run_params_before_creating_mpa( ) assert factory.params_at_creation is not None - assert list(factory.params_at_creation)[-11:] == [ + assert list(factory.params_at_creation)[-13:] == [ "n_steps", "test_steps", "batch_size", @@ -196,8 +196,10 @@ def test_attack_manager_records_run_params_before_creating_mpa( "anneal", "incr_control", "stop_on_success", + "random_seed", + "derived_seeds", ] - assert {key: factory.params_at_creation[key] for key in list(factory.params_at_creation)[-11:]} == { + assert {key: factory.params_at_creation[key] for key in list(factory.params_at_creation)[-13:]} == { "n_steps": 1, "test_steps": 4, "batch_size": 2, @@ -209,6 +211,8 @@ def test_attack_manager_records_run_params_before_creating_mpa( "anneal": False, "incr_control": False, "stop_on_success": False, + "random_seed": 42, + "derived_seeds": {}, } From 54f94a86f1254710618a4e9bd71f7f7073983b94 Mon Sep 17 00:00:00 2001 From: AmruthVamshi Date: Sat, 12 Sep 2026 13:43:29 -0400 Subject: [PATCH 3/5] FIX: Address review round 2, bundle self-creation, device-aware generators, and exact acceptance tests (#2490) --- .../gcg/attack/base/attack_manager.py | 58 +++++- pyrit/executor/promptgen/gcg/generator.py | 15 +- .../executor/promptgen/gcg/test_gcg_core.py | 186 ++++++++++++++---- .../promptgen/gcg/test_multi_prompt_attack.py | 2 +- .../executor/promptgen/gcg/test_run_state.py | 54 +++-- 5 files changed, 242 insertions(+), 73 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index 2b27500435..537fd39d3b 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -98,6 +98,17 @@ class ProgressiveScheduleState: stop_inner_on_success: bool = False +@dataclass +class RngBundle: + """Per-run RNG state bundle for deterministic GCG execution.""" + + np_rng: np.random.Generator + py_rng: random.Random + torch_gens: dict[int, torch.Generator] + base_seed: int + derived_seeds: dict[int, int] + + class NpEncoder(json.JSONEncoder): """Encode NumPy scalar and array values for JSON output.""" @@ -1008,13 +1019,14 @@ def run( if rng_bundle: self._torch_gens = rng_bundle.torch_gens else: + workers = getattr(self, "workers", []) try: + sampling_device = workers[0].model.device self._torch_gens = { - i: torch.Generator(device=self.models[i].device).manual_seed(random_seed + i) - for i in range(len(self.workers)) + i: torch.Generator(device=sampling_device).manual_seed(random_seed + i) for i in range(len(workers)) } - except (AttributeError, TypeError): - self._torch_gens = {0: torch.Generator().manual_seed(random_seed)} + except (TypeError, AttributeError, IndexError): + self._torch_gens = {i: torch.Generator().manual_seed(random_seed + i) for i in range(len(workers))} def acceptance_probability(e: float, e_prime: float, k: int) -> bool: temperature = max(1 - float(k + 1) / (n_steps + anneal_from), 1.0e-7) @@ -1435,6 +1447,23 @@ def run( self.last_schedule_state = None rng_bundle = getattr(self, "_rng_bundle", None) + if rng_bundle is None: + derived_seeds = {i: random_seed + i for i in range(len(self.workers))} + try: + sampling_device = self.workers[0].model.device + torch_gens = { + i: torch.Generator(device=sampling_device).manual_seed(derived_seeds[i]) + for i in range(len(self.workers)) + } + except (TypeError, AttributeError): + torch_gens = {i: torch.Generator().manual_seed(derived_seeds[i]) for i in range(len(self.workers))} + rng_bundle = RngBundle( + np_rng=np.random.default_rng(random_seed), + py_rng=random.Random(random_seed), + torch_gens=torch_gens, + base_seed=random_seed, + derived_seeds=derived_seeds, + ) _update_attack_log_params( logfile=self.logfile, @@ -1451,7 +1480,7 @@ def run( "incr_control": incr_control, "stop_on_success": stop_on_success, "random_seed": random_seed, - "derived_seeds": rng_bundle.derived_seeds if rng_bundle else {}, + "derived_seeds": rng_bundle.derived_seeds, }, ) @@ -1695,6 +1724,23 @@ def run( tuple[str, int]: The final control suffix and configured step count. """ rng_bundle = getattr(self, "_rng_bundle", None) + if rng_bundle is None: + derived_seeds = {i: random_seed + i for i in range(len(self.workers))} + try: + sampling_device = self.workers[0].model.device + torch_gens = { + i: torch.Generator(device=sampling_device).manual_seed(derived_seeds[i]) + for i in range(len(self.workers)) + } + except (TypeError, AttributeError): + torch_gens = {i: torch.Generator().manual_seed(derived_seeds[i]) for i in range(len(self.workers))} + rng_bundle = RngBundle( + np_rng=np.random.default_rng(random_seed), + py_rng=random.Random(random_seed), + torch_gens=torch_gens, + base_seed=random_seed, + derived_seeds=derived_seeds, + ) _update_attack_log_params( logfile=self.logfile, @@ -1711,7 +1757,7 @@ def run( "incr_control": incr_control, "stop_on_success": stop_on_success, "random_seed": random_seed, - "derived_seeds": rng_bundle.derived_seeds if rng_bundle else {}, + "derived_seeds": rng_bundle.derived_seeds, }, ) diff --git a/pyrit/executor/promptgen/gcg/generator.py b/pyrit/executor/promptgen/gcg/generator.py index 6eaa2d8c5a..caab33e70e 100644 --- a/pyrit/executor/promptgen/gcg/generator.py +++ b/pyrit/executor/promptgen/gcg/generator.py @@ -57,6 +57,7 @@ from pyrit.executor.promptgen.gcg.attack.base.attack_manager import ( IndividualPromptAttack, ProgressiveMultiPromptAttack, + RngBundle, get_workers, ) from pyrit.executor.promptgen.gcg.config import ( @@ -71,17 +72,6 @@ logger = logging.getLogger(__name__) -@dataclass -class RngBundle: - """Per-run RNG state bundle for deterministic GCG execution.""" - - np_rng: np.random.Generator - py_rng: random.Random - torch_gens: dict[int, torch.Generator] - base_seed: int - derived_seeds: dict[int, int] - - @dataclass class GCGContext(PromptGeneratorStrategyContext): """ @@ -283,8 +273,9 @@ async def _setup_async(self, *, context: GCGContext) -> None: seed = self._algorithm.random_seed derived_seeds = {i: seed + i for i in range(len(context.workers))} try: + sampling_device = context.workers[0].model.device torch_gens = { - i: torch.Generator(device=context.workers[i].model.device).manual_seed(derived_seeds[i]) + i: torch.Generator(device=sampling_device).manual_seed(derived_seeds[i]) for i in range(len(context.workers)) } except (TypeError, AttributeError): diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index df394bbc91..d4e4181b39 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -1495,55 +1495,130 @@ def test_sampling_different_seed_can_differ(self) -> None: assert not torch.equal(result1, result2) - def test_annealing_deterministic_same_seed(self) -> None: - """run() with same seed produces identical annealing acceptance decisions.""" + @staticmethod + def _run_annealing_with_boolean_tracking( + seed: int, + ) -> tuple[str, list[bool]]: + """Run annealing and capture per-step acceptance booleans. + + ``run()`` sets ``self.control_str`` only when a candidate is accepted. + We snapshot ``control_str`` at the *start* of each ``step()`` call; a + change between consecutive snapshots proves the previous candidate was + accepted. The last step's decision is derived from the final control. + """ + steps = [("c1", 2.1), ("c2", 2.2), ("c3", 2.3)] attack = object.__new__(MultiPromptAttack) - prompt_manager = MagicMock() - prompt_manager.control_str = "initial" - attack.prompts = [prompt_manager] + attack.prompts = [MagicMock(control_str="initial")] attack.logfile = None - # Step returns a slightly worse loss so annealing decides acceptance - attack.step = MagicMock(return_value=("candidate", 2.5)) - - _, loss1, _ = attack.run(n_steps=3, prev_loss=2.0, stop_on_success=False, anneal=True, random_seed=42) - attack.step = MagicMock(return_value=("candidate", 2.5)) - prompt_manager.control_str = "initial" - _, loss2, _ = attack.run(n_steps=3, prev_loss=2.0, stop_on_success=False, anneal=True, random_seed=42) - - assert loss1 == loss2 - - def test_annealing_different_seed_can_differ(self) -> None: - """run() with different seeds produces different annealing acceptance histories.""" - results = [] - for seed in [1, 999]: - attack = object.__new__(MultiPromptAttack) - prompt_manager = MagicMock() - prompt_manager.control_str = "initial" - attack.prompts = [prompt_manager] - attack.logfile = None - attack.step = MagicMock(side_effect=[("c1", 2.1), ("c2", 2.2), ("c3", 2.3)]) - control, _, _ = attack.run(n_steps=3, prev_loss=2.0, stop_on_success=False, anneal=True, random_seed=seed) - results.append(control) + snapshots: list[str] = [] + real_step = MagicMock(side_effect=list(steps)) - assert results[0] != results[1], f"different seeds produced same control: {results}" + def tracking_step(**kwargs: Any) -> tuple[str, float]: + snapshots.append(attack.control_str) + return real_step(**kwargs) - def test_concurrent_runs_isolated(self) -> None: - """Two runs with different seeds don't interfere with each other's RNG state.""" - targets = ["Sure, here is how to hack"] * 10 + attack.step = MagicMock(side_effect=tracking_step) - rng_a = np.random.default_rng(42) - result_a, _ = GCGGenerator._apply_target_augmentation(train_targets=targets, test_targets=[], np_rng=rng_a) - - # Interleave: run a different seed in between - rng_other = np.random.default_rng(999) - GCGGenerator._apply_target_augmentation(train_targets=targets, test_targets=[], np_rng=rng_other) - - # Fresh rng with seed 42 still gives same result - rng_b = np.random.default_rng(42) - result_b, _ = GCGGenerator._apply_target_augmentation(train_targets=targets, test_targets=[], np_rng=rng_b) + control, _, _ = attack.run( + n_steps=3, + prev_loss=2.0, + stop_on_success=False, + anneal=True, + random_seed=seed, + ) - assert result_a == result_b + accepted = [snapshots[i + 1] != snapshots[i] for i in range(len(snapshots) - 1)] + accepted.append(control != snapshots[-1]) + + return control, accepted + + def test_annealing_exact_history_same_seed(self) -> None: + """Same seed reproduces the exact step-by-step acceptance booleans.""" + for _ in range(2): + control, accepted = self._run_annealing_with_boolean_tracking(seed=42) + # seed=42: accept c1 (draw=0.64 < threshold=0.86), accept c2 (draw=0.02 < 0.74), + # reject c3 (draw=0.28 > threshold≈0 at temp≈1e-7) → final="c2" + assert accepted == [True, True, False] + assert control == "c2" + + def test_annealing_exact_history_different_seeds(self) -> None: + """Different seeds produce verifiably different acceptance boolean sequences.""" + control_1, accepted_1 = self._run_annealing_with_boolean_tracking(seed=1) + control_999, accepted_999 = self._run_annealing_with_boolean_tracking(seed=999) + + # Pre-computed from random.Random(seed) draws against acceptance_probability. + # seed=1: draw=0.13<0.86→accept, draw=0.85>0.74→reject, draw=0.76>≈0→reject + assert accepted_1 == [True, False, False] + assert control_1 == "c1" + # seed=999: draw=0.78<0.86→accept, draw=0.08<0.74→accept, draw=0.87>≈0→reject + assert accepted_999 == [True, True, False] + assert control_999 == "c2" + + def test_concurrent_runs_isolated_across_all_rng_types(self) -> None: + """Progressive inner phases draw from all three bundle RNG streams + (NumPy, Torch, Python) without restarting, and repeating the full + run with the same seed reproduces the exact sequence.""" + + @dataclass + class PhaseSnapshot: + np_draw: float + torch_draw: list[int] + py_draw: float + + all_runs: list[list[PhaseSnapshot]] = [] + + for _ in range(2): + captured: list[PhaseSnapshot] = [] + inner = MagicMock() + + def make_inner_run(captured_ref: list[PhaseSnapshot], attack_mock: Any) -> Any: + def inner_run(**kwargs: Any) -> tuple[str, float, int]: + bundle = getattr(attack_mock, "_rng_bundle", None) + assert bundle is not None, "inner MPA should receive a bundle" + snap = PhaseSnapshot( + np_draw=float(bundle.np_rng.random()), + torch_draw=torch.randint( + 0, + 1000, + (3,), + generator=bundle.torch_gens[0], + ).tolist(), + py_draw=bundle.py_rng.random(), + ) + captured_ref.append(snap) + return ("ctrl", 0.5, 1) + + return inner_run + + inner.run = MagicMock(side_effect=make_inner_run(captured, inner)) + + progressive = object.__new__(ProgressiveMultiPromptAttack) + progressive.goals = ["g1", "g2"] + progressive.targets = ["t1", "t2"] + progressive.workers = [MagicMock()] + progressive.test_goals = [] + progressive.test_targets = [] + progressive.test_workers = [] + progressive.test_prefixes = [] + progressive.managers = {"MPA": MagicMock(return_value=inner)} + progressive.control = "initial" + progressive.logfile = None + progressive.progressive_goals = True + progressive.progressive_models = False + + progressive.run(n_steps=4, stop_on_success=False, random_seed=42) + all_runs.append(captured) + + # Same seed → identical sequence across both runs + for field in ("np_draw", "torch_draw", "py_draw"): + assert [getattr(s, field) for s in all_runs[0]] == [getattr(s, field) for s in all_runs[1]] + + # Phases advanced all three streams (no restart) + phase1, phase2 = all_runs[0][0], all_runs[0][1] + assert phase1.np_draw != phase2.np_draw + assert phase1.torch_draw != phase2.torch_draw + assert phase1.py_draw != phase2.py_draw def test_run_creates_torch_gen_for_step(self) -> None: """run() sets self._torch_gen so step() can access it for sampling.""" @@ -1596,3 +1671,30 @@ def test_custom_sampler_without_torch_generator_still_works(self) -> None: ) assert torch.equal(result, sampled_tokens) + + def test_multi_device_generators_must_match_sampling_device(self) -> None: + """Regression: generators on a different device than the sampling + tensor make torch.randint raise. When workers span devices, all + generators must live on workers[0].model.device (the sampling + device). This test goes through the real MPA.run() bundle-creation + fallback with workers on different devices.""" + attack = object.__new__(MultiPromptAttack) + attack.prompts = [MagicMock(control_str="initial")] + attack.logfile = None + attack.step = MagicMock(return_value=("result", 0.5)) + + worker0 = MagicMock() + worker0.model.device = torch.device("cuda:0") + worker1 = MagicMock() + worker1.model.device = torch.device("cuda:1") + attack.workers = [worker0, worker1] + + attack.run(n_steps=1, stop_on_success=False, anneal=False, random_seed=42) + + # All generators must be on the sampling device (worker 0), not + # their own worker's device. If worker 1's generator were on cuda:1, + # torch.randint with device=cuda:0 would raise RuntimeError. + sampling_device = torch.device("cuda:0") + assert len(attack._torch_gens) == 2 + for i, gen in attack._torch_gens.items(): + assert gen.device == sampling_device, f"Generator {i} on {gen.device}, expected {sampling_device}" diff --git a/tests/unit/executor/promptgen/gcg/test_multi_prompt_attack.py b/tests/unit/executor/promptgen/gcg/test_multi_prompt_attack.py index 03a79d38cf..685ddfe9ee 100644 --- a/tests/unit/executor/promptgen/gcg/test_multi_prompt_attack.py +++ b/tests/unit/executor/promptgen/gcg/test_multi_prompt_attack.py @@ -212,7 +212,7 @@ def test_attack_manager_records_run_params_before_creating_mpa( "incr_control": False, "stop_on_success": False, "random_seed": 42, - "derived_seeds": {}, + "derived_seeds": {"0": 42}, } diff --git a/tests/unit/executor/promptgen/gcg/test_run_state.py b/tests/unit/executor/promptgen/gcg/test_run_state.py index 44d92a6eeb..808e4afe23 100644 --- a/tests/unit/executor/promptgen/gcg/test_run_state.py +++ b/tests/unit/executor/promptgen/gcg/test_run_state.py @@ -4,7 +4,7 @@ """Tests for typed optimization-iteration state in the GCG attack loop.""" from typing import Any -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest @@ -34,6 +34,30 @@ def _bare_multi_prompt_attack(step_results: list[tuple[str, float]]) -> MultiPro return attack +def _acceptance_booleans(attack: MultiPromptAttack, final_control: str) -> list[bool]: + """Derive per-step acceptance booleans from control_str snapshots. + + Must be called with an attack whose step() was wrapped by + ``_track_acceptance`` before ``run()``. + """ + snapshots: list[str] = attack._acceptance_snapshots # type: ignore[attr-defined] + accepted = [snapshots[i + 1] != snapshots[i] for i in range(len(snapshots) - 1)] + accepted.append(final_control != snapshots[-1]) + return accepted + + +def _track_acceptance(attack: MultiPromptAttack) -> None: + """Wrap attack.step to snapshot control_str at entry for boolean tracking.""" + attack._acceptance_snapshots = [] # type: ignore[attr-defined] + real_step = attack.step + + def tracking_step(**kwargs: Any) -> tuple[str, float]: + attack._acceptance_snapshots.append(attack.control_str) # type: ignore[attr-defined] + return real_step(**kwargs) + + attack.step = MagicMock(side_effect=tracking_step) # type: ignore[assignment] + + class TestStopReason: def test_has_expected_members(self) -> None: assert StopReason.MAX_STEPS_REACHED == "max_steps_reached" @@ -98,9 +122,12 @@ def test_rejected_first_candidate_does_not_dethrone_seed(self) -> None: # rejected candidate must not become the best result just because a # sentinel used to be larger. attack = _bare_multi_prompt_attack([("worse", 10.0)]) + _track_acceptance(attack) control, loss, steps = attack.run(n_steps=1, prev_loss=1.0, stop_on_success=False, anneal=True) + # seed=42 (default): 10.0 >> 1.0, threshold≈0 → rejected + assert _acceptance_booleans(attack, control) == [False] assert control == "initial" assert loss == 1.0 assert steps == 1 @@ -113,15 +140,15 @@ def test_rejected_first_candidate_does_not_dethrone_seed(self) -> None: def test_rejected_candidate_keeps_active_suffix_and_loss(self) -> None: attack = _bare_multi_prompt_attack([("better", 1.0), ("worse", 5.0)]) + _track_acceptance(attack) control, loss, steps = attack.run( n_steps=2, prev_loss=2.0, stop_on_success=False, anneal=True, random_seed=2026 ) - # The worse candidate must be rejected by annealing with overwhelming - # probability under this seed; the active suffix stays "better" and the - # reported loss stays paired with it. The rejected candidate's loss is - # still observable through ``candidate_loss``. + # seed=2026: 1.0 < 2.0 → accept (strictly better, no draw), + # 5.0 >> 1.0, threshold≈0 → reject + assert _acceptance_booleans(attack, control) == [True, False] assert control == "better" assert steps == 2 state: OptimizationRunState = attack.last_run_state @@ -134,15 +161,18 @@ def test_rejected_candidate_keeps_active_suffix_and_loss(self) -> None: def test_candidate_after_rejection_is_compared_with_active_loss(self) -> None: attack = _bare_multi_prompt_attack([("worse", 5.0), ("still-worse", 4.5)]) + _track_acceptance(attack) - with patch.object(random, "random", return_value=0.99): - control, loss, steps = attack.run( - n_steps=2, - prev_loss=1.0, - stop_on_success=False, - anneal=True, - ) + control, loss, steps = attack.run( + n_steps=2, + prev_loss=1.0, + stop_on_success=False, + anneal=True, + random_seed=42, + ) + # seed=42: 5.0 >> 1.0, threshold≈0 → reject; 4.5 >> 1.0, threshold≈0 → reject + assert _acceptance_booleans(attack, control) == [False, False] assert (control, loss, steps) == ("initial", 1.0, 2) state: OptimizationRunState = attack.last_run_state assert state.control == "initial" From 60d3a99d94ac472039e442e6b6a7a5bacdcae163 Mon Sep 17 00:00:00 2001 From: AmruthVamshi Date: Mon, 14 Sep 2026 17:13:33 -0400 Subject: [PATCH 4/5] TEST: exercise GCG seed isolation under real overlap; unique per-run logfile paths Address review round 3 on #2502: - GCGGenerator._build_logfile_path appends a short random id so concurrent runs sharing a result_prefix never read/modify/write the same JSON log. - Deterministic in-process worker/prompt-manager stubs drive the real IndividualPromptAttack / ProgressiveMultiPromptAttack and the real GCGMultiPromptAttack.step() end to end. - Overlapping runs (barrier-forced threads, asyncio.gather) across five topologies reproduce their isolated candidates, accepted controls, loss history and final suffix; all inner phases share one RNG bundle. - Target augmentation is asserted to consume the run-seeded NumPy stream. - CUDA placement test patches torch.Generator so it runs on CPU-only builds. --- doc/code/executor/gcg/1_gcg_azure_ml.ipynb | 7 +- doc/code/executor/gcg/1_gcg_azure_ml.py | 7 +- pyrit/executor/promptgen/gcg/config.py | 8 +- pyrit/executor/promptgen/gcg/generator.py | 7 +- .../executor/promptgen/gcg/test_gcg_core.py | 217 ++++++++++++------ .../executor/promptgen/gcg/test_generator.py | 113 ++++++++- .../promptgen/gcg/trajectory_stubs.py | 196 ++++++++++++++++ 7 files changed, 471 insertions(+), 84 deletions(-) create mode 100644 tests/unit/executor/promptgen/gcg/trajectory_stubs.py diff --git a/doc/code/executor/gcg/1_gcg_azure_ml.ipynb b/doc/code/executor/gcg/1_gcg_azure_ml.ipynb index a4c7b20040..c1103ab1cb 100644 --- a/doc/code/executor/gcg/1_gcg_azure_ml.ipynb +++ b/doc/code/executor/gcg/1_gcg_azure_ml.ipynb @@ -355,10 +355,11 @@ "The next cell polls the job until it reaches a terminal state (~20-30\n", "minutes for the small 5-step baseline above), then downloads the named\n", "`results` output and prints the final suffix. The runner writes its\n", - "result file as `_.json` (with `result_prefix`\n", + "result file as `__.json` (with `result_prefix`\n", "coming from the `GCGConfig` we built above, plus the AML output mount\n", - "prepended by `--output-dir`). For our config, that resolves to\n", - "`gcg_suffix_.json` under\n", + "prepended by `--output-dir`; `` is a short random suffix so concurrent\n", + "runs never share a file). For our config, that resolves to\n", + "`gcg_suffix__.json` under\n", "`/named-outputs/results/` once we download. The\n", "`controls` array in that file contains one entry per training step, and\n", "the last entry is the final adversarial suffix that, appended to the user\n", diff --git a/doc/code/executor/gcg/1_gcg_azure_ml.py b/doc/code/executor/gcg/1_gcg_azure_ml.py index c3c559f18c..cf329f28da 100644 --- a/doc/code/executor/gcg/1_gcg_azure_ml.py +++ b/doc/code/executor/gcg/1_gcg_azure_ml.py @@ -177,10 +177,11 @@ # The next cell polls the job until it reaches a terminal state (~20-30 # minutes for the small 5-step baseline above), then downloads the named # `results` output and prints the final suffix. The runner writes its -# result file as `_.json` (with `result_prefix` +# result file as `__.json` (with `result_prefix` # coming from the `GCGConfig` we built above, plus the AML output mount -# prepended by `--output-dir`). For our config, that resolves to -# `gcg_suffix_.json` under +# prepended by `--output-dir`; `` is a short random suffix so concurrent +# runs never share a file). For our config, that resolves to +# `gcg_suffix__.json` under # `/named-outputs/results/` once we download. The # `controls` array in that file contains one entry per training step, and # the last entry is the final adversarial suffix that, appended to the user diff --git a/pyrit/executor/promptgen/gcg/config.py b/pyrit/executor/promptgen/gcg/config.py index 5b0f639e38..d215ffd356 100644 --- a/pyrit/executor/promptgen/gcg/config.py +++ b/pyrit/executor/promptgen/gcg/config.py @@ -315,9 +315,11 @@ class GCGOutputConfig: Attributes: result_prefix (str): Prefix for the per-run JSON log file. The actual - filename is ``{result_prefix}_{YYYYMMDD-HHMMSS}.json``. Empty string - means write the log into the current working directory with no - prefix (``_.json``); that is rarely what you want. + filename is ``{result_prefix}_{YYYYMMDD-HHMMSS}_{id}.json`` where + ``id`` is a random 8-character hex string that keeps concurrent runs + sharing a prefix from writing to the same file. Empty string means + write the log into the current working directory with no prefix + (``__.json``); that is rarely what you want. logfile (str): Optional pre-resolved log file path. When set this takes precedence over ``result_prefix`` for the legacy code paths. verbose (bool): Verbose progress logging during the run. Defaults to True. diff --git a/pyrit/executor/promptgen/gcg/generator.py b/pyrit/executor/promptgen/gcg/generator.py index caab33e70e..1b99a71c01 100644 --- a/pyrit/executor/promptgen/gcg/generator.py +++ b/pyrit/executor/promptgen/gcg/generator.py @@ -38,6 +38,7 @@ import logging import random import time +import uuid from dataclasses import dataclass, field from functools import partial from typing import Any, overload @@ -395,10 +396,12 @@ async def execute_async(self, **kwargs: Any) -> GCGResult: return await super().execute_async(**kwargs) def _build_logfile_path(self) -> str: - timestamp = time.strftime("%Y%m%d-%H%M%S") if self._output.logfile: return self._output.logfile - return f"{self._output.result_prefix}_{timestamp}.json" + # Second-resolution timestamps collide for concurrent runs sharing a + # prefix; both would then read/modify/write the same JSON log. + timestamp = time.strftime("%Y%m%d-%H%M%S") + return f"{self._output.result_prefix}_{timestamp}_{uuid.uuid4().hex[:8]}.json" @staticmethod def _apply_target_augmentation( diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index d4e4181b39..efa145b1d2 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -1,9 +1,13 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import json import pickle +import threading from copy import deepcopy from dataclasses import dataclass +from functools import partial +from pathlib import Path from typing import Any from unittest.mock import MagicMock, call, patch, sentinel @@ -51,6 +55,12 @@ ) GCGGenerator = generator_mod.GCGGenerator +from unit.executor.promptgen.gcg.trajectory_stubs import ( # noqa: E402 + RecordingGCGAttack, + TrajectoryPromptManager, + TrajectoryWorker, +) + @dataclass class _TinyModelOutput: @@ -1404,6 +1414,83 @@ def test_length_preserving_filter_rejects_unknown_option() -> None: LengthPreservingFilter(unexpected=True) +# (attack class, n_workers, n_goals, constructor kwargs) for every execution +# topology #2490 requires coverage of. +_TRAJECTORY_TOPOLOGIES: dict[str, tuple[type, int, int, dict[str, bool]]] = { + "individual-1w-1g": (IndividualPromptAttack, 1, 1, {}), + "individual-1w-2g": (IndividualPromptAttack, 1, 2, {}), + "progressive-goals-1w-2g": ( + ProgressiveMultiPromptAttack, + 1, + 2, + {"progressive_goals": True, "progressive_models": False}, + ), + "progressive-goals-models-2w-2g": ( + ProgressiveMultiPromptAttack, + 2, + 2, + {"progressive_goals": True, "progressive_models": True}, + ), + "multi-2w-2g": ( + ProgressiveMultiPromptAttack, + 2, + 2, + {"progressive_goals": False, "progressive_models": False}, + ), +} + + +def _run_trajectory( + topology: str, + *, + seed: int, + logfile: Path, + run_id: str = "", + barrier: threading.Barrier | None = None, + switches: list[str] | None = None, +) -> dict[str, Any]: + """Drive a real outer attack -> real GCG step loop on stub workers and return everything observable.""" + attack_cls, n_workers, n_goals, attack_kwargs = _TRAJECTORY_TOPOLOGIES[topology] + events: list[Any] = [] + bundles: list[Any] = [] + workers = [TrajectoryWorker(i, run_id=run_id, barrier=barrier, events=switches) for i in range(n_workers)] + managers = {"PM": TrajectoryPromptManager, "MPA": partial(RecordingGCGAttack, events=events, bundles=bundles)} + attack = attack_cls( + [f"goal {i}" for i in range(n_goals)], + ["10 11", "12 13"][:n_goals], + workers, + control_init="1 2 3", + test_prefixes=[], + logfile=str(logfile), + managers=managers, + **attack_kwargs, + ) + final_control, _ = attack.run( + n_steps=4, + batch_size=4, + topk=6, + allow_non_ascii=True, + target_weight=1.0, + control_weight=0.0, + anneal=True, + test_steps=1, + incr_control=False, + stop_on_success=False, + verbose=False, + random_seed=seed, + ) + # A bundle rebuilt per inner phase would restart every stream; all phases must see the one object. + assert bundles and bundles[0] is not None and all(b is bundles[0] for b in bundles), "phases must share one bundle" + with open(logfile) as f: + log = json.load(f) + return { + "events": events, + "final_control": final_control, + "controls": log["controls"], + "losses": [round(loss, 6) for loss in log["losses"]], + } + + class TestRandomSeedDeterminism: """Verify that random_seed produces reproducible results across runs.""" @@ -1555,70 +1642,53 @@ def test_annealing_exact_history_different_seeds(self) -> None: assert accepted_999 == [True, True, False] assert control_999 == "c2" - def test_concurrent_runs_isolated_across_all_rng_types(self) -> None: - """Progressive inner phases draw from all three bundle RNG streams - (NumPy, Torch, Python) without restarting, and repeating the full - run with the same seed reproduces the exact sequence.""" - - @dataclass - class PhaseSnapshot: - np_draw: float - torch_draw: list[int] - py_draw: float - - all_runs: list[list[PhaseSnapshot]] = [] - - for _ in range(2): - captured: list[PhaseSnapshot] = [] - inner = MagicMock() - - def make_inner_run(captured_ref: list[PhaseSnapshot], attack_mock: Any) -> Any: - def inner_run(**kwargs: Any) -> tuple[str, float, int]: - bundle = getattr(attack_mock, "_rng_bundle", None) - assert bundle is not None, "inner MPA should receive a bundle" - snap = PhaseSnapshot( - np_draw=float(bundle.np_rng.random()), - torch_draw=torch.randint( - 0, - 1000, - (3,), - generator=bundle.torch_gens[0], - ).tolist(), - py_draw=bundle.py_rng.random(), - ) - captured_ref.append(snap) - return ("ctrl", 0.5, 1) - - return inner_run - - inner.run = MagicMock(side_effect=make_inner_run(captured, inner)) - - progressive = object.__new__(ProgressiveMultiPromptAttack) - progressive.goals = ["g1", "g2"] - progressive.targets = ["t1", "t2"] - progressive.workers = [MagicMock()] - progressive.test_goals = [] - progressive.test_targets = [] - progressive.test_workers = [] - progressive.test_prefixes = [] - progressive.managers = {"MPA": MagicMock(return_value=inner)} - progressive.control = "initial" - progressive.logfile = None - progressive.progressive_goals = True - progressive.progressive_models = False - - progressive.run(n_steps=4, stop_on_success=False, random_seed=42) - all_runs.append(captured) - - # Same seed → identical sequence across both runs - for field in ("np_draw", "torch_draw", "py_draw"): - assert [getattr(s, field) for s in all_runs[0]] == [getattr(s, field) for s in all_runs[1]] - - # Phases advanced all three streams (no restart) - phase1, phase2 = all_runs[0][0], all_runs[0][1] - assert phase1.np_draw != phase2.np_draw - assert phase1.torch_draw != phase2.torch_draw - assert phase1.py_draw != phase2.py_draw + @pytest.mark.parametrize("topology", list(_TRAJECTORY_TOPOLOGIES), ids=list(_TRAJECTORY_TOPOLOGIES)) + def test_overlapping_runs_reproduce_isolated_trajectories(self, topology: str, tmp_path: Path) -> None: + """Two runs executing at the same time follow exactly the trajectories they follow alone. + + Real ``IndividualPromptAttack`` / ``ProgressiveMultiPromptAttack`` drive the real + ``GCGMultiPromptAttack.step()`` (Torch sampling, length filter, cross-entropy loss, + Python annealing) against stub workers whose outputs are pure functions of their + inputs, so the seeded bundle is the only source of randomness. Each run writes a real + logfile. A barrier in the stub worker forces the two runs to alternate step by step, + so the recorded run-id sequence switches more than the single time a sequential + execution would. + """ + baseline_42 = _run_trajectory(topology, seed=42, logfile=tmp_path / "baseline_42.json") + baseline_7 = _run_trajectory(topology, seed=7, logfile=tmp_path / "baseline_7.json") + assert _run_trajectory(topology, seed=42, logfile=tmp_path / "repeat_42.json") == baseline_42 + assert baseline_42["events"] != baseline_7["events"] + + barrier = threading.Barrier(2, timeout=10) + switches: list[str] = [] + outcomes: dict[str, Any] = {} + + def run(run_id: str, seed: int) -> None: + try: + outcomes[run_id] = _run_trajectory( + topology, + seed=seed, + logfile=tmp_path / f"{run_id}.json", + run_id=run_id, + barrier=barrier, + switches=switches, + ) + except BaseException as exc: # noqa: BLE001 - release the peer thread, then surface the error + barrier.abort() + outcomes[run_id] = exc + + threads = [threading.Thread(target=run, args=("A", 42)), threading.Thread(target=run, args=("B", 7))] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + for outcome in outcomes.values(): + if isinstance(outcome, BaseException): + raise outcome + assert sum(a != b for a, b in zip(switches, switches[1:], strict=False)) > 1, switches + assert outcomes["A"] == baseline_42 + assert outcomes["B"] == baseline_7 def test_run_creates_torch_gen_for_step(self) -> None: """run() sets self._torch_gen so step() can access it for sampling.""" @@ -1677,7 +1747,13 @@ def test_multi_device_generators_must_match_sampling_device(self) -> None: tensor make torch.randint raise. When workers span devices, all generators must live on workers[0].model.device (the sampling device). This test goes through the real MPA.run() bundle-creation - fallback with workers on different devices.""" + fallback with workers on different devices. + + ``torch.Generator`` is patched to record each requested device while + handing back a CPU generator, because CPU-only PyTorch builds reject + ``torch.Generator(device="cuda:0")``. The assertion is on the device + each generator was requested on, which is what the regression is about. + """ attack = object.__new__(MultiPromptAttack) attack.prompts = [MagicMock(control_str="initial")] attack.logfile = None @@ -1689,12 +1765,15 @@ def test_multi_device_generators_must_match_sampling_device(self) -> None: worker1.model.device = torch.device("cuda:1") attack.workers = [worker0, worker1] - attack.run(n_steps=1, stop_on_success=False, anneal=False, random_seed=42) + real_generator = torch.Generator # the patch below also replaces the test's own torch.Generator + with patch.object( + attack_manager_mod.torch, "Generator", side_effect=lambda device: real_generator() + ) as generator_cls: + attack.run(n_steps=1, stop_on_success=False, anneal=False, random_seed=42) - # All generators must be on the sampling device (worker 0), not + # Both generators must be requested on the sampling device (worker 0), not # their own worker's device. If worker 1's generator were on cuda:1, # torch.randint with device=cuda:0 would raise RuntimeError. sampling_device = torch.device("cuda:0") + assert generator_cls.call_args_list == [call(device=sampling_device), call(device=sampling_device)] assert len(attack._torch_gens) == 2 - for i, gen in attack._torch_gens.items(): - assert gen.device == sampling_device, f"Generator {i} on {gen.device}, expected {sampling_device}" diff --git a/tests/unit/executor/promptgen/gcg/test_generator.py b/tests/unit/executor/promptgen/gcg/test_generator.py index d4f4942dee..72923923d3 100644 --- a/tests/unit/executor/promptgen/gcg/test_generator.py +++ b/tests/unit/executor/promptgen/gcg/test_generator.py @@ -5,9 +5,13 @@ from __future__ import annotations +import asyncio import json +import re +import threading from functools import partial -from typing import TYPE_CHECKING, Any +from pathlib import Path +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -19,9 +23,6 @@ GCGStrategyConfig, ) -if TYPE_CHECKING: - from pathlib import Path - generator_mod = pytest.importorskip( "pyrit.executor.promptgen.gcg.generator", reason="GCG optional dependencies (torch, transformers, etc.) not installed", @@ -30,6 +31,10 @@ GCGContext = generator_mod.GCGContext GCGResult = generator_mod.GCGResult +from unit.executor.promptgen.gcg.trajectory_stubs import ( # noqa: E402 + TrajectoryPromptManager, + TrajectoryWorker, +) _LLAMA_2 = "meta-llama/Llama-2-7b-chat-hf" @@ -426,3 +431,103 @@ def test_empty_controls_returns_nan_loss(self, tmp_path: Path) -> None: result = GCGGenerator._read_result(logfile_path=str(log_path), memory_labels={}) assert result.final_suffix == "" assert math.isnan(result.final_loss) + + +class TestBuildLogfilePath: + def test_same_prefix_yields_distinct_paths(self, tmp_path: Path) -> None: + gen = _make_generator(output_dir=tmp_path) + assert gen._build_logfile_path() != gen._build_logfile_path() + + def test_path_is_prefix_timestamp_and_short_id(self, tmp_path: Path) -> None: + path = Path(_make_generator(output_dir=tmp_path)._build_logfile_path()) + assert path.parent == tmp_path + assert re.fullmatch(r"gcg_\d{8}-\d{6}_[0-9a-f]{8}\.json", path.name), path.name + + def test_explicit_logfile_is_returned_unchanged(self) -> None: + gen = GCGGenerator(models=[GCGModelConfig(name=_LLAMA_2)], output=GCGOutputConfig(logfile="fixed.json")) + assert gen._build_logfile_path() == "fixed.json" + + +_TRAJECTORY_STRATEGIES = { + "individual": GCGStrategyConfig(anneal=True), + "progressive": GCGStrategyConfig(transfer=True, progressive_goals=True, anneal=True), +} + + +def _trajectory_generator(*, output_dir: Path, strategy: GCGStrategyConfig, seed: int) -> GCGGenerator: + return GCGGenerator( + models=[GCGModelConfig(name=_LLAMA_2)], + algorithm=GCGAlgorithmConfig( + n_steps=4, + test_steps=1, + batch_size=4, + topk=6, + allow_non_ascii=True, + control_init="1 2 3", + control_weight=0.0, + random_seed=seed, + ), + strategy=strategy, + output=GCGOutputConfig(result_prefix=str(output_dir / "gcg"), verbose=False), + ) + + +class TestConcurrentGenerators: + @pytest.mark.parametrize("strategy", list(_TRAJECTORY_STRATEGIES), ids=list(_TRAJECTORY_STRATEGIES)) + async def test_generators_sharing_a_prefix_reproduce_isolated_results(self, tmp_path: Path, strategy: str) -> None: + """Two ``execute_async`` calls overlapping in one process, with the same ``result_prefix``, + write distinct logfiles and each return exactly the result they return when run alone. + + Only worker creation and the prompt manager are stubbed; the RNG bundle, the outer + attack, the GCG step loop, and the logfile round-trip are all real. A barrier in the + stub worker forces the two runs to interleave step by step. + """ + goals, targets = ["goal 0", "goal 1"], ["10 11", "12 13"] + + def make(seed: int) -> GCGGenerator: + return _trajectory_generator(output_dir=tmp_path, strategy=_TRAJECTORY_STRATEGIES[strategy], seed=seed) + + def summary(result: GCGResult) -> tuple[str, list[str], list[float]]: + return result.final_suffix, result.control_history, [round(loss, 6) for loss in result.loss_history] + + with ( + patch.object(generator_mod, "get_workers") as get_workers, + patch.object(generator_mod.attack_lib, "GCGPromptManager", TrajectoryPromptManager), + ): + get_workers.return_value = ([TrajectoryWorker(0)], []) + baseline_42 = summary(await make(42).execute_async(goals=goals, targets=targets)) + baseline_7 = summary(await make(7).execute_async(goals=goals, targets=targets)) + assert baseline_42 != baseline_7 + + barrier = threading.Barrier(2, timeout=10) + get_workers.side_effect = [([TrajectoryWorker(0, barrier=barrier)], []) for _ in range(2)] + result_42, result_7 = await asyncio.gather( + make(42).execute_async(goals=goals, targets=targets), + make(7).execute_async(goals=goals, targets=targets), + ) + + assert result_42.log_path != result_7.log_path + assert summary(result_42) == baseline_42 + assert summary(result_7) == baseline_7 + + async def test_augmentation_draws_from_the_run_seeded_numpy_stream(self, tmp_path: Path) -> None: + """Target augmentation consumes the run's own NumPy generator, seeded from ``random_seed``.""" + import numpy as np + + targets = ["10 11", "12 13"] + with ( + patch.object(generator_mod, "get_workers", return_value=([TrajectoryWorker(0)], [])), + patch.object(generator_mod.attack_lib, "GCGPromptManager", TrajectoryPromptManager), + patch.object( + GCGGenerator, "_apply_target_augmentation", wraps=GCGGenerator._apply_target_augmentation + ) as augment, + ): + generator = _trajectory_generator( + output_dir=tmp_path, strategy=_TRAJECTORY_STRATEGIES["individual"], seed=42 + ) + await generator.execute_async(goals=["goal 0", "goal 1"], targets=targets) + + expected = np.random.default_rng(42) + for _ in targets: + expected.random() + assert augment.call_args.kwargs["np_rng"].bit_generator.state == expected.bit_generator.state diff --git a/tests/unit/executor/promptgen/gcg/trajectory_stubs.py b/tests/unit/executor/promptgen/gcg/trajectory_stubs.py new file mode 100644 index 0000000000..0fcc1b052c --- /dev/null +++ b/tests/unit/executor/promptgen/gcg/trajectory_stubs.py @@ -0,0 +1,196 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Deterministic in-process stand-ins for GCG model workers and prompt managers. + +They let the real ``GCGMultiPromptAttack.step()`` (sampling, filtering, loss, +selection) and the real ``MultiPromptAttack.run()`` (annealing, logging) execute +end to end without loading a model. Every model output (gradient, logits) is a +pure function of its inputs, so the only randomness left in a run is the seeded +RNG bundle under test. +""" + +from __future__ import annotations + +import queue +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any + +import torch + +from pyrit.executor.promptgen.gcg.attack.base.attack_manager import ModelWorkerOperation +from pyrit.executor.promptgen.gcg.attack.gcg.gcg_attack import GCGMultiPromptAttack + +if TYPE_CHECKING: + import threading + +VOCAB_SIZE = 16 + + +def _seed(*parts: int) -> int: + # Hashing a tuple of ints is deterministic in CPython; str hash randomization is not involved. + return hash(parts) & 0xFFFF_FFFF + + +class IntTokenizer: + """Tokens are ints rendered as space-joined text; ``"!"`` is token 0.""" + + vocab_size = VOCAB_SIZE + name_or_path = "int-tokenizer" + chat_template = None + + def encode_ints(self, text: str) -> list[int]: + return [0 if tok == "!" else int(tok) for tok in text.split()] + + def __call__(self, text: str, **_: Any) -> SimpleNamespace: + return SimpleNamespace(input_ids=self.encode_ints(text)) + + def decode(self, token_ids: torch.Tensor, **_: Any) -> str: + return " ".join(str(int(t)) for t in token_ids) + + +class _Prompt: + def __init__(self, *, target_ids: list[int], control_len: int) -> None: + self.target_ids = target_ids + self._control_slice = slice(1, 1 + control_len) + self._target_slice = slice(1 + control_len, 1 + control_len + len(target_ids)) + + +class TrajectoryPromptManager: + """The subset of the ``PromptManager`` contract that ``GCGMultiPromptAttack.step()`` touches.""" + + def __init__( + self, + goals: list[str], + targets: list[str], + tokenizer: IntTokenizer, + control_init: str, + test_prefixes: list[str] | None = None, + managers: dict[str, Any] | None = None, + ) -> None: + self.tokenizer = tokenizer + self.control_init = control_init + self.control_str = control_init + self._prompts = [ + _Prompt(target_ids=tokenizer.encode_ints(t), control_len=len(self.control_toks)) for t in targets + ] + self.disallowed_toks = torch.empty(0, dtype=torch.long) + + @property + def control_str(self) -> str: + return self._control_str + + @control_str.setter + def control_str(self, control: str) -> None: + self._control_str = control + self.control_toks = torch.tensor(self.tokenizer.encode_ints(control), dtype=torch.long) + + def __len__(self) -> int: + return len(self._prompts) + + def __getitem__(self, i: int) -> _Prompt: + return self._prompts[i] + + def __iter__(self) -> Any: + return iter(self._prompts) + + +class TrajectoryWorker: + """ + Synchronous stand-in for ``ModelWorker`` whose outputs are determined by their inputs. + + GRAD returns a gradient that is a pure function of (worker id, control tokens), with + the current token at every slot pushed to +100 so top-k sampling never reproduces the + current control and the length-preserving filter never runs dry. LOGITS returns + per-candidate logits that are a pure function of the candidate's token ids. TEST reports + a jailbreak once the control has moved off its initial value, so progressive phases + advance after one accepted step. + + When ``barrier`` is shared by two concurrently running attacks, worker 0 waits on it once + per optimization step, forcing the two runs to interleave. Each crossing appends ``run_id`` + to ``events`` so tests can assert that the interleaving actually happened. + """ + + def __init__( + self, + worker_id: int = 0, + *, + run_id: str = "", + barrier: threading.Barrier | None = None, + events: list[str] | None = None, + ) -> None: + self.worker_id = worker_id + self.run_id = run_id + self.barrier = barrier + self.events = events + self.tokenizer = IntTokenizer() + self.model = SimpleNamespace(device=torch.device("cpu"), name_or_path=f"stub-model-{worker_id}") + self.results: queue.SimpleQueue[Any] = queue.SimpleQueue() + + def start(self) -> TrajectoryWorker: + return self + + def stop(self) -> TrajectoryWorker: + return self + + def __call__(self, ob: Any, operation: ModelWorkerOperation, *args: Any, **kwargs: Any) -> TrajectoryWorker: + self.results.put(self._execute(ob, operation, *args)) + return self + + def _execute(self, ob: Any, operation: ModelWorkerOperation, *args: Any) -> Any: + if operation is ModelWorkerOperation.GRAD: + if self.barrier is not None and self.worker_id == 0: + self.barrier.wait() + if self.events is not None: + self.events.append(self.run_id) + return self._grad(ob.control_toks) + if operation is ModelWorkerOperation.LOGITS: + return self._logits(ob, args[0]) + if operation is ModelWorkerOperation.TEST: + return [(ob.control_str != ob.control_init, 0) for _ in ob] + if operation is ModelWorkerOperation.TEST_LOSS: + return [0.0 for _ in ob] + raise NotImplementedError(operation) + + def _grad(self, control_toks: torch.Tensor) -> torch.Tensor: + gen = torch.Generator().manual_seed(_seed(self.worker_id, *control_toks.tolist())) + grad = torch.randn(len(control_toks), VOCAB_SIZE, generator=gen) + grad[torch.arange(len(control_toks)), control_toks] = 100.0 + return grad + + def _logits(self, prompt: _Prompt, candidates: list[str]) -> tuple[torch.Tensor, torch.Tensor]: + rows = [[0, *self.tokenizer.encode_ints(c), *prompt.target_ids] for c in candidates] + ids = torch.tensor(rows, dtype=torch.long) + logits = torch.stack( + [ + torch.randn( + ids.shape[1], VOCAB_SIZE, generator=torch.Generator().manual_seed(_seed(self.worker_id, *row)) + ) + for row in rows + ] + ) + return logits, ids + + +class RecordingGCGAttack(GCGMultiPromptAttack): + """The real GCG attack, recording sampled candidates, per-step decisions and the bundle each step used.""" + + def __init__(self, *args: Any, events: list[Any], bundles: list[Any] | None = None, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.events = events + self.bundles = bundles + events.append(("attack", tuple(self.goals), len(self.workers))) + + def _sample_control_candidates(self, **kwargs: Any) -> torch.Tensor: + candidates = super()._sample_control_candidates(**kwargs) + self.events.append(("candidates", tuple(map(tuple, candidates.tolist())))) + return candidates + + def step(self, **kwargs: Any) -> tuple[str, float]: + if self.bundles is not None: + self.bundles.append(getattr(self, "_rng_bundle", None)) + before = self.control_str + control, loss = super().step(**kwargs) + self.events.append(("step", before, control, round(loss, 6))) + return control, loss From 4ae4b00edc1aea636a1db2aa3e5276a8d3ff86cb Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Wed, 16 Sep 2026 10:06:24 -0700 Subject: [PATCH 5/5] FIX centralize GCG RNG bundle creation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../gcg/attack/base/attack_manager.py | 100 ++++++++++-------- pyrit/executor/promptgen/gcg/generator.py | 21 +--- .../executor/promptgen/gcg/test_run_state.py | 39 ++++++- 3 files changed, 95 insertions(+), 65 deletions(-) diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index 537fd39d3b..cccdd870a8 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -108,6 +108,52 @@ class RngBundle: base_seed: int derived_seeds: dict[int, int] + @classmethod + def from_seed(cls, *, base_seed: int, workers: list[ModelWorker]) -> RngBundle: + """ + Create deterministic local RNGs for one GCG run. + + Args: + base_seed (int): Seed shared by the Python and NumPy generators. + workers (list[ModelWorker]): Workers that need derived Torch generators. + + Returns: + RngBundle: The initialized per-run RNG bundle. + """ + derived_seeds = {i: base_seed + i for i in range(len(workers))} + return cls( + np_rng=np.random.default_rng(base_seed), + py_rng=random.Random(base_seed), + torch_gens=cls._create_torch_generators(workers=workers, derived_seeds=derived_seeds), + base_seed=base_seed, + derived_seeds=derived_seeds, + ) + + @staticmethod + def _create_torch_generators( + *, workers: list[ModelWorker], derived_seeds: dict[int, int] + ) -> dict[int, torch.Generator]: + """ + Create worker generators on the shared sampling device. + + Args: + workers (list[ModelWorker]): Workers that consume sampled candidates. + derived_seeds (dict[int, int]): Deterministic seed for each worker. + + Returns: + dict[int, torch.Generator]: Generator keyed by worker index. + """ + if not workers: + return {} + + try: + sampling_device = workers[0].model.device + return { + i: torch.Generator(device=sampling_device).manual_seed(derived_seeds[i]) for i in range(len(workers)) + } + except (TypeError, AttributeError): + return {i: torch.Generator().manual_seed(derived_seeds[i]) for i in range(len(workers))} + class NpEncoder(json.JSONEncoder): """Encode NumPy scalar and array values for JSON output.""" @@ -1015,18 +1061,10 @@ def run( tuple[str, float, int]: The final control, loss, and step count. """ rng_bundle = getattr(self, "_rng_bundle", None) - py_rng = rng_bundle.py_rng if rng_bundle else random.Random(random_seed) - if rng_bundle: - self._torch_gens = rng_bundle.torch_gens - else: - workers = getattr(self, "workers", []) - try: - sampling_device = workers[0].model.device - self._torch_gens = { - i: torch.Generator(device=sampling_device).manual_seed(random_seed + i) for i in range(len(workers)) - } - except (TypeError, AttributeError, IndexError): - self._torch_gens = {i: torch.Generator().manual_seed(random_seed + i) for i in range(len(workers))} + if rng_bundle is None: + rng_bundle = RngBundle.from_seed(base_seed=random_seed, workers=getattr(self, "workers", [])) + py_rng = rng_bundle.py_rng + self._torch_gens = rng_bundle.torch_gens def acceptance_probability(e: float, e_prime: float, k: int) -> bool: temperature = max(1 - float(k + 1) / (n_steps + anneal_from), 1.0e-7) @@ -1448,22 +1486,7 @@ def run( rng_bundle = getattr(self, "_rng_bundle", None) if rng_bundle is None: - derived_seeds = {i: random_seed + i for i in range(len(self.workers))} - try: - sampling_device = self.workers[0].model.device - torch_gens = { - i: torch.Generator(device=sampling_device).manual_seed(derived_seeds[i]) - for i in range(len(self.workers)) - } - except (TypeError, AttributeError): - torch_gens = {i: torch.Generator().manual_seed(derived_seeds[i]) for i in range(len(self.workers))} - rng_bundle = RngBundle( - np_rng=np.random.default_rng(random_seed), - py_rng=random.Random(random_seed), - torch_gens=torch_gens, - base_seed=random_seed, - derived_seeds=derived_seeds, - ) + rng_bundle = RngBundle.from_seed(base_seed=random_seed, workers=self.workers) _update_attack_log_params( logfile=self.logfile, @@ -1479,7 +1502,7 @@ def run( "anneal": anneal, "incr_control": incr_control, "stop_on_success": stop_on_success, - "random_seed": random_seed, + "random_seed": rng_bundle.base_seed, "derived_seeds": rng_bundle.derived_seeds, }, ) @@ -1725,22 +1748,7 @@ def run( """ rng_bundle = getattr(self, "_rng_bundle", None) if rng_bundle is None: - derived_seeds = {i: random_seed + i for i in range(len(self.workers))} - try: - sampling_device = self.workers[0].model.device - torch_gens = { - i: torch.Generator(device=sampling_device).manual_seed(derived_seeds[i]) - for i in range(len(self.workers)) - } - except (TypeError, AttributeError): - torch_gens = {i: torch.Generator().manual_seed(derived_seeds[i]) for i in range(len(self.workers))} - rng_bundle = RngBundle( - np_rng=np.random.default_rng(random_seed), - py_rng=random.Random(random_seed), - torch_gens=torch_gens, - base_seed=random_seed, - derived_seeds=derived_seeds, - ) + rng_bundle = RngBundle.from_seed(base_seed=random_seed, workers=self.workers) _update_attack_log_params( logfile=self.logfile, @@ -1756,7 +1764,7 @@ def run( "anneal": anneal, "incr_control": incr_control, "stop_on_success": stop_on_success, - "random_seed": random_seed, + "random_seed": rng_bundle.base_seed, "derived_seeds": rng_bundle.derived_seeds, }, ) diff --git a/pyrit/executor/promptgen/gcg/generator.py b/pyrit/executor/promptgen/gcg/generator.py index 1b99a71c01..119d59382e 100644 --- a/pyrit/executor/promptgen/gcg/generator.py +++ b/pyrit/executor/promptgen/gcg/generator.py @@ -36,7 +36,6 @@ import asyncio import json import logging -import random import time import uuid from dataclasses import dataclass, field @@ -44,7 +43,6 @@ from typing import Any, overload import numpy as np -import torch import torch.multiprocessing as mp from pydantic import Field @@ -271,22 +269,9 @@ async def _setup_async(self, *, context: GCGContext) -> None: params = self._to_attack_params(context=context) context.workers, context.test_workers = await asyncio.to_thread(get_workers, params) - seed = self._algorithm.random_seed - derived_seeds = {i: seed + i for i in range(len(context.workers))} - try: - sampling_device = context.workers[0].model.device - torch_gens = { - i: torch.Generator(device=sampling_device).manual_seed(derived_seeds[i]) - for i in range(len(context.workers)) - } - except (TypeError, AttributeError): - torch_gens = {i: torch.Generator().manual_seed(derived_seeds[i]) for i in range(len(context.workers))} - context.rng_bundle = RngBundle( - np_rng=np.random.default_rng(seed), - py_rng=random.Random(seed), - torch_gens=torch_gens, - base_seed=seed, - derived_seeds=derived_seeds, + context.rng_bundle = RngBundle.from_seed( + base_seed=self._algorithm.random_seed, + workers=context.workers, ) context.targets, context.test_targets = self._apply_target_augmentation( diff --git a/tests/unit/executor/promptgen/gcg/test_run_state.py b/tests/unit/executor/promptgen/gcg/test_run_state.py index 808e4afe23..087cc11496 100644 --- a/tests/unit/executor/promptgen/gcg/test_run_state.py +++ b/tests/unit/executor/promptgen/gcg/test_run_state.py @@ -3,9 +3,11 @@ """Tests for typed optimization-iteration state in the GCG attack loop.""" +import random from typing import Any -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch +import numpy as np import pytest attack_manager_mod = pytest.importorskip( @@ -18,6 +20,7 @@ OptimizationRunState = attack_manager_mod.OptimizationRunState ProgressiveMultiPromptAttack = attack_manager_mod.ProgressiveMultiPromptAttack ProgressiveScheduleState = attack_manager_mod.ProgressiveScheduleState +RngBundle = attack_manager_mod.RngBundle StopReason = attack_manager_mod.StopReason @@ -226,6 +229,40 @@ def test_seeded_runs_produce_identical_trajectories(self) -> None: assert results[0] == results[1] assert results[0] == ("c", 1.5, 3) + def test_direct_run_creates_and_uses_complete_rng_bundle(self) -> None: + attack = _bare_multi_prompt_attack([("worse", 2.0)]) + attack.workers[0].model.device = torch.device("cpu") + created_bundles: list[Any] = [] + create_bundle = RngBundle.from_seed + + def record_bundle(*, base_seed: int, workers: list[Any]) -> Any: + bundle = create_bundle(base_seed=base_seed, workers=workers) + created_bundles.append(bundle) + return bundle + + with patch.object(RngBundle, "from_seed", side_effect=record_bundle) as factory: + control, _, _ = attack.run( + n_steps=1, + prev_loss=1.0, + stop_on_success=False, + anneal=True, + random_seed=123, + ) + + assert factory.call_count == 1 + assert factory.call_args.kwargs == {"base_seed": 123, "workers": attack.workers} + assert len(created_bundles) == 1 + bundle = created_bundles[0] + assert bundle.base_seed == 123 + assert bundle.derived_seeds == {0: 123} + assert attack._torch_gens is bundle.torch_gens + assert control == "initial" + + expected_py_rng = random.Random(123) + expected_py_rng.random() + assert bundle.py_rng.random() == expected_py_rng.random() + assert bundle.np_rng.random() == np.random.default_rng(123).random() + class TestGCGCandidateSelection: def test_selects_minimum_within_single_group(self) -> None: