fix(sampling): sampling span missing for strategies that override sam… by cptnm3 · Pull Request #1584 · generative-computing/mellea

cptnm3 · 2026-08-25T16:42:00Z

Pull Request

Issue

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.

…ple() 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 Signed-off-by: Vishal V <VishalV@ibm.com>

github-actions · 2026-08-25T16:42:49Z

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 · 2026-08-26T14:57:01Z

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 · 2026-08-27T19:31:51Z

    """

-    @abc.abstractmethod
    async def sample(


Consider marking this @final, matching the reference method Backend.generate_from_raw. It enforces that subclasses override _sample_impl rather than sample() and prevents a subclass from silently bypassing the whole sampling lifecycle — the exact gap this PR closes. (Would also need final added to the typing import.)

Method has been marked @final

ajbozarth · 2026-08-27T19:31:51Z

+                )
+
+    @abc.abstractmethod
+    async def _sample_impl(


Naming nit: the reference pattern this mirrors uses a bare underscore — generate_from_raw / _generate_from_raw. For consistency with that convention, consider _sample here rather than _sample_impl. Not blocking, but it'd read as the same twin-method pattern already used elsewhere in core. (Touches all five strategy files if changed.)

Renamed to _sample

ajbozarth · 2026-08-27T19:31:51Z

+            reqs += call_requirements
+        return list(set(reqs))
+
+    def _get_loop_budget(self) -> int:


These getattr(self, ...) shims exist only because the base SamplingStrategy doesn't declare loop_budget/requirements, yet the wrapper already assumes every strategy has them (defaulting to 1/None). Cleaner to declare them on the base class directly — then _get_loop_budget goes away entirely (just self.loop_budget), and _merge_requirements reads self.requirements without the getattr. SOFAI, which has no global requirements today, simply inherits the None default, so no behavior change.

Added to base class

ajbozarth · 2026-08-27T19:31:51Z

-                    sampled_results,
-                    sampled_scores,
-                )
+                if loop_count < effective_loop_budget:


Re your question about the repair after budget exhaustion: yes, this is a real behavior change and it's the correct one — it aligns budget-forcing with the base loop, which already skips the final-iteration repair (base.py:505-508). It's safe because repair() here is pure (inherited RejectionSamplingStrategy.repair returns the last action unedited): on the final iteration the only things removed are a discarded return value and a repair event base never emitted anyway.

ajbozarth · 2026-08-27T19:31:51Z

        for i in range(self.number_of_samples):
            task = asyncio.create_task(
-                super().sample(
+                super()._sample_impl(


Re your question on the iteration/repair numbering collisions — good catch, and worth fixing rather than leaving as-is. The ideal is to mirror how base already handles fan-out: it keeps concurrent attempts distinct under one sampling_id via the subsample_index offset (base.py:426), so the majority-vote fan-out should do the same rather than let branches collide.

Add an optional sample_index: int | None to SamplingIterationPayload/SamplingRepairPayload, thread it through _sample_impl → the _emit_* helpers, and pass sample_index=i for the i-th fan-out sample. Then (sampling_id, sample_index, iteration) is unique and iteration keeps its single meaning; non-fan-out strategies leave it None. (sample_index mirrors number_of_samples; _index since it's an ordinal, not a UUID like sampling_id.)

Then surface it on the tracing side too — one extra event attribute (mellea.sampling.sample_index) in SamplingTracingPlugin's iteration/repair handlers — so the whole change lands together in this PR. Happy to walk through the threading if any of it's unclear.

I have made the changes based on your suggested plan. Please take a look

…e hook payloads BaseMBRDSampling._sample fans out number_of_samples concurrent calls to BaseSamplingStrategy._sample, each receiving the same sampling_id and starting their own _subsample_iteration loop from subsample_index=0. This caused all branches to emit the same iteration numbers under one sampling_id, making (sampling_id, iteration) non-unique for consumers of SAMPLING_ITERATION and SAMPLING_REPAIR hooks. Fix by introducing sample_index: int | None = None on both payload classes and threading it through the emit helpers and _subsample_iteration so each fan-out branch carries a distinct 0-based ordinal. Non-fan-out strategies leave sample_index=None; no existing call sites change. - Renamed _sample_impl to _sample - Marked sample method as @Final enforces that subclasses override _sample rather than sample() - Add sample_index field to SamplingIterationPayload and SamplingRepairPayload (mellea/plugins/hooks/sampling.py) - Add sample_index kwarg to _emit_sampling_iteration and _emit_sampling_repair, forwarded to the payload (mellea/core/sampling.py) - Add sample_index param to BaseSamplingStrategy._sample and _subsample_iteration; forward to both _emit_* calls (mellea/stdlib/sampling/base.py) - Pass sample_index=i in the BaseMBRDSampling fan-out loop (mellea/stdlib/sampling/majority_voting.py) - Emit mellea.sampling.sample_index span-event attribute in SamplingTracingPlugin.on_iteration and on_repair when not None (mellea/telemetry/tracing_plugins.py) Tests added: - test_majority_vote_iteration_sample_index_is_unique: e2e regression proving (sampling_id, sample_index, iteration) is unique across all branches with number_of_samples=3, loop_budget=2 - test_majority_vote_repair_sample_index_matches_branch: repair events carry the same sample_index as the failed iteration that triggered them - test_sample_index_defaults_to_none / test_sample_index_construction on both payload classes - test_sampling_iteration_includes/omits_sample_index_when_set/none - test_sampling_repair_includes/omits_sample_index_when_set/none Assisted-by: IBM Bob Signed-off-by: Vishal V <VishalV@ibm.com>

github-actions Bot added the bug Something isn't working label Aug 25, 2026

cptnm3 mentioned this pull request Aug 25, 2026

fix(sampling): sampling span missing for strategies that override sample() #1487

Open

4 tasks

ajbozarth requested changes Aug 27, 2026

View reviewed changes

cptnm3 marked this pull request as ready for review September 1, 2026 06:15

cptnm3 requested a review from a team as a code owner September 1, 2026 06:15

cptnm3 requested review from AngeloDanducci, ajbozarth and akihikokuroda September 1, 2026 06:15

Navigation Menu

Uh oh!

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

fix(sampling): sampling span missing for strategies that override sam…#1584
cptnm3 wants to merge 2 commits into
generative-computing:mainfrom
cptnm3:add-missing-sampling-span

cptnm3 commented Aug 25, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Aug 25, 2026

Uh oh!

ajbozarth commented Aug 26, 2026

Uh oh!

ajbozarth left a comment

Uh oh!

ajbozarth Aug 27, 2026

Uh oh!

cptnm3 Sep 1, 2026

Uh oh!

ajbozarth Aug 27, 2026

Uh oh!

cptnm3 Sep 1, 2026

Uh oh!

ajbozarth Aug 27, 2026

Uh oh!

cptnm3 Sep 1, 2026

Uh oh!

ajbozarth Aug 27, 2026

Uh oh!

ajbozarth Aug 27, 2026

Uh oh!

cptnm3 Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

cptnm3 commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request

Issue

Description

Testing

Attribution

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

Uh oh!

github-actions Bot commented Aug 25, 2026

Sampling Strategy PR Checklist

Base Class

Return Value

Integration

Uh oh!

ajbozarth commented Aug 26, 2026

Uh oh!

ajbozarth left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

cptnm3 commented Aug 25, 2026 •

edited

Loading