Skip to content

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

Closed

Description

@ajbozarth

Background

The sampling span added under #1050 is emitted by SamplingTracingPlugin subscribing to the sampling_loop_start / sampling_iteration / sampling_repair / sampling_loop_end hooks. Those hooks are dispatched inline inside BaseSamplingStrategy.sample() (mellea/stdlib/sampling/base.py).

Any strategy that overrides sample() rather than reusing the base loop therefore emits no sampling span, no iteration/repair events, and no error-path closure:

  • BudgetForcingSamplingStrategy.sample() (budget_forcing.py:101) — own loop, no hook dispatch.
  • SOFAISamplingStrategy.sample() (sofai.py:565) — own S1/S2 loop, no hook dispatch.
  • BaseMBRDSampling.sample() / majority-voting (majority_voting.py:102) — fans out to super().sample() per inner sample, so it emits one sampling span per inner sample with no enclosing span for the majority-vote operation. The nesting is incidental, not designed.

Why this matters

sample() is a public, override-designed method, but sampling telemetry is coupled to one specific method body rather than to the strategy boundary. A strategy silently drops out of tracing with no signal, and the majority-voting case shows the coupling already produces surprising span shapes. This is the same class of gap PR #1264 / #1183 fixed for generate_from_raw, where inline hook firing duplicated across backends was moved into a base-class wrapper (generate_from_raw fires hooks, _generate_from_raw_impl is the override).

Scope

Mirror the generate_from_raw fix:

  1. SamplingStrategy.sample() becomes a concrete base method that mints the sampling_id and fires the sampling lifecycle hooks (start/iteration/repair/end) with error-path closure.
  2. Add an abstract _sample_impl (or equivalent) that strategies override for their actual loop.
  3. Move BaseSamplingStrategy's current loop body into its _sample_impl; migrate BudgetForcing, SOFAI, and majority-voting to the same pattern so each emits exactly one enclosing sampling span.
  4. Add firing-site tests using a mock strategy that overrides only _sample_impl.

Acceptance criteria

  • Every SamplingStrategy (base loop, budget-forcing, SOFAI, majority-voting) emits exactly one enclosing sampling span per top-level sample() call.
  • Iteration/repair events and error-path closure fire regardless of which strategy is used.
  • No inline sampling_loop_* hook dispatch remains in individual strategy sample() bodies.
  • Firing-site tests cover the hooks against a mock strategy.

Activity

  1. added
    enhancementNew feature or request
    area/samplingSamplingStrategy, SamplingResult, ModelOption, generation options
    on Jul 31, 2026
  2. ajbozarth commented on Jul 31, 2026

    @ajbozarth
    MemberAuthor

    @jakelorocco thoughts on this? If you agree it's a good improvement it'd probably be a good issue for newer contributors.

  3. jakelorocco commented on Jul 31, 2026

    @jakelorocco
    Contributor

    Yes. I agree with this issue / improvement. It was flagged when we added hooks initially but I guess we never opened an issue for it / fixed it. Thanks for catching it!

  4. cptnm3 commented on Aug 10, 2026

    @cptnm3
    Contributor

    Hi @ajbozarth,
    If you don't mind, I'd like to take a shot at this.

  5. ajbozarth commented on Aug 10, 2026

    @ajbozarth
    MemberAuthor

    @cptnm3 sure, I believe I outlined the design decently to match what was in my head, but feel free to reach out if you hit any questions and add me as a reviewer on the PR if I'm not auto tagged :)

  6. cptnm3 commented on Aug 25, 2026

    @cptnm3
    Contributor

    Hey @ajbozarth, I have opened up a draft PR #1584 with the changes to address this issue. Does this match the design you had?
    I have couple of questions on the implementation -

    • Budget forcing previously performed an unused repair after budget exhaustion. Currently, this change removes that call to align repair execution and telemetry semantics with BaseSamplingStrategy. Is this acceptable?
    • majority voting now has one lifecycle, but its iteration/repair numbering can collide - each inner call to BaseSamplingStrategy._sample_impl() starts its own iteration numbering. For number_of_samples=3, loop_budget=1, all three iteration events are effectively iteration 1. Larger loop budget would result in concurrent duplicate sequence. Is this behaviour acceptable? What would be the ideal behaviour in this situation?
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Labels

area/samplingSamplingStrategy, SamplingResult, ModelOption, generation optionsenhancementNew feature or request

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions