Fix UsageTracker AttributeError when using ParallelExecutor with dspy.context (#9095)

* Initial plan

* Fix UsageTracker AttributeError when using ParallelExecutor with dspy.context

Co-authored-by: chenmoneygithub <22925031+chenmoneygithub@users.noreply.github.com>

* Address code review: check new_overrides for usage_tracker instead of parent_overrides

Co-authored-by: chenmoneygithub <22925031+chenmoneygithub@users.noreply.github.com>

* better tests

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: chenmoneygithub <22925031+chenmoneygithub@users.noreply.github.com>
Co-authored-by: chenmoneygithub <chen.qian@databricks.com>

Author: Copilot

dspy/utils/parallelizer.py

@@ -89,10 +89,11 @@ def worker(parent_overrides, submission_id, index, item):
             from dspy.dsp.utils.settings import thread_local_overrides
 
             original = thread_local_overrides.get()
-            token = thread_local_overrides.set({**original, **parent_overrides.copy()})
-            if parent_overrides.get("usage_tracker"):
+            new_overrides = {**original, **parent_overrides.copy()}
+            if new_overrides.get("usage_tracker"):
                 # Usage tracker needs to be deep copied across threads so that each thread tracks its own usage
-                thread_local_overrides.overrides["usage_tracker"] = copy.deepcopy(parent_overrides["usage_tracker"])
+                new_overrides["usage_tracker"] = copy.deepcopy(new_overrides["usage_tracker"])
+            token = thread_local_overrides.set(new_overrides)
 
             try:
                 return index, function(item)

tests/utils/test_usage_tracker.py

@@ -1,3 +1,5 @@
+from unittest import mock
+
 from pydantic import BaseModel
 
 import dspy
@@ -325,3 +327,58 @@ class PromptTokensDetailsWrapper(BaseModel):
     assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["audio_tokens"] == 1
     assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["accepted_prediction_tokens"] == 1
     assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["rejected_prediction_tokens"] == 1
+
+
+def test_parallel_executor_with_usage_tracker():
+    """Test that usage tracking works correctly with ParallelExecutor and mocked LM calls."""
+
+    parent_tracker = UsageTracker()
+
+    # Mock LM with different responses
+    mock_lm = mock.MagicMock(spec=dspy.LM)
+    mock_lm.return_value = ['{"answer": "Mocked answer"}']
+    mock_lm.kwargs = {}
+    mock_lm.model = "openai/gpt-4o-mini"
+
+    dspy.configure(lm=mock_lm, adapter=dspy.JSONAdapter())
+
+    def task1():
+        # Simulate LM usage tracking for task 1
+        dspy.settings.usage_tracker.add_usage(
+            "openai/gpt-4o-mini",
+            {
+                "prompt_tokens": 50,
+                "completion_tokens": 10,
+                "total_tokens": 60,
+            },
+        )
+        return dspy.settings.usage_tracker.get_total_tokens()
+
+    def task2():
+        # Simulate LM usage tracking for task 2 with different values
+        dspy.settings.usage_tracker.add_usage(
+            "openai/gpt-4o-mini",
+            {
+                "prompt_tokens": 80,
+                "completion_tokens": 15,
+                "total_tokens": 95,
+            },
+        )
+        return dspy.settings.usage_tracker.get_total_tokens()
+
+    # Execute tasks in parallel
+    with dspy.context(track_usage=True, usage_tracker=parent_tracker):
+        executor = dspy.Parallel()
+        results = executor([(task1, {}), (task2, {})])
+    # Verify that the two workers had different usage
+    usage1 = results[0]
+    usage2 = results[1]
+
+    # Task 1 should have 50 prompt tokens, task 2 should have 80
+    assert usage1["openai/gpt-4o-mini"]["prompt_tokens"] == 50
+    assert usage1["openai/gpt-4o-mini"]["completion_tokens"] == 10
+    assert usage2["openai/gpt-4o-mini"]["prompt_tokens"] == 80
+    assert usage2["openai/gpt-4o-mini"]["completion_tokens"] == 15
+
+    # Parent tracker should remain unchanged (workers have independent copies)
+    assert len(parent_tracker.usage_data) == 0