[cptnm3]

Pull Request

Issue

Fixes #1487

Description

Move the sampling lifecycle boundary from BaseSamplingStrategy.sample() into the SamplingStrategy base class.

SamplingStrategy.sample() now owns:

• requirement merging and deduplication
• sampling_id creation
• sampling_loop_start dispatch and effective loop budget resolution
• validation of hook-modified loop budgets
• exception handling and error-path lifecycle closure
• sampling_loop_end dispatch

Concrete strategies now implement _sample_impl() for their sampling algorithm. Shared helpers centralize sampling iteration and repair payload construction and hook dispatch.

Migrate BaseSamplingStrategy, BudgetForcingSamplingStrategy, SOFAISamplingStrategy, and majority-voting to the new contract.

This ensures every top-level sample() call emits exactly one enclosing sampling lifecycle, regardless of whether a strategy uses the base sampling loop, implements its own loop, or fans out into multiple inner samples.

Intentional behavior changes:

• Majority voting emits one enclosing lifecycle for the top-level majority-vote operation instead of one lifecycle per inner sample.

• Budget Forcing and SOFAI emit sampling iteration and repair events for their strategy-specific attempts and repairs.

• Budget Forcing and SOFAI do not perform repairs after the final allowed failed iteration, since the repaired action/context cannot be consumed, aligning them with BaseSamplingStrategy.

• sampling_loop_end fires for failures during lifecycle setup, including requirement merging, start-hook execution, and invalid hook-modified loop budgets, allowing lifecycle consumers to close error paths.

Add firing-site and regression coverage for:

• strategies overriding only _sample_impl()
• sampling_id propagation and correlation
• start-hook exceptions and lifecycle setup failures
• hook-modified effective loop budgets
• invalid effective loop budgets and error-path closure
• Budget Forcing iteration and repair events
• SOFAI S1/S2 iteration and repair behavior
• majority-voting iteration and repair events
• a single enclosing lifecycle for majority voting

Assisted-by: IBM Bob

Testing

• Tests added to the respective file if code was changed
• New code has 100% coverage if code was added
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

Attribution

• AI coding assistants used

---

Adding a new component, requirement, sampling strategy, or tool?

If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.

• Component
• Requirement
• Sampling Strategy
• Tool

NOTE: Please ensure you have an issue that has been acknowledged by a core contributor and routed you to open a pull request against this repository. Otherwise, please open an issue before continuing with this pull request.

[github-actions]

This comment is managed by a bot. Editing it is fine — checking off boxes, adding notes — but please leave the HTML comment marker on the first line alone, otherwise checklist updates will break.

Sampling Strategy PR Checklist

Use this checklist when adding or modifying sampling strategies in mellea/stdlib/sampling/.

Base Class

• Extends appropriate base class:
  • BaseSamplingStrategy if your changes are mostly modifying the repair and/or select_from_failure functions
  • SamplingStrategy if your changes involve a new sample method
  • Other defined sampling strategies if your implementation is similar to existing implementations

Return Value

• Returns a properly typed SamplingResult. Specifically, this means:
  • ModelOutputThunks in sample_generations are properly typed from the Component and the parsed_repr is the expected type.

Integration

• Strategy exported in mellea/stdlib/sampling/__init__.py

[ajbozarth]

Just got back from vacation, I'll deep dive review this by EOW. Thank you for the contribution

[ajbozarth]

Thanks for taking this on. Solid PR — cleanly mirrors the generate_from_raw split, and firing sampling_loop_end on setup failures (with tests) is a nice touch beyond the issue scope.

One item that has no diff anchor: docs/docs/community/building-extensions.md still teaches extension authors to override sample(). After this change a strategy overriding only sample() is missing _sample_impl and can't be instantiated — please update that page to override _sample_impl instead (leave docs/versioned_docs/** alone, it's a frozen version snapshot).

Inline: answers to your two questions, plus suggestions on @final, method naming, and declaring loop_budget/requirements on the base class.

[ajbozarth]

one docs nit otherwise this looks good, but I'll defer to @jakelorocco for final approval as he knows this code better than me and might catch things I missed

[planetf1]

Nicely structured refactor — the @final sample() / _sample() split is a clean way to guarantee the lifecycle wrapper can't be bypassed, and the new tests for setup-failure lifecycle closure are a real improvement. Two things below are worth fixing before merge; a couple of smaller ones are optional.

One doc note that doesn't have a line in this diff to anchor to: docs/docs/concepts/plugins.mdx documents the sampling_iteration/sampling_repair payload fields but doesn't list the new sample_index field yet — worth a follow-up so plugin authors know it exists.

[cptnm3]

I did not get a chance to work on requested changes last week. I'll make the changes by end of this week. Sorry for the delay!

[planetf1]

Two regressions from this refactor, plus a one-liner, below. Prior rounds' items checked against this head and are resolved.

[planetf1]

Not part of the diff, but same gap sample_index was meant to close:

• mellea/plugins/builtin_debug/sampling.py (~lines 88, 123): iteration/repair logs don't include sample_index — a majority-vote run logs [SAMPLING-ITER 1] per branch with no way to tell them apart
• docs/docs/observability/metrics.md: sampling-counters table wasn't updated for the same cardinality shift tracing.md already documents (successes/failures now fire once per call, not once per fan-out branch)

Neither blocks merge.

[cptnm3]

Hi @planetf1, I have addressed your comments, please let me know if there are any open issues that need to be addressed. Thanks!

[ajbozarth]

Some feedback from Claude:

Re-reviewed at a081fc0. My earlier nits (@final, _sample naming, base-class loop_budget/requirements, majority-vote sample_index, docs) and @planetf1's **kwargs and list(set()) ordering items are all addressed. Requesting changes on one item: the majority-voting cancellation path is now correct in code but still untested (inline). Three smaller nits inline too — none blocking on their own.

[planetf1]

CI's code-checks / quality is failing on 3.11/3.12/3.13 with a real error, not a flake:

test/stdlib/test_context_type_enforcement.py:429: error: Cannot override final attribute "sample" (previously declared in base class "SamplingStrategy")  [misc]
test/stdlib/test_context_type_enforcement.py:429: error: Signature of "sample" incompatible with supertype "mellea.core.sampling.SamplingStrategy"  [override]

_StubStrategy in that file (added on main after this branch forked, so it won't show in this PR's diff) subclasses SamplingStrategy and overrides sample() directly — now illegal now that sample() is @final. Needs renaming to override _sample/_sample_impl instead, matching the new template-method contract the rest of this PR migrates to.

---

Re-checked against a081fc0d: all three items from my last round are addressed —

• majority_voting.py gather now cancels and drains sibling tasks on failure, matching the base.py pattern.
• sample() dropped **kwargs for an explicit show_progress param, so typo'd kwargs raise again.
• _merge_requirements now preserves order (dict.fromkeys) instead of list(set(...)).

The plugins.mdx doc gap from round 1 is fixed too (sample_index/sampling_id now listed).

I'm happy with the state of my own findings. Once @ajbozarth's outstanding item (missing test for the gather cancellation path) is addressed and the mypy CI failure above is fixed, this looks good to merge from my side.

[cptnm3]

Somehow the pre-commit did not catch the reasons for CI failure (again). I have made the requested changes. Thank you @planetf1 @ajbozarth for the through review and suggestions!

[ajbozarth]

Some feedback from Claude:

Re-reviewed fresh at 1c4c5bb. The prior rounds' blocking items are resolved at this head — the majority-voting cancellation path now has a test (test_branch_exception_cancels_siblings_and_propagates), the _StubStrategy @final mypy failure is fixed, and planetf1's **kwargs / requirement-ordering / sample_index items all check out. mypy and the targeted sampling/telemetry/hook suites pass locally. One trivial consistency nit inline; not blocking.

[ajbozarth]

@planetf1 If you want to take another pass to confirm your review has been addressed and update your review to approve this LGTM now

[planetf1]

LGTM — the success/all_results/all_validations docstring caveat resolves my last item, matches tracing.md. Content is ready to merge; confirm the pending action_required CI runs at this head go green first.

[planetf1]

LGTM — the success/all_results/all_validations docstring caveat resolves my last item, matches tracing.md. Content is ready to merge; confirm the pending action_required CI runs at this head go green first.

[ajbozarth]

@cptnm3 Thank you for you contribution :)