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/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index c50e72f406..cccdd870a8 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -98,6 +98,63 @@ 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] + + @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.""" @@ -995,6 +1052,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. @@ -1002,10 +1060,15 @@ def run( Returns: tuple[str, float, int]: The final control, loss, and step count. """ + rng_bundle = getattr(self, "_rng_bundle", None) + 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) - 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: @@ -1378,6 +1441,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. @@ -1409,6 +1473,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. @@ -1418,6 +1484,10 @@ def run( # not keep looking current. self.last_schedule_state = None + rng_bundle = getattr(self, "_rng_bundle", None) + if rng_bundle is None: + rng_bundle = RngBundle.from_seed(base_seed=random_seed, workers=self.workers) + _update_attack_log_params( logfile=self.logfile, params={ @@ -1432,6 +1502,8 @@ def run( "anneal": anneal, "incr_control": incr_control, "stop_on_success": stop_on_success, + "random_seed": rng_bundle.base_seed, + "derived_seeds": rng_bundle.derived_seeds, }, ) @@ -1462,6 +1534,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, @@ -1477,6 +1550,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 @@ -1634,6 +1708,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. @@ -1665,10 +1740,16 @@ 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. """ + rng_bundle = getattr(self, "_rng_bundle", None) + if rng_bundle is None: + rng_bundle = RngBundle.from_seed(base_seed=random_seed, workers=self.workers) + _update_attack_log_params( logfile=self.logfile, params={ @@ -1683,6 +1764,8 @@ def run( "anneal": anneal, "incr_control": incr_control, "stop_on_success": stop_on_success, + "random_seed": rng_bundle.base_seed, + "derived_seeds": rng_bundle.derived_seeds, }, ) @@ -1703,6 +1786,7 @@ def run( self.test_targets, self.test_workers, ) + attack._rng_bundle = rng_bundle attack.run( n_steps=n_steps, batch_size=batch_size, @@ -1719,6 +1803,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..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 @@ -104,6 +105,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 +116,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 +130,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 +204,22 @@ 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_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, + "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: + 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( self, 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/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..119d59382e 100644 --- a/pyrit/executor/promptgen/gcg/generator.py +++ b/pyrit/executor/promptgen/gcg/generator.py @@ -37,6 +37,7 @@ import json import logging import time +import uuid from dataclasses import dataclass, field from functools import partial from typing import Any, overload @@ -55,6 +56,7 @@ from pyrit.executor.promptgen.gcg.attack.base.attack_manager import ( IndividualPromptAttack, ProgressiveMultiPromptAttack, + RngBundle, get_workers, ) from pyrit.executor.promptgen.gcg.config import ( @@ -97,6 +99,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,17 +263,21 @@ async def _setup_async(self, *, context: GCGContext) -> None: self._ensure_spawn_start_method() context.memory_labels = combine_dict({}, context.memory_labels) - context.targets, context.test_targets = self._apply_target_augmentation( - train_targets=context.targets, - test_targets=context.test_targets, - ) - 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) + context.rng_bundle = RngBundle.from_seed( + base_seed=self._algorithm.random_seed, + workers=context.workers, + ) + + 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. @@ -303,6 +310,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, @@ -318,6 +327,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) @@ -371,16 +381,19 @@ 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( *, 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 +401,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 +410,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 ded5749cff..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 @@ -41,6 +45,21 @@ 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 + +from unit.executor.promptgen.gcg.trajectory_stubs import ( # noqa: E402 + RecordingGCGAttack, + TrajectoryPromptManager, + TrajectoryWorker, +) @dataclass @@ -1393,3 +1412,368 @@ 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) + + +# (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.""" + + 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) + + @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) + attack.prompts = [MagicMock(control_str="initial")] + attack.logfile = None + + snapshots: list[str] = [] + real_step = MagicMock(side_effect=list(steps)) + + def tracking_step(**kwargs: Any) -> tuple[str, float]: + snapshots.append(attack.control_str) + return real_step(**kwargs) + + attack.step = MagicMock(side_effect=tracking_step) + + control, _, _ = attack.run( + n_steps=3, + prev_loss=2.0, + stop_on_success=False, + anneal=True, + random_seed=seed, + ) + + 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" + + @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.""" + 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_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 + 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)) + 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_manager = MagicMock() + prompt_manager.control_toks = control_tokens + prompt_manager.disallowed_toks = disallowed_tokens + + sampled_tokens = torch.tensor([[8, 8, 8]], dtype=torch.long) + 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" + attack._torch_gens = {0: torch.Generator(device=torch.device("cpu")).manual_seed(42)} + + 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) + + 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. + + ``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 + 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] + + 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) + + # 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 diff --git a/tests/unit/executor/promptgen/gcg/test_generator.py b/tests/unit/executor/promptgen/gcg/test_generator.py index af19732a06..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" @@ -253,9 +258,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 @@ -425,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/test_multi_prompt_attack.py b/tests/unit/executor/promptgen/gcg/test_multi_prompt_attack.py index 43bb84c229..685ddfe9ee 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": {"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 1b6f865b1f..087cc11496 100644 --- a/tests/unit/executor/promptgen/gcg/test_run_state.py +++ b/tests/unit/executor/promptgen/gcg/test_run_state.py @@ -7,6 +7,7 @@ from typing import Any from unittest.mock import MagicMock, patch +import numpy as np import pytest attack_manager_mod = pytest.importorskip( @@ -19,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 @@ -35,6 +37,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" @@ -99,9 +125,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 @@ -114,14 +143,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)]) - random.seed(2026) + _track_acceptance(attack) - 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 - # 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 +164,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" @@ -190,13 +223,46 @@ 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) + 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: @@ -289,6 +355,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: 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