MAINT GCG: make optimization iteration state explicit and typed

[fei (feiiiiii5)]

Purpose

Implements the structural refactor requested in #2416: a typed optimization-iteration state for the GCG loop, with the iteration lifecycle made explicit — while keeping public attack behavior and extension protocols unchanged.

Closes #2416

What changed

Typed state (attack_manager.py)

• StopReason enum: MAX_STEPS_REACHED, ALL_PROMPTS_JAILBROKEN
• OptimizationRunState dataclass: current suffix, best suffix, current/best loss, step counter, step runtime, stop reason. Exposed as MultiPromptAttack.last_run_state after run(), so callers and tests can programmatically inspect why/when optimization stopped
• ProgressiveScheduleState dataclass: goals/workers admitted so far, shared step counter, loss carry-over, inner stop flag. Exposed as ProgressiveMultiPromptAttack.last_schedule_state

Explicit phases with stable contracts

• MultiPromptAttack.run: stopping phase extracted to _all_training_prompts_jailbroken(); periodic logging phase extracted to _log_best_checkpoint() (same best-suffix swap-and-restore semantics as before, now with a finally guarantee)
• ProgressiveMultiPromptAttack.run: scheduling counters moved into ProgressiveScheduleState; final evaluation extracted to _finalize_progressive_run()
• GCGMultiPromptAttack.step: selection phase (flat argmin → group/in-batch decomposition) extracted to _select_best_candidate(). Candidate batches intentionally remain step-local to preserve the existing VRAM-bounding behavior noted in the loop comment

Compatibility

• All public signatures and return values are unchanged (run() still returns (control, loss, steps) / (control, steps)); the new attributes are additive
• No algorithm changes: acceptance/annealing, admission schedule, control-weight bumping, and logging cadence are byte-for-byte the same decisions as before, just routed through typed state

Validation

• New tests/unit/executor/promptgen/gcg/test_run_state.py (12 tests): stop-reason assignment (incl. that the final success check does not count as an executed step), best tracking under annealing rejection, periodic checkpoint restore contract, seeded identical trajectories, argmin decomposition across worker groups, progressive finalize path with call-through assertions on the inner attack log
• Full GCG unit suite: 211 passed (199 pre-existing + 12 new)
• ruff check + ruff format --check clean with the repo-pinned ruff v0.16.0

AI usage disclosure

• This PR was developed with AI assistance (Claude code agent). All code was human-reviewed, and every behavioral claim above is backed by the tests in this PR.

Human verification: I reviewed the diff hunk-by-hunk against the pre-refactor control flow, confirmed the annealing acceptance rule, admission ordering, weight bump threshold (<= 0.09), and return semantics are preserved, and re-derived the argmin index decomposition used in the selection tests.

[fei (feiiiiii5)]

Rebased onto current main (now includes #2364 logging consolidation and #2434 ty adoption):

• The extracted stopping/logging/finalize phases are adapted to the new module-level log helpers (_initialize_attack_log/_update_attack_log_params); no behavioral overlap remains with MAINT: Consolidate GCG attack manager logging #2364.
• While rebasing I caught and fixed a bug in my own branch: my initial extraction had renamed the inner-run loss unpack target, which silently broke the progressive schedule's annealing continuity (prev_loss was pinned to inf across admissions instead of feeding the inner result back). The seeded regression test now asserts the correct feed-back semantics.
• Validation on the rebased commit: GCG unit suite 217 passed (199 pre-existing incl. MAINT: Adopt ty 0.0.73 #2434's new tests + 12 new state tests), ty check clean under the repo's all = "error" config, ruff check + ruff format --check clean with the pinned ruff.

[Roman Lutz (romanlutz)]

Thanks for the submission! A few things need adjusting, though.

[fei (feiiiiii5)]

Addressed — thank you for the precise repro. The root cause is exactly as you described: when an inner run consumes the entire remaining step budget while a further goal/worker is still waiting to be admitted, the admission transition resets schedule.loss to inf right before the budget-exhausted loop exits, so the old guard ("loss was never updated") crashed a valid run instead of returning.

Fix: the guard now only requires a measurable loss while budget remains (steps_completed < n_steps). A budget-exhausted exit legitimately carries an inf loss for the goal that was admitted but never scored — returning is the correct outcome there.

Regression test reproduces the exact scenario (two progressive goals, n_steps=3, inner result ("ctrl", 0.75, 3)): goals_admitted=2, steps_completed=3, loss=inf, and the run returns ("ctrl", 3) instead of asserting. Pre-fix this test crashes; post-fix it passes. test_run_state.py + test_multi_prompt_attack.py: 27 passed; rebased onto current main.

[fei (feiiiiii5)]

Follow-up to the comment above: while preparing these fixes we consolidated two parallel work attempts on this branch. The budget-exhaustion case is now handled by gating the phase admissions and their sentinel resets on remaining step budget (the direction suggested in review), rather than relaxing the carried-loss assertion — see c8fd9a4. Gating keeps last_schedule_state honest (no admitted-but-never-scored phase recorded with an inf loss) and preserves the assertion's value as a carry-over regression guard. Finalize-on-success behavior is unchanged, and exact-budget regression tests now cover goal, worker, and control-weight transitions (tests/unit/executor/promptgen/gcg/test_run_state.py). All 226 tests in the GCG unit suite pass.

[Roman Lutz (romanlutz)]

This assertion uses inf to mean "the inner run never updated the loss," but a completed inner run can also report inf because of a non-finite model loss or numeric overflow. That case previously returned normally and now raises AssertionError. Could we track whether an inner result was received explicitly, or validate non-finite losses with a deliberate error at the source, rather than infer update state from the numeric value?