esh-pfi-infrastructure/services/semif-serve/tests/test_engine.py

"""TorchEngine's OOM path against a fake torch (INV-4). Found on the card 2026-09-27: after a
503 the failed call's tensors stayed alive (11.92 GiB allocated) because the raised
OutOfMemory chained back to the torch exception, whose traceback held the scorer's frames."""
import weakref

import pytest

from semif_serve.config import Settings
from semif_serve.engine import TorchEngine
from semif_serve.errors import OutOfMemory, ScoringFailed


class FakeTorch:
    class cuda:
        class OutOfMemoryError(RuntimeError):
            pass

        empties = []           # for each empty_cache() call: was the failed call's tensor already freed?
        watched = []

        @classmethod
        def empty_cache(cls):
            cls.empties.append(all(ref() is None for ref in cls.watched))


class Tensor:
    pass


def failing_scorer(*_args):
    kv_cache = Tensor()                                   # stands in for the replicated prefix cache
    FakeTorch.cuda.watched.append(weakref.ref(kv_cache))
    raise FakeTorch.cuda.OutOfMemoryError("CUDA out of memory. Tried to allocate 490.00 MiB.\nGPU 0 has ...")


@pytest.mark.parametrize("call", ["direct", "shared"])
def test_oom_frees_the_failed_call_before_emptying_the_cache_and_keeps_nothing_alive(call):
    FakeTorch.cuda.empties.clear(), FakeTorch.cuda.watched.clear()
    engine = TorchEngine(FakeTorch, model=None, tokenizer=None, metadata={}, settings=Settings(api_token="t" * 40),
                         direct_fn=failing_scorer, shared_fn=failing_scorer)
    with pytest.raises(OutOfMemory) as info:
        getattr(engine, call)({})
    assert str(info.value) == "CUDA out of memory. Tried to allocate 490.00 MiB."
    assert info.value.__cause__ is None and info.value.__context__ is None   # no chain back to the frames
    assert FakeTorch.cuda.watched[0]() is None                               # nothing keeps the tensor alive
    assert FakeTorch.cuda.empties == [True]                                  # emptied once, after it was freed


def test_other_scorer_errors_pass_through_unchanged():
    def bad(*_args):
        raise ValueError("Row r: 5000 input tokens exceed limit 4096")
    engine = TorchEngine(FakeTorch, None, None, {}, Settings(api_token="t" * 40), direct_fn=bad, shared_fn=bad)
    with pytest.raises(ValueError, match="exceed limit"):
        engine.direct({})


class ReservingTorch(FakeTorch):
    class cuda(FakeTorch.cuda):
        reserved = 0
        emptied = 0

        @classmethod
        def memory_reserved(cls, _device=0):
            return cls.reserved

        @classmethod
        def empty_cache(cls):
            cls.emptied += 1


@pytest.mark.parametrize("reserved_after, released", [(8 * 2**30, 0), (8 * 2**30 + 512 * 2**20, 0),
                                                      (8 * 2**30 + 513 * 2**20, 1), (12 * 2**30, 1)])
def test_a_burst_is_returned_to_the_driver_after_the_call(reserved_after, released):
    ReservingTorch.cuda.emptied = 0

    def scorer(*_args):
        ReservingTorch.cuda.reserved = reserved_after
        return {"ok": True}

    engine = TorchEngine(ReservingTorch, None, None, {}, Settings(api_token="t" * 40),
                         direct_fn=scorer, shared_fn=scorer, release_above_bytes=8 * 2**30 + 512 * 2**20)
    assert engine.direct({}) == {"ok": True}
    assert ReservingTorch.cuda.emptied == released


def cyclic_tensor():
    """A tensor held in a reference cycle, as real frames and tensors often are: only gc frees it."""
    t = Tensor()
    t.self_ref = t
    FakeTorch.cuda.watched.append(weakref.ref(t))
    return t


@pytest.mark.parametrize("raised, expected_type, expected_message", [
    (lambda: FakeTorch.cuda.OutOfMemoryError(""), OutOfMemory, "CUDA out of memory"),
    (lambda: RuntimeError("CUBLAS_STATUS_ALLOC_FAILED: CUDA error: out of memory"), OutOfMemory,
     "CUBLAS_STATUS_ALLOC_FAILED: CUDA error: out of memory"),
    (lambda: RuntimeError("Invalid native prefix cache"), ScoringFailed, "RuntimeError: Invalid native prefix cache"),
    (lambda: KeyError("option_logits"), ScoringFailed, "KeyError: 'option_logits'"),
])
def test_every_non_validation_failure_is_released_unchained_after_gc(raised, expected_type, expected_message):
    FakeTorch.cuda.empties.clear(), FakeTorch.cuda.watched.clear()

    def scorer(*_args):
        kv_cache = cyclic_tensor()   # noqa: F841 — alive in this frame when it raises
        raise raised()

    engine = TorchEngine(FakeTorch, None, None, {}, Settings(api_token="t" * 40), direct_fn=scorer, shared_fn=scorer)
    with pytest.raises(expected_type) as info:
        engine.shared([])
    assert str(info.value) == expected_message
    assert info.value.__cause__ is None and info.value.__context__ is None
    assert FakeTorch.cuda.empties == [True]          # gc freed the cycle BEFORE the cache was emptied