agentic-kvc/tests/test_proxy_pick.py

"""Minimal coverage for scripts/cache_aware_proxy pick_instance + cache LRU (S1)."""

from __future__ import annotations

import importlib.util
import sys
import types
from pathlib import Path

import pytest

PROXY_PATH = Path(__file__).resolve().parent.parent / "scripts" / "cache_aware_proxy.py"


def _install_stub_modules() -> None:
    """Provide minimal stand-ins for fastapi/uvicorn/httpx so the proxy
    module imports cleanly without the full server deps."""
    if "uvicorn" not in sys.modules:
        sys.modules["uvicorn"] = types.ModuleType("uvicorn")

    if "fastapi" not in sys.modules:
        fastapi_mod = types.ModuleType("fastapi")

        class _FastAPI:
            def __init__(self, *a, **kw):
                self.state = types.SimpleNamespace()

            def post(self, *a, **kw):
                def deco(fn): return fn
                return deco

            def get(self, *a, **kw):
                def deco(fn): return fn
                return deco

        class _HTTPException(Exception):
            def __init__(self, status_code=500, detail=""):
                self.status_code = status_code
                self.detail = detail

        class _Request:  # not actually instantiated by the routing tests
            pass

        fastapi_mod.FastAPI = _FastAPI
        fastapi_mod.HTTPException = _HTTPException
        fastapi_mod.Request = _Request
        sys.modules["fastapi"] = fastapi_mod

        responses_mod = types.ModuleType("fastapi.responses")

        class _StreamingResponse:
            def __init__(self, *a, **kw): pass

        responses_mod.StreamingResponse = _StreamingResponse
        sys.modules["fastapi.responses"] = responses_mod

    if "httpx" not in sys.modules:
        httpx_mod = types.ModuleType("httpx")

        class _AsyncClient:
            def __init__(self, *a, **kw): pass
            async def aclose(self): pass

        class _Limits:
            def __init__(self, *a, **kw): pass

        httpx_mod.AsyncClient = _AsyncClient
        httpx_mod.Limits = _Limits
        sys.modules["httpx"] = httpx_mod


@pytest.fixture(scope="module")
def proxy():
    _install_stub_modules()
    spec = importlib.util.spec_from_file_location("cache_aware_proxy", PROXY_PATH)
    if spec is None or spec.loader is None:
        pytest.skip(f"cannot load proxy module at {PROXY_PATH}")
    mod = importlib.util.module_from_spec(spec)
    sys.modules["cache_aware_proxy"] = mod
    try:
        spec.loader.exec_module(mod)
    except ModuleNotFoundError as exc:
        pytest.skip(f"proxy dependency missing: {exc}")
    return mod


def _make_inst(proxy, url: str, ongoing_tokens: int = 0,
               active_p_offloads: int = 0):
    inst = proxy.InstanceState(url)
    inst.ongoing_tokens = ongoing_tokens
    inst.active_p_offloads = active_p_offloads
    return inst


def test_record_prefix_evicts_oldest_block(proxy):
    """LRU bound on cached_blocks must evict the oldest entry once full."""
    inst = proxy.InstanceState("http://x")
    saved = proxy.CACHE_CAPACITY_BLOCKS
    proxy.CACHE_CAPACITY_BLOCKS = 2
    try:
        block_size = proxy.BLOCK_SIZE
        # Three distinct one-block prefixes; first must be evicted.
        prefix_a = [1] * block_size
        prefix_b = [2] * block_size
        prefix_c = [3] * block_size
        inst.record_prefix(prefix_a)
        inst.record_prefix(prefix_b)
        inst.record_prefix(prefix_c)
        assert len(inst.cached_blocks) == 2
        # A should have been evicted.
        assert inst.estimate_cache_hit(prefix_a) == 0
        assert inst.estimate_cache_hit(prefix_b) == block_size
        assert inst.estimate_cache_hit(prefix_c) == block_size
    finally:
        proxy.CACHE_CAPACITY_BLOCKS = saved


def test_estimate_cache_hit_touches_lru(proxy):
    """A cache hit must move the block to the MRU position."""
    inst = proxy.InstanceState("http://x")
    saved = proxy.CACHE_CAPACITY_BLOCKS
    proxy.CACHE_CAPACITY_BLOCKS = 2
    try:
        block_size = proxy.BLOCK_SIZE
        a = [1] * block_size
        b = [2] * block_size
        c = [3] * block_size
        inst.record_prefix(a)
        inst.record_prefix(b)
        # Touch A so it becomes MRU; B is now LRU.
        assert inst.estimate_cache_hit(a) == block_size
        # Insert C: B should be evicted, A should remain.
        inst.record_prefix(c)
        assert inst.estimate_cache_hit(a) == block_size
        assert inst.estimate_cache_hit(b) == 0
    finally:
        proxy.CACHE_CAPACITY_BLOCKS = saved


def test_pick_instance_session_affinity_sticks(proxy):
    insts = [_make_inst(proxy, "http://a"), _make_inst(proxy, "http://b")]
    affinity = {"sess1": 1}
    chosen, idx = proxy.pick_instance(insts, None, "sess1", 1000, affinity)
    assert idx == 1 and chosen is insts[1]


def test_pick_instance_session_affinity_breaks_on_overload(proxy):
    """When the pinned instance is heavily overloaded, fallback to load-aware pick."""
    insts = [
        _make_inst(proxy, "http://a", ongoing_tokens=100),
        _make_inst(proxy, "http://b", ongoing_tokens=1_000_000),
        _make_inst(proxy, "http://c", ongoing_tokens=100),
    ]
    affinity = {"sess1": 1}
    chosen, idx = proxy.pick_instance(insts, None, "sess1", 1000, affinity)
    # avg ~333k; B at 1M is ~3x avg, well above OVERLOAD_FACTOR=2.0 -> fallback.
    assert idx != 1
    assert chosen is not insts[1]


def test_pick_instance_p_offload_penalty_steers_away(proxy):
    """Instances actively running offloaded HEAVY prefills get penalized."""
    insts = [
        _make_inst(proxy, "http://a", ongoing_tokens=0, active_p_offloads=2),
        _make_inst(proxy, "http://b", ongoing_tokens=1000),
    ]
    chosen, idx = proxy.pick_instance(insts, None, None, 5000, {})
    # B's 1000-token load is much smaller than A's 2 * HEAVY_THRESHOLD penalty.
    assert idx == 1 and chosen is insts[1]


def test_pick_instance_lmetric_picks_lowest_score(proxy):
    insts = [_make_inst(proxy, "http://a"), _make_inst(proxy, "http://b")]
    insts[0].pending_prefill_tokens = 0
    insts[0].num_requests = 0
    insts[1].pending_prefill_tokens = 5000
    insts[1].num_requests = 4
    chosen, idx = proxy.pick_instance_lmetric(insts, None, None, 1000, {})
    # Empty instance has score = 1000 * 0 = 0; busy one has (5000+1000)*4.
    assert idx == 0 and chosen is insts[0]