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

[ajbozarth]

nit: I see no reason to add all this whitespace and dashes

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

[planetf1]

Before this PR, the whole sampling loop — including the start/end hook dispatch — ran inside log_context(strategy=..., loop_budget=...). In this wrapper, log_context is only opened inside _sample, which runs after sampling_loop_start and before sampling_loop_end fires in the finally. So a plugin that logs from inside a sampling_loop_start/sampling_loop_end handler (the built-in debug plugin does this) loses the strategy/loop_budget fields on those specific log lines.

Not a functional break — just missing context for anyone debugging via logs on those two hook types. Suggest wrapping this method's try body in with log_context(strategy=type(self).__name__): — nesting is safe per log_context's own docs, and _sample's inner call still shadows loop_budget with the post-hook value.

[planetf1]

This test's name and docstring promise it checks that the S2 escalation fires as the final iteration number, but the assertion only checks the count — a regression to the wrong iteration number (e.g. loop_count instead of loop_count + 1) would still pass this test.

Suggested change

	assert len(observed) == 3
	assert len(observed) == 3
	assert observed[-1].iteration == 3

[planetf1]

Minor: sample() now accepts **kwargs and forwards it straight through to every _sample(), so a misspelled keyword (e.g. show_progres instead of show_progress) silently gets swallowed instead of raising TypeError like the old explicit signatures did. Not urgent, but worth naming show_progress explicitly if you want to keep that safety net.

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