← back to Exo
Add usage stats (#1333)
d9eca7589556f745eee2b2476552f27e29ddb15b · 2026-01-30 10:23:08 +0000 · rltakashige
## Motivation
(Probably) the final missing piece of the Chat Completions API
## Changes
Add UsageStats
## Why It Works
OpenCode reviewed my PR and gave me stats:
<img width="1150" height="802" alt="image"
src="https://github.com/user-attachments/assets/ebc06bae-797f-4087-87d5-2f26cf60fc48"
/>
## Test Plan
### Automated Testing
No tests were broken.
Files touched
M src/exo/master/api.pyM src/exo/shared/types/api.pyM src/exo/shared/types/chunks.pyM src/exo/shared/types/worker/runner_response.pyM src/exo/worker/engines/mlx/generator/generate.pyM src/exo/worker/runner/runner.pyM src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
Diff
commit d9eca7589556f745eee2b2476552f27e29ddb15b
Author: rltakashige <rl.takashige@gmail.com>
Date: Fri Jan 30 10:23:08 2026 +0000
Add usage stats (#1333)
## Motivation
(Probably) the final missing piece of the Chat Completions API
## Changes
Add UsageStats
## Why It Works
OpenCode reviewed my PR and gave me stats:
<img width="1150" height="802" alt="image"
src="https://github.com/user-attachments/assets/ebc06bae-797f-4087-87d5-2f26cf60fc48"
/>
## Test Plan
### Automated Testing
No tests were broken.
---
src/exo/master/api.py | 21 +++++++--
src/exo/shared/types/api.py | 13 +++--
src/exo/shared/types/chunks.py | 4 +-
src/exo/shared/types/worker/runner_response.py | 3 ++
src/exo/worker/engines/mlx/generator/generate.py | 55 +++++++++++++++++-----
src/exo/worker/runner/runner.py | 9 +++-
.../unittests/test_runner/test_event_ordering.py | 4 +-
7 files changed, 85 insertions(+), 24 deletions(-)
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index 3d56887f..de2b3a3e 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -65,7 +65,9 @@ from exo.shared.types.api import (
StartDownloadParams,
StartDownloadResponse,
StreamingChoiceResponse,
+ StreamOptions,
ToolCall,
+ Usage,
)
from exo.shared.types.chunks import (
ErrorChunk,
@@ -113,7 +115,9 @@ def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None)
def chunk_to_response(
- chunk: TokenChunk | ToolCallChunk, command_id: CommandId
+ chunk: TokenChunk | ToolCallChunk,
+ command_id: CommandId,
+ usage: Usage | None,
) -> ChatCompletionResponse:
return ChatCompletionResponse(
id=command_id,
@@ -138,6 +142,7 @@ def chunk_to_response(
finish_reason=chunk.finish_reason,
)
],
+ usage=usage,
)
@@ -522,9 +527,10 @@ class API:
del self._chat_completion_queues[command_id]
async def _generate_chat_stream(
- self, command_id: CommandId
+ self, command_id: CommandId, stream_options: StreamOptions | None = None
) -> AsyncGenerator[str, None]:
"""Generate chat completion stream as JSON strings."""
+ include_usage = stream_options.include_usage if stream_options else False
async for chunk in self._chat_chunk_stream(command_id):
assert not isinstance(chunk, ImageChunk)
@@ -540,8 +546,10 @@ class API:
yield "data: [DONE]\n\n"
return
+ usage = chunk.usage if include_usage else None
+
chunk_response: ChatCompletionResponse = chunk_to_response(
- chunk, command_id
+ chunk, command_id, usage=usage
)
logger.debug(f"chunk_response: {chunk_response}")
@@ -559,6 +567,7 @@ class API:
tool_calls: list[ToolCall] = []
model: str | None = None
finish_reason: FinishReason | None = None
+ usage: Usage | None = None
async for chunk in self._chat_chunk_stream(command_id):
if isinstance(chunk, ErrorChunk):
@@ -583,6 +592,9 @@ class API:
for i, tool in enumerate(chunk.tool_calls)
)
+ if chunk.usage is not None:
+ usage = chunk.usage
+
if chunk.finish_reason is not None:
finish_reason = chunk.finish_reason
@@ -604,6 +616,7 @@ class API:
finish_reason=finish_reason,
)
],
+ usage=usage,
)
async def _collect_chat_completion_with_stats(
@@ -691,7 +704,7 @@ class API:
await self._send(command)
if payload.stream:
return StreamingResponse(
- self._generate_chat_stream(command.command_id),
+ self._generate_chat_stream(command.command_id, payload.stream_options),
media_type="text/event-stream",
)
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index 1a758a3c..1dac848f 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -3,7 +3,7 @@ from collections.abc import Generator
from typing import Annotated, Any, Literal
from fastapi import UploadFile
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import BaseModel, Field, field_validator
from pydantic_core import PydanticUseDefault
from exo.shared.models.model_cards import ModelCard, ModelId
@@ -11,7 +11,7 @@ from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.memory import Memory
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding, ShardMetadata
-from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
+from exo.utils.pydantic_ext import CamelCaseModel, ConfigDict, TaggedModel
FinishReason = Literal[
"stop", "length", "tool_calls", "content_filter", "function_call", "error"
@@ -116,8 +116,8 @@ class Usage(BaseModel):
prompt_tokens: int
completion_tokens: int
total_tokens: int
- prompt_tokens_details: PromptTokensDetails | None = None
- completion_tokens_details: CompletionTokensDetails | None = None
+ prompt_tokens_details: PromptTokensDetails
+ completion_tokens_details: CompletionTokensDetails
class StreamingChoiceResponse(BaseModel):
@@ -170,6 +170,10 @@ class BenchChatCompletionResponse(ChatCompletionResponse):
generation_stats: GenerationStats | None = None
+class StreamOptions(BaseModel):
+ include_usage: bool = False
+
+
class ChatCompletionTaskParams(TaggedModel):
model_config = ConfigDict(extra="ignore")
@@ -186,6 +190,7 @@ class ChatCompletionTaskParams(TaggedModel):
seed: int | None = None
stop: str | list[str] | None = None
stream: bool = False
+ stream_options: StreamOptions | None = None
temperature: float | None = None
top_p: float | None = None
tools: list[dict[str, Any]] | None = None
diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py
index 235ef70d..e96dbc9d 100644
--- a/src/exo/shared/types/chunks.py
+++ b/src/exo/shared/types/chunks.py
@@ -2,7 +2,7 @@ from collections.abc import Generator
from typing import Any, Literal
from exo.shared.models.model_cards import ModelId
-from exo.shared.types.api import GenerationStats, ImageGenerationStats
+from exo.shared.types.api import GenerationStats, ImageGenerationStats, Usage
from exo.utils.pydantic_ext import TaggedModel
from .api import FinishReason
@@ -17,6 +17,7 @@ class BaseChunk(TaggedModel):
class TokenChunk(BaseChunk):
text: str
token_id: int
+ usage: Usage | None
finish_reason: Literal["stop", "length", "content_filter"] | None = None
stats: GenerationStats | None = None
@@ -28,6 +29,7 @@ class ErrorChunk(BaseChunk):
class ToolCallChunk(BaseChunk):
tool_calls: list[ToolCallItem]
+ usage: Usage | None
finish_reason: Literal["tool_calls"] = "tool_calls"
stats: GenerationStats | None = None
diff --git a/src/exo/shared/types/worker/runner_response.py b/src/exo/shared/types/worker/runner_response.py
index 9f867413..5dfbe547 100644
--- a/src/exo/shared/types/worker/runner_response.py
+++ b/src/exo/shared/types/worker/runner_response.py
@@ -6,6 +6,7 @@ from exo.shared.types.api import (
GenerationStats,
ImageGenerationStats,
ToolCallItem,
+ Usage,
)
from exo.utils.pydantic_ext import TaggedModel
@@ -24,6 +25,7 @@ class GenerationResponse(BaseRunnerResponse):
# logprobs: list[float] | None = None # too big. we can change to be top-k
finish_reason: FinishReason | None = None
stats: GenerationStats | None = None
+ usage: Usage | None
class ImageGenerationResponse(BaseRunnerResponse):
@@ -57,6 +59,7 @@ class PartialImageResponse(BaseRunnerResponse):
class ToolCallResponse(BaseRunnerResponse):
tool_calls: list[ToolCallItem]
+ usage: Usage | None
class FinishedResponse(BaseRunnerResponse):
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index 56327a0f..c392f6a8 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -10,8 +10,11 @@ from mlx_lm.tokenizer_utils import TokenizerWrapper
from exo.shared.types.api import (
BenchChatCompletionTaskParams,
ChatCompletionMessage,
+ CompletionTokensDetails,
FinishReason,
GenerationStats,
+ PromptTokensDetails,
+ Usage,
)
from exo.shared.types.memory import Memory
from exo.shared.types.mlx import KVCacheType
@@ -216,22 +219,37 @@ def mlx_generate(
max_tokens = task.max_tokens or MAX_TOKENS
generated_text_parts: list[str] = []
generation_start_time = time.perf_counter()
- for out in stream_generate(
- model=model,
- tokenizer=tokenizer,
- prompt=last_token,
- max_tokens=max_tokens,
- sampler=sampler,
- logits_processors=logits_processors,
- prompt_cache=caches,
- # TODO: Dynamically change prefill step size to be the maximum possible without timing out.
- prefill_step_size=2048,
- kv_group_size=KV_GROUP_SIZE,
- kv_bits=KV_BITS,
+ usage: Usage | None = None
+ in_thinking = False
+ reasoning_tokens = 0
+ think_start = tokenizer.think_start
+ think_end = tokenizer.think_end
+ for completion_tokens, out in enumerate(
+ stream_generate(
+ model=model,
+ tokenizer=tokenizer,
+ prompt=last_token,
+ max_tokens=max_tokens,
+ sampler=sampler,
+ logits_processors=logits_processors,
+ prompt_cache=caches,
+ # TODO: Dynamically change prefill step size to be the maximum possible without timing out.
+ prefill_step_size=2048,
+ kv_group_size=KV_GROUP_SIZE,
+ kv_bits=KV_BITS,
+ ),
+ start=1,
):
generated_text_parts.append(out.text)
logger.info(out.text)
+ if think_start is not None and out.text == think_start:
+ in_thinking = True
+ elif think_end is not None and out.text == think_end:
+ in_thinking = False
+ if in_thinking:
+ reasoning_tokens += 1
+
stats: GenerationStats | None = None
if out.finish_reason is not None:
stats = GenerationStats(
@@ -249,11 +267,24 @@ def mlx_generate(
f"Model generated unexpected finish_reason: {out.finish_reason}"
)
+ usage = Usage(
+ prompt_tokens=int(out.prompt_tokens),
+ completion_tokens=completion_tokens,
+ total_tokens=int(out.prompt_tokens) + completion_tokens,
+ prompt_tokens_details=PromptTokensDetails(
+ cached_tokens=prefix_hit_length
+ ),
+ completion_tokens_details=CompletionTokensDetails(
+ reasoning_tokens=reasoning_tokens
+ ),
+ )
+
yield GenerationResponse(
text=out.text,
token=out.token,
finish_reason=cast(FinishReason | None, out.finish_reason),
stats=stats,
+ usage=usage,
)
if out.finish_reason is not None:
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 937aa7a9..98705eea 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -277,9 +277,11 @@ def main(
tokenizer.tool_parser, # pyright: ignore[reportAny]
)
+ completion_tokens = 0
for response in mlx_generator:
match response:
case GenerationResponse():
+ completion_tokens += 1
if (
device_rank == 0
and response.finish_reason == "error"
@@ -307,6 +309,7 @@ def main(
model=shard_metadata.model_card.model_id,
text=response.text,
token_id=response.token,
+ usage=response.usage,
finish_reason=response.finish_reason,
stats=response.stats,
),
@@ -320,6 +323,7 @@ def main(
chunk=ToolCallChunk(
tool_calls=response.tool_calls,
model=shard_metadata.model_card.model_id,
+ usage=response.usage,
),
)
)
@@ -535,7 +539,8 @@ def parse_gpt_oss(
name=current_tool_name,
arguments="".join(tool_arg_parts).strip(),
)
- ]
+ ],
+ usage=response.usage,
)
tool_arg_parts = []
current_tool_name = recipient
@@ -683,7 +688,7 @@ def parse_tool_calls(
tools = [_validate_single_tool(tool) for tool in parsed]
else:
tools = [_validate_single_tool(parsed)]
- yield ToolCallResponse(tool_calls=tools)
+ yield ToolCallResponse(tool_calls=tools, usage=response.usage)
except (
json.JSONDecodeError,
diff --git a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
index 045bb072..d92a4a83 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
@@ -120,7 +120,7 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(mlx_runner, "detect_thinking_prompt_suffix", make_nothin(False))
def fake_generate(*_1: object, **_2: object):
- yield GenerationResponse(token=0, text="hi", finish_reason="stop")
+ yield GenerationResponse(token=0, text="hi", finish_reason="stop", usage=None)
monkeypatch.setattr(mlx_runner, "mlx_generate", fake_generate)
@@ -182,6 +182,8 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
text="hi",
token_id=0,
finish_reason="stop",
+ usage=None,
+ stats=None,
),
)
← 9dabde7e Fix bench after recent updates (#1331)
·
back to Exo
·
improve log message in shard downloader 9ba61f37 →