Team Ai
Modelpublic

Felipe97/llama-cpp-compiled

sourceHugging Faceupdated 21d agoView on Hugging Face
0likes1.2kdownloads
test_router.py622 linesDownload Raw Back to unit
1import threading2import pytest3from utils import *4 5server: ServerProcess6 7@pytest.fixture(autouse=True)8def create_server():9    global server10    server = ServerPreset.router()11 12 13def test_router_props():14    global server15    server.models_max = 216    server.no_models_autoload = True17    server.start()18    res = server.make_request("GET", "/props")19    assert res.status_code == 20020    assert res.body["role"] == "router"21    assert res.body["max_instances"] == 222    assert res.body["models_autoload"] is False23    assert res.body["build_info"].startswith("b")24 25 26@pytest.mark.parametrize(27    "model,success",28    [29        ("ggml-org/tinygemma3-GGUF:Q8_0", True),30        ("non-existent/model", False),31    ]32)33def test_router_chat_completion_stream(model: str, success: bool):34    global server35    server.start()36    content = ""37    ex: ServerError | None = None38    try:39        res = server.make_stream_request("POST", "/chat/completions", data={40            "model": model,41            "max_tokens": 16,42            "messages": [43                {"role": "user", "content": "hello"},44            ],45            "stream": True,46        })47        for data in res:48            if data["choices"]:49                choice = data["choices"][0]50                if choice["finish_reason"] in ["stop", "length"]:51                    assert "content" not in choice["delta"]52                else:53                    assert choice["finish_reason"] is None54                    content += choice["delta"]["content"] or ''55    except ServerError as e:56        ex = e57 58    if success:59        assert ex is None60        assert len(content) > 061    else:62        assert ex is not None63        assert content == ""64 65 66def _get_model_ids(is_reload: bool, headers: dict | None = None) -> set[str]:67    res = server.make_request(68        "GET", "/models" + ("?reload=1" if is_reload else ""), headers=headers69    )70    assert res.status_code == 20071    return {item["id"] for item in res.body.get("data", [])}72 73 74def _get_model_status(model_id: str, headers: dict | None = None) -> str:75    res = server.make_request("GET", "/models", headers=headers)76    assert res.status_code == 20077    for item in res.body.get("data", []):78        if item.get("id") == model_id or item.get("model") == model_id:79            return item["status"]["value"]80    raise AssertionError(f"Model {model_id} not found in /models response")81 82 83def _wait_for_model_status(model_id: str, desired: set[str], timeout: int = 60, headers: dict | None = None) -> str:84    deadline = time.time() + timeout85    last_status = None86    while time.time() < deadline:87        last_status = _get_model_status(model_id, headers=headers)88        if last_status in desired:89            return last_status90        time.sleep(0.01)91    raise AssertionError(92        f"Timed out waiting for {model_id} to reach {desired}, last status: {last_status}"93    )94 95 96def _load_model_and_wait(97    model_id: str, timeout: int = 60, headers: dict | None = None98) -> None:99    load_res = server.make_request(100        "POST", "/models/load", data={"model": model_id}, headers=headers101    )102    assert load_res.status_code == 200103    assert isinstance(load_res.body, dict)104    assert load_res.body.get("success") is True105    _wait_for_model_status(model_id, {"loaded"}, timeout=timeout, headers=headers)106 107 108def test_router_unload_model():109    global server110    server.start()111    model_id = "ggml-org/tinygemma3-GGUF:Q8_0"112 113    _load_model_and_wait(model_id)114 115    unload_res = server.make_request("POST", "/models/unload", data={"model": model_id})116    assert unload_res.status_code == 200117    assert unload_res.body.get("success") is True118    _wait_for_model_status(model_id, {"unloaded"})119 120 121def test_router_models_max_evicts_lru():122    global server123    server.models_max = 2124    server.start()125 126    candidate_models = [127        "ggml-org/tinygemma3-GGUF:Q8_0",128        "ggml-org/test-model-stories260K:F32",129        "ggml-org/test-model-stories260K-infill:F32",130    ]131 132    # Load only the first 2 models to fill the cache133    first, second, third = candidate_models[:3]134 135    _load_model_and_wait(first, timeout=120)136    _load_model_and_wait(second, timeout=120)137 138    # Verify both models are loaded139    assert _get_model_status(first) == "loaded"140    assert _get_model_status(second) == "loaded"141 142    # Load the third model - this should trigger LRU eviction of the first model143    _load_model_and_wait(third, timeout=120)144 145    # Verify eviction: third is loaded, first was evicted146    assert _get_model_status(third) == "loaded"147    assert _get_model_status(first) == "unloaded"148 149 150# server_lru_sched tests (relying on LLAMA_SERVER_DEBUG_FAKE_TIMING)151 152MODEL_A = "ggml-org/tinygemma3-GGUF:Q8_0"153MODEL_B = "ggml-org/test-model-stories260K:F32"154MODEL_C = "ggml-org/test-model-stories260K-infill:F32"155 156 157def _tokenize(model_id: str, timeout: float | None = DEFAULT_REQUEST_TIMEOUT) -> ServerResponse:158    return server.make_request(159        "POST", "/tokenize", data={"model": model_id, "content": "hello world"}, timeout=timeout160    )161 162 163class _Bg:164    """runs one request in a thread, keeps its result, error and finish time"""165 166    def __init__(self, fn):167        self.result = None168        self.error: Exception | None = None169        self.done_at: float = 0.0170        self._thread = threading.Thread(target=self._run, args=(fn,), daemon=True)171 172    def _run(self, fn):173        try:174            self.result = fn()175        except Exception as e:176            self.error = e177        self.done_at = time.time()178 179    def start(self):180        self._thread.start()181        return self182 183    def join(self, timeout: int = 180):184        self._thread.join(timeout)185        assert not self._thread.is_alive(), "background request did not finish in time"186        return self187 188    def assert_ok(self, what: str):189        assert self.error is None, f"{what} raised {self.error!r}"190        assert self.result is not None and self.result.status_code == 200, \191            f"{what} failed: {self.result.status_code if self.result else None} {self.result.body if self.result else None}"192 193 194def test_router_queue_does_not_evict_busy_model():195    """a request that finds no free slot waits, and the model serving a request survives it"""196    global server197    server.models_max = 1198    server.start()199 200    _load_model_and_wait(MODEL_A, timeout=120)201 202    busy = _Bg(lambda: _tokenize(MODEL_A)).start()203    time.sleep(0.5)  # let the request reach the child and take the only slot204 205    # no slot free and MODEL_A is busy, so this queues instead of evicting mid-request206    queued = _Bg(lambda: _tokenize(MODEL_B)).start()207 208    busy.join()209    queued.join()210 211    # had MODEL_A been evicted while serving, its own request would have died212    busy.assert_ok("request against the busy model")213    queued.assert_ok("queued request")214 215    _wait_for_model_status(MODEL_B, {"loaded"}, timeout=120)216    assert _get_model_status(MODEL_A) == "unloaded"217 218 219def test_router_queue_coalesces_requests_for_same_model():220    """many requests for one missing model share a slot, so only one model is given up"""221    global server222    server.models_max = 2223    server.start()224 225    _load_model_and_wait(MODEL_A, timeout=120)226    _load_model_and_wait(MODEL_B, timeout=120)227 228    # keep MODEL_A busy so MODEL_B is the only model that can be given up229    busy = _Bg(lambda: _tokenize(MODEL_A)).start()230    time.sleep(0.5)231 232    waiters = [_Bg(lambda: _tokenize(MODEL_C)).start() for _ in range(3)]233 234    busy.join()235    for w in waiters:236        w.join()237 238    busy.assert_ok("request against the busy model")239    for i, w in enumerate(waiters):240        w.assert_ok(f"queued request {i}")241 242    _wait_for_model_status(MODEL_C, {"loaded"}, timeout=120)243    # one entry for 3 requests means one eviction: MODEL_B goes, MODEL_A is left alone.244    # without coalescing the leftover entries still ask for a slot,245    # and MODEL_A is taken too as soon as it goes idle246    assert _get_model_status(MODEL_A) == "loaded"247    assert _get_model_status(MODEL_B) == "unloaded"248 249 250def test_router_queue_client_disconnect_keeps_model():251    """a client that leaves while queued must not cost a running model its slot"""252    global server253    server.models_max = 1254    server.start()255 256    _load_model_and_wait(MODEL_A, timeout=120)257 258    busy = _Bg(lambda: _tokenize(MODEL_A)).start()259    time.sleep(0.5)260 261    # queues behind MODEL_A, then gives up long before MODEL_A goes idle262    with pytest.raises(requests.exceptions.RequestException):263        _tokenize(MODEL_B, timeout=1)264 265    busy.join()266    busy.assert_ok("request against the busy model")267 268    # nobody is waiting anymore, so MODEL_A keeps its slot269    time.sleep(3)270    assert _get_model_status(MODEL_A) == "loaded"271    assert _get_model_status(MODEL_B) == "unloaded"272 273 274def test_router_queue_is_fifo():275    """the queue is served in arrival order"""276    global server277    server.models_max = 1278    server.start()279 280    _load_model_and_wait(MODEL_A, timeout=120)281 282    busy = _Bg(lambda: _tokenize(MODEL_A)).start()283    time.sleep(0.5)284 285    first = _Bg(lambda: _tokenize(MODEL_B)).start()286    time.sleep(1)  # keep the arrival order unambiguous287    second = _Bg(lambda: _tokenize(MODEL_C)).start()288 289    busy.join()290    first.join()291    second.join()292 293    busy.assert_ok("request against the busy model")294    first.assert_ok("first queued request")295    second.assert_ok("second queued request")296 297    assert first.done_at < second.done_at, "queue was not served in arrival order"298 299 300def test_router_queue_two_waiters_share_one_eviction():301    """two requests that both find the same idle model must both be served in the end"""302    global server303    server.models_max = 1304    server.start()305 306    _load_model_and_wait(MODEL_A, timeout=120)307 308    # both arrive while MODEL_A is idle, so both want its slot; only one eviction can happen309    first = _Bg(lambda: _tokenize(MODEL_B)).start()310    second = _Bg(lambda: _tokenize(MODEL_C)).start()311 312    first.join(90)313    second.join(90)314 315    first.assert_ok("first queued request")316    second.assert_ok("second queued request")317    assert _get_model_status(MODEL_A) == "unloaded"318 319 320def test_router_no_models_autoload():321    global server322    server.no_models_autoload = True323    server.start()324    model_id = "ggml-org/tinygemma3-GGUF:Q8_0"325 326    res = server.make_request(327        "POST",328        "/v1/chat/completions",329        data={330            "model": model_id,331            "messages": [{"role": "user", "content": "hello"}],332            "max_tokens": 4,333        },334    )335    assert res.status_code == 400336    assert "error" in res.body337 338    _load_model_and_wait(model_id)339 340    success_res = server.make_request(341        "POST",342        "/v1/chat/completions",343        data={344            "model": model_id,345            "messages": [{"role": "user", "content": "hello"}],346            "max_tokens": 4,347        },348    )349    assert success_res.status_code == 200350    assert "error" not in success_res.body351 352 353def test_router_api_key_required():354    global server355    server.api_key = "sk-router-secret"356    server.start()357 358    model_id = "ggml-org/tinygemma3-GGUF:Q8_0"359    auth_headers = {"Authorization": f"Bearer {server.api_key}"}360 361    res = server.make_request(362        "POST",363        "/v1/chat/completions",364        data={365            "model": model_id,366            "messages": [{"role": "user", "content": "hello"}],367            "max_tokens": 4,368        },369    )370    assert res.status_code == 401371    assert res.body.get("error", {}).get("type") == "authentication_error"372 373    _load_model_and_wait(model_id, headers=auth_headers)374 375    authed = server.make_request(376        "POST",377        "/v1/chat/completions",378        headers=auth_headers,379        data={380            "model": model_id,381            "messages": [{"role": "user", "content": "hello"}],382            "max_tokens": 4,383        },384    )385    assert authed.status_code == 200386    assert "error" not in authed.body387 388 389def test_router_reload_models():390    """POST /models/reload re-reads the INI preset and updates the model list."""391    global server392 393    preset_path = os.path.join(TMP_DIR, "test_reload.ini")394 395    # Initial preset: two models396    with open(preset_path, "w") as f:397        f.write(398            "[model-reload-a]\n"399            "hf-repo = ggml-org/test-model-stories260K\n"400            "\n"401            "[model-reload-b]\n"402            "hf-repo = ggml-org/test-model-stories260K-infill\n"403        )404 405    server.models_preset = preset_path406    server.start()407 408    ids = _get_model_ids(is_reload=False)409    assert "model-reload-a" in ids410    assert "model-reload-b" in ids411 412    # Updated preset: remove a, keep b unchanged, add c413    with open(preset_path, "w") as f:414        f.write(415            "[model-reload-b]\n"416            "hf-repo = ggml-org/test-model-stories260K-infill\n"417            "\n"418            "[model-reload-c]\n"419            "hf-repo = ggml-org/test-model-stories260K\n"420        )421 422    try:423        ids = _get_model_ids(is_reload=True)424        assert "model-reload-a" not in ids, "removed model should no longer appear"425        assert "model-reload-b" in ids, "unchanged model should still appear"426        assert "model-reload-c" in ids, "newly added model should appear"427    finally:428        os.remove(preset_path)429 430 431def test_router_dedup_cache_models():432    """dedup-cache-models hides the cache entry backing a preset from GET /models"""433    global server434 435    preset_path = os.path.join(TMP_DIR, "test_dedup.ini")436    cache_id = "ggml-org/test-model-stories260K:F32"437 438    with open(preset_path, "w") as f:439        f.write(440            "[model-dedup]\n"441            "hf-repo = ggml-org/test-model-stories260K\n"442            "dedup-cache-models = 1\n"443        )444 445    server.models_preset = preset_path446    server.start()447 448    try:449        ids = _get_model_ids(is_reload=False)450        assert "model-dedup" in ids451        assert cache_id not in ids, "cache model should be hidden by dedup"452        # other cache models are unaffected453        assert "ggml-org/tinygemma3-GGUF:Q8_0" in ids454 455        # the hidden model is only hidden from the listing, it can still be used456        res = server.make_request("POST", "/tokenize", data={"model": cache_id, "content": "hello"})457        assert res.status_code == 200458 459        # disabling the flag brings the cache entry back on reload460        with open(preset_path, "w") as f:461            f.write(462                "[model-dedup]\n"463                "hf-repo = ggml-org/test-model-stories260K\n"464            )465        ids = _get_model_ids(is_reload=True)466        assert cache_id in ids467 468        # the flag also works from the global section469        with open(preset_path, "w") as f:470            f.write(471                "[*]\n"472                "dedup-cache-models = 1\n"473                "\n"474                "[model-dedup]\n"475                "hf-repo = ggml-org/test-model-stories260K\n"476            )477        ids = _get_model_ids(is_reload=True)478        assert "model-dedup" in ids479        assert cache_id not in ids, "cache model should be hidden by global dedup"480    finally:481        os.remove(preset_path)482 483 484def test_router_remote_preset():485    global server486    server.model_hf_repo = "ggml-org/test-preset-ci"487    server.model_hf_file = None488    server.offline = False489    server.start()490 491    # Should see preset models in GET /models492    res = server.make_request("GET", "/models")493    assert res.status_code == 200494    ids = {item["id"] for item in res.body.get("data", [])}495    assert "tinygemma3-preset" in ids496    assert "stories260K-test" in ids497 498    # Should be able to load a preset model499    model_id = "tinygemma3-preset"500    _load_model_and_wait(model_id)501 502 503MODEL_DOWNLOAD_ID = "ggml-org/test-model-router-download:F16"504MODEL_DOWNLOAD_TIMEOUT = 30505 506 507def _listen_sse(508    server: ServerProcess, collected: list, stop: threading.Event, ready: threading.Event | None = None509):510    """Collect /models/sse events into `collected` until `stop` is set.511 512    When `ready` is provided, it is set once the streaming response is open,513    i.e. the server has accepted the connection and registered us as a514    subscriber. Callers that trigger one-shot events (e.g. download_finished)515    must wait on `ready` before acting, otherwise the event can be broadcast516    before this client is subscribed and be lost.517    """518    url = f"http://{server.server_host}:{server.server_port}/models/sse"519    try:520        with requests.get(url, stream=True, timeout=MODEL_DOWNLOAD_TIMEOUT) as resp:521            if ready is not None:522                ready.set()523            for line_bytes in resp.iter_lines():524                if stop.is_set():525                    break526                line = line_bytes.decode("utf-8")527                if line.startswith("data: "):528                    collected.append(json.loads(line[6:]))529    except Exception:530        pass531 532 533def _wait_for_sse_event(collected: list, event_type: str, model: str, timeout: int) -> bool:534    deadline = time.time() + timeout535    while time.time() < deadline:536        if any(e.get("event") == event_type and e.get("model") == model for e in collected):537            return True538        time.sleep(0.01)539    return False540 541 542def test_router_download_model():543    """Case 1: download a model at the model limit, verify SSE events and GET /models."""544    global server545    server.models_max = 1546    server.start()547 548    # Ensure the model is not present before we start549    server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}")550 551    # A download worker must not consume or evict a model slot552    _load_model_and_wait(MODEL_B, timeout=120)553 554    sse_events: list = []555    stop = threading.Event()556    sse_ready = threading.Event()557    sse_thread = threading.Thread(558        target=_listen_sse, args=(server, sse_events, stop, sse_ready), daemon=True559    )560    sse_thread.start()561 562    # wait for the SSE client to be subscribed before triggering the download,563    # otherwise the one-shot download_finished event can be broadcast before564    # this client is registered and be lost565    assert sse_ready.wait(10), "SSE client failed to connect"566 567    # Trigger the download568    res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID})569    assert res.status_code == 200570    assert res.body.get("success") is True571 572    # Wait for download_finished SSE event573    finished = _wait_for_sse_event(574        sse_events, "download_finished", MODEL_DOWNLOAD_ID, MODEL_DOWNLOAD_TIMEOUT575    )576    stop.set()577 578    assert finished, "Never received download_finished SSE event"579    assert any(580        e.get("event") == "download_progress" and e.get("model") == MODEL_DOWNLOAD_ID581        for e in sse_events582    ), "No download_progress events received"583 584    # Model should now appear in GET /models585    ids = _get_model_ids(is_reload=False)586    assert MODEL_DOWNLOAD_ID in ids, f"{MODEL_DOWNLOAD_ID} not found in /models after download"587    assert _get_model_status(MODEL_B) == "loaded"588 589 590def test_router_delete_model():591    """Case 2: delete the downloaded model, verify it disappears from GET /models."""592    global server593    server.start()594 595    # Ensure the model exists (download it if needed)596    if MODEL_DOWNLOAD_ID not in _get_model_ids(is_reload=False):597        sse_events: list = []598        stop = threading.Event()599        sse_ready = threading.Event()600        threading.Thread(601            target=_listen_sse, args=(server, sse_events, stop, sse_ready), daemon=True602        ).start()603        # subscribe before triggering the download so the one-shot604        # download_finished event is not lost (see test_router_download_model)605        assert sse_ready.wait(10), "SSE client failed to connect"606        res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID})607        assert res.status_code == 200608        finished = _wait_for_sse_event(609            sse_events, "download_finished", MODEL_DOWNLOAD_ID, MODEL_DOWNLOAD_TIMEOUT610        )611        stop.set()612        assert finished, "Model did not finish downloading before delete test"613 614    # Delete the model615    del_res = server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}")616    assert del_res.status_code == 200617    assert del_res.body.get("success") is True618 619    # Model should no longer appear in GET /models620    ids = _get_model_ids(is_reload=False)621    assert MODEL_DOWNLOAD_ID not in ids, f"{MODEL_DOWNLOAD_ID} still present after deletion"622