optimize further and get usage stats

bench/exo_eval.py

@@ -9,8 +9,8 @@ Supports multiple evaluation frameworks via TOML configuration:
 - custom: Custom evaluation scripts
 
 Usage:
-    uv run python -m bench.exo_eval --config bench/eval_config.toml --model llama-3.2-1b
-    uv run python -m bench.exo_eval --config bench/eval_config.toml --model llama-3.2-1b --dry-run
+    uv run python -m bench.exo_eval --config bench/eval_config.toml --model Llama-3.2-1b-Instruct-4bit
+    uv run python -m bench.exo_eval --config bench/eval_config.toml --model Llama-3.2-1b-Instruct-4bit --dry-run
 """
 
 from __future__ import annotations
@@ -238,6 +238,8 @@ def build_lm_eval_args(
         model_type,
         "--model_args",
         model_args,
+        "--verbosity",
+        "WARNING",
     ]
 
     # Tasks
@@ -317,6 +319,30 @@ def run_lm_eval(
 
     try:
         result = subprocess.run(args, check=False)
+
+        # Print token usage summary from exo
+        try:
+            import httpx
+            usage_resp = httpx.get(f"{exo_base_url}/v1/usage", timeout=5)
+            if usage_resp.status_code == 200:
+                usage = usage_resp.json()
+                logger.info("--- Token Usage (Total) ---")
+                logger.info(f"  Requests:          {usage.get('total_requests', 0)}")
+                logger.info(f"  Prompt tokens:     {usage.get('total_prompt_tokens', 0)}")
+                logger.info(f"  Completion tokens: {usage.get('total_completion_tokens', 0)}")
+                logger.info(f"  Reasoning tokens:  {usage.get('total_reasoning_tokens', 0)}")
+                logger.info(f"  Total tokens:      {usage.get('total_tokens', 0)}")
+                by_model = usage.get("by_model", {})
+                if by_model:
+                    for model_name, counters in by_model.items():
+                        logger.info(f"--- Token Usage ({model_name}) ---")
+                        logger.info(f"  Requests:          {counters.get('requests', 0)}")
+                        logger.info(f"  Prompt tokens:     {counters.get('prompt_tokens', 0)}")
+                        logger.info(f"  Completion tokens: {counters.get('completion_tokens', 0)}")
+                        logger.info(f"  Reasoning tokens:  {counters.get('reasoning_tokens', 0)}")
+        except Exception:
+            pass  # Usage endpoint not available
+
         return result.returncode
     except FileNotFoundError:
         logger.error("lm_eval not found. Install with: uv sync --extra eval")

src/exo/master/api.py

@@ -4,7 +4,7 @@ import re
 import time
 from collections.abc import AsyncGenerator
 from http import HTTPStatus
-from typing import Literal, cast
+from typing import Any, Literal, cast
 from uuid import uuid4
 
 import anyio
@@ -233,11 +233,8 @@ class API:
         self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
         self._tg: TaskGroup | None = None
 
-        # Accumulated usage stats across all requests
-        self._total_prompt_tokens: int = 0
-        self._total_completion_tokens: int = 0
-        self._total_reasoning_tokens: int = 0
-        self._total_requests: int = 0
+        # Accumulated usage stats per instance (keyed by model id)
+        self._usage_by_model: dict[str, dict[str, int]] = {}
 
     def reset(self, new_session_id: SessionId, result_clock: int):
         logger.info("Resetting API State")
@@ -306,16 +303,41 @@ class API:
         self.app.get("/events")(lambda: self._event_log)
         self.app.get("/v1/usage")(self.get_usage)
 
-    def get_usage(self) -> dict[str, int]:
-        """Return accumulated token usage across all requests."""
+    def get_usage(self) -> dict[str, Any]:
+        """Return accumulated token usage per model instance."""
+        total_requests = 0
+        total_prompt = 0
+        total_completion = 0
+        total_reasoning = 0
+        for counters in self._usage_by_model.values():
+            total_requests += counters.get("requests", 0)
+            total_prompt += counters.get("prompt_tokens", 0)
+            total_completion += counters.get("completion_tokens", 0)
+            total_reasoning += counters.get("reasoning_tokens", 0)
         return {
-            "total_requests": self._total_requests,
-            "total_prompt_tokens": self._total_prompt_tokens,
-            "total_completion_tokens": self._total_completion_tokens,
-            "total_reasoning_tokens": self._total_reasoning_tokens,
-            "total_tokens": self._total_prompt_tokens + self._total_completion_tokens,
+            "total_requests": total_requests,
+            "total_prompt_tokens": total_prompt,
+            "total_completion_tokens": total_completion,
+            "total_reasoning_tokens": total_reasoning,
+            "total_tokens": total_prompt + total_completion,
+            "by_model": self._usage_by_model,
         }
 
+    def _accumulate_usage(self, model: str, prompt_tokens: int, completion_tokens: int, reasoning_tokens: int) -> None:
+        """Accumulate usage stats for a model instance."""
+        if model not in self._usage_by_model:
+            self._usage_by_model[model] = {
+                "requests": 0,
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "reasoning_tokens": 0,
+            }
+        counters = self._usage_by_model[model]
+        counters["requests"] += 1
+        counters["prompt_tokens"] += prompt_tokens
+        counters["completion_tokens"] += completion_tokens
+        counters["reasoning_tokens"] += reasoning_tokens
+
     async def place_instance(self, payload: PlaceInstanceParams):
         command = PlaceInstance(
             model_card=await resolve_model_card(payload.model_id),
@@ -585,10 +607,12 @@ class API:
                 # Accumulate usage stats from the final chunk
                 if isinstance(chunk, TokenChunk) and chunk.stats is not None:
                     s = chunk.stats
-                    self._total_prompt_tokens += s.prompt_tokens
-                    self._total_completion_tokens += s.generation_tokens
-                    self._total_reasoning_tokens += s.reasoning_tokens
-                    self._total_requests += 1
+                    self._accumulate_usage(
+                        model=chunk.model,
+                        prompt_tokens=s.prompt_tokens,
+                        completion_tokens=s.generation_tokens,
+                        reasoning_tokens=s.reasoning_tokens,
+                    )
                 yield "data: [DONE]\n\n"
 
     async def _collect_chat_completion(
@@ -661,10 +685,12 @@ class API:
                     reasoning_tokens=stats.reasoning_tokens,
                 ) if stats.reasoning_tokens > 0 else None,
             )
-            self._total_prompt_tokens += stats.prompt_tokens
-            self._total_completion_tokens += completion_tokens
-            self._total_reasoning_tokens += stats.reasoning_tokens
-            self._total_requests += 1
+            self._accumulate_usage(
+                model=model or "unknown",
+                prompt_tokens=stats.prompt_tokens,
+                completion_tokens=completion_tokens,
+                reasoning_tokens=stats.reasoning_tokens,
+            )
 
         return ChatCompletionResponse(
             id=command_id,