diff --git a/src/ometer/api.py b/src/ometer/api.py index 0cd13e47..8997ffcc 100644 --- a/src/ometer/api.py +++ b/src/ometer/api.py @@ -1,5 +1,7 @@ from __future__ import annotations +import asyncio +import functools import json import time from dataclasses import dataclass, field @@ -32,6 +34,45 @@ def _key(model: dict[str, Any]) -> datetime: return sorted(models, key=_key, reverse=True) +class StreamEndedError(RuntimeError): + pass + + +def retry( + status_codes: list[int] | None = None, + max_attempts: int = 4, + return_error_dict: bool = False, +): + codes = status_codes if status_codes is not None else [429, 500, 502, 503, 504] + + def decorator(func): + @functools.wraps(func) + async def wrapper(*args, **kwargs): + for attempt in range(1, max_attempts + 1): + try: + return await func(*args, **kwargs) + except Exception as e: + should_retry = False + if isinstance(e, httpx.HTTPStatusError): + if e.response.status_code in codes: + should_retry = True + elif isinstance(e, (httpx.RequestError, StreamEndedError)): + should_retry = True + + if should_retry and attempt < max_attempts: + await asyncio.sleep(2 ** (attempt - 1)) + continue + + if return_error_dict: + return {"ttft": None, "tps": None, "error": str(e)} + raise + + return wrapper + + return decorator + + +@retry() async def fetch_tags(client: httpx.AsyncClient, base_url: str) -> list[dict[str, Any]]: resp = await client.get(f"{base_url}/api/tags") resp.raise_for_status() @@ -39,6 +80,7 @@ async def fetch_tags(client: httpx.AsyncClient, base_url: str) -> list[dict[str, return data.get("models", []) +@retry() async def fetch_model_show( client: httpx.AsyncClient, base_url: str, model_name: str ) -> dict[str, Any]: @@ -55,6 +97,7 @@ def is_embedding_model(show_data: dict[str, Any]) -> bool: return "embedding" in show_data.get("capabilities", []) +@retry(return_error_dict=True) async def benchmark_chat_single_run( client: httpx.AsyncClient, base_url: str, @@ -72,69 +115,62 @@ async def benchmark_chat_single_run( if num_predict is not None: payload["options"] = {"num_predict": num_predict} + is_thinking = bool(show_data) and "thinking" in show_data.get("capabilities", []) + start = time.perf_counter() - first_token_time: float = -1.0 + first_token_time: float | None = None eval_count = 0 eval_duration = 0 total_duration = 0 - error: str | None = None seen_done = False - is_thinking = bool(show_data) and "thinking" in show_data.get("capabilities", []) - - try: - async with client.stream( - "POST", - f"{base_url}/api/chat", - json=payload, - headers=headers, - timeout=300.0, - ) as response: - response.raise_for_status() - async for line in response.aiter_lines(): - if not line.strip(): - continue - try: - chunk = json.loads(line) - except json.JSONDecodeError: - continue - - if chunk.get("error"): - error = chunk["error"] - break - - msg = chunk.get("message") or {} - if is_thinking: - if first_token_time < 0 and ( - msg.get("thinking") or msg.get("content") - ): - first_token_time = time.perf_counter() - start - else: - if first_token_time < 0 and msg.get("content"): - first_token_time = time.perf_counter() - start - - if chunk.get("done"): - eval_count = chunk.get("eval_count", 0) - eval_duration = chunk.get("eval_duration", 0) - total_duration = chunk.get("total_duration", 0) - seen_done = True - break - except Exception as e: - error = str(e) - - if not seen_done and not error: - error = "Stream ended without completion" - - if first_token_time >= 0: - ttft = first_token_time - else: - ttft = None + async with client.stream( + "POST", + f"{base_url}/api/chat", + json=payload, + headers=headers, + timeout=300.0, + ) as response: + response.raise_for_status() + async for line in response.aiter_lines(): + if not line.strip(): + continue + try: + chunk = json.loads(line) + except json.JSONDecodeError: + continue + + if chunk.get("error"): + raise RuntimeError(chunk["error"]) + + msg = chunk.get("message") or {} + if is_thinking: + if first_token_time is None and ( + msg.get("thinking") or msg.get("content") + ): + first_token_time = time.perf_counter() - start + else: + if first_token_time is None and msg.get("content"): + first_token_time = time.perf_counter() - start + + if chunk.get("done"): + eval_count = chunk.get("eval_count", 0) + eval_duration = chunk.get("eval_duration", 0) + total_duration = chunk.get("total_duration", 0) + seen_done = True + break + + if not seen_done: + raise StreamEndedError("Stream ended without completion") + + ttft = first_token_time duration = eval_duration or total_duration tps = eval_count / (duration / 1e9) if duration else None - return {"ttft": ttft, "tps": tps, "error": error} + return {"ttft": ttft, "tps": tps, "error": None} +@retry(return_error_dict=True) async def benchmark_embed_single_run( client: httpx.AsyncClient, base_url: str, @@ -148,28 +184,21 @@ async def benchmark_embed_single_run( } start = time.perf_counter() - prompt_eval_count = 0 - total_duration = 0 - error: str | None = None - - try: - resp = await client.post( - f"{base_url}/api/embed", - json=payload, - headers=headers, - timeout=300.0, - ) - resp.raise_for_status() - data = resp.json() - prompt_eval_count = data.get("prompt_eval_count", 0) - total_duration = data.get("total_duration", 0) - except Exception as e: - error = str(e) + resp = await client.post( + f"{base_url}/api/embed", + json=payload, + headers=headers, + timeout=300.0, + ) + resp.raise_for_status() + data = resp.json() + prompt_eval_count = data.get("prompt_eval_count", 0) + total_duration = data.get("total_duration", 0) ttft = time.perf_counter() - start tps = prompt_eval_count / (total_duration / 1e9) if total_duration else None - return {"ttft": ttft, "tps": tps, "error": error} + return {"ttft": ttft, "tps": tps, "error": None} async def benchmark_model( diff --git a/tests/test_api.py b/tests/test_api.py index 9ad5d1e4..3ab41509 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -20,6 +20,13 @@ from ometer.config import Config +@pytest.fixture(autouse=True) +def mock_sleep(monkeypatch): + async def dummy_sleep(*args, **kwargs): + pass + monkeypatch.setattr("asyncio.sleep", dummy_sleep) + + class TestSortByModified: def test_sorts_newest_first(self): models = [ @@ -87,6 +94,56 @@ async def test_fetch_tags_empty(self, httpx_mock: pytest_httpx.HTTPXMock): result = await fetch_tags(client, "http://localhost:11434") assert result == [] + @pytest.mark.asyncio + async def test_fetch_tags_retry_success(self, httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_response( + url="http://localhost:11434/api/tags", + status_code=500, + ) + httpx_mock.add_response( + url="http://localhost:11434/api/tags", + json={"models": [{"name": "llama3"}]}, + ) + async with httpx.AsyncClient() as client: + result = await fetch_tags(client, "http://localhost:11434") + assert result == [{"name": "llama3"}] + + @pytest.mark.asyncio + async def test_fetch_tags_exhausted_fail(self, httpx_mock: pytest_httpx.HTTPXMock): + for _ in range(4): + httpx_mock.add_response( + url="http://localhost:11434/api/tags", + status_code=500, + ) + async with httpx.AsyncClient() as client: + with pytest.raises(httpx.HTTPStatusError): + await fetch_tags(client, "http://localhost:11434") + + @pytest.mark.asyncio + async def test_fetch_tags_network_retry_success(self, httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_exception( + httpx.ConnectError("Connection failed"), + url="http://localhost:11434/api/tags", + ) + httpx_mock.add_response( + url="http://localhost:11434/api/tags", + json={"models": [{"name": "llama3"}]}, + ) + async with httpx.AsyncClient() as client: + result = await fetch_tags(client, "http://localhost:11434") + assert result == [{"name": "llama3"}] + + @pytest.mark.asyncio + async def test_fetch_tags_network_exhausted_fail(self, httpx_mock: pytest_httpx.HTTPXMock): + for _ in range(4): + httpx_mock.add_exception( + httpx.ConnectError("Connection failed"), + url="http://localhost:11434/api/tags", + ) + async with httpx.AsyncClient() as client: + with pytest.raises(httpx.ConnectError): + await fetch_tags(client, "http://localhost:11434") + class TestFetchModelShow: @pytest.mark.asyncio @@ -103,6 +160,56 @@ async def test_fetch_model_show(self, httpx_mock: pytest_httpx.HTTPXMock): result = await fetch_model_show(client, "http://localhost:11434", "llama3") assert result == show_data + @pytest.mark.asyncio + async def test_fetch_model_show_retry_success(self, httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_response( + url="http://localhost:11434/api/show", + status_code=500, + ) + httpx_mock.add_response( + url="http://localhost:11434/api/show", + json={"details": {"parameter_size": "7B"}}, + ) + async with httpx.AsyncClient() as client: + result = await fetch_model_show(client, "http://localhost:11434", "llama3") + assert result == {"details": {"parameter_size": "7B"}} + + @pytest.mark.asyncio + async def test_fetch_model_show_exhausted_fail(self, httpx_mock: pytest_httpx.HTTPXMock): + for _ in range(4): + httpx_mock.add_response( + url="http://localhost:11434/api/show", + status_code=500, + ) + async with httpx.AsyncClient() as client: + with pytest.raises(httpx.HTTPStatusError): + await fetch_model_show(client, "http://localhost:11434", "llama3") + + @pytest.mark.asyncio + async def test_fetch_model_show_network_retry_success(self, httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_exception( + httpx.ConnectError("Connection failed"), + url="http://localhost:11434/api/show", + ) + httpx_mock.add_response( + url="http://localhost:11434/api/show", + json={"details": {"parameter_size": "7B"}}, + ) + async with httpx.AsyncClient() as client: + result = await fetch_model_show(client, "http://localhost:11434", "llama3") + assert result == {"details": {"parameter_size": "7B"}} + + @pytest.mark.asyncio + async def test_fetch_model_show_network_exhausted_fail(self, httpx_mock: pytest_httpx.HTTPXMock): + for _ in range(4): + httpx_mock.add_exception( + httpx.ConnectError("Connection failed"), + url="http://localhost:11434/api/show", + ) + async with httpx.AsyncClient() as client: + with pytest.raises(httpx.ConnectError): + await fetch_model_show(client, "http://localhost:11434", "llama3") + class TestBenchmarkChatSingleRun: @pytest.mark.asyncio @@ -187,15 +294,98 @@ async def test_stream_without_done(self, httpx_mock: pytest_httpx.HTTPXMock): json.dumps({"message": {"content": "Hello"}, "done": False}), ] body = "\n".join(chunks) + for _ in range(4): + httpx_mock.add_response( + url="http://localhost:11434/api/chat", + content=body.encode(), + ) + async with httpx.AsyncClient() as client: + result = await benchmark_chat_single_run( + client, "http://localhost:11434", "llama3", "hi" + ) + assert result["error"] == "Stream ended without completion" + + @pytest.mark.asyncio + async def test_retry_success(self, httpx_mock: pytest_httpx.HTTPXMock): httpx_mock.add_response( url="http://localhost:11434/api/chat", - content=body.encode(), + status_code=500, + ) + chunks = [ + json.dumps( + { + "message": {"content": "resolved"}, + "done": True, + "eval_count": 5, + "eval_duration": 500_000_000, + "total_duration": 1_000_000_000, + } + ), + ] + httpx_mock.add_response( + url="http://localhost:11434/api/chat", + content="\n".join(chunks).encode(), ) async with httpx.AsyncClient() as client: result = await benchmark_chat_single_run( client, "http://localhost:11434", "llama3", "hi" ) - assert result["error"] == "Stream ended without completion" + assert result["error"] is None + assert result["tps"] == 10.0 + + @pytest.mark.asyncio + async def test_network_retry_success(self, httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_exception( + httpx.ConnectError("Connection failed"), + url="http://localhost:11434/api/chat", + ) + chunks = [ + json.dumps( + { + "message": {"content": "resolved"}, + "done": True, + "eval_count": 5, + "eval_duration": 500_000_000, + "total_duration": 1_000_000_000, + } + ), + ] + httpx_mock.add_response( + url="http://localhost:11434/api/chat", + content="\n".join(chunks).encode(), + ) + async with httpx.AsyncClient() as client: + result = await benchmark_chat_single_run( + client, "http://localhost:11434", "llama3", "hi" + ) + assert result["error"] is None + assert result["tps"] == 10.0 + + @pytest.mark.asyncio + async def test_network_exhausted_fail(self, httpx_mock: pytest_httpx.HTTPXMock): + for _ in range(4): + httpx_mock.add_exception( + httpx.ConnectError("Connection failed"), + url="http://localhost:11434/api/chat", + ) + async with httpx.AsyncClient() as client: + result = await benchmark_chat_single_run( + client, "http://localhost:11434", "llama3", "hi" + ) + assert result["error"] is not None + assert "Connection failed" in result["error"] + + @pytest.mark.asyncio + async def test_unexpected_exception(self, httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_exception( + ValueError("Unexpected exception"), + url="http://localhost:11434/api/chat", + ) + async with httpx.AsyncClient() as client: + result = await benchmark_chat_single_run( + client, "http://localhost:11434", "llama3", "hi" + ) + assert result["error"] == "Unexpected exception" @pytest.mark.asyncio async def test_thinking_model_first_token(self, httpx_mock: pytest_httpx.HTTPXMock): @@ -295,15 +485,76 @@ async def test_embedding_successful_run(self, httpx_mock: pytest_httpx.HTTPXMock @pytest.mark.asyncio async def test_http_error(self, httpx_mock: pytest_httpx.HTTPXMock): + for _ in range(4): + httpx_mock.add_response( + url="http://localhost:11434/api/embed", + status_code=500, + ) + async with httpx.AsyncClient() as client: + result = await benchmark_embed_single_run( + client, "http://localhost:11434", "nomic", "hello" + ) + assert result["error"] is not None + + @pytest.mark.asyncio + async def test_retry_success(self, httpx_mock: pytest_httpx.HTTPXMock): httpx_mock.add_response( url="http://localhost:11434/api/embed", status_code=500, ) + httpx_mock.add_response( + url="http://localhost:11434/api/embed", + json={"prompt_eval_count": 8, "total_duration": 500_000_000}, + ) + async with httpx.AsyncClient() as client: + result = await benchmark_embed_single_run( + client, "http://localhost:11434", "nomic", "hello" + ) + assert result["error"] is None + assert result["tps"] == 16.0 + + @pytest.mark.asyncio + async def test_network_retry_success(self, httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_exception( + httpx.ConnectError("Connection failed"), + url="http://localhost:11434/api/embed", + ) + httpx_mock.add_response( + url="http://localhost:11434/api/embed", + json={"prompt_eval_count": 8, "total_duration": 500_000_000}, + ) + async with httpx.AsyncClient() as client: + result = await benchmark_embed_single_run( + client, "http://localhost:11434", "nomic", "hello" + ) + assert result["error"] is None + assert result["tps"] == 16.0 + + @pytest.mark.asyncio + async def test_network_exhausted_fail(self, httpx_mock: pytest_httpx.HTTPXMock): + for _ in range(4): + httpx_mock.add_exception( + httpx.ConnectError("Connection failed"), + url="http://localhost:11434/api/embed", + ) async with httpx.AsyncClient() as client: result = await benchmark_embed_single_run( client, "http://localhost:11434", "nomic", "hello" ) assert result["error"] is not None + assert "Connection failed" in result["error"] + + @pytest.mark.asyncio + async def test_unexpected_exception(self, httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_exception( + ValueError("Unexpected exception"), + url="http://localhost:11434/api/embed", + ) + async with httpx.AsyncClient() as client: + result = await benchmark_embed_single_run( + client, "http://localhost:11434", "nomic", "hello" + ) + assert result["error"] == "Unexpected exception" class TestBenchmarkModel: @@ -356,18 +607,11 @@ async def test_embedding_model_local(self, httpx_mock: pytest_httpx.HTTPXMock): @pytest.mark.asyncio async def test_all_runs_fail(self, httpx_mock: pytest_httpx.HTTPXMock): - httpx_mock.add_response( - url="http://localhost:11434/api/chat", - status_code=500, - ) - httpx_mock.add_response( - url="http://localhost:11434/api/chat", - status_code=500, - ) - httpx_mock.add_response( - url="http://localhost:11434/api/chat", - status_code=500, - ) + for _ in range(12): + httpx_mock.add_response( + url="http://localhost:11434/api/chat", + status_code=500, + ) cfg = Config("http://localhost:11434", "https://ollama.com", "", 3, 1) show_data = {"capabilities": ["completion"]} async with httpx.AsyncClient() as client: @@ -394,10 +638,11 @@ async def test_partial_runs_fail(self, httpx_mock: pytest_httpx.HTTPXMock): url="http://localhost:11434/api/chat", content="\n".join(chunks).encode(), ) - httpx_mock.add_response( - url="http://localhost:11434/api/chat", - status_code=500, - ) + for _ in range(4): + httpx_mock.add_response( + url="http://localhost:11434/api/chat", + status_code=500, + ) cfg = Config("http://localhost:11434", "https://ollama.com", "", 3, 1) show_data = {"capabilities": ["completion"]} async with httpx.AsyncClient() as client: