Merge pull request lightspeed-core#159 from bsatapat-jpg/dev

[LEADS-208] Fix TokenTracker double-counting in multi-thread evaluation

Author: asamal4

src/lightspeed_evaluation/core/llm/custom.py

@@ -2,6 +2,7 @@
 
 import os
 import logging
+import threading
 from typing import Any, Optional, Union
 
 import litellm
@@ -11,62 +12,76 @@
 
 logger = logging.getLogger(__name__)
 
+# Thread-local storage for active TokenTracker
+_active_tracker: threading.local = threading.local()
+
 
 class TokenTracker:
-    """Tracks token usage from LiteLLM calls via callbacks.
+    """Tracks token usage from LLM calls using direct response extraction.
+
+    Uses thread-local storage to track the active tracker. Tokens are captured
+    directly from litellm response in BaseCustomLLM.call() - no callbacks,
+    no timeouts, no race conditions.
 
     Usage:
         tracker = TokenTracker()
-        tracker.start()  # Register callback
-        # ... make LLM calls ...
-        tracker.stop()   # Unregister callback
+        tracker.start()  # Set as active tracker for this thread
+        # ... make LLM calls (tokens captured automatically) ...
+        tracker.stop()   # Unset as active tracker
         input_tokens, output_tokens = tracker.get_counts()
     """
 
     def __init__(self) -> None:
         """Initialize token tracker."""
         self.input_tokens = 0
         self.output_tokens = 0
-        self._callback_registered = False
+        self._lock = threading.Lock()  # Instance lock for token counter updates
 
-    def _token_callback(
-        self,
-        _kwargs: dict[str, Any],
-        completion_response: Any,
-        _start_time: float,
-        _end_time: float,
-    ) -> None:
-        """Capture token usage from LiteLLM completion response."""
-        if hasattr(completion_response, "usage") and completion_response.usage:
-            usage = completion_response.usage
-            self.input_tokens += getattr(usage, "prompt_tokens", 0)
-            self.output_tokens += getattr(usage, "completion_tokens", 0)
+    def add_tokens(self, prompt_tokens: int, completion_tokens: int) -> None:
+        """Add token counts (thread-safe).
+
+        Called by BaseCustomLLM.call() to record tokens from LLM response.
+
+        Args:
+            prompt_tokens: Number of input/prompt tokens.
+            completion_tokens: Number of output/completion tokens.
+        """
+        with self._lock:
+            self.input_tokens += prompt_tokens
+            self.output_tokens += completion_tokens
 
     def start(self) -> None:
-        """Register the token tracking callback."""
-        if self._callback_registered:
-            return
-        if not hasattr(litellm, "success_callback") or litellm.success_callback is None:
-            litellm.success_callback = []
-        litellm.success_callback.append(self._token_callback)
-        self._callback_registered = True
+        """Set this tracker as active for the current thread."""
+        _active_tracker.tracker = self
 
     def stop(self) -> None:
-        """Unregister the token tracking callback."""
-        if not self._callback_registered:
-            return
-        if self._token_callback in litellm.success_callback:
-            litellm.success_callback.remove(self._token_callback)
-        self._callback_registered = False
+        """Unset this tracker as active for the current thread."""
+        if getattr(_active_tracker, "tracker", None) is self:
+            _active_tracker.tracker = None
 
     def get_counts(self) -> tuple[int, int]:
-        """Get accumulated token counts."""
-        return self.input_tokens, self.output_tokens
+        """Get accumulated token counts.
+
+        Returns:
+            Tuple of (input_tokens, output_tokens)
+        """
+        with self._lock:
+            return self.input_tokens, self.output_tokens
 
     def reset(self) -> None:
         """Reset token counts to zero."""
-        self.input_tokens = 0
-        self.output_tokens = 0
+        with self._lock:
+            self.input_tokens = 0
+            self.output_tokens = 0
+
+    @staticmethod
+    def get_active() -> Optional["TokenTracker"]:
+        """Get the active tracker for the current thread.
+
+        Returns:
+            The active TokenTracker, or None if no tracker is active.
+        """
+        return getattr(_active_tracker, "tracker", None)
 
 
 class BaseCustomLLM:  # pylint: disable=too-few-public-methods
@@ -133,6 +148,14 @@ def call(
         try:
             response = litellm.completion(**call_params)
 
+            # Direct token extraction - capture tokens synchronously from response
+            tracker = TokenTracker.get_active()
+            if tracker and hasattr(response, "usage") and response.usage:
+                tracker.add_tokens(
+                    getattr(response.usage, "prompt_tokens", 0),
+                    getattr(response.usage, "completion_tokens", 0),
+                )
+
             # Extract content from all choices
             results = []
             for choice in response.choices:  # type: ignore

tests/unit/core/llm/test_custom.py

@@ -2,6 +2,8 @@
 
 """Unit tests for custom LLM classes."""
 
+import threading
+
 import pytest
 from pytest_mock import MockerFixture
 
@@ -12,21 +14,102 @@
 class TestTokenTracker:
     """Tests for TokenTracker."""
 
-    def test_token_callback_accumulates_tokens(self, mocker: MockerFixture) -> None:
-        """Test that token callback accumulates token counts."""
+    def test_add_tokens_accumulates(self) -> None:
+        """Test that add_tokens accumulates token counts."""
         tracker = TokenTracker()
 
-        # Mock completion response with usage
-        mock_response = mocker.Mock()
-        mock_response.usage = mocker.Mock()
-        mock_response.usage.prompt_tokens = 10
-        mock_response.usage.completion_tokens = 20
+        tracker.add_tokens(10, 20)
+        tracker.add_tokens(5, 15)
+
+        input_tokens, output_tokens = tracker.get_counts()
+        assert input_tokens == 15
+        assert output_tokens == 35
+
+    def test_reset_clears_counts(self) -> None:
+        """Test that reset clears token counts."""
+        tracker = TokenTracker()
+        tracker.add_tokens(100, 200)
 
-        tracker._token_callback({}, mock_response, 0.0, 0.0)
+        tracker.reset()
 
         input_tokens, output_tokens = tracker.get_counts()
-        assert input_tokens == 10
-        assert output_tokens == 20
+        assert input_tokens == 0
+        assert output_tokens == 0
+
+    def test_start_sets_active_tracker(self) -> None:
+        """Test that start sets the tracker as active for current thread."""
+        tracker = TokenTracker()
+        tracker.start()
+
+        try:
+            assert TokenTracker.get_active() is tracker
+        finally:
+            tracker.stop()
+
+    def test_stop_clears_active_tracker(self) -> None:
+        """Test that stop clears the active tracker."""
+        tracker = TokenTracker()
+        tracker.start()
+        tracker.stop()
+
+        assert TokenTracker.get_active() is None
+
+    def test_get_active_returns_none_when_no_tracker(self) -> None:
+        """Test that get_active returns None when no tracker is active."""
+        # Ensure clean state by starting and stopping a tracker
+        temp = TokenTracker()
+        temp.start()
+        temp.stop()
+
+        assert TokenTracker.get_active() is None
+
+    def test_thread_local_isolation(self) -> None:
+        """Test that each thread has its own active tracker."""
+        tracker1 = TokenTracker()
+        tracker2 = TokenTracker()
+        results: dict[str, TokenTracker | None] = {}
+
+        def thread_work(name: str, tracker: TokenTracker) -> None:
+            tracker.start()
+            results[name] = TokenTracker.get_active()
+            # Deliberately don't stop to check isolation
+
+        # Start tracker1 in main thread
+        tracker1.start()
+
+        # Start tracker2 in another thread
+        thread = threading.Thread(target=thread_work, args=("thread2", tracker2))
+        thread.start()
+        thread.join()
+
+        # Main thread should still have tracker1
+        assert TokenTracker.get_active() is tracker1
+        # Other thread had tracker2
+        assert results["thread2"] is tracker2
+
+        tracker1.stop()
+
+    def test_add_tokens_thread_safe(self) -> None:
+        """Test that add_tokens is thread-safe under concurrent access."""
+        tracker = TokenTracker()
+        num_threads = 10
+        tokens_per_thread = 100
+
+        def add_tokens_worker() -> None:
+            for _ in range(tokens_per_thread):
+                tracker.add_tokens(1, 2)
+
+        threads = [
+            threading.Thread(target=add_tokens_worker) for _ in range(num_threads)
+        ]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        input_tokens, output_tokens = tracker.get_counts()
+        assert input_tokens == num_threads * tokens_per_thread
+        assert output_tokens == num_threads * tokens_per_thread * 2
 
 
 class TestBaseCustomLLM:
@@ -103,3 +186,63 @@ def test_call_raises_llm_error_on_failure(self, mocker: MockerFixture) -> None:
 
         with pytest.raises(LLMError, match="LLM call failed"):
             llm.call("test prompt")
+
+    def test_call_captures_tokens_with_active_tracker(
+        self, mocker: MockerFixture
+    ) -> None:
+        """Test call captures tokens when a TokenTracker is active."""
+        mock_litellm = mocker.patch("lightspeed_evaluation.core.llm.custom.litellm")
+        mocker.patch.dict("os.environ", {})
+
+        # Mock response with usage
+        mock_choice = mocker.Mock()
+        mock_choice.message.content = "Test response"
+        mock_response = mocker.Mock()
+        mock_response.choices = [mock_choice]
+        mock_response.usage = mocker.Mock()
+        mock_response.usage.prompt_tokens = 50
+        mock_response.usage.completion_tokens = 100
+        mock_litellm.completion.return_value = mock_response
+
+        # Start a tracker
+        tracker = TokenTracker()
+        tracker.start()
+
+        try:
+            llm = BaseCustomLLM("gpt-4", {"temperature": 0.0})
+            llm.call("test prompt")
+
+            # Tokens should be captured
+            input_tokens, output_tokens = tracker.get_counts()
+            assert input_tokens == 50
+            assert output_tokens == 100
+        finally:
+            tracker.stop()
+
+    def test_call_does_not_capture_tokens_without_active_tracker(
+        self, mocker: MockerFixture
+    ) -> None:
+        """Test call does not fail when no TokenTracker is active."""
+        mock_litellm = mocker.patch("lightspeed_evaluation.core.llm.custom.litellm")
+        mocker.patch.dict("os.environ", {})
+
+        # Mock response with usage
+        mock_choice = mocker.Mock()
+        mock_choice.message.content = "Test response"
+        mock_response = mocker.Mock()
+        mock_response.choices = [mock_choice]
+        mock_response.usage = mocker.Mock()
+        mock_response.usage.prompt_tokens = 50
+        mock_response.usage.completion_tokens = 100
+        mock_litellm.completion.return_value = mock_response
+
+        # Ensure no tracker is active
+        temp = TokenTracker()
+        temp.start()
+        temp.stop()
+
+        llm = BaseCustomLLM("gpt-4", {"temperature": 0.0})
+        result = llm.call("test prompt")
+
+        # Should succeed without error
+        assert result == "Test response"

tests/unit/pipeline/evaluation/test_evaluator.py

@@ -882,16 +882,16 @@ def test_token_tracker_start_stop(self) -> None:
         """Test start and stop methods."""
         tracker = TokenTracker()
         tracker.start()
-        assert tracker._callback_registered is True
+        assert TokenTracker.get_active() is tracker
         tracker.stop()
-        assert tracker._callback_registered is False
+        assert TokenTracker.get_active() is None
 
     def test_token_tracker_double_start(self) -> None:
-        """Test calling start twice doesn't register callback twice."""
+        """Test calling start twice doesn't fail."""
         tracker = TokenTracker()
         tracker.start()
         tracker.start()  # Should not fail
-        assert tracker._callback_registered is True
+        assert TokenTracker.get_active() is tracker
         tracker.stop()
 
     def test_token_tracker_double_stop(self) -> None:
@@ -900,7 +900,7 @@ def test_token_tracker_double_stop(self) -> None:
         tracker.start()
         tracker.stop()
         tracker.stop()  # Should not fail
-        assert tracker._callback_registered is False
+        assert TokenTracker.get_active() is None
 
     def test_token_tracker_independent_instances(self) -> None:
         """Test multiple TokenTracker instances are independent."""