From f0e946e0df85755bfb33426358863f177f79e509 Mon Sep 17 00:00:00 2001 From: Lewis Cowles Date: Sun, 28 Jun 2026 20:09:26 +0100 Subject: [PATCH 1/3] feat: retry on failures --- src/ometer/api.py | 250 ++++++++++++++++++++++++++--------------- tests/test_api.py | 281 +++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 421 insertions(+), 110 deletions(-) diff --git a/src/ometer/api.py b/src/ometer/api.py index 0cd13e47..4930ae16 100644 --- a/src/ometer/api.py +++ b/src/ometer/api.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import json import time from dataclasses import dataclass, field @@ -33,22 +34,54 @@ def _key(model: dict[str, Any]) -> datetime: 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() - data = resp.json() - return data.get("models", []) + max_attempts = 4 + for attempt in range(1, max_attempts + 1): + should_retry = False + try: + resp = await client.get(f"{base_url}/api/tags") + resp.raise_for_status() + data = resp.json() + return data.get("models", []) + except httpx.HTTPStatusError as e: + if e.response.status_code in (429, 500, 502, 503, 504) and attempt < max_attempts: + should_retry = True + else: + raise + except httpx.RequestError as e: + if attempt < max_attempts: + should_retry = True + else: + raise + if should_retry: + await asyncio.sleep(2 ** (attempt - 1)) async def fetch_model_show( client: httpx.AsyncClient, base_url: str, model_name: str ) -> dict[str, Any]: - resp = await client.post( - f"{base_url}/api/show", - json={"model": model_name}, - timeout=60.0, - ) - resp.raise_for_status() - return resp.json() + max_attempts = 4 + for attempt in range(1, max_attempts + 1): + should_retry = False + try: + resp = await client.post( + f"{base_url}/api/show", + json={"model": model_name}, + timeout=60.0, + ) + resp.raise_for_status() + return resp.json() + except httpx.HTTPStatusError as e: + if e.response.status_code in (429, 500, 502, 503, 504) and attempt < max_attempts: + should_retry = True + else: + raise + except httpx.RequestError as e: + if attempt < max_attempts: + should_retry = True + else: + raise + if should_retry: + await asyncio.sleep(2 ** (attempt - 1)) def is_embedding_model(show_data: dict[str, Any]) -> bool: @@ -72,67 +105,84 @@ async def benchmark_chat_single_run( if num_predict is not None: payload["options"] = {"num_predict": num_predict} - start = time.perf_counter() - first_token_time: float = -1.0 - 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 - duration = eval_duration or total_duration - tps = eval_count / (duration / 1e9) if duration else None - - return {"ttft": ttft, "tps": tps, "error": error} + max_attempts = 4 + for attempt in range(1, max_attempts + 1): + start = time.perf_counter() + first_token_time: float = -1.0 + eval_count = 0 + eval_duration = 0 + total_duration = 0 + error: str | None = None + seen_done = False + should_retry = False + + 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 httpx.HTTPStatusError as e: + error = str(e) + if e.response.status_code in (429, 500, 502, 503, 504): + should_retry = True + except httpx.RequestError as e: + error = str(e) + should_retry = True + except Exception as e: + error = str(e) + + if not seen_done and not error: + error = "Stream ended without completion" + should_retry = True + + if error and should_retry and attempt < max_attempts: + await asyncio.sleep(2 ** (attempt - 1)) + continue + + if first_token_time >= 0: + ttft = first_token_time + else: + ttft = None + duration = eval_duration or total_duration + tps = eval_count / (duration / 1e9) if duration else None + + return {"ttft": ttft, "tps": tps, "error": error} + + return {"ttft": None, "tps": None, "error": error or "Unknown failure"} async def benchmark_embed_single_run( @@ -147,29 +197,45 @@ async def benchmark_embed_single_run( "input": prompt, } - 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) + max_attempts = 4 + for attempt in range(1, max_attempts + 1): + start = time.perf_counter() + prompt_eval_count = 0 + total_duration = 0 + error: str | None = None + should_retry = False - ttft = time.perf_counter() - start - tps = prompt_eval_count / (total_duration / 1e9) if total_duration else None - - return {"ttft": ttft, "tps": tps, "error": error} + 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 httpx.HTTPStatusError as e: + error = str(e) + if e.response.status_code in (429, 500, 502, 503, 504): + should_retry = True + except httpx.RequestError as e: + error = str(e) + should_retry = True + except Exception as e: + error = str(e) + + if error and should_retry and attempt < max_attempts: + await asyncio.sleep(2 ** (attempt - 1)) + continue + + 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": None, "tps": None, "error": error or "Unknown failure"} 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: From 11b3c2fd4a782eca5b3b5c72ee94b8c220d54ccf Mon Sep 17 00:00:00 2001 From: Lewis Cowles Date: Sun, 28 Jun 2026 20:15:17 +0100 Subject: [PATCH 2/3] refactor: retry wrapped functools decorator --- src/ometer/api.py | 272 ++++++++++++++++++++-------------------------- 1 file changed, 119 insertions(+), 153 deletions(-) diff --git a/src/ometer/api.py b/src/ometer/api.py index 4930ae16..2bc93af1 100644 --- a/src/ometer/api.py +++ b/src/ometer/api.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import functools import json import time from dataclasses import dataclass, field @@ -33,61 +34,70 @@ 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]]: - max_attempts = 4 - for attempt in range(1, max_attempts + 1): - should_retry = False - try: - resp = await client.get(f"{base_url}/api/tags") - resp.raise_for_status() - data = resp.json() - return data.get("models", []) - except httpx.HTTPStatusError as e: - if e.response.status_code in (429, 500, 502, 503, 504) and attempt < max_attempts: - should_retry = True - else: - raise - except httpx.RequestError as e: - if attempt < max_attempts: - should_retry = True - else: - raise - if should_retry: - await asyncio.sleep(2 ** (attempt - 1)) + resp = await client.get(f"{base_url}/api/tags") + resp.raise_for_status() + data = resp.json() + return data.get("models", []) +@retry() async def fetch_model_show( client: httpx.AsyncClient, base_url: str, model_name: str ) -> dict[str, Any]: - max_attempts = 4 - for attempt in range(1, max_attempts + 1): - should_retry = False - try: - resp = await client.post( - f"{base_url}/api/show", - json={"model": model_name}, - timeout=60.0, - ) - resp.raise_for_status() - return resp.json() - except httpx.HTTPStatusError as e: - if e.response.status_code in (429, 500, 502, 503, 504) and attempt < max_attempts: - should_retry = True - else: - raise - except httpx.RequestError as e: - if attempt < max_attempts: - should_retry = True - else: - raise - if should_retry: - await asyncio.sleep(2 ** (attempt - 1)) + resp = await client.post( + f"{base_url}/api/show", + json={"model": model_name}, + timeout=60.0, + ) + resp.raise_for_status() + return resp.json() 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, @@ -107,84 +117,63 @@ async def benchmark_chat_single_run( is_thinking = bool(show_data) and "thinking" in show_data.get("capabilities", []) - max_attempts = 4 - for attempt in range(1, max_attempts + 1): - start = time.perf_counter() - first_token_time: float = -1.0 - eval_count = 0 - eval_duration = 0 - total_duration = 0 - error: str | None = None - seen_done = False - should_retry = False + start = time.perf_counter() + first_token_time: float = -1.0 + eval_count = 0 + eval_duration = 0 + total_duration = 0 + seen_done = False + + 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 < 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 - 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("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 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 httpx.HTTPStatusError as e: - error = str(e) - if e.response.status_code in (429, 500, 502, 503, 504): - should_retry = True - except httpx.RequestError as e: - error = str(e) - should_retry = True - except Exception as e: - error = str(e) - - if not seen_done and not error: - error = "Stream ended without completion" - should_retry = True - - if error and should_retry and attempt < max_attempts: - await asyncio.sleep(2 ** (attempt - 1)) - continue - - if first_token_time >= 0: - ttft = first_token_time - else: - ttft = None - duration = eval_duration or total_duration - tps = eval_count / (duration / 1e9) if duration else None + if not seen_done: + raise StreamEndedError("Stream ended without completion") - return {"ttft": ttft, "tps": tps, "error": error} + if first_token_time >= 0: + ttft = first_token_time + else: + ttft = None + duration = eval_duration or total_duration + tps = eval_count / (duration / 1e9) if duration else None - return {"ttft": None, "tps": None, "error": error or "Unknown failure"} + return {"ttft": ttft, "tps": tps, "error": None} +@retry(return_error_dict=True) async def benchmark_embed_single_run( client: httpx.AsyncClient, base_url: str, @@ -197,45 +186,22 @@ async def benchmark_embed_single_run( "input": prompt, } - max_attempts = 4 - for attempt in range(1, max_attempts + 1): - start = time.perf_counter() - prompt_eval_count = 0 - total_duration = 0 - error: str | None = None - should_retry = False + start = time.perf_counter() + 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) - 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 httpx.HTTPStatusError as e: - error = str(e) - if e.response.status_code in (429, 500, 502, 503, 504): - should_retry = True - except httpx.RequestError as e: - error = str(e) - should_retry = True - except Exception as e: - error = str(e) - - if error and should_retry and attempt < max_attempts: - await asyncio.sleep(2 ** (attempt - 1)) - continue - - 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": None, "tps": None, "error": error or "Unknown failure"} + ttft = time.perf_counter() - start + tps = prompt_eval_count / (total_duration / 1e9) if total_duration else None + + return {"ttft": ttft, "tps": tps, "error": None} async def benchmark_model( From 88b2c41b8636c0f711e8caec5da64082f7037611 Mon Sep 17 00:00:00 2001 From: Lewis Cowles Date: Sun, 28 Jun 2026 20:23:32 +0100 Subject: [PATCH 3/3] chore: collapse conditional to restore 100% coverage --- src/ometer/api.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/ometer/api.py b/src/ometer/api.py index 2bc93af1..8997ffcc 100644 --- a/src/ometer/api.py +++ b/src/ometer/api.py @@ -118,7 +118,7 @@ async def benchmark_chat_single_run( 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 @@ -145,12 +145,12 @@ async def benchmark_chat_single_run( msg = chunk.get("message") or {} if is_thinking: - if first_token_time < 0 and ( + 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 < 0 and msg.get("content"): + if first_token_time is None and msg.get("content"): first_token_time = time.perf_counter() - start if chunk.get("done"): @@ -163,10 +163,7 @@ async def benchmark_chat_single_run( if not seen_done: raise StreamEndedError("Stream ended without completion") - if first_token_time >= 0: - ttft = first_token_time - else: - ttft = None + ttft = first_token_time duration = eval_duration or total_duration tps = eval_count / (duration / 1e9) if duration else None