lightseekorg--tokenspeed/test/runtime/test_grammar_cache_timeout.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Regression tests for BaseGrammarBackend's compile-timeout cache.

These tests pin the contract that GrammarManager relies on:

1. Permanent failures (``cache_invalid``) stay sticky forever.
2. Transient timeouts (``record_compile_timeout``) cache a TTL'd marker;
   while it's still valid, a still-running compile that finishes late must
   NOT be allowed to clobber the marker.
3. After the TTL expires, a new request triggers a fresh compile and the
   successful result becomes the new cached value (timeout marker is
   replaced, not left to poison the key).
4. Repeated timeouts escalate: after ``max_retries`` consecutive timeouts
   the marker becomes permanent (no TTL) — at that point the compiler is
   consistently broken for this key.
5. ``compile_started_at`` measures from when the worker actually begins
   running, not from when the future was submitted, so executor queueing
   delay doesn't eat into the timeout budget.
6. The base API exposes ``is_invalid`` / ``expired`` attributes on every
   value so callers don't need isinstance checks.

Pure CPU; no model, tokenizer, or GPU required.
"""

from __future__ import annotations

import os
import sys
import threading
import time
import unittest

# CI registration (AST-parsed, runtime no-op).
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from ci_system.ci_register import register_cuda_ci  # noqa: E402

register_cuda_ci(est_time=20, suite="runtime-1gpu")

from tokenspeed.runtime.grammar.base_grammar_backend import (  # noqa: E402
    BaseGrammarBackend,
    BaseGrammarObject,
    InvalidGrammarObject,
)


class _FakeGrammar(BaseGrammarObject):
    """Trivial successful compile result; copy() returns self for cheap equality."""

    def copy(self):
        return self


class _GatedBackend(BaseGrammarBackend):
    """Backend whose ``init_value_impl`` blocks on an externally-controlled
    Event so tests can drive the worker thread's lifecycle precisely."""

    def __init__(self):
        super().__init__()
        self.compile_started = threading.Event()
        self.compile_can_finish = threading.Event()
        self.compile_count = 0

    def init_value_impl(self, key, require_reasoning: bool = False):
        self.compile_count += 1
        self.compile_started.set()
        self.compile_can_finish.wait()
        return _FakeGrammar()


class GrammarBackendCacheTests(unittest.TestCase):
    KEY = ("json", "schema-A")

    # --- attribute interface (no isinstance) ---

    def test_is_invalid_attribute_split(self):
        """Successful and failed grammars share the is_invalid attribute so
        callers don't need isinstance to branch."""
        ok = _FakeGrammar()
        bad = InvalidGrammarObject("nope")
        self.assertFalse(ok.is_invalid)
        self.assertTrue(bad.is_invalid)
        self.assertFalse(ok.expired)
        self.assertFalse(bad.expired)  # no expires_at -> permanent

    # --- permanent failures ---

    def test_cache_invalid_is_permanent(self):
        b = _GatedBackend()
        b.cache_invalid(self.KEY, "bad regex")
        v, hit = b.get_cached_or_future_value(self.KEY)
        self.assertTrue(hit)
        self.assertTrue(v.is_invalid)
        self.assertEqual(v.error_message, "bad regex")
        self.assertIsNone(v.expires_at)
        # Even after a long sleep it stays invalid.
        time.sleep(0.05)
        self.assertFalse(v.expired)

    # --- transient timeouts (TTL) ---

    def test_timeout_marker_sticks_within_ttl(self):
        """While the TTL'd marker is fresh, a slow compile that completes
        late must not be able to overwrite it. Otherwise the abort message
        we returned to the user would be silently contradicted by the next
        request seeing a real grammar."""
        b = _GatedBackend()
        # Start a compile and let it block.
        fut = b.executor.submit(b.init_value, self.KEY)
        self.assertTrue(b.compile_started.wait(timeout=1.0))

        # Mid-compile: cache a transient timeout marker.
        b.record_compile_timeout(self.KEY, "timed out", ttl_secs=2.0, max_retries=2)

        v, hit = b.get_cached_or_future_value(self.KEY)
        self.assertTrue(hit)
        self.assertTrue(v.is_invalid)
        self.assertFalse(v.expired)
        self.assertEqual(v.error_message, "timed out")

        # Now let the slow compile finish — it must NOT clobber the marker.
        b.compile_can_finish.set()
        fut.result(timeout=1.0)

        v, hit = b.get_cached_or_future_value(self.KEY)
        self.assertTrue(hit, "expected cache hit, got miss")
        self.assertTrue(
            v.is_invalid, "late finish wrongly un-invalidated the cache marker"
        )
        self.assertEqual(v.error_message, "timed out")

    def test_marker_expires_and_recovers(self):
        """After the TTL, a fresh compile is dispatched and the success path
        replaces the marker — the key is not poisoned forever."""
        b = _GatedBackend()
        # Cache a short-lived timeout marker (no in-flight compile this time).
        b.record_compile_timeout(self.KEY, "timed out", ttl_secs=0.1, max_retries=2)
        time.sleep(0.15)

        v, hit = b.get_cached_or_future_value(self.KEY)
        self.assertFalse(hit, "expired marker should evict on lookup")
        # The returned value is a Future from the executor; the worker is
        # already running and blocked on compile_can_finish.
        self.assertTrue(b.compile_started.wait(timeout=1.0))
        b.compile_can_finish.set()
        result = v.result(timeout=1.0)
        self.assertFalse(result.is_invalid)
        self.assertEqual(b.compile_count, 1, "exactly one fresh compile expected")

        # Subsequent lookups see the real grammar.
        v, hit = b.get_cached_or_future_value(self.KEY)
        self.assertTrue(hit)
        self.assertFalse(v.is_invalid)

    def test_repeated_timeouts_escalate_to_permanent(self):
        """After max_retries+1 consecutive timeouts the marker becomes
        permanent (no TTL) so we stop hitting the compiler in a loop."""
        b = _GatedBackend()
        # max_retries=2 → first 2 timeouts get a TTL, the 3rd is permanent.
        for _ in range(2):
            b.record_compile_timeout(
                self.KEY, "timed out", ttl_secs=0.01, max_retries=2
            )
            v, hit = b.get_cached_or_future_value(self.KEY)
            self.assertTrue(hit)
            self.assertTrue(v.is_invalid)
            self.assertIsNotNone(v.expires_at, "first 2 timeouts must have a TTL")
            time.sleep(0.02)  # let TTL expire so the next call sees a miss path

        # 3rd timeout escalates to permanent.
        b.record_compile_timeout(self.KEY, "timed out", ttl_secs=0.01, max_retries=2)
        v, hit = b.get_cached_or_future_value(self.KEY)
        self.assertTrue(hit)
        self.assertTrue(v.is_invalid)
        self.assertIsNone(v.expires_at, "post-escalation marker must be permanent")
        self.assertIn("gave up after", v.error_message)

    def test_successful_compile_clearstimeout_history(self):
        """A successful compile after some transient timeouts resets the
        retry counter, so a future timeout cycle gets the full retry budget
        again."""
        b = _GatedBackend()
        b.record_compile_timeout(self.KEY, "t1", ttl_secs=0.01, max_retries=2)
        time.sleep(0.02)
        b.record_compile_timeout(self.KEY, "t2", ttl_secs=0.01, max_retries=2)
        self.assertEqual(b.timeout_history[self.KEY].count, 2)
        time.sleep(0.02)

        # Trigger a fresh compile that succeeds.
        v, hit = b.get_cached_or_future_value(self.KEY)
        self.assertFalse(hit)
        b.compile_can_finish.set()
        v.result(timeout=1.0)

        # Successful compile must have cleared the per-key history; a
        # subsequent timeout is also a no-op here (cache holds a valid
        # grammar, so record_compile_timeout drops the spurious mark).
        self.assertNotIn(self.KEY, b.timeout_history)
        b.record_compile_timeout(self.KEY, "t3", ttl_secs=0.01, max_retries=2)
        self.assertNotIn(
            self.KEY,
            b.timeout_history,
            "spurious timeout against a valid cached grammar must not increment",
        )
        v, hit = b.get_cached_or_future_value(self.KEY)
        self.assertTrue(hit)
        self.assertFalse(v.is_invalid, "valid grammar must not be clobbered by timeout")

    # --- compile_started_at excludes executor queueing ---

    def test_started_at_records_worker_start_not_submit(self):
        """compile_started_at must reflect when the worker actually begins
        running init_value_impl, so an arbitrary executor queueing delay
        doesn't count toward the timeout budget."""
        # Use a single-thread executor and pre-fill it with a blocker so the
        # second submission must wait in the queue.
        from concurrent.futures import ThreadPoolExecutor

        b = _GatedBackend()
        b.executor.shutdown(wait=False)
        b.executor = ThreadPoolExecutor(max_workers=1)

        block_release = threading.Event()
        b.executor.submit(block_release.wait)  # occupy the only worker

        submit_ts = time.monotonic()
        future = b.executor.submit(b.init_value, self.KEY)

        # While queued, no entry exists yet -> compile_started_at is None.
        for _ in range(5):
            if b.compile_started.is_set():
                self.fail("worker must not have started while queue is blocked")
            self.assertIsNone(b.compile_started_at(self.KEY))
            time.sleep(0.02)

        queue_delay = time.monotonic() - submit_ts
        self.assertGreater(
            queue_delay,
            0.05,
            "test setup failed: expected measurable queue delay before unblocking",
        )

        # Release the blocker so the worker can pick up the real compile.
        block_release.set()
        self.assertTrue(b.compile_started.wait(timeout=1.0))

        started_at = b.compile_started_at(self.KEY)
        self.assertIsNotNone(started_at)
        # started_at must reflect the worker start, well after submit_ts +
        # queue_delay. Slack the bound a bit but require it to be later than
        # the submit instant by at least the observed queue delay.
        self.assertGreaterEqual(
            started_at,
            submit_ts + queue_delay - 0.02,
            "started_at should be past submit_ts + queue delay",
        )

        # After the compile finishes, started_at must REMAIN set — callers
        # use it to compute compile-only elapsed time, and a worker that
        # finishes between a future.done() check and the elapsed check must
        # not silently revert callers to a wall-clock-from-submit fallback
        # (that was a real bug — see GrammarManager's done()-recheck).
        b.compile_can_finish.set()
        future.result(timeout=1.0)
        self.assertEqual(b.compile_started_at(self.KEY), started_at)


class GrammarManagerTimeoutTests(unittest.TestCase):
    """End-to-end timeout behavior in GrammarManager.get_ready_grammar_requests:

    queued-but-not-yet-started compiles must still time out (otherwise a
    wedged or saturated executor could let a request wait forever), and
    once the worker starts the per-compile budget begins fresh.
    """

    KEY = ("json", "queued-forever")

    def _make_manager(self, *, compile_timeout: float = 0.05, max_retries: int = 2):
        # Build a GrammarManager without going through ServerArgs (the parent
        # imports too much). Subclass it to plug in our gated backend.
        from tokenspeed.runtime.grammar.grammar_manager import (
            GrammarManager as RealGrammarManager,
        )

        backend = _GatedBackend()
        # Saturate the executor so no worker is available — submitted
        # init_value calls will sit in the queue, never running.
        n_workers = backend.executor._max_workers
        block_release = threading.Event()
        for _ in range(n_workers):
            backend.executor.submit(block_release.wait)

        mgr = RealGrammarManager.__new__(RealGrammarManager)
        mgr.server_args = None
        mgr.grammar_queue = []
        mgr.compile_timeout_secs = compile_timeout
        mgr.compile_max_retries = max_retries
        mgr.grammar_backend = backend
        mgr.grammar_sync_group = None
        mgr.grammar_sync_size = 1
        return mgr, backend, block_release

    def _fake_state(self, key=KEY):
        from tokenspeed.runtime.engine.generation_output_processor import RequestState
        from tokenspeed.runtime.sampling.sampling_params import SamplingParams

        sp = SamplingParams(json_schema='{"type":"object"}')
        s = RequestState(
            prompt_input_ids=[0],
            sampling_params=sp,
            stream=False,
            tokenizer=None,
            eos_token_ids=[1],
        )
        s.rid = "test-rid"
        return s

    def test_queued_compile_times_out_via_queued_ts(self):
        """A request whose worker NEVER starts (executor saturated) must
        still time out within ~compile_timeout_secs of being queued."""
        mgr, backend, block_release = self._make_manager(compile_timeout=0.05)
        try:
            state = self._fake_state()
            value, hit = backend.get_cached_or_future_value(self.KEY)
            self.assertFalse(hit)
            state.grammar = value  # the Future
            state.grammar_key = self.KEY
            mgr.add_to_queue(spec=object(), state=state, bootstrap_info=object())

            # compile_started_at must be None (worker is queued, not running).
            self.assertIsNone(backend.compile_started_at(self.KEY))

            # First poll — within the budget, no failure.
            promoted = mgr.get_ready_grammar_requests()
            self.assertEqual(promoted, [])

            # Wait past the compile timeout; worker still hasn't started.
            time.sleep(0.07)
            self.assertIsNone(backend.compile_started_at(self.KEY))

            promoted = mgr.get_ready_grammar_requests()
            self.assertEqual(len(promoted), 1)
            _, promoted_state, _ = promoted[0]
            self.assertIs(promoted_state, state)
            self.assertTrue(state.finished)
            self.assertIn("timed out", state.to_abort_message)
        finally:
            block_release.set()

    def test_completion_at_timeout_boundary_admits_not_aborts(self):
        """Race: the worker finishes the compile in the same poll iteration
        where the elapsed check would otherwise mark it timed out. The
        request must be admitted (compile succeeded) — not aborted."""
        mgr, backend, block_release = self._make_manager(compile_timeout=0.05)
        try:
            state = self._fake_state()
            value, hit = backend.get_cached_or_future_value(self.KEY)
            self.assertFalse(hit)
            state.grammar = value
            state.grammar_key = self.KEY
            mgr.add_to_queue(spec=object(), state=state, bootstrap_info=object())

            # Let the worker start compiling, then sleep past the timeout
            # budget WITHOUT completing the compile yet.
            block_release.set()
            self.assertTrue(backend.compile_started.wait(timeout=1.0))
            time.sleep(0.07)  # past compile_timeout_secs from started_at

            # NOW finish the compile, *before* the next poll. From the poll's
            # perspective, the future is done at the moment we'd otherwise
            # decide "timed out" — must admit, not abort.
            backend.compile_can_finish.set()
            value.result(timeout=1.0)

            promoted = mgr.get_ready_grammar_requests()
            self.assertEqual(len(promoted), 1)
            _, promoted_state, _ = promoted[0]
            self.assertIs(promoted_state, state)
            self.assertFalse(
                state.finished,
                f"expected admission, got abort: {state.to_abort_message!r}",
            )
            self.assertFalse(state.grammar.is_invalid)
        finally:
            block_release.set()

    def test_compile_budget_resets_on_worker_start(self):
        """If the request waits in the queue almost to the timeout and only
        then the worker starts, we get a fresh compile budget — total max
        wait is ~2 * compile_timeout_secs, not 1."""
        mgr, backend, block_release = self._make_manager(compile_timeout=0.10)
        try:
            state = self._fake_state()
            value, hit = backend.get_cached_or_future_value(self.KEY)
            self.assertFalse(hit)
            state.grammar = value
            state.grammar_key = self.KEY
            mgr.add_to_queue(spec=object(), state=state, bootstrap_info=object())

            # Sit just under the queue-wait budget, then unblock the executor
            # so the worker starts running our compile.
            time.sleep(0.07)
            block_release.set()
            self.assertTrue(backend.compile_started.wait(timeout=1.0))

            # We're past the queue budget, but compile just started → must
            # NOT be marked failed yet.
            promoted = mgr.get_ready_grammar_requests()
            self.assertEqual(promoted, [], "compile budget should reset on start")
            self.assertFalse(state.finished)

            # Within the fresh compile budget, complete the compile.
            backend.compile_can_finish.set()
            value.result(timeout=1.0)
            promoted = mgr.get_ready_grammar_requests()
            self.assertEqual(len(promoted), 1)
            self.assertFalse(state.finished, "successful compile should not abort")
            self.assertFalse(state.grammar.is_invalid)
        finally:
            block_release.set()


if __name__ == "__main__":
    unittest.main()