fix(dspy): Fix Module.__call__ to enable nested module usage tracking with track_usage

dspy/primitives/module.py

@@ -72,14 +72,39 @@ def __call__(self, *args, **kwargs) -> Prediction:
 
         with settings.context(caller_modules=caller_modules):
             if settings.track_usage and thread_local_overrides.get().get("usage_tracker") is None:
+                # No existing tracker - create a new one for this module call
                 with track_usage() as usage_tracker:
                     output = self.forward(*args, **kwargs)
                 tokens = usage_tracker.get_total_tokens()
                 self._set_lm_usage(tokens, output)
 
                 return output
-
-            return self.forward(*args, **kwargs)
+            elif settings.track_usage:
+                # Tracker exists (nested call) - capture usage for this specific module call
+                usage_tracker = thread_local_overrides.get().get("usage_tracker")
+                # Get current token count before calling forward
+                tokens_before = usage_tracker.get_total_tokens()
+                output = self.forward(*args, **kwargs)
+                # Get tokens used by this module call
+                tokens_after = usage_tracker.get_total_tokens()
+
+                # Calculate the delta (tokens used by this module call only)
+                tokens_delta = {}
+                # Process all LM names from both before and after
+                all_lm_names = set(tokens_before.keys()) | set(tokens_after.keys())
+                for lm_name in all_lm_names:
+                    if lm_name in tokens_before and lm_name in tokens_after:
+                        # Compute difference for each token type
+                        tokens_delta[lm_name] = self._compute_token_delta(tokens_before[lm_name], tokens_after[lm_name])
+                    elif lm_name in tokens_after:
+                        # New LM used in this call
+                        tokens_delta[lm_name] = tokens_after[lm_name]
+                    # If lm_name only in tokens_before, it means no usage in this call, so skip
+
+                self._set_lm_usage(tokens_delta, output)
+                return output
+            else:
+                return self.forward(*args, **kwargs)
 
     @with_callbacks
     async def acall(self, *args, **kwargs) -> Prediction:
@@ -181,18 +206,81 @@ def batch(
             results = parallel_executor.forward(exec_pairs)
             return results
 
+    def _compute_token_delta(self, tokens_before: dict[str, Any], tokens_after: dict[str, Any]) -> dict[str, Any]:
+        """Compute the difference in token usage between two snapshots.
+
+        Args:
+            tokens_before: Token usage before the operation
+            tokens_after: Token usage after the operation
+
+        Returns:
+            Dictionary with the delta (difference) in token usage
+        """
+        delta = {}
+        # Process all keys from both dicts
+        all_keys = set(tokens_before.keys()) | set(tokens_after.keys())
+        for key in all_keys:
+            value_before = tokens_before.get(key)
+            value_after = tokens_after.get(key)
+
+            if isinstance(value_after, dict) or isinstance(value_before, dict):
+                # Recursively compute delta for nested dicts
+                delta[key] = self._compute_token_delta(
+                    value_before if isinstance(value_before, dict) else {},
+                    value_after if isinstance(value_after, dict) else {}
+                )
+            elif isinstance(value_after, (int, float)) or isinstance(value_before, (int, float)):
+                # Compute numeric difference only for numeric types
+                delta[key] = (value_after or 0) - (value_before or 0)
+            elif value_after is not None:
+                # For non-numeric types (strings, etc.), just use the after value
+                delta[key] = value_after
+        return delta
+
     def _set_lm_usage(self, tokens: dict[str, Any], output: Any):
-        # Some optimizers (e.g., GEPA bootstrap tracing) temporarily patch
-        # module.forward to return a tuple: (prediction, trace).
-        # When usage tracking is enabled, ensure we attach usage to the
-        # prediction object if present.
-        prediction_in_output = None
-        if isinstance(output, Prediction):
-            prediction_in_output = output
-        elif isinstance(output, tuple) and len(output) > 0 and isinstance(output[0], Prediction):
-            prediction_in_output = output[0]
-        if prediction_in_output:
-            prediction_in_output.set_lm_usage(tokens)
+        # When usage tracking is enabled, recursively find all Prediction objects
+        # in the output and set usage on them. This handles any structure:
+        # - Single Prediction
+        # - List/tuple of Predictions
+        # - Nested structures (e.g., [[Prediction, ...], ...])
+        # - Tuples from optimizers (e.g., (prediction, trace))
+        # - Dicts with Prediction values
+        # - Any combination of the above
+        predictions_in_output = []
+
+        def collect_predictions(obj, seen=None):
+            """Recursively collect all Prediction objects from any structure.
+
+            Args:
+                obj: The object to search for Predictions (can be any type)
+                seen: Set of object ids already visited (prevents infinite recursion on circular refs)
+
+            Side effects:
+                Appends found Prediction objects to predictions_in_output list in outer scope
+            """
+            if seen is None:
+                seen = set()
+
+            # Avoid infinite recursion on circular references
+            obj_id = id(obj)
+            if obj_id in seen:
+                return
+            seen.add(obj_id)
+
+            if isinstance(obj, Prediction):
+                predictions_in_output.append(obj)
+            elif isinstance(obj, (list, tuple)):
+                for item in obj:
+                    collect_predictions(item, seen)
+            elif isinstance(obj, dict):
+                for value in obj.values():
+                    collect_predictions(value, seen)
+
+        collect_predictions(output)
+
+        if predictions_in_output:
+            for prediction in predictions_in_output:
+                prediction.set_lm_usage(tokens)
         else:
             logger.warning("Failed to set LM usage. Please return `dspy.Prediction` object from dspy.Module to enable usage tracking.")
 

tests/predict/test_parallel.py

@@ -194,3 +194,156 @@ def forward(self, value: int) -> str:
     assert len(exceptions) == 1
     assert isinstance(exceptions[0], ValueError)
     assert str(exceptions[0]) == "test error"
+
+
+def test_parallel_with_track_usage():
+    """Test that usage tracking works correctly when using Parallel inside a Module."""
+    lm = DummyLM(
+        [
+            {"output": "test output 1"},
+            {"output": "test output 2"},
+        ]
+    )
+
+    class MyModule(dspy.Module):
+        def forward(self):
+            predictor = dspy.Predict("input -> output")
+            parallel = dspy.Parallel(num_threads=2)
+            return parallel(
+                [
+                    (predictor, {"input": "test input 1"}),
+                    (predictor, {"input": "test input 2"}),
+                ]
+            )
+
+    # Test with track_usage enabled
+    with dspy.settings.context(lm=lm, track_usage=True):
+        results = MyModule()()
+
+    # Verify that results are Prediction objects with usage data set
+    assert len(results) == 2
+    assert all(isinstance(r, dspy.Prediction) for r in results)
+    assert all(r.get_lm_usage() is not None for r in results)
+
+    # Verify predictions have the expected outputs
+    expected_outputs = {"test output 1", "test output 2"}
+    assert {r.output for r in results} == expected_outputs
+
+
+def test_parallel_with_track_usage_and_tuple_return():
+    """Test that usage tracking works when Module returns a tuple with list of Predictions."""
+    lm = DummyLM(
+        [
+            {"output": "test output 1"},
+            {"output": "test output 2"},
+        ]
+    )
+
+    class MyModule(dspy.Module):
+        def forward(self):
+            predictor = dspy.Predict("input -> output")
+            parallel = dspy.Parallel(num_threads=2)
+            results = parallel(
+                [
+                    (predictor, {"input": "test input 1"}),
+                    (predictor, {"input": "test input 2"}),
+                ]
+            )
+            # Return tuple mimicking optimizer behavior (e.g., GEPA bootstrap tracing)
+            return (results, {"trace": "some_trace_data"})
+
+    # Test with track_usage enabled
+    with dspy.settings.context(lm=lm, track_usage=True):
+        output = MyModule()()
+
+    # Verify that output is a tuple
+    assert isinstance(output, tuple)
+    assert len(output) == 2
+    results, trace = output
+
+    # Verify that results are Prediction objects with usage data set
+    assert len(results) == 2
+    assert all(isinstance(r, dspy.Prediction) for r in results)
+    assert all(r.get_lm_usage() is not None for r in results)
+
+    # Verify predictions have the expected outputs
+    expected_outputs = {"test output 1", "test output 2"}
+    assert {r.output for r in results} == expected_outputs
+
+    # Verify trace is preserved
+    assert trace == {"trace": "some_trace_data"}
+
+
+def test_track_usage_with_nested_and_dict_structures():
+    """Test that usage tracking works with nested lists, dicts, and mixed structures."""
+    lm = DummyLM([{"output": "test1"}, {"output": "test2"}, {"output": "test3"}])
+
+    class NestedStructureModule(dspy.Module):
+        def forward(self):
+            predictor = dspy.Predict("input -> output")
+            p1 = predictor(input="input1")
+            p2 = predictor(input="input2")
+            p3 = predictor(input="input3")
+            # Return a complex nested structure
+            return {
+                "nested_list": [[p1, p2]],
+                "single": p3,
+            }
+
+    # Test with track_usage enabled
+    with dspy.settings.context(lm=lm, track_usage=True):
+        result = NestedStructureModule()()
+
+    # Verify structure
+    assert isinstance(result, dict)
+    assert "nested_list" in result
+    assert "single" in result
+
+    # Verify all predictions have usage data
+    all_predictions = result["nested_list"][0] + [result["single"]]
+    assert len(all_predictions) == 3
+    assert all(isinstance(p, dspy.Prediction) for p in all_predictions)
+    assert all(p.get_lm_usage() is not None for p in all_predictions)
+
+    # Verify predictions have the expected outputs
+    expected_outputs = {"test1", "test2", "test3"}
+    assert {p.output for p in all_predictions} == expected_outputs
+
+
+def test_track_usage_with_nested_modules():
+    """Test that usage tracking works correctly for nested module calls."""
+    lm = DummyLM([{"output": "inner1"}, {"output": "inner2"}])
+
+    class InnerModule(dspy.Module):
+        def forward(self, input):
+            predictor = dspy.Predict("input -> output")
+            result = predictor(input=input)
+            # Verify usage is set inside the nested module
+            assert result.get_lm_usage() is not None, "Nested module should have usage set"
+            return result
+
+    class OuterModule(dspy.Module):
+        def __init__(self):
+            super().__init__()
+            self.inner = InnerModule()
+
+        def forward(self):
+            r1 = self.inner(input="test1")
+            r2 = self.inner(input="test2")
+            # Verify usage is set on nested module results
+            assert r1.get_lm_usage() is not None, "Nested call result should have usage set"
+            assert r2.get_lm_usage() is not None, "Nested call result should have usage set"
+            return [r1, r2]
+
+    # Test with track_usage enabled
+    with dspy.settings.context(lm=lm, track_usage=True):
+        results = OuterModule()()
+
+    # Verify all results have usage data
+    assert len(results) == 2
+    assert all(isinstance(r, dspy.Prediction) for r in results)
+    assert all(r.get_lm_usage() is not None for r in results)
+
+    # Verify predictions have the expected outputs
+    expected_outputs = {"inner1", "inner2"}
+    assert {r.output for r in results} == expected_outputs