STT-Assistant/tests/test_laya_lane.py

"""Tests for the in-process Laya triage lane (todos 5, 7, 8).

All lane tests inject a fake agent factory; the real `laya` package is never
required. The single measurement test (todo 5) is decorated with
`pytest.importorskip("laya")` so the default suite skips it green.
"""

import sys
import threading
import time
from types import SimpleNamespace

import pytest

from stt_daemon.assistant.laya_lane import LayaLane, device_or_none

LAYA_PIN_SHA = "970dc8c5f63d7b886a68409493f37d569424f933"


def make_cfg(**overrides):
    cfg = {
        "device": "",
        "preload": ["multilingual"],
        "default_checkpoint": "multilingual",
        "predict_max_tokens": 1024,
    }
    cfg.update(overrides)
    return SimpleNamespace(**cfg)


class FakeRouter:
    """Fake laya Router: records ctor kwargs, preload list, predict calls."""

    def __init__(self, calls, **kwargs):
        self._calls = calls
        self.ctor_kwargs = kwargs
        self.preload_names = None
        self.predict_log = []
        self.predict_impl = None
        self.fail_warmup = False
        self.current = 0
        self.max_concurrent = 0

    def preload(self, names):
        self.preload_names = list(names)
        return self

    def predict(self, state, questions, model=None, lang=None, **kwargs):
        self._calls["count"] += 1
        self.current += 1
        self.max_concurrent = max(self.max_concurrent, self.current)
        try:
            self.predict_log.append(
                {"state": state, "questions": questions, "model": model, "lang": lang,
                 "kwargs": kwargs}
            )
            if "_warmup" in questions:
                # loader warmup: never gated, returns immediately
                if self.fail_warmup:
                    raise RuntimeError("warmup boom")
                return {"model": model, "answers": {}, "usage": {}, "routing": {}}
            if self.predict_impl is not None:
                return self.predict_impl(state, questions, model=model, lang=lang)
            return {"model": model, "answers": {}, "usage": {}, "routing": {}}
        finally:
            self.current -= 1


class FakeFactory:
    """Callable recording ctor kwargs and spawning FakeRouter instances."""

    def __init__(self, calls, predict_impl=None):
        self._calls = calls
        self._predict_impl = predict_impl
        self.routers = []

    def __call__(self, **kwargs):
        router = FakeRouter(self._calls, **kwargs)
        router.predict_impl = self._predict_impl
        self.routers.append(router)
        self._calls["ctor_kwargs"] = kwargs
        return router


def make_factory(calls, predict_impl=None):
    return FakeFactory(calls, predict_impl=predict_impl)


def wait_until(predicate, timeout=5.0):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return True
        time.sleep(0.005)
    return predicate()


def start_ready_lane(calls=None, **kwargs):
    """Start a lane whose loader succeeds immediately; wait for ready."""
    calls = calls if calls is not None else {"count": 0}
    results = []
    states = []
    factory = make_factory(calls)
    lane = LayaLane(
        make_cfg(),
        agent_factory=factory,
        on_result=lambda gen, payload: results.append((gen, payload)),
        on_state=lambda s, r: states.append((s, r)),
        **kwargs,
    )
    lane.start()
    assert wait_until(lane.ready), "lane did not become ready"
    return lane, factory, results, states


# -- todo 7: lazy import + loader --------------------------------------

def test_no_top_level_laya_import():
    assert "laya" not in sys.modules
    assert "torch" not in sys.modules


def test_device_or_none():
    assert device_or_none("") is None
    assert device_or_none("cpu") == "cpu"
    assert device_or_none("cuda") == "cuda"


def test_loader_builds_router_preloads_and_warms_up():
    calls = {"count": 0}
    lane, factory, results, states = start_ready_lane(calls)
    try:
        assert lane.state() == ("ok", "")
        assert states == [("ok", "")]
        assert calls["ctor_kwargs"] == {"default": "multilingual", "device": None}
        router = factory.routers[0]
        assert router.preload_names == ["multilingual"]
        # one throwaway warmup predict with model="multilingual"
        assert len(router.predict_log) == 1
        warmup = router.predict_log[0]
        assert warmup["model"] == "multilingual"
        assert warmup["questions"] == {
            "_warmup": {"type": "noul", "instructions": "warmup"}
        }
        assert results == []  # warmup never emits on_result
    finally:
        lane.stop()


def test_loader_device_passthrough():
    calls = {"count": 0}
    factory = make_factory(calls)
    lane = LayaLane(make_cfg(device="cpu"), agent_factory=factory)
    lane.start()
    try:
        assert wait_until(lane.ready)
        assert calls["ctor_kwargs"] == {"default": "multilingual", "device": "cpu"}
    finally:
        lane.stop()


def test_loader_import_error_gives_laya_missing():
    def factory(**kwargs):
        raise ImportError("No module named 'laya'")

    states = []
    lane = LayaLane(make_cfg(), agent_factory=factory,
                    on_state=lambda s, r: states.append((s, r)))
    lane.start()  # must not raise
    try:
        assert wait_until(lambda: lane.state()[0] == "laya_missing")
        assert lane.ready() is False
        status, hint = lane.state()
        assert status == "laya_missing"
        assert "uv sync --extra triage" in hint
        assert states == [("laya_missing", hint)]
    finally:
        lane.stop()


def test_loader_generic_exception_gives_error():
    def factory(**kwargs):
        raise RuntimeError("CUDA OOM")

    lane = LayaLane(make_cfg(), agent_factory=factory)
    lane.start()  # must not raise
    try:
        assert wait_until(lambda: lane.state()[0] == "error")
        assert lane.ready() is False
        assert "CUDA OOM" in lane.state()[1]
    finally:
        lane.stop()


def test_loader_warmup_crash_gives_error():
    calls = {"count": 0}
    real_factory = make_factory(calls)

    def arming_factory(**kwargs):
        router = real_factory(**kwargs)
        router.fail_warmup = True
        return router

    lane = LayaLane(make_cfg(), agent_factory=arming_factory)
    lane.start()
    try:
        assert wait_until(lambda: lane.state()[0] == "error")
        assert lane.ready() is False
    finally:
        lane.stop()


def test_start_never_raises_even_when_spawn_fails(monkeypatch):
    calls = {"count": 0}
    lane = LayaLane(make_cfg(), agent_factory=make_factory(calls))

    def boom(*args, **kwargs):
        raise RuntimeError("no threads")

    monkeypatch.setattr(threading, "Thread", boom)
    lane.start()  # must not raise


def test_submit_false_when_failed():
    def factory(**kwargs):
        raise ImportError("No module named 'laya'")

    lane = LayaLane(make_cfg(), agent_factory=factory)
    lane.start()
    try:
        assert wait_until(lambda: lane.state()[0] == "laya_missing")
        assert lane.submit("s", {}, None, 1) is False
    finally:
        lane.stop()


def test_submit_false_when_error():
    def factory(**kwargs):
        raise RuntimeError("hub down")

    lane = LayaLane(make_cfg(), agent_factory=factory)
    lane.start()
    try:
        assert wait_until(lambda: lane.state()[0] == "error")
        assert lane.submit("s", {}, None, 1) is False
    finally:
        lane.stop()


# -- todo 8: single-flight / coalesce / shutdown ------------------------

def test_pre_ready_submit_drops_silently():
    calls = {"count": 0}
    results = []
    gate = threading.Event()

    def predict_impl(state, questions, model=None, lang=None):
        gate.wait(timeout=5.0)
        return {"answers": {}}

    factory = make_factory(calls, predict_impl=predict_impl)
    lane = LayaLane(make_cfg(), agent_factory=factory,
                    on_result=lambda gen, payload: results.append((gen, payload)))
    # submit before start(): not ready -> False, nothing recorded
    assert lane.submit("early", {"q": 1}, None, 0) is False
    assert calls["count"] == 0
    assert results == []
    lane.stop()
    assert calls["count"] == 0
    assert results == []


def test_submit_newest_wins_while_blocked():
    calls = {"count": 0}
    entered = threading.Event()
    release = threading.Event()
    first = {"done": False}

    def predict_impl(state, questions, model=None, lang=None):
        if not first["done"]:
            first["done"] = True
            entered.set()
            assert release.wait(timeout=5.0)
            return {"answers": {"gen": 0}}
        return {"answers": {"state": state}}

    factory = make_factory(calls, predict_impl=predict_impl)
    results = []
    done = threading.Event()

    def on_result(gen, payload):
        results.append((gen, payload))
        if len(results) >= 2:
            done.set()

    lane = LayaLane(
        make_cfg(), agent_factory=factory, on_result=on_result,
    )
    lane.start()
    try:
        assert wait_until(lane.ready)
        router = factory.routers[0]
        router.predict_log.clear()
        calls["count"] = 0
        assert lane.submit("state-0", {"q": 0}, None, 0) is True
        assert entered.wait(timeout=5.0)
        # 10 rapid submits while the first predict is blocked
        for i in range(1, 11):
            assert lane.submit("state-%d" % i, {"q": i}, "es", i) is True
        release.set()
        assert wait_until(lambda: len(results) >= 2, timeout=5.0)
        # exactly one coalesced predict for the newest window
        assert len(router.predict_log) == 2
        assert router.predict_log[0]["state"] == "state-0"
        assert router.predict_log[1]["state"] == "state-10"
        assert router.predict_log[1]["questions"] == {"q": 10}
        assert router.predict_log[1]["lang"] == "es"
        assert router.predict_log[1]["model"] == "multilingual"
        assert [gen for gen, _ in results] == [0, 10]
    finally:
        lane.stop()


def test_max_concurrent_predicts_is_one():
    calls = {"count": 0}
    started = threading.Barrier(2, timeout=5.0)

    def predict_impl(state, questions, model=None, lang=None):
        try:
            started.wait(timeout=5.0)
        except threading.BrokenBarrierError:
            pass
        time.sleep(0.02)
        return {"answers": {}}

    factory = make_factory(calls, predict_impl=predict_impl)
    results = []
    results_lock = threading.Lock()

    def on_hammer_result(gen, payload):
        with results_lock:
            results.append(gen)

    lane = LayaLane(
        make_cfg(), agent_factory=factory, on_result=on_hammer_result,
    )
    lane.start()
    try:
        assert wait_until(lane.ready)
        router = factory.routers[0]

        def hammer(n):
            for i in range(n):
                lane.submit("s-%d-%d" % (n, i), {}, None, n * 100 + i)

        threads = [threading.Thread(target=hammer, args=(n,)) for n in range(4)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        assert wait_until(lambda: calls["count"] >= 2, timeout=5.0)
        # drain: wait until no more predicts happen
        before = calls["count"]
        time.sleep(0.15)
        # single worker thread => concurrency counter never exceeded 1
        assert router.max_concurrent == 1
        assert calls["count"] >= before
    finally:
        lane.stop()


def test_on_result_exactly_once_per_predict_none_on_exception():
    calls = {"count": 0}

    def predict_impl(state, questions, model=None, lang=None):
        if state == "bad":
            raise ValueError("infer failed")
        return {"answers": {"ok": True}, "model": model}

    factory = make_factory(calls, predict_impl=predict_impl)
    results = []
    lane = LayaLane(
        make_cfg(), agent_factory=factory,
        on_result=lambda gen, payload: results.append((gen, payload)),
    )
    lane.start()
    try:
        assert wait_until(lane.ready)
        router = factory.routers[0]
        router.predict_log.clear()
        assert lane.submit("good", {}, None, 7) is True
        assert wait_until(lambda: len(results) >= 1, timeout=5.0)
        assert lane.submit("bad", {}, None, 8) is True
        assert wait_until(lambda: len(results) >= 2, timeout=5.0)
        assert results[0][0] == 7
        assert results[0][1]["answers"] == {"ok": True}
        assert results[1] == (8, None)
        assert len(router.predict_log) == 2
    finally:
        lane.stop()


def test_drop_pending_clears_slot():
    calls = {"count": 0}
    entered = threading.Event()
    release = threading.Event()
    gate_calls = {"n": 0}

    def predict_impl(state, questions, model=None, lang=None):
        gate_calls["n"] += 1
        if gate_calls["n"] == 1:
            entered.set()
            assert release.wait(timeout=5.0)
        return {"answers": {"state": state}}

    factory = make_factory(calls, predict_impl=predict_impl)
    results = []
    lane = LayaLane(
        make_cfg(), agent_factory=factory,
        on_result=lambda gen, payload: results.append((gen, payload)),
    )
    lane.start()
    try:
        assert wait_until(lane.ready)
        assert lane.submit("first", {}, None, 1) is True
        assert entered.wait(timeout=5.0)
        assert lane.submit("dropped", {}, None, 2) is True
        lane.drop_pending()
        release.set()
        assert wait_until(lambda: len(results) >= 1, timeout=5.0)
        time.sleep(0.1)  # allow any stray coalesced predict to (not) run
        states = [entry["state"] for entry in factory.routers[0].predict_log[1:]]
        assert "dropped" not in states
        assert [gen for gen, _ in results] == [1]
    finally:
        lane.stop()


def test_stop_during_blocked_predict_bounded_and_silent():
    calls = {"count": 0}
    entered = threading.Event()
    release = threading.Event()

    def predict_impl(state, questions, model=None, lang=None):
        entered.set()
        release.wait(timeout=10.0)
        return {"answers": {}}

    factory = make_factory(calls, predict_impl=predict_impl)
    results = []
    lane = LayaLane(
        make_cfg(), agent_factory=factory,
        on_result=lambda gen, payload: results.append((gen, payload)),
    )
    lane.start()
    assert wait_until(lane.ready)
    router = factory.routers[0]
    router.predict_log.clear()
    calls["count"] = 0  # exclude the loader warmup predict
    assert lane.submit("blocked", {}, None, 3) is True
    assert entered.wait(timeout=5.0)
    start = time.monotonic()
    lane.stop()
    elapsed = time.monotonic() - start
    assert elapsed < 2.5  # 1.0 s join bound per thread (loader already done)
    assert lane.submit("after", {}, None, 4) is False
    release.set()
    time.sleep(0.3)  # let the stranded predict finish if it can
    assert results == []  # no callback fires after stop() returns
    assert calls["count"] == 1


def test_repeated_submit_drop_stop_leaves_no_thread_alive():
    before = set(threading.enumerate())
    for _ in range(3):
        calls = {"count": 0}
        lane = LayaLane(make_cfg(), agent_factory=make_factory(calls))
        lane.start()
        assert wait_until(lane.ready)
        for i in range(5):
            lane.submit("s-%d" % i, {}, None, i)
        lane.drop_pending()
        lane.stop()
        lane.stop()  # idempotent
    # give daemon threads a beat to exit, then check none linger
    assert wait_until(
        lambda: not ({t.name for t in threading.enumerate() if t.is_alive()}
                     & {"laya-loader", "laya-worker"}),
        timeout=5.0,
    )
    lingering = [t for t in threading.enumerate()
                 if t.is_alive() and t.name in ("laya-loader", "laya-worker")
                 and t not in before]
    assert lingering == []


def test_predict_passes_max_len_from_cfg():
    calls = {"count": 0}
    factory = make_factory(calls)
    lane = LayaLane(make_cfg(predict_max_tokens=768), agent_factory=factory)
    lane.start()
    try:
        assert wait_until(lane.ready)
        router = factory.routers[0]
        # loader warmup carries the budget too
        assert router.predict_log[0]["kwargs"].get("max_len") == 768
        router.predict_log.clear()
        assert lane.submit("hello world", {"q": 1}, "es", 1) is True
        assert wait_until(lambda: len(router.predict_log) >= 1, timeout=5.0)
        call = router.predict_log[0]
        assert call["kwargs"].get("max_len") == 768
        assert call["model"] == "multilingual"
        assert call["lang"] == "es"
    finally:
        lane.stop()


# -- todo 5: multilingual pin + lang + resource measurement ------------

def test_multilingual_pin_lang_and_resources():
    laya = pytest.importorskip("laya")
    import resource

    router = laya.Router(default="multilingual")
    assert router.default == "multilingual"

    questions = {"_probe": {"type": "noul", "instructions": "warmup probe"}}
    rss_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    start = time.monotonic()
    result = router.predict("hola mundo", questions, model="multilingual", lang="es")
    latency_s = time.monotonic() - start
    rss_after = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss

    assert set(result) >= {"model", "answers", "usage", "routing"}
    assert "multilingual" in str(result["routing"])
    assert "_probe" in result["answers"]
    assert latency_s >= 0.0
    assert isinstance(router.loaded, list) and "multilingual" in router.loaded
    assert rss_after >= 0 and rss_before >= 0
    # pin verified against the todo 2 laya revision
    assert LAYA_PIN_SHA == "970dc8c5f63d7b886a68409493f37d569424f933"
    assert laya.__version__ == "0.3.20"