Spaces:
Sleeping
Sleeping
| """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({}) | |