Spaces:
Sleeping
Sleeping
| """HTTP adapter tests: no native model, backend, or model download required.""" | |
| import asyncio | |
| from concurrent.futures import ThreadPoolExecutor | |
| import json | |
| import threading | |
| import time | |
| from types import SimpleNamespace | |
| from fastapi.testclient import TestClient | |
| import pytest | |
| from cism.engine import Engine, GenerationChunk as Chunk | |
| from cism.server import _Chunk, _CompletionRequest, _Job, _Worker, create_app | |
| class FakeEngine: | |
| model_id = "test-model" | |
| context_length = 512 | |
| def __init__(self): | |
| self.calls = [] | |
| self.chat_messages = [] | |
| self.thread_ids = [] | |
| self.started = threading.Event() | |
| self.cancelled = threading.Event() | |
| self.closed = threading.Event() | |
| self.gate = None | |
| self.error = None | |
| self.fail_after_first = False | |
| self.active = 0 | |
| self.max_active = 0 | |
| self.tokenizer = self | |
| self._generation_lock = threading.Lock() | |
| def encode(self, prompt): | |
| self.thread_ids.append(threading.get_ident()) | |
| if prompt == "invalid": | |
| raise ValueError("PRIVATE PROMPT in tokenizer error") | |
| if isinstance(prompt, list): | |
| if any(token > 1000 for token in prompt): | |
| raise ValueError("PRIVATE TOKEN") | |
| return prompt.copy() | |
| return list(range(len(prompt))) | |
| def format_chat(self, messages): | |
| self.thread_ids.append(threading.get_ident()) | |
| self.chat_messages.append(messages) | |
| return [100, 101, 102] | |
| def stream(self, prompt, **kwargs): | |
| self.thread_ids.append(threading.get_ident()) | |
| self.calls.append((prompt, kwargs)) | |
| self.active += 1 | |
| self.max_active = max(self.max_active, self.active) | |
| self.started.set() | |
| try: | |
| while self.gate is not None and not self.gate.wait(0.01): | |
| if kwargs["cancel"].is_set(): | |
| self.cancelled.set() | |
| return | |
| if self.error: | |
| raise self.error | |
| yield Chunk("hello", [7], None, len(prompt), 1) | |
| if self.fail_after_first: | |
| raise RuntimeError("PRIVATE PROMPT in generation error") | |
| yield Chunk(" world", [8], "stop", len(prompt), 2) | |
| finally: | |
| if kwargs["cancel"].is_set(): | |
| self.cancelled.set() | |
| self.active -= 1 | |
| self.closed.set() | |
| def no_environment_key(monkeypatch): | |
| monkeypatch.delenv("CISM_API_KEY", raising=False) | |
| def completion(**overrides): | |
| return {"model": "test-model", "prompt": "hi", **overrides} | |
| def chat(**overrides): | |
| return {"model": "test-model", "messages": [{"role": "user", "content": "hi"}], **overrides} | |
| def wait_until(predicate, timeout=2): | |
| deadline = time.monotonic() + timeout | |
| while not predicate(): | |
| if time.monotonic() >= deadline: | |
| pytest.fail("Timed out waiting for worker state") | |
| time.sleep(0.005) | |
| def events(response): | |
| assert response.headers["content-type"].startswith("text/event-stream") | |
| frames = [line[6:] for line in response.text.splitlines() if line.startswith("data: ")] | |
| assert frames[-1] == "[DONE]" | |
| return [json.loads(frame) for frame in frames[:-1]] | |
| def assert_error(response, status): | |
| assert response.status_code == status | |
| error = response.json()["error"] | |
| assert set(error) == {"message", "type", "param", "code"} | |
| assert "PRIVATE" not in response.text | |
| assert "traceback" not in response.text.lower() | |
| return error | |
| def test_models_probes_and_lifecycle(): | |
| app = create_app(FakeEngine()) | |
| with TestClient(app) as client: | |
| assert client.get("/health").json() == {"status": "ok"} | |
| assert client.get("/ready").json() == {"status": "ready"} | |
| model = client.get("/v1/models").json()["data"][0] | |
| assert model == {"id": "test-model", "object": "model", "created": 0, "owned_by": "cism", | |
| "context_length": 512} | |
| assert_error(client.get("/missing"), 404) | |
| worker = app.state.worker | |
| assert not worker.thread.is_alive() | |
| assert not worker.jobs | |
| with TestClient(app) as client: | |
| assert client.get("/ready").status_code == 200 | |
| assert app.state.worker is not worker | |
| def test_not_ready_without_lifespan(): | |
| client = TestClient(create_app(FakeEngine())) | |
| assert client.get("/health").status_code == 200 | |
| assert_error(client.get("/ready"), 503) | |
| assert_error(client.post("/v1/completions", json=completion()), 503) | |
| client.close() | |
| def test_auth_and_environment_snapshot(monkeypatch): | |
| monkeypatch.setenv("CISM_API_KEY", "original-secret") | |
| app = create_app(FakeEngine()) | |
| monkeypatch.setenv("CISM_API_KEY", "changed-secret") | |
| with TestClient(app) as client: | |
| for authorization in (None, "Basic original-secret", "Bearer changed-secret", "Bearer"): | |
| headers = {"Authorization": authorization} if authorization else {} | |
| response = client.get("/v1/models", headers=headers) | |
| assert_error(response, 401) | |
| assert response.headers["www-authenticate"] == "Bearer" | |
| assert client.get("/v1/models", headers={"Authorization": "bearer original-secret"}).status_code == 200 | |
| assert client.get("/health").status_code == 200 | |
| assert client.get("/ready").status_code == 200 | |
| with TestClient(create_app(FakeEngine(), api_key="")) as client: | |
| assert client.get("/v1/models").status_code == 200 | |
| with TestClient(create_app(FakeEngine(), api_key="explicit")) as client: | |
| assert_error(client.post("/v1/completions", json=completion()), 401) | |
| assert client.get("/v1/models", headers={"Authorization": "Bearer explicit"}).status_code == 200 | |
| def test_invalid_capacity(size): | |
| with pytest.raises(ValueError): | |
| create_app(FakeEngine(), queue_size=size) | |
| def test_completion_defaults_usage_and_no_state_leak(): | |
| engine = FakeEngine() | |
| with TestClient(create_app(engine)) as client: | |
| first = client.post("/v1/completions", json=completion()) | |
| second = client.post("/v1/completions", json=completion(prompt=[42])) | |
| assert first.status_code == second.status_code == 200 | |
| result = first.json() | |
| assert result["object"] == "text_completion" | |
| assert result["choices"] == [{"index": 0, "text": "hello world", "finish_reason": "stop", "logprobs": None}] | |
| assert result["usage"] == {"prompt_tokens": 2, "completion_tokens": 2, "total_tokens": 4, | |
| "spec_proposed": 0, "spec_accepted": 0, "spec_accept_rate": None} | |
| assert second.json()["usage"]["prompt_tokens"] == 1 | |
| assert result["id"] != second.json()["id"] | |
| assert engine.calls[0][0] == [0, 1] | |
| assert engine.calls[1][0] == [42] | |
| kwargs = engine.calls[0][1] | |
| assert {key: value for key, value in kwargs.items() if key != "cancel"} == { | |
| "max_new_tokens": 128, "temperature": 0.0, "top_p": 1.0, | |
| "top_k": 0, "seed": 0, "stop": None, "chunk_size": 16, "spec_k": 0, | |
| } | |
| assert engine.calls[0][1]["cancel"] is not engine.calls[1][1]["cancel"] | |
| assert len(set(engine.thread_ids)) == 1 | |
| assert engine.thread_ids[0] != threading.get_ident() | |
| assert engine.max_active == 1 | |
| def test_chat_format_and_sampling_options(): | |
| engine = FakeEngine() | |
| messages = [{"role": role, "content": role} for role in ("system", "developer", "user", "assistant")] | |
| with TestClient(create_app(engine)) as client: | |
| response = client.post("/v1/chat/completions", json=chat( | |
| messages=messages, max_completion_tokens=3, temperature=0.7, top_p=0.8, | |
| top_k=9, seed=17, stop="END", response_format={"type": "text"}, | |
| )) | |
| assert response.status_code == 200 | |
| assert response.json()["object"] == "chat.completion" | |
| assert response.json()["choices"][0]["message"] == {"role": "assistant", "content": "hello world"} | |
| assert engine.chat_messages == [messages] | |
| assert engine.calls[0][0] == [100, 101, 102] | |
| assert {key: value for key, value in engine.calls[0][1].items() if key != "cancel"} == { | |
| "max_new_tokens": 3, "temperature": 0.7, "top_p": 0.8, | |
| "top_k": 9, "seed": 17, "stop": ["END"], "chunk_size": 16, "spec_k": 0, | |
| } | |
| def test_stream_shape_and_usage(is_chat, include_usage): | |
| path = "/v1/chat/completions" if is_chat else "/v1/completions" | |
| body = chat if is_chat else completion | |
| with TestClient(create_app(FakeEngine())) as client: | |
| response = client.post(path, json=body(stream=True, stream_options={"include_usage": include_usage})) | |
| frames = events(response) | |
| assert len({frame["id"] for frame in frames}) == 1 | |
| assert len({frame["created"] for frame in frames}) == 1 | |
| assert all(frame["object"] == ("chat.completion.chunk" if is_chat else "text_completion") for frame in frames) | |
| if include_usage: | |
| assert frames[-1]["choices"] == [] | |
| assert frames[-1]["usage"]["completion_tokens"] == 2 | |
| assert all(frame["usage"] is None for frame in frames[:-1]) | |
| frames = frames[:-1] | |
| else: | |
| assert all("usage" not in frame for frame in frames) | |
| assert frames[-1]["choices"][0]["finish_reason"] == "stop" | |
| assert all(frame["choices"][0]["finish_reason"] is None for frame in frames[:-1]) | |
| if is_chat: | |
| assert frames[0]["choices"][0]["delta"] == {"role": "assistant", "content": ""} | |
| assert frames[-1]["choices"][0]["delta"] == {} | |
| assert "".join(frame["choices"][0]["delta"].get("content", "") for frame in frames) == "hello world" | |
| else: | |
| assert "".join(frame["choices"][0]["text"] for frame in frames) == "hello world" | |
| def test_strict_completion_validation_before_sse(overrides): | |
| engine = FakeEngine() | |
| with TestClient(create_app(engine)) as client: | |
| response = client.post("/v1/completions", json=completion(**overrides)) | |
| assert_error(response, 400) | |
| assert response.headers["content-type"].startswith("application/json") | |
| assert not engine.calls | |
| def test_unsupported_chat_fields(overrides): | |
| engine = FakeEngine() | |
| with TestClient(create_app(engine)) as client: | |
| response = client.post("/v1/chat/completions", json=chat(stream=True, **overrides)) | |
| assert_error(response, 400) | |
| assert not engine.calls | |
| def test_malformed_missing_model_and_nonfinite_json(): | |
| with TestClient(create_app(FakeEngine())) as client: | |
| for raw in ('{', 'null', '[]', '{"prompt":"PRIVATE"}', | |
| '{"model":"test-model","prompt":"PRIVATE","temperature":NaN}', | |
| '{"model":"test-model","prompt":"PRIVATE","temperature":Infinity}'): | |
| assert_error(client.post("/v1/completions", content=raw, headers={"Content-Type": "application/json"}), 400) | |
| assert_error(client.post("/v1/completions", json=completion(model="PRIVATE")), 404) | |
| def test_context_and_token_validation_before_sse(prompt, max_tokens): | |
| engine = FakeEngine() | |
| with TestClient(create_app(engine)) as client: | |
| assert_error(client.post("/v1/completions", json=completion(prompt=prompt, max_tokens=max_tokens, stream=True)), 400) | |
| assert client.post("/v1/completions", json=completion()).status_code == 200 | |
| assert len(engine.calls) == 1 | |
| def test_over_budget_max_tokens_is_clamped_not_rejected(): | |
| engine = FakeEngine() | |
| with TestClient(create_app(engine)) as client: | |
| response = client.post("/v1/completions", json=completion(prompt="hi", max_tokens=4096)) | |
| assert response.status_code == 200 | |
| response = client.post("/v1/completions", json=completion(prompt="hi", max_tokens=4096, stream=True)) | |
| assert response.status_code == 200 | |
| # The engine, not the server, clamps the budget to the model context. | |
| assert len(engine.calls) == 2 | |
| def test_generation_errors_are_sanitized_and_worker_recovers(error_type): | |
| engine = FakeEngine() | |
| engine.error = error_type("PRIVATE PROMPT and stack") | |
| with TestClient(create_app(engine)) as client: | |
| assert_error(client.post("/v1/completions", json=completion(stream=True)), 503) | |
| engine.error = None | |
| engine.fail_after_first = True | |
| response = client.post("/v1/chat/completions", json=chat(stream=True)) | |
| frames = events(response) | |
| assert frames[-1]["error"]["code"] == "generation_failed" | |
| assert "PRIVATE" not in response.text | |
| engine.fail_after_first = False | |
| assert client.post("/v1/completions", json=completion()).status_code == 200 | |
| def test_fifo_total_admission_bound_and_responsive_event_loop(): | |
| engine = FakeEngine() | |
| engine.gate = threading.Event() | |
| app = create_app(engine, queue_size=2) | |
| with TestClient(app) as client, ThreadPoolExecutor(max_workers=2) as pool: | |
| first = pool.submit(client.post, "/v1/completions", json=completion(prompt=[1])) | |
| try: | |
| assert engine.started.wait(2) | |
| second = pool.submit(client.post, "/v1/completions", json=completion(prompt=[2])) | |
| wait_until(lambda: len(app.state.worker.jobs) == 2) | |
| assert app.state.worker.pending.qsize() == 1 | |
| started = time.monotonic() | |
| assert client.get("/health").status_code == 200 | |
| assert client.get("/ready").status_code == 200 | |
| response = client.post("/v1/completions", json=completion(prompt=[3])) | |
| assert_error(response, 429) | |
| assert response.headers["retry-after"] == "1" | |
| assert time.monotonic() - started < 1 | |
| finally: | |
| engine.gate.set() | |
| assert first.result(timeout=2).status_code == 200 | |
| assert second.result(timeout=2).status_code == 200 | |
| wait_until(lambda: not app.state.worker.jobs) | |
| assert [prompt for prompt, _ in engine.calls] == [[1], [2]] | |
| assert engine.max_active == 1 | |
| def test_bounded_output_cancel_and_shutdown(): | |
| class EndlessEngine(FakeEngine): | |
| def stream(self, prompt, **kwargs): | |
| self.started.set() | |
| try: | |
| index = 0 | |
| while not kwargs["cancel"].is_set(): | |
| index += 1 | |
| yield Chunk("x", [1], None, len(prompt), index) | |
| finally: | |
| self.cancelled.set() | |
| engine = EndlessEngine() | |
| app = create_app(engine, queue_size=2) | |
| with TestClient(app): | |
| worker = app.state.worker | |
| first = worker.submit(_CompletionRequest(**completion())) | |
| assert engine.started.wait(2) | |
| second = worker.submit(_CompletionRequest(**completion())) | |
| wait_until(lambda: first.output.full()) | |
| assert first.output.qsize() == first.output.maxsize == 8 | |
| worker.release(second) | |
| worker.release(first) | |
| assert engine.cancelled.wait(2) | |
| wait_until(lambda: not worker.jobs) | |
| assert not worker.thread.is_alive() | |
| def test_shutdown_cancels_active_and_queued_jobs(): | |
| engine = FakeEngine() | |
| engine.gate = threading.Event() | |
| app = create_app(engine, queue_size=2) | |
| with TestClient(app): | |
| worker = app.state.worker | |
| first = worker.submit(_CompletionRequest(**completion())) | |
| assert engine.started.wait(2) | |
| second = worker.submit(_CompletionRequest(**completion())) | |
| assert first.cancel.is_set() and second.cancel.is_set() | |
| assert engine.cancelled.is_set() | |
| assert not worker.thread.is_alive() | |
| assert not worker.jobs | |
| assert len(engine.calls) == 1 | |
| def test_shutdown_does_not_wait_forever_for_uncooperative_engine(): | |
| class StuckEngine(FakeEngine): | |
| def stream(self, prompt, **kwargs): | |
| self.started.set() | |
| self.gate.wait() | |
| yield Chunk("x", [1], "stop", len(prompt), 1) | |
| engine = StuckEngine() | |
| engine.gate = threading.Event() | |
| app = create_app(engine) | |
| try: | |
| with TestClient(app): | |
| worker = app.state.worker | |
| job = worker.submit(_CompletionRequest(**completion())) | |
| assert engine.started.wait(2) | |
| worker.release(job) | |
| start = time.monotonic() | |
| assert time.monotonic() - start < 2 | |
| assert worker.thread.is_alive() | |
| assert not worker.ready | |
| with pytest.raises(RuntimeError, match="previous inference worker"): | |
| with TestClient(app): | |
| pass | |
| finally: | |
| engine.gate.set() | |
| worker.thread.join(timeout=2) | |
| assert not worker.thread.is_alive() | |
| def test_actual_asgi_disconnect_cancels_before_first_chunk(stream): | |
| # TestClient buffers streaming responses, so drive ASGI directly for a real | |
| # disconnect while prefill/generation has not produced any response yet. | |
| engine = FakeEngine() | |
| engine.gate = threading.Event() | |
| app = create_app(engine) | |
| async def run(): | |
| async with app.router.lifespan_context(app): | |
| messages = asyncio.Queue() | |
| await messages.put({"type": "http.request", "body": json.dumps(completion(stream=stream)).encode(), "more_body": False}) | |
| sent = [] | |
| async def send(message): | |
| sent.append(message) | |
| task = asyncio.create_task(app({ | |
| "type": "http", "asgi": {"version": "3.0", "spec_version": "2.3"}, | |
| "http_version": "1.1", "method": "POST", "scheme": "http", | |
| "path": "/v1/completions", "raw_path": b"/v1/completions", | |
| "query_string": b"", "root_path": "", "headers": [(b"content-type", b"application/json")], | |
| "client": ("test", 123), "server": ("test", 80), | |
| }, messages.get, send)) | |
| try: | |
| deadline = asyncio.get_running_loop().time() + 2 | |
| while not engine.started.is_set(): | |
| assert asyncio.get_running_loop().time() < deadline | |
| await asyncio.sleep(0.01) | |
| await messages.put({"type": "http.disconnect"}) | |
| await asyncio.wait_for(task, timeout=2) | |
| finally: | |
| task.cancel() | |
| assert engine.cancelled.is_set() | |
| assert not app.state.worker.thread.is_alive() | |
| assert not app.state.worker.jobs | |
| asyncio.run(run()) | |
| def test_stream_disconnect_releases_job_after_headers(disconnect): | |
| class StreamingEngine(FakeEngine): | |
| def stream(self, prompt, **kwargs): | |
| self.started.set() | |
| try: | |
| yield Chunk("first", [1], None, len(prompt), 1) | |
| while not kwargs["cancel"].wait(0.01): | |
| pass | |
| finally: | |
| self.cancelled.set() | |
| engine = StreamingEngine() | |
| app = create_app(engine, queue_size=1) | |
| async def run(): | |
| async with app.router.lifespan_context(app): | |
| messages = asyncio.Queue() | |
| await messages.put({"type": "http.request", "body": json.dumps(chat(stream=True)).encode(), "more_body": False}) | |
| first_sent = asyncio.Event() | |
| async def send(message): | |
| if message["type"] == "http.response.body": | |
| first_sent.set() | |
| if disconnect == "send_error": | |
| raise OSError("Connection closed") | |
| if disconnect in ("event", "event_2_4"): | |
| await messages.put({"type": "http.disconnect"}) | |
| else: | |
| await asyncio.Event().wait() | |
| task = asyncio.create_task(app({ | |
| "type": "http", "asgi": {"version": "3.0", "spec_version": "2.4" if disconnect in ("send_error", "event_2_4") else "2.3"}, | |
| "http_version": "1.1", "method": "POST", "scheme": "http", | |
| "path": "/v1/chat/completions", "raw_path": b"/v1/chat/completions", | |
| "query_string": b"", "root_path": "", "headers": [(b"content-type", b"application/json")], | |
| "client": ("test", 123), "server": ("test", 80), | |
| }, messages.get, send)) | |
| await asyncio.wait_for(first_sent.wait(), 2) | |
| if disconnect == "task_cancel": | |
| task.cancel() | |
| with pytest.raises(asyncio.CancelledError): | |
| await task | |
| elif disconnect == "send_error": | |
| from starlette.requests import ClientDisconnect | |
| with pytest.raises(ClientDisconnect): | |
| await asyncio.wait_for(task, 2) | |
| else: | |
| await asyncio.wait_for(task, 2) | |
| deadline = asyncio.get_running_loop().time() + 2 | |
| while app.state.worker.jobs: | |
| assert asyncio.get_running_loop().time() < deadline | |
| await asyncio.sleep(0.01) | |
| assert engine.cancelled.is_set() | |
| asyncio.run(run()) | |
| def test_slow_tokenization_runs_off_event_loop(): | |
| class SlowTokenizer(FakeEngine): | |
| def encode(self, prompt): | |
| self.started.set() | |
| self.gate.wait(2) | |
| return super().encode(prompt) | |
| engine = SlowTokenizer() | |
| engine.gate = threading.Event() | |
| with TestClient(create_app(engine, queue_size=1)) as client, ThreadPoolExecutor(max_workers=1) as pool: | |
| result = pool.submit(client.post, "/v1/completions", json=completion()) | |
| try: | |
| assert engine.started.wait(2) | |
| start = time.monotonic() | |
| assert client.get("/health").status_code == 200 | |
| assert_error(client.post("/v1/completions", json=completion()), 429) | |
| assert time.monotonic() - start < 1 | |
| finally: | |
| engine.gate.set() | |
| assert result.result(timeout=2).status_code == 200 | |
| def test_shutdown_releases_waiting_http_requests(): | |
| engine = FakeEngine() | |
| engine.gate = threading.Event() | |
| app = create_app(engine, queue_size=2) | |
| with TestClient(app) as client, ThreadPoolExecutor(max_workers=2) as pool: | |
| first = pool.submit(client.post, "/v1/completions", json=completion()) | |
| try: | |
| assert engine.started.wait(2) | |
| second = pool.submit(client.post, "/v1/chat/completions", json=chat(stream=True)) | |
| wait_until(lambda: len(app.state.worker.jobs) == 2) | |
| app.state.worker.stop() | |
| assert_error(first.result(timeout=2), 503) | |
| assert_error(second.result(timeout=2), 503) | |
| assert_error(client.get("/ready"), 503) | |
| assert_error(client.post("/v1/completions", json=completion()), 503) | |
| finally: | |
| engine.gate.set() | |
| assert not app.state.worker.jobs | |
| assert not app.state.worker.thread.is_alive() | |
| def test_context_exact_boundary_and_chat_overflow(): | |
| engine = FakeEngine() | |
| with TestClient(create_app(engine)) as client: | |
| assert client.post("/v1/completions", json=completion(prompt=[1] * 510, max_tokens=2)).status_code == 200 | |
| # Chat tokens [100, 101, 102] clamp over-budget requests instead of failing. | |
| assert client.post("/v1/chat/completions", json=chat(max_tokens=510, stream=True)).status_code == 200 | |
| class TinyContext(FakeEngine): | |
| context_length = 2 | |
| with TestClient(create_app(TinyContext())) as client: | |
| assert_error(client.post("/v1/chat/completions", json=chat(stream=True)), 400) | |
| assert len(engine.calls) == 2 | |
| def test_missing_terminal_chunk_is_an_error(): | |
| class EmptyEngine(FakeEngine): | |
| def stream(self, prompt, **kwargs): | |
| return iter(()) | |
| with TestClient(create_app(EmptyEngine())) as client: | |
| assert_error(client.post("/v1/completions", json=completion()), 503) | |
| assert_error(client.post("/v1/completions", json=completion(stream=True)), 503) | |
| def test_seed_uint64_validation_before_admission(is_chat, stream): | |
| engine = FakeEngine() | |
| path = "/v1/chat/completions" if is_chat else "/v1/completions" | |
| body = chat if is_chat else completion | |
| with TestClient(create_app(engine)) as client: | |
| for seed in (2**64, 2**128, -1, True): | |
| response = client.post(path, json=body(seed=seed, stream=stream)) | |
| assert_error(response, 400) | |
| assert response.headers["content-type"].startswith("application/json") | |
| assert not engine.thread_ids | |
| assert not engine.calls | |
| response = client.post(path, json=body(seed=2**64 - 1, stream=stream)) | |
| assert response.status_code == 200 | |
| assert engine.calls[0][1]["seed"] == 2**64 - 1 | |
| def test_top_k_above_vocabulary_rejected_before_admission(is_chat, stream, monkeypatch): | |
| engine = FakeEngine() | |
| engine.config = {"vocab_size": 32} | |
| app = create_app(engine) | |
| path = "/v1/chat/completions" if is_chat else "/v1/completions" | |
| body = chat if is_chat else completion | |
| with TestClient(app) as client: | |
| def unexpected_submit(body): | |
| pytest.fail("Invalid top_k must not enter the generation queue") | |
| monkeypatch.setattr(app.state.worker, "submit", unexpected_submit) | |
| for top_k in (33, 2**64): | |
| response = client.post(path, json=body(top_k=top_k, stream=stream)) | |
| error = assert_error(response, 400) | |
| assert error["param"] == "top_k" | |
| assert response.headers["content-type"].startswith("application/json") | |
| assert not engine.thread_ids | |
| assert not engine.calls | |
| assert not app.state.worker.jobs | |
| def test_output_notifications_are_coalesced_and_buffered_reads_do_not_wait(): | |
| async def run(): | |
| worker = _Worker(FakeEngine(), 1) | |
| job = _Job(_CompletionRequest(**completion())) | |
| callbacks = [] | |
| worker.loop = SimpleNamespace(call_soon_threadsafe=callbacks.append) | |
| async def no_disconnect_poll(): | |
| pytest.fail("Buffered data must not poll disconnect or wait on a timer") | |
| request = SimpleNamespace(scope={}, is_disconnected=no_disconnect_poll) | |
| value = _Chunk("x", None, 2, 1) | |
| for _ in range(10): | |
| results = [] | |
| def produce(): | |
| results.extend(worker._put(job, value) for _ in range(8)) | |
| producer = threading.Thread(target=produce) | |
| producer.start() | |
| producer.join(timeout=2) | |
| assert not producer.is_alive() | |
| assert results == [True] * 8 | |
| assert job.output.full() | |
| # Even across empty/nonempty transitions, a callback that has not | |
| # run yet must not be duplicated by subsequent producer bursts. | |
| assert len(callbacks) == 1 | |
| for _ in range(8): | |
| assert await worker.receive(job, request) is value | |
| callbacks.pop()() | |
| assert job.output_ready.is_set() | |
| assert worker._put(job, value) | |
| assert not callbacks | |
| assert await worker.receive(job, request) is value | |
| asyncio.run(run()) | |
| def test_output_notification_wakes_idle_consumer_without_timeout(monkeypatch): | |
| async def run(): | |
| worker = _Worker(FakeEngine(), 1) | |
| worker.thread.start() | |
| job = _Job(_CompletionRequest(**completion())) | |
| request = SimpleNamespace(scope={}) | |
| waiting = threading.Event() | |
| original_wait = job.output_ready.wait | |
| async def wait(): | |
| waiting.set() | |
| await original_wait() | |
| async def no_disconnect_poll(): | |
| pytest.fail("Producer notification should wake the receiver before timeout") | |
| request.is_disconnected = no_disconnect_poll | |
| monkeypatch.setattr(job.output_ready, "wait", wait) | |
| # Disable the idle timeout entirely: successful receipt must be driven | |
| # by a notification, not by timing-dependent periodic queue checks. | |
| timeout = asyncio.timeout | |
| monkeypatch.setattr(asyncio, "timeout", lambda delay: timeout(None)) | |
| value = _Chunk("x", None, 2, 1) | |
| def produce(): | |
| if waiting.wait(2): | |
| worker._put(job, value) | |
| try: | |
| for _ in range(20): | |
| waiting.clear() | |
| producer = threading.Thread(target=produce) | |
| producer.start() | |
| assert await asyncio.wait_for(worker.receive(job, request), 2) is value | |
| producer.join(timeout=2) | |
| assert not producer.is_alive() | |
| finally: | |
| worker.stop() | |
| worker.thread.join(timeout=2) | |
| asyncio.run(run()) | |
| def test_real_engine_contract_with_python_backend(is_chat, stream): | |
| class Tokenizer: | |
| chat_template = "test" | |
| bos_token_id = 1 | |
| eos_token_id = None | |
| def encode(self, prompt, *, add_special_tokens): | |
| assert add_special_tokens is True | |
| return [1, 2] | |
| def apply_chat_template(self, messages, *, tokenize, add_generation_prompt): | |
| assert messages == [{"role": "user", "content": "hi"}] | |
| assert tokenize is True and add_generation_prompt is True | |
| return [1, 2, 3] | |
| def decode(self, tokens, *, skip_special_tokens, clean_up_tokenization_spaces): | |
| assert skip_special_tokens is True | |
| assert clean_up_tokenization_spaces is False | |
| return "".join({7: "hello", 8: " world"}[token] for token in tokens) | |
| class Session: | |
| finish_reason = "" | |
| def __init__(self): | |
| self.remaining = [7, 8] | |
| self.cancelled = False | |
| def next_tokens(self, count, spec_k=0): | |
| batch, self.remaining = self.remaining[:count], self.remaining[count:] | |
| if not self.remaining: | |
| self.finish_reason = "length" | |
| return batch | |
| def cancel(self): | |
| self.cancelled = True | |
| class Backend: | |
| def __init__(self): | |
| self.calls = [] | |
| self.sessions = [] | |
| def create_session(self, tokens, **kwargs): | |
| self.calls.append((tokens, kwargs)) | |
| session = Session() | |
| self.sessions.append(session) | |
| return session | |
| backend = Backend() | |
| engine = Engine(backend, Tokenizer(), {"max_position_embeddings": 512, "vocab_size": 32}, "test-model") | |
| app = create_app(engine) | |
| path = "/v1/chat/completions" if is_chat else "/v1/completions" | |
| body = chat if is_chat else completion | |
| with TestClient(app) as client: | |
| assert_error(client.post(path, json=body(top_k=33, stream=stream)), 400) | |
| assert not backend.calls | |
| response = client.post(path, json=body(seed=2**64 - 1, top_k=32, max_tokens=2, stream=stream, | |
| **({"stream_options": {"include_usage": True}} if stream else {}))) | |
| assert response.status_code == 200 | |
| if stream: | |
| frames = events(response) | |
| assert frames[-2]["choices"][0]["finish_reason"] == "length" | |
| usage = frames[-1]["usage"] | |
| text = "".join( | |
| frame["choices"][0]["delta"].get("content", "") if is_chat else frame["choices"][0]["text"] | |
| for frame in frames[:-1] | |
| ) | |
| else: | |
| result = response.json() | |
| usage = result["usage"] | |
| choice = result["choices"][0] | |
| text = choice["message"]["content"] if is_chat else choice["text"] | |
| assert choice["finish_reason"] == "length" | |
| assert text == "hello world" | |
| assert usage == {"prompt_tokens": 3 if is_chat else 2, "completion_tokens": 2, | |
| "total_tokens": 5 if is_chat else 4, | |
| "spec_proposed": 0, "spec_accepted": 0, "spec_accept_rate": None} | |
| wait_until(lambda: not app.state.worker.jobs) | |
| assert backend.sessions[0].cancelled | |
| assert backend.calls == [([1, 2, 3] if is_chat else [1, 2], { | |
| "max_new_tokens": 2, "temperature": 0.0, "top_p": 1.0, "top_k": 32, | |
| "seed": 2**64 - 1, "eos_token_ids": [], | |
| })] | |
| assert engine._generation_lock.acquire(blocking=False) | |
| engine._generation_lock.release() | |
| engine.tokenizer.chat_template = None | |
| assert_error(client.post("/v1/chat/completions", json=chat(stream=True)), 400) | |
| def test_official_openai_client(): | |
| openai = pytest.importorskip("openai") | |
| with TestClient(create_app(FakeEngine(), api_key="test-key")) as transport: | |
| sdk = openai.OpenAI(api_key="test-key", base_url="http://testserver/v1", http_client=transport, max_retries=0) | |
| assert sdk.models.list().data[0].id == "test-model" | |
| result = sdk.completions.create(model="test-model", prompt="hi", max_tokens=2) | |
| assert result.choices[0].text == "hello world" | |
| assert result.usage.total_tokens == 4 | |
| result = sdk.chat.completions.create(model="test-model", messages=[{"role": "user", "content": "hi"}]) | |
| assert result.choices[0].message.content == "hello world" | |
| stream = sdk.chat.completions.create(model="test-model", messages=[{"role": "user", "content": "hi"}], | |
| stream=True, stream_options={"include_usage": True}) | |
| chunks = list(stream) | |
| assert chunks[0].choices[0].delta.role == "assistant" | |
| assert chunks[-1].choices == [] | |
| assert chunks[-1].usage.completion_tokens == 2 | |
| assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "hello world" | |
| chunks = list(sdk.completions.create(model="test-model", prompt=[1, 2], stream=True)) | |
| assert all(chunk.object == "text_completion" for chunk in chunks) | |
| assert "".join(chunk.choices[0].text for chunk in chunks) == "hello world" | |
| assert chunks[-1].choices[0].finish_reason == "stop" | |
| with pytest.raises(openai.BadRequestError): | |
| sdk.chat.completions.create(model="test-model", messages=[{"role": "user", "content": "hi"}], | |
| response_format={"type": "json_object"}) | |
| def autobench_client(monkeypatch, run): | |
| import cism.autobench as module | |
| monkeypatch.setattr(module, "run_benchmark", run) | |
| return TestClient(create_app(FakeEngine())) | |
| def wait_bench(client, want): | |
| state = client.get("/v1/bench/auto").json() | |
| for _ in range(200): | |
| state = client.get("/v1/bench/auto").json() | |
| if state["status"] == want: | |
| return state | |
| time.sleep(0.02) | |
| return state | |
| def test_autobench_runs_to_completion(monkeypatch): | |
| calls = [] | |
| def run(model_id, **kwargs): | |
| calls.append((model_id, kwargs)) | |
| kwargs["progress"]("fake: decoding") | |
| return {"cism": {"tokens_per_second": 9.0}, "speedup_vs_eager": 1.5} | |
| with autobench_client(monkeypatch, run) as client: | |
| assert client.get("/v1/bench/auto").json()["status"] == "idle" | |
| response = client.post("/v1/bench/auto", json={"model": "test-model", "precision": "hybrid-int4"}) | |
| assert response.status_code == 202 | |
| assert response.json() == {"status": "running", "model": "test-model", "precision": "hybrid-int4"} | |
| state = wait_bench(client, "done") | |
| assert state["status"] == "done" | |
| assert state["error"] is None | |
| assert state["result"]["speedup_vs_eager"] == 1.5 | |
| assert calls[0][0] == "test-model" | |
| assert calls[0][1]["precision"] == "hybrid-int4" | |
| assert calls[0][1]["threads"] == 2 # request default | |
| assert client.get("/v1/bench/auto").json()["status"] == "done" | |
| def test_autobench_rejects_concurrent_and_reports_failure(monkeypatch): | |
| release = threading.Event() | |
| def run(model_id, **kwargs): | |
| assert release.wait(5) | |
| raise RuntimeError("synthetic bench failure") | |
| with autobench_client(monkeypatch, run) as client: | |
| assert client.post("/v1/bench/auto", json={"model": "test-model"}).status_code == 202 | |
| # While the first benchmark runs, a second one is refused. | |
| assert client.post("/v1/bench/auto", json={"model": "test-model"}).status_code == 409 | |
| release.set() | |
| state = wait_bench(client, "failed") | |
| assert state["status"] == "failed" | |
| assert "synthetic bench failure" in state["error"] | |
| assert state["result"] is None | |
| # After failure the slot is free again. | |
| assert client.post("/v1/bench/auto", json={"model": "test-model"}).status_code == 202 | |
| def test_autobench_validates_request(monkeypatch): | |
| with autobench_client(monkeypatch, lambda *a, **k: {}) as client: | |
| assert client.post("/v1/bench/auto", json={"model": "test-model", "precision": "int5"}).status_code == 400 | |
| assert client.post("/v1/bench/auto", json={"model": "test-model", "threads": 0}).status_code == 400 | |
| assert client.post("/v1/bench/auto", json={"model": "test-model", "max_tokens": 0}).status_code == 400 | |