Team Ai
Modelpublic

Felipe97/llama-cpp-compiled

sourceHugging Faceupdated 21d agoView on Hugging Face
0likes1.2kdownloads
test_sleep.py128 linesDownload Raw Back to unit
1import pytest2import time3from utils import *4 5server = ServerPreset.tinyllama2()6 7 8@pytest.fixture(autouse=True)9def create_server():10    global server11    server = ServerPreset.tinyllama2()12 13 14def is_sleeping(server: ServerProcess) -> bool:15    res = server.make_request("GET", "/props")16    assert res.status_code == 20017    return res.body["is_sleeping"]18 19 20def wait_for_sleep(server: ServerProcess, timeout: float = 10.0):21    start = time.time()22    while time.time() - start < timeout:23        if is_sleeping(server):24            return25        time.sleep(0.1)26    raise TimeoutError("server did not go to sleep")27 28 29def fetch_metrics(server: ServerProcess) -> str:30    res = server.make_request("GET", "/metrics")31    assert res.status_code == 20032    assert isinstance(res.body, str)33    return res.body34 35 36def get_metric(text: str, name: str) -> float:37    prefix = f"llamacpp:{name} "38    values = [ln for ln in text.splitlines() if ln.startswith(prefix)]39    assert len(values) == 1, f"{name} not found in metrics"40    return float(values[0][len(prefix):])41 42 43def test_server_sleep():44    global server45    server.sleep_idle_seconds = 146    server.start()47 48    # wait a bit so that server can go to sleep49    time.sleep(2)50 51    # make sure these endpoints are still responsive after sleep52    res = server.make_request("GET", "/health")53    assert res.status_code == 20054    res = server.make_request("GET", "/props")55    assert res.status_code == 20056    assert res.body["is_sleeping"] == True57    res = server.make_request("GET", "/models")58    assert res.status_code == 20059    assert len(res.body["data"]) == 160    assert res.body["data"][0]["id"] == server.model_alias61 62    # make a generation request to wake up the server63    res = server.make_request("POST", "/completion", data={64        "n_predict": 1,65        "prompt": "Hello",66    })67    assert res.status_code == 20068 69    # it should no longer be sleeping70    res = server.make_request("GET", "/props")71    assert res.status_code == 20072    assert res.body["is_sleeping"] == False73 74 75def test_server_sleep_read_only_endpoints():76    global server77    server.sleep_idle_seconds = 178    server.server_metrics = True79    server.start()80 81    res = server.make_request("POST", "/completion", data={82        "n_predict": 4,83        "prompt": "Hello",84    })85    assert res.status_code == 20086 87    # the first scrape resets the throughput buckets, so that the second one reports88    # the same zero rates as the snapshot taken on entering sleep89    fetch_metrics(server)90    metrics_awake = fetch_metrics(server)91    assert get_metric(metrics_awake, "tokens_predicted_total") > 092 93    wait_for_sleep(server)94 95    # during sleep, metrics are served from the snapshot taken right before sleeping96    assert fetch_metrics(server) == metrics_awake97 98    # scraping /metrics must not wake the server up99    assert is_sleeping(server)100 101 102def test_server_sleep_metrics_buckets():103    global server104    server.sleep_idle_seconds = 1105    server.server_metrics = True106    server.start()107 108    res = server.make_request("POST", "/completion", data={109        "n_predict": 8,110        "prompt": "Hello",111    })112    assert res.status_code == 200113 114    wait_for_sleep(server)115 116    # the first scrape reports the throughput of the last generation117    assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") > 0118 119    # nothing runs while sleeping, so the next scrapes report an empty window120    assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") == 0121    assert is_sleeping(server)122 123    # waking up must not report the buckets again124    res = server.make_request("POST", "/tokenize", data={"content": "Hello"})125    assert res.status_code == 200126    assert is_sleeping(server) == False127    assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") == 0128