Spaces:
Sleeping
Sleeping
File size: 3,316 Bytes
421b8c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | """Multi-model routing: several loaded engines behind one shared FIFO worker."""
from concurrent.futures import ThreadPoolExecutor
import threading
from fastapi.testclient import TestClient
import pytest
from cism.server import create_app
from test_server import FakeEngine, chat, completion, wait_until
def two_engines():
first, second = FakeEngine(), FakeEngine()
first.model_id, second.model_id = "model-one", "model-two"
second.context_length = 256
return first, second
def test_models_lists_every_engine():
first, second = two_engines()
with TestClient(create_app({"model-one": first, "model-two": second})) as client:
data = client.get("/v1/models").json()["data"]
assert [model["id"] for model in data] == ["model-one", "model-two"]
assert [model["context_length"] for model in data] == [512, 256]
def test_requests_route_to_named_engine():
first, second = two_engines()
with TestClient(create_app({"model-one": first, "model-two": second})) as client:
assert client.post("/v1/completions",
json=completion(model="model-one", prompt=[1])).status_code == 200
assert client.post("/v1/chat/completions",
json=chat(model="model-two")).status_code == 200
assert len(first.calls) == 1 and len(second.calls) == 1
assert first.calls[0][0] == [1]
assert second.chat_messages
def test_unknown_model_is_404():
first, _ = two_engines()
with TestClient(create_app({"model-one": first})) as client:
response = client.post("/v1/completions", json=completion(model="model-two"))
assert response.status_code == 404
assert response.json()["error"]["code"] == "model_not_found"
def test_context_limits_apply_per_engine():
_, second = two_engines()
with TestClient(create_app({"model-one": FakeEngine(), "model-two": second})) as client:
response = client.post("/v1/completions",
json=completion(model="model-two", prompt=[1] * 257))
assert response.status_code == 400
assert client.post("/v1/completions",
json=completion(model="model-two", max_tokens=9999)).status_code == 200
assert second.calls[0][1]["max_new_tokens"] == 9999
def test_single_fifo_queue_is_shared_across_models():
first, second = two_engines()
first.gate = threading.Event()
app = create_app({"model-one": first, "model-two": second}, queue_size=2)
with TestClient(app) as client, ThreadPoolExecutor(max_workers=2) as pool:
blocked = pool.submit(client.post, "/v1/completions",
json=completion(model="model-one", prompt=[1], stream=True))
try:
assert first.started.wait(2)
queued = pool.submit(client.post, "/v1/completions",
json=completion(model="model-two", prompt=[1]))
wait_until(lambda: len(app.state.worker.jobs) == 2)
assert not second.calls
finally:
first.gate.set()
assert blocked.result(5).status_code == 200
assert queued.result(5).status_code == 200
assert second.calls
def test_empty_model_map_is_rejected():
with pytest.raises(ValueError, match="At least one model"):
create_app({})
|