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

[Copilot]

📝 Changes Description

dspy.Parallel returns a list of Prediction objects, but Module._set_lm_usage() only handled single Prediction or tuple cases. Additionally, nested module calls did not track usage - only the root module would set usage on its output.

This PR fixes both issues with complementary changes that address different concerns:

1. Refactored _set_lm_usage() to use a recursive approach that finds all Prediction objects in any output structure, making it robust and future-proof for user-defined custom return types
2. Fixed Module.__call__() to properly handle nested module calls by computing token usage deltas for each module invocation

Why both fixes are needed:

• The Module.__call__ fix ensures each module call (root or nested) tracks its usage
• The recursive _set_lm_usage ensures all Prediction objects within a single module's output get usage set (e.g., when returning lists/dicts/nested structures)
• Example: A nested module returning a list from dspy.Parallel requires both fixes to work correctly

class InnerModule(dspy.Module):
    def forward(self, input):
        pred = dspy.Predict("input -> output")
        return pred(input=input)  # Now gets usage set

class OuterModule(dspy.Module):
    def __init__(self):
        self.inner = InnerModule()
    
    def forward(self):
        results = dspy.Parallel(num_threads=2)(
            [(self.inner, {"input": "test1"}),
             (self.inner, {"input": "test2"})]
        )
        # All predictions have usage set, including nested module calls
        return {"results": results, "metadata": {...}}

# Before: nested module predictions had None for usage
# After: All Prediction objects have usage set, including nested modules

Changes:

• Module.call: Modified to always call _set_lm_usage when track_usage is enabled:
  • For root modules (no existing tracker): creates tracker and sets usage
  • For nested modules (tracker exists): captures token delta before/after forward() and sets usage
• _set_lm_usage: Replaced type-specific checks with recursive collect_predictions() function that traverses any data structure
• _compute_token_delta: Added helper method to calculate token usage differences with robust handling of numeric/non-numeric values and missing keys
• Handles single Prediction, lists/tuples, nested structures (e.g., [[Prediction, ...], ...]), dicts with Prediction values, and any combination
• Includes circular reference prevention to avoid infinite recursion
• Added comprehensive tests for nested modules, nested structures, dicts, and mixed combinations

Benefits:

• Nested modules now properly track usage - fixes the core issue where only root modules tracked usage
• Lists/dicts/complex structures from Parallel are handled - fixes the original bug where lists weren't traversed
• Works with any custom return structure users might define
• More maintainable than adding type checks for each new pattern
• Handles edge cases (circular references, non-numeric values, missing keys) gracefully

✅ Contributor Checklist

• Pre-Commit checks are passing (locally and remotely)
• Title of your PR / MR corresponds to the required format
• Commit message follows required format {label}(dspy): {message}

⚠️ Warnings

None. The implementation properly handles both root and nested module usage tracking. All existing tests pass, plus new tests validate nested modules, nested structures, and edge cases. Both fixes (Module.call and recursive _set_lm_usage) are necessary and address orthogonal concerns.

Original prompt

---

This section details on the original issue you should resolve

<issue_title>[Bug] dspy.Parallel inside dspy.Module not working when track_usage is enabled</issue_title>
<issue_description>### What happened?

When I enable track_usage and use dspy.Parallel for parallel predictions inside the dspy.Module forward method, only a list of None is returned without any error.

It works when I call this code without dspy.Module or use the forward method of my module directly. Therefore, I assume, there is a bug in the dspy.Module.__call__ method.

Thank you

Steps to reproduce

import dspy
import os

lm = dspy.LM(
    "azure/gpt-4o",
    api_key=os.environ["LM_API_KEY"],
    api_base=os.environ["LM_API_ENDPOINT"],
    api_version="2024-12-01-preview",
    cache=False,
    temperature=1,
)


def code():
    pred = dspy.Predict("q -> a")
    r = dspy.Parallel(num_threads=1, max_errors=0)(
        [
            (pred, {"q": "President of Czechia"}),
            (pred, {"q": "Capital of Vietnam"}),
        ],
    )
    print(r)
    for a in r:
        print(a.get_lm_usage())
    return r


class Mod(dspy.Module):
    def forward(self):
        return code()


print("Without module")
with dspy.settings.context(lm=lm, track_usage=True):
    code()


print("\nWith module, using __call__")
with dspy.settings.context(lm=lm, track_usage=True):
    r = Mod()()

Output:

Without module
  0%|          | 0/2 [00:00<?, ?it/s]
Processed 2 [/](https://file+.vscode-resource.vscode-cdn.net/) 2 examples: 100%|██████████| 2/2 [00:01<00:00,  1.77it/s]
[Prediction(
    a='The current President of Czechia is Petr Pavel. He assumed office on March 9, 2023.'
), Prediction(
    a='Hanoi'
)]
{'azure/gpt-4o': {'completion_tokens': 33, 'prompt_tokens': 141, 'total_tokens': 174, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0, 'text_tokens': None, 'image_tokens': None}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0, 'text_tokens': None, 'image_tokens': None}}}
{'azure/gpt-4o': {'completion_tokens': 14, 'prompt_tokens': 140, 'total_tokens': 154, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0, 'text_tokens': None, 'image_tokens': None}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0, 'text_tokens': None, 'image_tokens': None}}}

With module, using __call__
  0%|          | 0/2 [00:00<?, ?it/s]
[None, None]

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
...
AttributeError: 'NoneType' object has no attribute 'get_lm_usage'

DSPy version

3.1.0</issue_description>

Comments on the Issue (you are @copilot in this section)

• Fixes [Bug] dspy.Parallel inside dspy.Module not working when track_usage is enabled #9201

---

💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.

[TomeHirata]

@copilot I think this fix is fragile as there are many other types that users' custom modules may return. It it difficult to set usage for all module output when it returns dspy.Prediction? The usage is the total token usage during the module.__call__ execution.

[Copilot]

@copilot I think this fix is fragile as there are many other types that users' custom modules may return. It it difficult to set usage for all module output when it returns dspy.Prediction? The usag...

Good point! I've refactored the implementation to use a recursive approach that finds all Prediction objects in any structure (lists, tuples, dicts, nested combinations). This makes it robust for any custom return types users might define.

The new implementation traverses the entire output structure recursively and sets usage on all Prediction objects found, regardless of how they're nested. I've added tests for nested lists, dicts, and mixed structures to verify it works correctly.

Fixed in commit adc30e3.

[TomeHirata]

@copilot I think we should fix Module.__call__ instead of _set_lm_usage. Currently, _set_lm_usage is called only on the root module call and no usage is added to the nested module call.

[Copilot]

You're absolutely right! I've fixed Module.__call__ to handle nested module calls. Previously, only the root module would call _set_lm_usage. Now when track_usage is enabled:

1. If no tracker exists: creates one and sets usage (root module case)
2. If tracker exists: captures token delta before/after forward() and sets usage on output (nested module case)

This ensures all module calls (root and nested) properly track and set usage on their Prediction outputs.

Fixed in commit e8c0ee0.

[TomeHirata]

@copilot do we still need the change for _set_lm_usage?

[Copilot]

Yes, we still need the recursive _set_lm_usage because both fixes address different concerns:

1. Module.call fix (commit e8c0ee0): Ensures each module call (root or nested) gets its usage tracked

• Handles: OuterModule() → calls inner_module() → both get usage set

2. Recursive _set_lm_usage (commits 7f381e8-adc30e3): Ensures all Predictions in a single module's output get usage set

• Handles: Module returns [Prediction, Prediction] from dspy.Parallel → all get usage set
• Original only handled single Prediction or tuple, so lists would get the warning and no usage

The original bug was a module returning a list of Predictions from Parallel. Without the recursive _set_lm_usage, even with the call fix, the list wouldn't be traversed and individual Predictions wouldn't get usage set.

Both fixes are needed for the complete solution.