vllm-frontdoor/router/tests/test_idle.py

"""Tiered idle management: thresholds, lock races, active-request guard."""

from __future__ import annotations

import asyncio
import time

from config import DEPTH_AWAKE, DEPTH_OFFLOADED, DEPTH_SLEEPING


def _age(svc, seconds: float) -> None:
    svc.last_activity = time.monotonic() - seconds


async def test_idle_level1_after_threshold(stack, backend):
    stack.cfg.idle_sleep_min = 15 / 60.0     # 15 s in "minutes"
    stack.cfg.idle_offload_min = 180 / 60.0
    text = stack.manager.services["text"]
    text.depth = DEPTH_AWAKE
    _age(text, 20)

    await stack.manager.idle_tick()

    assert backend.count("POST vllm-text/sleep") == 1
    assert "level=1" in backend.calls[backend.last("POST vllm-text/sleep")]
    assert text.depth == DEPTH_SLEEPING


async def test_idle_escalates_to_level2(stack, backend):
    stack.cfg.idle_sleep_min = 15 / 60.0
    stack.cfg.idle_offload_min = 180 / 60.0
    text = stack.manager.services["text"]
    text.depth = DEPTH_AWAKE
    _age(text, 400)

    await stack.manager.idle_tick()
    assert "level=2" in backend.calls[backend.last("POST vllm-text/sleep")]
    assert text.depth == DEPTH_OFFLOADED


async def test_level1_service_escalates_to_level2(stack, backend):
    """Already napping (level 1) and still idle -> escalate to offload.

    A direct POST /sleep?level=2 on a level-1-sleeping backend is a
    well-behaved NO-OP that retains the host-RAM copy (calibration
    2026-08-17), so the escalation must wake into RAM first and then offload.
    """
    stack.cfg.idle_sleep_min = 1 / 60.0
    stack.cfg.idle_offload_min = 5 / 60.0
    text = stack.manager.services["text"]
    assert (await stack.manager.sleep_service("text", 1))["ok"] is True
    backend.calls.clear()
    svc = backend.services["vllm-text"]
    svc["wake_up_calls"] = svc["sleep_calls"] = 0
    _age(text, 400)

    await stack.manager.idle_tick()
    assert text.depth == DEPTH_OFFLOADED
    assert svc["wake_up_calls"] == 1                     # wake into RAM ...
    assert svc["sleep_calls"] == 1                       # ... then offload
    assert backend.last("POST vllm-text/sleep?level=2") > backend.last("POST vllm-text/wake_up")


async def test_offload_from_awake_is_direct(stack, backend):
    """Only from depth 'awake' can level 2 be entered directly."""
    backend.services["vllm-embed"]["wake_up_calls"] = 0
    result = await stack.manager.sleep_service("embed", 2)
    assert result["ok"] is True
    assert backend.services["vllm-embed"]["wake_up_calls"] == 0
    assert backend.count("POST vllm-embed/sleep?level=2") == 1
    assert stack.manager.services["embed"].depth == DEPTH_OFFLOADED


async def test_escalation_failure_is_reported(stack, backend):
    assert (await stack.manager.sleep_service("ocr", 1))["ok"] is True
    backend.services["vllm-ocr"]["wake_fails"] = 1
    backend.calls.clear()
    result = await stack.manager.sleep_service("ocr", 2)
    assert result["ok"] is False
    assert result["reason"] == "escalation_failed"
    assert backend.count("POST vllm-ocr/sleep") == 0
    assert stack.manager.services["ocr"].depth == DEPTH_SLEEPING


async def test_offloaded_service_is_left_alone(stack, backend):
    stack.cfg.idle_offload_min = 1 / 60.0
    text = stack.manager.services["text"]
    text.depth = DEPTH_OFFLOADED
    _age(text, 10_000)
    await stack.manager.idle_tick()
    assert backend.calls == []


async def test_active_requests_block_idle_sleep(stack, backend):
    stack.cfg.idle_sleep_min = 1 / 60.0
    text = stack.manager.services["text"]
    text.depth = DEPTH_AWAKE
    _age(text, 600)

    manager = stack.manager
    manager.begin_request(text)                 # long generation in flight
    try:
        await manager.idle_tick()
        assert backend.calls == []
    finally:
        manager.end_request(text)


async def test_stream_completes_then_idle_can_sleep(stack, backend, pub):
    """end_request (background task of the streamed response) re-opens sleep."""
    stack.cfg.idle_sleep_min = 1 / 60.0
    text = stack.manager.services["text"]
    r = await pub.post("/v1/chat/completions", json={"model": "text"})
    assert r.status_code == 200
    assert text.active_requests == 0           # released after the body drained

    _age(text, 600)
    await stack.manager.idle_tick()
    assert backend.count("POST vllm-text/sleep") == 1


async def test_race_last_activity_refreshed_under_lock(stack, backend):
    """A request landed between the threshold check and the locked re-check."""
    stack.cfg.idle_sleep_min = 1 / 60.0        # threshold = 60 s
    text = stack.manager.services["text"]
    text.depth = DEPTH_AWAKE
    _age(text, 600)                            # what the idle scan saw

    # ... and then a request arrived, refreshing last_activity *now*
    text.last_activity = time.monotonic()
    result = await stack.manager.sleep_service("text", 1, reason="idle", min_idle_s=60.0)
    assert result["ok"] is False
    assert result["reason"] == "activity_resumed"
    assert backend.count("POST vllm-text/sleep") == 0


async def test_race_active_request_seen_under_lock(stack, backend):
    text = stack.manager.services["text"]
    text.depth = DEPTH_AWAKE
    _age(text, 600)
    text.active_requests = 2                   # arrived while we were scanning
    result = await stack.manager.sleep_service("text", 1, reason="idle", min_idle_s=60.0)
    assert result["ok"] is False
    assert result["reason"] == "active_requests"
    assert backend.calls == []
    text.active_requests = 0


async def test_request_at_the_moment_the_timer_expires(pub, backend, stack):
    """E2E case 11 in miniature: the request wins, no sleep mid-flight."""
    stack.cfg.idle_sleep_min = 1 / 60.0
    stack.cfg.idle_offload_min = 5 / 60.0
    text = stack.manager.services["text"]
    text.depth = DEPTH_AWAKE
    _age(text, 600)

    async def tick_soon():
        await asyncio.sleep(0)                 # run after the request started
        return await stack.manager.idle_tick()

    tick, response = await asyncio.gather(tick_soon(), pub.post(
        "/v1/chat/completions", json={"model": "text"}))
    assert response.status_code == 200
    assert backend.count("POST vllm-text/sleep") == 0


async def test_admin_sleep_refuses_when_active(adm, stack, backend):
    text = stack.manager.services["text"]
    text.depth = DEPTH_AWAKE
    stack.manager.begin_request(text)
    r = await adm.post("/admin/sleep/text?level=1")
    assert r.status_code == 409
    assert r.json()["reason"] == "active_requests"
    assert backend.count("/sleep") == 0
    stack.manager.end_request(text)


async def test_idle_loop_runs_in_background_when_started(stack, backend):
    stack.cfg.idle_enabled = True
    stack.cfg.idle_poll_s = 0.01
    stack.cfg.idle_sleep_min = 1 / 60.0      # 60 s
    stack.cfg.idle_offload_min = 1000.0      # far away: expect the level-1 tier
    text = stack.manager.services["text"]
    text.depth = DEPTH_AWAKE
    _age(text, 600)
    stack.manager.start()
    try:
        for _ in range(100):
            if backend.count("POST vllm-text/sleep"):
                break
            await asyncio.sleep(0.01)
    finally:
        await stack.manager.stop()
    assert backend.count("POST vllm-text/sleep") == 1
    assert "level=1" in backend.calls[backend.last("POST vllm-text/sleep")]