esh-pfi-infrastructure/stacks/chatterbox-fast/test_scheduler.py

"""GPU-free simulation of the adaptive-chunk scheduler.

Validates the two acceptance properties from the plan (§6) without a GPU:
  * NO STARVATION — the stream stays ahead of 1× playback.
  * RATCHET — chunks grow after the latency-critical first one.

Run directly (``python test_scheduler.py``) or under pytest.
"""

from __future__ import annotations

from scheduler import (
    ChunkConfig,
    plan_chunk,
    protect_first_audio,
    split_sentences,
    stream_chunks,
)

# A multi-sentence paragraph with varied lengths. The long sentences carry
# commas (as natural prose does), so the clause-split starvation relief has
# boundaries to work with.
PARAGRAPH = (
    "The cake is a lie. "
    "I am being entirely sincere, without a trace of sarcasm, when I say "
    "that this is the single most important scientific breakthrough in the "
    "entire history of this facility. "
    "You will be baked, and then there will be cake. "
    "It is delicious and moist, assuming you survive the testing protocol, "
    "which the available data suggests you almost certainly will not. "
    "Goodbye."
)


class FakeClock:
    """A clock the fake generator advances, so gen_time reflects simulated work."""

    def __init__(self) -> None:
        self.t = 0.0

    def __call__(self) -> float:
        return self.t


def make_generator(clock: FakeClock, *, true_rtf: float, sec_per_char: float):
    """A fake generate() that costs realistic wall-clock and returns audio_sec.

    Audio duration is proportional to text length; generation costs
    ``audio_sec / true_rtf`` of (simulated) wall-clock, advancing the clock.
    """

    def generate(text: str):
        audio_sec = len(text) * sec_per_char
        clock.t += audio_sec / true_rtf
        return None, audio_sec

    return generate


def run(true_rtf: float = 3.8, sec_per_char: float = 0.060, cfg: ChunkConfig | None = None):
    clock = FakeClock()
    gen = make_generator(clock, true_rtf=true_rtf, sec_per_char=sec_per_char)
    return list(stream_chunks(PARAGRAPH, generate=gen, clock=clock, cfg=cfg))


# ── splitting ─────────────────────────────────────────────────────────────


def test_split_sentences_count():
    units = split_sentences(PARAGRAPH)
    assert units[0] == "The cake is a lie."
    assert units[-1] == "Goodbye."
    assert len(units) == 5


def test_split_preserves_inline_tags():
    units = split_sentences("Hello there. [laugh] That was funny. Bye.")
    # The tag must not be torn from its sentence nor split on.
    assert any("[laugh]" in u for u in units)
    assert len(units) == 3


def test_protect_first_audio_clause_splits_long_opener():
    cfg = ChunkConfig(max_first_sec=0.5)  # force the guard
    long_open = ["This is a long opener, with a clause, and another clause.", "Short."]
    out = protect_first_audio(long_open, cfg)
    assert len(out) > len(long_open)
    assert out[0] == "This is a long opener,"


def test_protect_first_audio_keeps_unsplittable_opener_whole():
    cfg = ChunkConfig(max_first_sec=0.1)
    out = protect_first_audio(["No clause boundaries here at all."], cfg)
    assert out == ["No clause boundaries here at all."]


# ── chunk planning ────────────────────────────────────────────────────────


def test_first_chunk_is_single_unit():
    # Zero buffer ⇒ exactly the first unit, regardless of margin/rtf.
    units = split_sentences(PARAGRAPH)
    chunk, rest = plan_chunk(units, 0.0, margin=0.65, rtf=3.8, sec_per_char=0.06)
    assert chunk == units[0]
    assert len(rest) == len(units) - 1


def test_plan_never_returns_empty():
    chunk, rest = plan_chunk(["Only one."], 0.0, margin=0.8, rtf=3.8, sec_per_char=0.06)
    assert chunk == "Only one."
    assert rest == []


# ── the acceptance properties ─────────────────────────────────────────────


def test_no_starvation_nominal():
    """Stream never starves when the RTF prior matches reality."""
    results = run(true_rtf=3.8)
    assert all(not r.starved for r in results), [
        (r.index, r.drained) for r in results if r.starved
    ]


def test_no_starvation_when_rtf_overestimated():
    """Prior says 3.8× but the GPU only delivers 3.0× — live correction + margin
    must still keep the stream fed. This is the property that matters: the
    scheduler must be robust to an optimistic prior, not just a matching one."""
    results = run(true_rtf=3.0)
    assert all(not r.starved for r in results), [
        (r.index, r.drained) for r in results if r.starved
    ]


def test_no_starvation_slow_gpu():
    """Even at 2.5× (well below prior) the margin absorbs it for this text."""
    results = run(true_rtf=2.5)
    assert all(not r.starved for r in results), [
        (r.index, r.drained) for r in results if r.starved
    ]


def test_chunks_ratchet_up():
    """After the latency-critical first chunk, chunks grow (buffer-ratchet)."""
    results = run(true_rtf=3.8)
    assert len(results) >= 3, "paragraph should not collapse to one chunk"
    # Chunk 1 is the smallest (single sentence); the second chunk is larger.
    assert len(results[1].text) > len(results[0].text)
    # The paragraph consolidates: the last chunk carries multiple sentences.
    assert results[-1].audio_sec >= results[0].audio_sec


def test_first_chunk_low_latency():
    """First chunk is one short sentence ⇒ smallest gen time ⇒ fast first audio."""
    results = run(true_rtf=3.8)
    assert results[0].text == "The cake is a lie."
    # In the sim, chunk 0's gen_time IS the time-to-first-audio, and it must be
    # the smallest of all chunks (everything after it is ≥ a clause).
    assert results[0].gen_time == min(r.gen_time for r in results)
    assert results[0].gen_time < 0.4


def test_unsplittable_long_sentence_flags_starvation():
    """KNOWN LIMITATION (surfaced, not hidden): a long COMMA-LESS sentence right
    after a short opener cannot be clause-split, so we honor the never-split-mid-
    sentence rule and accept a brief gap — which MUST show up as drained>0 in
    telemetry so it is measurable, never silent."""
    clock = FakeClock()
    gen = make_generator(clock, true_rtf=3.8, sec_per_char=0.060)
    text = (
        "Hi. "
        "I am now going to speak one extremely long sentence with no clause "
        "boundaries at all so that nothing in here can ever be split apart by "
        "the scheduler no matter how hard it tries to find a comma."
    )
    results = list(stream_chunks(text, generate=gen, clock=clock))
    assert any(r.starved for r in results), "expected the gap to be flagged"
    assert any(r.drained > 0 for r in results)


def test_full_text_reconstructed():
    """Every sentence is emitted exactly once, in order."""
    results = run(true_rtf=3.8)
    joined = " ".join(r.text for r in results)
    assert joined == " ".join(split_sentences(PARAGRAPH))


def _main():
    results = run(true_rtf=3.8)
    print(f"{'idx':>3} {'chars':>5} {'audio_s':>8} {'gen_s':>7} "
          f"{'est_s':>6} {'buf_before':>10} {'buf_after':>9} {'rtf':>5} {'drain':>6}")
    for r in results:
        print(f"{r.index:>3} {len(r.text):>5} {r.audio_sec:>8.2f} {r.gen_time:>7.3f} "
              f"{r.est_gen:>6.3f} {r.buffer_before:>10.2f} {r.buffer_after:>9.2f} "
              f"{r.rtf:>5.2f} {r.drained:>6.3f}")
    starved = [r.index for r in results if r.starved]
    print(f"\nchunks={len(results)}  starved={starved or 'none'}  "
          f"ttfa≈{results[0].gen_time:.3f}s")


if __name__ == "__main__":
    _main()