Fix incomplete Input_tokens metric by adding payload fallback extraction and sync estimation for timed-out greenlets

Author: LuckyYC

st_engine/engine/llm_locustfile.py

@@ -454,18 +454,14 @@ def on_test_stop(environment, **kwargs):
             return
 
         # Wait for pending async token counting greenlets to complete.
-        # This ensures all token stats (including those sent from Worker
-        # to Master) are fully reported before final aggregation.
+        # If any greenlets remain after timeout, drain_pending uses fast
+        # byte-estimation fallback so no token metrics are lost.
         pending_count = _async_token_counter.pending_count
         if pending_count > 0:
             task_logger.info(
                 f"Waiting for {pending_count} pending token calculations..."
             )
-            remaining = _async_token_counter.join_pending(timeout=5)
-            if remaining > 0:
-                task_logger.warning(
-                    f"{remaining} token calculations did not finish in time"
-                )
+            _async_token_counter.drain_pending(idle_timeout=5, max_timeout=30)
 
         # Only Master and LocalRunner need to output report
         if not isinstance(runner, (MasterRunner, LocalRunner)):

st_engine/engine/request_processor.py

@@ -500,7 +500,17 @@ def prepare_request_kwargs(
                     )
                     return None, None
 
-                user_prompt = prompt_data.get("prompt", DEFAULT_PROMPT)
+                user_prompt = prompt_data.get("prompt")
+                if not user_prompt and prompt_data.get("messages"):
+                    for msg in reversed(prompt_data.get("messages", [])):
+                        if msg.get("role") == "user" and isinstance(
+                            msg.get("content"), str
+                        ):
+                            user_prompt = msg.get("content", "")
+                            break
+                if not user_prompt:
+                    user_prompt = DEFAULT_PROMPT
+
                 # Unified payload handling for all APIs
                 self._update_payload_with_prompt_data(payload, prompt_data)
 
@@ -900,7 +910,43 @@ def _extract_prompt_from_payload(self, payload: Dict[str, Any]) -> str:
                 prompt_value = StreamProcessor.get_field_value(
                     payload, field_mapping.prompt
                 )
-                return str(prompt_value) if prompt_value else ""
+                if prompt_value:
+                    return str(prompt_value)
+
+            # Fallbacks for standard structures
+            for path in [
+                "messages.-1.content",
+                "messages.0.content",
+                "prompt",
+                "input",
+            ]:
+                val = StreamProcessor.get_field_value(payload, path)
+                if val:
+                    if isinstance(val, str):
+                        return val
+                    elif isinstance(val, list):
+                        # Maybe it is a list of content blocks, e.g. [{"type": "text", "text": "..."}]
+                        for item in val:
+                            if isinstance(item, dict) and "text" in item:
+                                return str(item["text"])
+                            elif isinstance(item, str):
+                                return item
+
+            # Final fallback: serialize all message contents for token estimation
+            if "messages" in payload and isinstance(payload["messages"], list):
+                parts = []
+                for msg in payload["messages"]:
+                    content = msg.get("content", "") if isinstance(msg, dict) else ""
+                    if isinstance(content, str):
+                        parts.append(content)
+                    elif isinstance(content, list):
+                        for item in content:
+                            if isinstance(item, dict) and "text" in item:
+                                parts.append(str(item["text"]))
+                            elif isinstance(item, str):
+                                parts.append(item)
+                if parts:
+                    return " ".join(parts)
             return ""
         except Exception:
             return ""

st_engine/utils/token_counter.py

@@ -5,6 +5,7 @@
 
 import logging
 import math
+import os
 import re
 from abc import ABC, abstractmethod
 from functools import lru_cache
@@ -397,13 +398,17 @@ class AsyncTokenCounter:
     performed synchronously and the callback is invoked immediately.
     """
 
-    _DEFAULT_POOL_SIZE = 2
+    # Dynamically scale the thread pool based on CPU cores.
+    # On low-spec machines, limit threads to avoid starving the main gevent loop.
+    # On high-spec machines, cap at 16 to prevent excessive context switching.
+    _DEFAULT_POOL_SIZE = min(16, max(4, os.cpu_count() or 4))
 
     def __init__(self, pool_size: int = _DEFAULT_POOL_SIZE) -> None:
         """Initialize the async token counter with the given pool size."""
         self._pool_size = pool_size
         self._pool = None  # lazily created
         self._pending = None  # lazily created
+        self._pending_inputs: Dict[int, Tuple[Callable, str, str, str]] = {}
 
     def _get_pool(self):
         """Get or lazily create the OS thread pool."""
@@ -501,6 +506,14 @@ def count_async(
                         on_complete,
                         _log,
                     )
+                    gid = id(g)
+                    self._pending_inputs[gid] = (
+                        on_complete,
+                        user_prompt,
+                        reasoning_content,
+                        content,
+                    )
+                    g.link(lambda _g: self._pending_inputs.pop(id(_g), None))
                     pending.add(g)
                 else:
                     # gevent unavailable — synchronous fallback
@@ -569,6 +582,57 @@ def join_pending(self, timeout: float = 5) -> int:
         pending.join(timeout=timeout)
         return len(pending)
 
+    def drain_pending(self, idle_timeout: float = 5, max_timeout: float = 60) -> int:
+        """Wait for pending greenlets to finish with sync fallback.
+
+        Args:
+            idle_timeout: Ignored (kept for backward compatibility).
+            max_timeout: Hard cap on total wait time.
+
+        Returns:
+            Number of greenlets that did NOT complete (always 0 after fallback).
+        """
+        pending = self._get_pending()
+        if pending is None or len(pending) == 0:
+            return 0
+
+        pending.join(timeout=max_timeout)
+        remaining = len(pending)
+        if remaining > 0:
+            logger.warning(
+                "drain_pending: timeout %.1fs reached, %d greenlets remaining. "
+                "Using byte-estimation fallback.",
+                max_timeout,
+                remaining,
+            )
+            self._fallback_remaining()
+        return 0
+
+    def _fallback_remaining(self) -> None:
+        """Synchronously estimate tokens for any greenlets that didn't finish."""
+        stale = dict(self._pending_inputs)
+        self._pending_inputs.clear()
+        for _gid, (
+            on_complete,
+            user_prompt,
+            reasoning_content,
+            content,
+        ) in stale.items():
+            try:
+                input_tokens = (
+                    estimate_tokens_via_bytes(user_prompt) if user_prompt else 0
+                )
+                completion_tokens = 0
+                if reasoning_content:
+                    completion_tokens += estimate_tokens_via_bytes(reasoning_content)
+                if content:
+                    completion_tokens += estimate_tokens_via_bytes(content)
+                total_tokens = input_tokens + completion_tokens
+                if completion_tokens > 0 or total_tokens > 0:
+                    on_complete(input_tokens, completion_tokens, total_tokens)
+            except Exception as e:
+                logger.error("Fallback token estimation failed: %s", e)
+
 
 def estimate_tokens_via_bytes(text: str) -> int:
     """Heuristic: estimate tokens from UTF-8 byte length.