fix(sampling): sampling span missing for strategies that override sam…

[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.

[planetf1]

With majority voting, iterations_used/all_results/all_validations here only reflect the winning branch's SamplingResult (BaseMBRDSampling._sample returns results[maxR][1]) — the other number_of_samples - 1 branches' generations aren't included anywhere. That's a real gap against what SamplingLoopEndPayload promises (all_results: List of ModelOutputThunk from every iteration), and it'll skew anything downstream that counts on this — e.g. the sampling-failures metric moves from "per branch" to "per top-level call" for fan-out strategies.

Worth either aggregating all branches into the result _sample returns, or updating the docstrings here (and in docs/docs/observability/tracing.md) to say "of the selected branch" so plugin authors aren't misled.

[cptnm3]

Hi @planetf1, this behaviour predates this change: BaseMBRDSampling._sample() has always fanned out number_of_samples branches and returned only results[maxR][1]. I have documented this in the docs and docstrings of the payload. Please let me know if you want to change this behaviour to aggregate all the branches into the returned result.

[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]

Two lines below (sampling_results = await asyncio.gather(*tasks), not shown in this diff so no inline anchor for it): branches now share one sampling_id. On failure:

• asyncio.gather doesn't cancel the other branches — they keep running
• sample()'s finally closes the shared span immediately
• surviving branches' events get silently dropped (tracing.py:815), and they keep burning backend calls for nothing
• no test covers this path

Fix — wrap the gather call:

try:
    sampling_results = await asyncio.gather(*tasks)
except BaseException:
    for t in tasks:
        t.cancel()
    await asyncio.gather(*tasks, return_exceptions=True)
    raise

(same pattern as base.py:336)

[planetf1]

• Pre-PR: SOFAI's sample() took no **kwargs → show_progress=False raised TypeError if unsupported
• Post-PR: sample() forwards show_progress through **kwargs, but SOFAI's _sample() has no such param — silently absorbed and ignored (recomputed from log level instead, sofai.py:667)
• same silent-swallow applies to any typo'd kwarg, on any strategy

Suggest dropping **kwargs from sample() (it's @final, only caller passes no extras) so bad keywords raise; give SOFAI an explicit show_progress param if it should honour it.

[cptnm3]

Dropped **kwargs.

[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.

[ajbozarth]

The cancel-and-drain guard is correct and matches the base.py producer-task pattern. The one item still open from @planetf1's review is that this path has no test. A branch raising should cancel the siblings and propagate, rather than leaving them running against the shared sampling_id (whose span the wrapper's finally has already closed). Since our review guidelines treat a missing test as blocking, requesting changes to add one: make a single branch's backend raise, then assert the exception propagates and the sibling tasks end up cancelled.

[cptnm3]

Sorry I missed this. I have added test_branch_exception_cancels_siblings_and_propagates test for this path.

[ajbozarth]

nit: sample_index is accepted here but never passed to the _emit_sampling_iteration / _emit_sampling_repair calls below (and majority_voting._sample accepts it as the fan-out parent, where it's always None) — so it's inert in both today. Budget forcing isn't fanned-out-into, but if it ever were, its iteration/repair events would silently drop the branch index. Either drop the param or thread it through the _emit_* calls.

[cptnm3]

Threaded sample_index in Buget forcing through _emit_* calls. But it can not be dropped in majority_voting._sample as it causes mypy to raise

error: Signature of "_sample" incompatible with supertype "mellea.stdlib.sampling.base.BaseSamplingStrategy" [override].

Do you want me to override it instead?

[ajbozarth]

nit (pre-existing, optional): this is the same getattr(self, ...) shim you removed for loop_budget / requirements. Base rejection sampling defines no _get_repair_type, so its repair events carry repair_type="unknown". Declaring _get_repair_type(self) -> str: return "unknown" on SamplingStrategy would let this read self._get_repair_type() and finish that cleanup.

[cptnm3]

Added

[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!